| 引用本文: | 舒善富,刘超,孙毓忠,张洪宇,高翠芸,张小洪.基于自适应知识蒸馏的代码大模型轻量化.软件学报,2026,37(6):2431-2454 |
| |
|
| |
|
|
| 本文已被：浏览 1898次 下载 1065次 |
码上扫一扫！ |
|
|
| 基于自适应知识蒸馏的代码大模型轻量化 |
|
舒善富1, 刘超1, 孙毓忠2, 张洪宇1, 高翠芸3, 张小洪1
|
|
1.重庆大学 大数据与软件学院, 重庆 401331;2.中国科学院 计算技术研究所, 北京 100190;3.哈尔滨工业大学(深圳) 计算机科学与技术学院, 广东 深圳 518055
|
|
| 摘要: |
| 以大语言模型(large language model, LLM)为基座的软件编程助手(如Copilot), 能够显著提升程序员开发效率, 但LLM的计算和存储需求大、本地化部署难. 构建轻量化小参数LLM能够满足计算、存储、部署需求, 但其代码生成的精度损失比大参数LLM 大. 知识蒸馏(knowledge distillation, KD)技术, 让小参数LLM (学生模型)在目标训练数据集上拟合大参数LLM (教师模型)的生成分布, 降低代码生成精度损失. 人工智能领域前沿的KD技术基于Kullback-Leibler (KL)散度损失函数, 度量并缩小因学生/教师模型的生成分布差异导致的精度损失, 但学生模型难以学习教师模型的趋零分布区域. 随后, 学者利用反向KL散度损失(RKL)函数解决该趋零分布区域的学习问题. 研究发现, RKL在高概率分布区域存在学习问题, 与KL散度损失函数存在互补性; 对于一些数据, 教师模型生成质量低, 导致学生模型学习效果差. 提出一种自适应知识蒸馏 (adaptive knowledge distillation, AKD)方法, 通过prompt提升教师模型的生成质量, 并构造自适应损失函数, 根据学生/教师模型之间的生成分布差异自适应调整学习的优先级, 确保学生模型在主要概率区域和趋零概率区域均具备学习能力. 基于AKD方法, 利用StarCoder-1B/7B (学生/教师模型)和CodeAlpaca数据, 训练了轻量化代码大模型, 并评估代码大模型的精度损失及生成代码的质量问题. 实验结果显示, 轻量化代码大模型规模降低85.7%, 在HumanEval和MBPP数据集上, 任务提示明确的prompt可以提高教师模型的代码生成质量, 使训练的学生模型降低6%的平均精度损失; AKD方法训练的模型较教师模型(StarCoder-7B)的平均精度损失为17.14%, 较原始学生模型平均降低30.6%; AKD方法训练的模型较前沿的KD和RKD方法的精度损失平均降低19.9%; 关于推理显存需求情况, KD和RKD方法需要54.7 GB, 而AKD方法仅增加3 GB. 关于训练时间方面, AKD方法所需训练时间增加30%; 相较而言, 即使KD和RKD方法训练至相同时长, 它们的平均效果仅提升3%, 相比AKD方法低16.9%. 因此, AKD方法增加的训练成本是值得的. 此外, 将AKD方法应用到Code Llama和CodeGen系列模型, 相较前沿的KD及RKD方法的精度损失平均降低19.2%, 证明了AKD方法的泛化能力. |
| 关键词: 代码生成 大语言模型 知识蒸馏 |
| DOI：10.13328/j.cnki.jos.007462 |
| 分类号:TP311 |
| 基金项目:国家自然科学基金(62202074, 62372071); 中国博士后科学基金(2022M710519); 重庆市技术创新与应用发展专项重点项目(CSTB2023TIAD-STX0015, CSTB2022TIAD-KPX0068); 重庆市出站留(来)渝博士后择优资助项目(2021LY23) |
|
| Adaptive Knowledge Distillation for Lightweight Large Code Models |
|
SHU Shan-Fu1, LIU Chao1, SUN Yu-Zhong2, ZHANG Hong-Yu1, GAO Cui-Yun3, ZHANG Xiao-Hong1
|
|
1.School of Big Data & Software Engineering, Chongqing University, Chongqing 401331, China;2.Institute of Computing Technology, Chinese Academy of Sciences, Beijing 100190, China;3.School of Computer Science and Technology, Harbin Institute of Technology, Shenzhen, Shenzhen 518055, China
|
| Abstract: |
| Software programming assistants based on large language models (LLMs), such as Copilot, significantly enhance programmer productivity. However, LLMs have large computing and storage requirements and are difficult to deploy locally. Building a lightweight, small LLM can meet computing, storage, and deployment requirements, but it leads to a greater accuracy loss in code generation compared to large LLMs. Knowledge distillation (KD) techniques allow small LLMs (student models) to approximate the output distributions of large LLMs (teacher models) on target training datasets, thus reducing accuracy loss in code generation. Cutting-edge KD techniques in artificial intelligence are based on the Kullback-Leibler (KL) divergence loss function, which measures and reduces accuracy loss due to discrepancies in the output distributions between student and teacher models. However, student models struggle to learn in the near-zero distribution regions of teacher models. Consequently, researchers have employed the reverse KL (RKL) divergence loss function to address this issue in near-zero distribution regions. This study finds that RKL faces learning challenges in high-probability distribution regions and complements the KL divergence loss function. For some datasets, low-quality outputs from teacher models lead to poor learning outcomes for the student models. This study proposes an adaptive knowledge distillation (AKD) method that uses prompts to enhance teacher model output quality and constructs an adaptive loss function to adjust learning priorities based on the distributional differences between student and teacher models. This ensures the student model effectively learns in both primary and near-zero probability regions. Using the AKD method, this study trains a lightweight large code model based on StarCoder-1B/7B (student/teacher models) and the CodeAlpaca dataset, evaluating accuracy loss and code quality issues. Experimental results show that the lightweight model size is reduced by 85.7%. On the HumanEval and MBPP data sets, prompts with clear instructions improve teacher model code generation quality, reducing the average accuracy loss of the trained student model by 6%. The AKD-trained model’s average accuracy loss compared to the teacher model (StarCoder-7B) is 17.14%, a 30.6% reduction over the original student model. The AKD-trained model’s accuracy loss is reduced by an average of 19.9% compared to state-of-the-art KD and RKD methods. Regarding inference memory requirements, the KD and RKD methods require 54.7 GB, while the AKD method only adds 3 GB. In terms of training time, the AKD method incurs a 30% increase. However, even when the KD and RKD methods are trained for the same duration, their average performance improves by only 3%, which is 16.9% lower than that of the AKD method. Therefore, the additional training cost of the AKD method is justified. Moreover, applying the AKD method to the Code Llama and CodeGen model series reduces accuracy loss by an average of 19.2% compared to state-of-the-art KD and RKD methods, demonstrating the generalizability of the AKD method. |
| Key words: code generation large language model (LLM) knowledge distillation (KD) |
|
|
|
|
