arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2602.02891cs.LGcs.CL
收藏

TraceNAS: 通过梯度轨迹相关性实现零样本LLM剪枝

TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

AI总结:

TraceNAS通过梯度轨迹相关性实现零样本LLM剪枝,高效发现高保真剪枝模型,减少GPU计算成本并提升模型性能。

摘要:

结构化剪枝对于高效部署大语言模型(LLM)至关重要。LLM子块对剪枝的敏感性差异需要识别最优的非均匀剪枝模型。现有方法单独评估层、注意力头或权重通道的重要性。这种局部聚焦忽略了模型中复杂的全局结构依赖关系。训练感知的结构化剪枝解决了全局依赖关系,但其计算成本可能与后剪枝训练一样高。为减轻训练感知剪枝的计算负担并捕捉全局结构依赖关系,我们提出TraceNAS,一种无需训练的神经架构搜索(NAS)框架,联合探索LLM深度和宽度的结构化剪枝。TraceNAS通过缩放不变的零样本代理识别保持与预训练模型高度损失景观一致的剪枝模型,有效选择在后剪枝训练中表现潜力最大的模型。TraceNAS效率极高,能够在单块GPU上8.5小时内发现高保真的剪枝模型,GPU小时数比训练感知方法减少10倍。在Llama和Qwen系列上的评估表明,TraceNAS在常识和推理基准上与训练感知基线具有竞争力。

英文摘要:

Structured pruning is essential for efficient deployment of Large Language Models (LLMs). The varying sensitivity of LLM sub-blocks to pruning necessitates the identification of optimal non-uniformly pruned models. Existing methods evaluate the importance of layers, attention heads, or weight channels in isolation. Such localized focus ignores the complex global structural dependencies that exist across the model. Training-aware structured pruning addresses global dependencies, but its computational cost can be just as expensive as post-pruning training. To alleviate the computational burden of training-aware pruning and capture global structural dependencies, we propose TraceNAS, a training-free Neural Architecture Search (NAS) framework that jointly explores structured pruning of LLM depth and width. TraceNAS identifies pruned models that maintain a high degree of loss landscape alignment with the pretrained model using a scale-invariant zero-shot proxy, effectively selecting models that exhibit maximal performance potential during post-pruning training. TraceNAS is highly efficient, enabling high-fidelity discovery of pruned models on a single GPU in 8.5 hours, yielding a 10$\times$ reduction in GPU-hours compared to training-aware methods. Evaluations on the Llama and Qwen families demonstrate that TraceNAS is competitive with training-aware baselines across commonsense and reasoning benchmarks.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑