NIRVANA通过新方法实现大模型结构化压缩,不降性能还能直接微调。
NIRVANA: Structured Pruning Reimagined for Large Language Model Compression
- 用函数空间敏感度替代传统损失,保留训练动态。
- 在高压缩率下仍保持性能,优于现有方法。
- 适合需要高效推理和快速微调的部署场景。
结构化剪枝是加速大语言模型推理的有效途径,但现有方法常导致性能显著下降,且需重新训练恢复能力。为此,我们提出NIRVANA,一种新型硬件感知的结构化剪枝框架,可同时保持零样本性能和下游微调的优化空间。不同于传统的基于损失的启发式方法,我们的方法借鉴神经切线核(NTK)思想,通过一阶函数空间灵敏度评估结构重要性,有效保护模型关键训练动态。为防止高压缩率下的结构坍塌,引入全局单元排序策略与解析导出的分配机制,合理平衡注意力头与MLP神经元的剪枝强度。此外,通过轻量级KL散度驱动的数据校准选择,消除随机采样带来的不稳定性。在Llama3、Qwen和T5架构上的广泛实验表明,NIRVANA在多个基准上持续取得新最优结果,提供了一种理论严谨且实用的大模型压缩方案。代码已开源:https://github.com/iDEA-iSAIL-Lab-UIUC/NIRVANA。
原文摘要 · Abstract (English)
While structured pruning presents a highly effective pathway for accelerating Large Language Model (LLM) inference, existing methods frequently suffer from significant performance degradation and demand computationally retraining to recover capabilities. To overcome these barriers, we present NIRVANA, a novel, hardware-aware structured pruning framework designed to preserve both zero-shot performance and the optimization landscape for downstream fine-tuning. Departing from traditional loss-based heuristics, our approach evaluates structural importance through a first-order function-space saliency inspired by the Neural Tangent Kernel (NTK), effectively safeguarding the model's critical training dynamics. To prevent structural collapse at high compression rates, we introduce a global unit-ranking strategy coupled with an analytically derived allocation mechanism, which optimally balances the pruning aggressiveness between attention heads and MLP neurons. Furthermore, we eliminate the instability typically associated with random data sampling by employing a lightweight, KL-divergence-driven calibration data selection process. Extensive evaluations across Llama3, Qwen, and T5 architectures demonstrate that NIRVANA consistently establishes new state-of-the-art results on different benchmarks, providing a theoretically sound and practical approach to LLM compression. The code is available at https://github.com/iDEA-iSAIL-Lab-UIUC/NIRVANA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。