LEANCODE: Understanding Models Better for Code Simplification of Pre-trained Large Language Models

📄 arXiv: 2505.14759v3 📥 PDF

作者: Yan Wang, Ling Ding, Tien N Nguyen, Shaohua Wang, Yanan Zheng

分类: cs.SE, cs.LG

发布日期: 2025-05-20 (更新: 2025-06-08)

备注: Accepted to ACL 2025 main conference

💡 一句话要点

LeanCode：利用注意力机制进行代码简化，加速预训练大语言模型

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 代码简化 注意力机制 预训练语言模型 代码搜索 代码摘要 模型压缩 计算效率

📋 核心要点

  1. 代码大语言模型计算复杂度高，尤其在处理长代码序列时，效率成为瓶颈。
  2. LeanCode利用代码上下文的注意力分数，选择性移除不重要的token，从而简化代码。
  3. 实验表明，LeanCode在代码搜索和摘要任务上显著优于现有SOTA方法DietCode和Slimcode。

📝 摘要（中文）

针对代码的大语言模型通常具有显著的计算复杂度，并且随着输入代码序列长度的增加而显著增长。我们提出了LeanCode用于代码简化，以减少训练和预测时间，它利用代码上下文中的注意力分数来表示token的重要性。我们提倡基于上下文感知的平均注意力分数选择性地移除token，而不是使用跨所有输入的平均分数。LeanCode使用编码器中CLS token的注意力分数来处理分类任务，例如代码搜索。它还采用编码器-解码器注意力分数来确定序列到序列任务（如代码摘要）中token的重要性。我们的评估表明，LeanCode优于SOTA方法DietCode和Slimcode，在代码搜索方面分别提高了60%和16%，在代码摘要方面分别提高了29%和27%。

🔬 方法详解

问题定义：现有代码大语言模型计算复杂度高，训练和预测耗时。尤其对于长代码序列，计算成本显著增加。DietCode和Slimcode等现有代码简化方法效果有限，仍有提升空间。

核心思路：LeanCode的核心在于利用注意力机制，更准确地评估代码token的重要性。通过上下文感知的注意力分数，选择性地移除对模型性能影响较小的token，从而在不显著降低模型性能的前提下，减少计算量。

技术框架：LeanCode主要针对两种任务：代码搜索（分类任务）和代码摘要（序列到序列任务）。对于代码搜索，LeanCode使用编码器中CLS token的注意力分数来判断token的重要性。对于代码摘要，LeanCode使用编码器-解码器注意力分数来确定token的重要性。整体流程包括：1) 计算注意力分数；2) 根据注意力分数对token进行排序；3) 选择性移除低重要性token；4) 使用简化后的代码进行训练或预测。

关键创新：LeanCode的关键创新在于使用上下文感知的注意力分数。与现有方法（如使用跨所有输入的平均注意力分数）不同，LeanCode考虑了特定代码上下文中的token重要性，从而更准确地识别可以安全移除的token。这种方法能够更好地保留代码的关键信息，避免过度简化导致性能下降。

关键设计：LeanCode的关键设计包括：1) 使用预训练语言模型的注意力头输出作为token重要性度量；2) 针对不同任务（分类和序列到序列）采用不同的注意力分数计算方式；3) 设计了token移除比例的超参数，用于控制简化程度。具体的损失函数和网络结构沿用了预训练模型的设计，没有进行显著修改。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LeanCode在代码搜索任务上比DietCode和Slimcode分别提高了60%和16%，在代码摘要任务上分别提高了29%和27%。这些显著的性能提升证明了LeanCode在代码简化方面的有效性。此外，实验还验证了LeanCode在不同代码数据集上的泛化能力。

🎯 应用场景

LeanCode具有广泛的应用前景，可用于加速代码搜索、代码摘要、代码生成等任务。通过降低计算复杂度，LeanCode可以使大型代码语言模型更容易部署在资源受限的设备上，并提高开发效率。此外，LeanCode还可以用于代码混淆和代码安全等领域，通过移除不重要的token来隐藏代码的真实意图。

📄 摘要（原文）

Large Language Models for code often entail significant computational complexity, which grows significantly with the length of the input code sequence. We propose LeanCode for code simplification to reduce training and prediction time, leveraging code contexts in utilizing attention scores to represent the tokens' importance. We advocate for the selective removal of tokens based on the average context-aware attention scores rather than average scores across all inputs. LeanCode uses the attention scores of `CLS' tokens within the encoder for classification tasks, such as code search. It also employs the encoder-decoder attention scores to determine token significance for sequence-to-sequence tasks like code summarization. Our evaluation shows LeanCode's superiority over the SOTAs DietCode and Slimcode, with improvements of 60% and 16% for code search, and 29% and 27% for code summarization, respectively.