本文提出Flab-Pruner，一种结合词汇、层和FFN剪枝的统一结构剪枝方法，通过KL散度优化和自定义微调策略，在减少代码LLM参数的同时保持高性能和效率。
Large Language Model, Efficiency, Generative AI, Code Generation, Fine-tuning, Pre-training
Guang Yang, Yu Zhou, Xiangyu Zhang, Wei Cheng, Ke Liu, Xiang Chen, Terry Yue Zhuo, Taolue Chen
Nanjing University of Aeronautics and Astronautics, Singapore Management University, National University of Defense Technology, Nantong University, Monash University, CSIRO’s Data61, Birkbeck University of London
Generated by grok-3-mini-latest
Background Problem
大型语言模型（LLMs）在代码生成任务中的广泛应用引发了对高计算需求和能源消耗的担忧，这与绿色软件工程的目标相悖。现有的结构剪枝方法主要针对分类模型设计，处理低维分类对数，而生成式代码LLMs处理高维令牌对数序列，导致传统剪枝目标不匹配。此外，现有方法通常采用单组件剪枝策略，忽略了多颗粒度剪枝的协同效应，且缺乏针对代码任务的特定微调策略，限制了剪枝后性能的恢复效率。本文通过分析这些问题，提出一种针对生成式代码LLMs的统一剪枝方法，以减少参数数量、降低计算开销并维持性能。
Method
核心思想：Flab-Pruner是一种统一结构剪枝方法，旨在最小化模型参数同时保持代码生成性能，通过定义KL散度最小化作为剪枝目标，确保剪枝后模型的令牌生成概率分布与原模型相似。具体实现包括三个组件：
- 词汇剪枝：基于代码语料中令牌的使用频率，移除低频令牌，减少嵌入矩阵和输出层的大小。
- 层剪枝：迭代评估并移除对KL散度影响最小的层，使用贪婪算法逐步优化模型深度。
- FFN剪枝：针对前馈网络的中间层神经元，使用启发式规则（如Top-K、Bottom-K等）选择移除神经元，以减少计算复杂度。 性能恢复策略：引入自定义代码指令数据，通过原模型生成高质量代码替换训练数据，并使用LoRA技术进行高效微调。
Experiment
实验在三个先进的代码LLM（CodeQwen-1.5、NxCode和CodeSlerp）上进行，涵盖代码生成、链式思考（CoT）生成和代码输出预测三项任务。数据集包括HumanEval、OpenEval、CodeHarmony（自建数据集，包含15,800个训练样本）和Crux-O等，实验设置旨在全面评估剪枝方法的有效性、效率和鲁棒性。剪枝比例为22%，结果显示性能保留约97%，微调后性能达到或超过原模型（如HumanEval上的Pass@1从77.44%提升至78.05%）。实验设计合理，考虑了不同任务和模型的多样性，验证了KL散度目标的优越性（与基线方法如ShortGPT相比，性能提升显著）。效率分析显示GPU使用减少约21%、FLOPs降低20%、CO2排放减少14%，并与量化方法兼容。鲁棒性测试（ReCode和EvoEval扰动）表明剪枝后模型在部分场景下性能略降，但微调后恢复或提升，符合预期。
Further Thoughts
本文的Flab-Pruner方法启发我们思考如何将结构剪枝扩展到其他AI领域，如结合量化或知识蒸馏进一步提升模型效率；同时，针对代码任务的自定义微调策略可借鉴到其他特定领域中，以提高剪枝后模型的泛化能力；未来，可以探索多语言支持或动态剪枝机制，以推动AI可持续发展，并减少碳排放；此外，与其他研究如DeepSeek的推理数据蒸馏相比，本文强调了任务特定优化的重要性，可能在跨领域应用中提供新思路。