arXiv:2503.03008cs.CLcs.AI2025-03AAAI被引 2

通过自蒸馏提升模型浅层表征,实现代码检索的高效灵活部署。

MoSE: Hierarchical Self-Distillation Enhances Early Layer Embeddings

  • 用自蒸馏让深层指导浅层,增强早期特征表达。
  • 在代码检索任务中,不同层级退出头均保持高精度。
  • 适合对延迟敏感的代码理解场景,如实时搜索工具。

部署语言模型常面临准确率与性能间的权衡,需在延迟约束下保持实用价值。传统模型蒸馏虽能缩小规模,但需训练独立模型,成本高昂。我们提出模块化多出口编码器MoSE(10亿参数),采用新型自蒸馏机制,显著提升低层表征质量,支持按需部署不同模型部分,实现性能与成本的灵活平衡。该架构通过在训练中由高层引导低层,优化中间表示,附加开销极小。进一步引入仓库级上下文损失,最大化训练上下文窗口利用率。此外,我们发布一个通过代码翻译构建的新数据集,扩展了文本到代码基准,包含跨语言代码对。实验表明,自蒸馏是一种在多种代码理解任务中有效权衡推理成本与准确性的原则性方法。

原文摘要 · Abstract (English)

Deploying language models often requires navigating accuracy vs. performance trade-offs to meet latency constraints while preserving utility. Traditional model distillation reduces size but incurs substantial costs through training separate models. We introduce ModularStarEncoder (MoSE), a 1-billion-parameter multi-exit encoder for code retrieval and classification that employs a novel Self-Distillation mechanism. This approach significantly enhances lower-layer representations, enabling flexible deployment of different model portions with favorable performance trade-offs. Our architecture improves text-to-code and code-to-code search by targeting specific encoder layers as exit heads, where higher layers guide earlier ones during training, thereby improving intermediate representations at minimal additional cost. We further enhance MoSE with a repository-level contextual loss that maximizes training context window utilization. Additionally, we release a new dataset created through code translation that extends text-to-code benchmarks with cross-language code-to-code pairs. Evaluations demonstrate the effectiveness of Self-Distillation as a principled approach to trading inference cost for accuracy across various code understanding tasks.

代码检索自蒸馏多出口模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。