C2LLM Technical Report: A New Frontier in Code Retrieval via Adaptive Cross-Attention Pooling
作者: Jin Qin, Zihan Liao, Ziyin Zhang, Hang Yu, Peng Di, Rui Wang
分类: cs.CL, cs.AI
发布日期: 2025-12-24
💡 一句话要点
C2LLM:通过自适应跨注意力池化实现代码检索的新突破
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 代码检索 代码大语言模型 多头注意力池化 对比学习 代码嵌入
📋 核心要点
- 现有基于EOS token的序列embedding方法存在信息瓶颈,无法有效利用所有token信息。
- C2LLM采用多头注意力池化(PMA)模块,聚合所有token信息,打破信息瓶颈,并灵活调整embedding维度。
- C2LLM在MTEB-Code上取得了显著成果,C2LLM-7B在同等规模模型中排名第一,验证了方法的有效性。
📝 摘要(中文)
本文介绍了C2LLM,一个对比代码大语言模型家族,包含0.5B和7B两种规模。C2LLM基于Qwen-2.5-Coder作为backbone,采用多头注意力池化(Pooling by Multihead Attention, PMA)模块从token embeddings生成序列embedding,有效地:1)利用了LLM在预训练期间获得的因果表示;2)能够聚合序列中所有token的信息,打破了基于EOS的序列embedding的信息瓶颈;3)支持embedding维度的灵活调整,作为MRL的替代方案。C2LLM模型在三百万公开数据上进行训练,在同等规模的模型中,于MTEB-Code上创造了新的记录,其中C2LLM-7B在总排行榜上排名第一。
🔬 方法详解
问题定义:代码检索任务旨在根据给定的查询语句,从代码库中检索出相关的代码片段。现有方法,特别是基于大语言模型的方法,通常依赖于EOS (End-of-Sequence) token的embedding来表示整个代码序列,这导致了严重的信息瓶颈,因为EOS token的embedding需要压缩整个序列的信息,无法充分利用序列中所有token的信息。此外,现有方法在embedding维度调整方面不够灵活。
核心思路:C2LLM的核心思路是利用多头注意力池化(PMA)模块,将序列中所有token的embedding进行加权聚合,从而生成更具代表性的序列embedding。PMA模块通过学习不同token的重要性,能够有效地提取序列中的关键信息,打破了EOS token的信息瓶颈。同时,PMA模块的设计允许灵活调整输出embedding的维度,以适应不同的应用场景。
技术框架:C2LLM的整体框架包括以下几个主要模块:1) Backbone LLM:使用Qwen-2.5-Coder作为backbone,负责生成token embeddings。2) PMA模块:将token embeddings作为输入,通过多头注意力机制进行加权聚合,生成序列embedding。3) 对比学习目标:使用对比学习目标函数训练模型,使得相似的代码片段具有相似的embedding,不相似的代码片段具有不同的embedding。
关键创新:C2LLM的关键创新在于引入了PMA模块来生成序列embedding,这与传统的基于EOS token的方法有本质区别。PMA模块能够利用序列中所有token的信息,避免了信息瓶颈,并且支持灵活的embedding维度调整。此外,C2LLM还充分利用了LLM在预训练期间学习到的因果表示,提升了代码embedding的质量。
关键设计:PMA模块的关键设计包括:1) 多头注意力机制:使用多个注意力头来学习不同token的重要性,从而更好地提取序列中的关键信息。2) 可学习的query vectors:PMA模块使用可学习的query vectors来查询token embeddings,从而实现自适应的加权聚合。3) 对比学习损失函数:使用InfoNCE损失函数来训练模型,鼓励相似的代码片段具有相似的embedding,不相似的代码片段具有不同的embedding。训练数据包含三百万公开数据。
🖼️ 关键图片
📊 实验亮点
C2LLM在MTEB-Code基准测试中取得了显著的成果。C2LLM-7B在总排行榜上排名第一,超越了同等规模的其他模型。具体而言,C2LLM在代码检索任务上的性能提升尤为明显,表明PMA模块能够有效地提升代码embedding的质量,从而提高代码检索的准确率。
🎯 应用场景
C2LLM在代码检索、代码克隆检测、代码摘要生成等领域具有广泛的应用前景。高质量的代码embedding可以提升代码检索的准确率和效率,帮助开发者快速找到所需的代码片段。此外,C2LLM还可以用于代码相似性分析,检测代码克隆,以及生成代码摘要,提高软件开发的效率和质量。未来,C2LLM有望成为软件开发工具链中的重要组成部分。
📄 摘要(原文)
We present C2LLM - Contrastive Code Large Language Models, a family of code embedding models in both 0.5B and 7B sizes. Building upon Qwen-2.5-Coder backbones, C2LLM adopts a Pooling by Multihead Attention (PMA) module for generating sequence embedding from token embeddings, effectively 1) utilizing the LLM's causal representations acquired during pretraining, while also 2) being able to aggregate information from all tokens in the sequence, breaking the information bottleneck in EOS-based sequence embeddings, and 3) supporting flexible adaptation of embedding dimension, serving as an alternative to MRL. Trained on three million publicly available data, C2LLM models set new records on MTEB-Code among models of similar sizes, with C2LLM-7B ranking 1st on the overall leaderboard.