Skip to main content

3 posts tagged with "embedding"

View All Tags

从现代Coding Agent视角回看代码搜索与嵌入

· 24 min read
ayanami

应该如何说起代码搜索呢,先说代码搜索的几个小的子流派吧,这方面可能略微和其他领域不同

代码的检索我认为是可以分解成以下几种的:

  • 搜索引擎式
  • grep传统搜
  • 向量embedding搜
  • 码仓index

每一种都是什么意思呢,举个例子

搜索引擎式是复用传统的搜索引擎,如elasticsearch, meilisearch等;也有一些变体,比如github的代码搜索,主要是弥补传统搜索引擎在代码搜索的不足

grep如其名字,基于linux grep或者riggrep,在agent内使用最广

向量embedding搜则有多个子任务,如NL2Code, Code2Code, Code2NL等,每个还有一些细微差别,这个很有意思,待会讲;

码仓index呢,则大致有两种:一种基于传统的AST分析,希望构建整个码仓的符号树或者CKG来辅助搜,另一种则寄希望于新型的生成式大模型,希望让LLM读了仓库之后能生成”索引“,典型的做法比如deepwiki。

搜索引擎式

传统的搜索引擎其实在代码搜上有很多问题,参见

https://github.blog/engineering/architecture-optimization/the-technology-behind-githubs-new-code-search/

最严重的问题甚至不是想象的NL2Code的问题,而是传统搜索引擎的部分优化不够、部分优化又不足,总之不与代码适配

  1. 索引的开销, 恐怖的内存消耗

当我们第一次部署 Elasticsearch 时,需要几个月的时间才能索引 GitHub 上的所有代码(当时大约有 800 万个存储库)。今天,这个数字已经超过 2 亿,而且代码不是静态的:它不断变化,这对搜索引擎来说是相当具有挑战性的。

  1. 可不可以不要索引?对大规模库高并发是不可能的:

首先,让我们探讨一下解决问题的蛮力方法。我们经常收到这样的问题:“你为什么不直接使用 grep?为了回答这个问题,让我们使用 ripgrep 对这 115 TB 的内容进行一些数学计算。在具有八核 Intel CPU 的机器上,ripgrep 可以在 2.769 秒内对缓存在内存中的 13 GB 文件运行详尽的正则表达式查询 ,即大约 0.6 GB/秒/内核。 我们很快就会发现,这对于我们拥有的大量数据来说确实不起作用。代码搜索在 64 个核心、32 个机器集群上运行。即使我们设法将 115 TB 的代码放入内存中并假设我们可以完美并行化工作,我们也会在 96 秒内使 2,048 个 CPU 内核饱和,以处理单个查询!只能运行一个查询。其他人都必须排队。结果是每秒 0.01 次查询

  1. 分词器不需要了:搜索引擎依赖分词来减少构建倒排的成本和作为基础搜索单元,但就和Tokenizer的引入一样,这个开销的减少从来不是免费的,接下来的问题就是:代码如何分词?另一个问题是:停用词全部没用了!无论是!@?/#$:{}()[]还是什么别的乱七八糟的符号,都是编程语言的最爱, 这使得传统的分词系统更加雪上加霜

    1. 如果提前跑一遍AST分析呢?我们的CPU要算爆炸了:), 并且你如何统一没有LSP小众语言,不同的语法版本(py2->3)... 工程量立刻爆炸了
    2. 能不能在char level倒排?太爆炸了索引,所以github是使用3-char的倒排的 argument -> arg、rgu、gum、...
    3. 代码里面还有注释,注释还有多语言,甚至单个仓库都很常见中英两种语言的注释...看起来朴素的分词器比如jieba要全面阵亡了...
  2. 基于git的增量更新:增量更新本身可以用merkel tree,这也是cursor在用的技术,但结合git版本?你发现事情变得复杂起来了,这是纯工程的复杂性。一次git commit涉及到十几个文件的几行变化,需要触发至少十几个chunk的embedding更新?如何处理多分支呢?我的每一个存储代码块是否还得加一个tag标识它的branch,然后在搜索引擎里面支持完备地按tag过滤,省的不同branch的代码在同一个搜索引擎中返回?(然后发现branch name作为tag简直太烂了,它是一个完全动态的无限的集合)

鉴于恐怖的存储代码数量,github采用的搜索引擎相当简陋,甚至是弱化版本的完全匹配:3-grams索引

对于代码搜索,我们需要一种特殊类型的倒排索引,称为 ngram 索引,它对于查找内容的子字符串很有用。ngram 是长度为 n 的字符序列。例如,如果我们选择 n=3,则构成内容“limits”的 ngram 是 limimimitits(二元组的选择性不够,四元组占用了太多空间)

这个索引显然是非常大无法放入内存的,所以github采用了一些传统数据库里面的懒加载和流式优化技术,使得可以仅读取一个小子集完成搜索

而关于构建索引本身,github还有很多特殊设计,但这其实属于system/后端任务了,不细讲:

  • 用Git blob object ID来分片,kafka分区

  • 用path, branch, repository + 元信息(owner, visibility, etc.) 来构建增量索引key

  • commit-level的一致性

  • Github相当多的blob是相同的,使用增量编码很有吸引力, 这里用到了概率上的近似数据结构和一些分布式图(近似)算法

    To determine the optimal ingest order, we need a way to tell how similar one repository is to another (similar in terms of their content), so we invented a new probabilistic data structure to do this in the same class of data structures as MinHash and HyperLogLog. This data structure, which we call a geometric filter, allows computing set similarity and the symmetric difference between sets with logarithmic space. In this case, the sets we’re comparing are the contents of each repository as represented by (path, blob_sha) tuples. Armed with that knowledge, we can construct a graph where the vertices are repositories and edges are weighted with this similarity metric. Calculating a minimum spanning tree of this graph (with similarity as cost) and then doing a level order traversal of the tree gives us an ingest order where we can make best use of delta encoding. Really though, this graph is enormous (millions of nodes, trillions of edges), so our MST algorithm computes an approximation that only takes a few minutes to calculate and provides 90% of the delta compression benefits we’re going for.


Grep

grep属实是在Coding Agent时代焕发了第二春,由于其系统级别自带+完美匹配Bash工具和Unix文本管道的特性,在现代的LLM之中都大量训练了如何写出各种米奇妙妙grep的数据

claude code这种经过更多优化的grep会更过分一点,它会有几个细节优化:

  • 使用更现代的rg(riggrep)代替原始的grep
  • 逆向cc源码可知,它的grep有七八个参数,分别对应grep里面的不同参数比如 -A -E -C , 除了一些呈现格式(比如带不带行号和文件名)之类的差别,主要的几个参数就是在匹配行前保留多少行、匹配行后保留多少行、和上下保留多少行
    • 如果读者熟悉coding agent的工作的话,其实早在swe-agent就已经探究过这个context window开多少的问题,原始论文的实验结论是50行
❯ tldr grep
grep
Find patterns in files using regular expressions.More information: https://www.gnu.org/software/grep/manual/grep.html.

- Search for a pattern within a file:
grep "{{search_pattern}}" {{path/to/file}}

- Search for an exact string (disables regular expressions):
grep {{[-F|--fixed-strings]}} "{{exact_string}}" {{path/to/file}}

- Search for a pattern in all files recursively in a directory, showing line numbers of matches, ignoring binary files:
grep {{[-r|--recursive]}} {{[-n|--line-number]}} --binary-files {{without-match}} "{{search_pattern}}" {{path/to/directory}}

- Use extended regular expressions (supports ?, +, {}, (), and |), in case-insensitive mode:
grep {{[-E|--extended-regexp]}} {{[-i|--ignore-case]}} "{{search_pattern}}" {{path/to/file}}

- Print 3 lines of [C]ontext around, [B]efore or [A]fter each match:
grep --{{context|before-context|after-context}} 3 "{{search_pattern}}" {{path/to/file}}

- Print file name and line number for each match with color output:
grep {{[-H|--with-filename]}} {{[-n|--line-number]}} --color=always "{{search_pattern}}" {{path/to/file}}

- Search for lines matching a pattern, printing only the matched text:
grep {{[-o|--only-matching]}} "{{search_pattern}}" {{path/to/file}}

- Search stdin for lines that do not match a pattern:
cat {{path/to/file}} | grep {{[-v|--invert-match]}} "{{search_pattern}}"

另一个有趣的事情是,现在的coding agent不约而同地使用了grep而不是rag作为其系统原生的工具,我觉得理由也是非常清晰的:

  1. grep的输出是标准可预测的,而rag的输出依赖于 {基础模型, 分块方法,召回topk,重排模型} 等多个配置参数,一个标准的输出带来的好处是 可强化学习, 如果对一个 code llm + rag 的系统做RL,最后的搜索策略一定会是拟合到和rag的embedding模型和具体策略相匹配,丧失了可迁移性
  2. 除此之外的好处也有很多,比如RL环境不需要embedding的额外开销(存储和计算上甚至编码成本上),整体轨迹可解释,精确匹配效果好,速度快...

rag的index开销其实相当大,学界不在乎这个,为了提升精度每个token一个embedding的方法也有,但一个embedding是一个1024维的向量,光存储开销就是4KB,对于百万行级别的代码仓库,其chunk可能在数万,达到了GB级别的存储开销

而工业项目有百万码仓,在TB级别的存储上进行高效地索引和查询着实压力很大,可以参考 美团和milvus/lancedb的相关文章 ,索引优化也有相当多的新实践,但这是做DB的人考虑的(雾

但Grep就是万灵药吗?并非如此

Grep提供了一个切面,能够让模型Agentic Search,根据搜索到的局部反馈调整搜索方法,从部分开始探索整个代码仓库——大部分需求的完成不需要对全仓的理解

——吗?

一个Grep的bad case是高阶语义的需求

  • 哪里导致了这个bug?
  • 某个模块的核心逻辑是什么?
  • 整体的代码结构?
  • ...

模型要么老实cat,要么就只能在log中见到它尝试“猜”你的变量名字,比如你问“...的实现”就会开始Grep impl,如果你把所有变量换成abcde,它立刻就GG了

比起失败的搜索浪费的上下文更糟地是浪费的交互轮数 —— 长达几百轮的agent轨迹是相当稀少的训练数据,如果再配合没那么好的历史压缩方法,或是没有精心设计的防止模型死循环的额外环境反馈,连续失败的grep会让agent的性能迅速地劣化

基于这方面的需求,在推理阶段Grep还是得配合别的工具,比如deepwiki,比如CKG(代码知识图谱,例如每个函数的caller和callee),

比如Code RAG

这方面也有一些新的探索,例如 https://cognition.ai/blog/swe-grep 的RL并行工具调用(关键不在速度,关键在减少交互轮数!),比如在工程侧融合Grep和RAG如https://github.com/daimh/mgrephttps://github.com/zilliztech/claude-context ,以及我们将要发的一篇文章(自吹自擂一下,关注主包后续的工作谢谢喵)


code embedding

主包主包,code embedding和文本的embedding有什么区别呢?为什么要强调code?

非常好问题,爱来自AI4SE。我认为code其实和图像比较像,某种意义上算是一种特殊的模态,不完全是文本——code某种意义上是“反语言常理”的,例如大部分语言的上下文有限,一句话很难和1000个字之前的某个东西形成强烈的联系,而这种长程交互在code之中非常常见——甚至有跨文件、跨模块、跨仓库的交互

而另一个很有趣的事情是,当我们在讲“某段代码的语义”的时候,这件事本身是模糊的,文本没有那么强的二义性,太阳就是太阳,月亮就是月亮,但一个递归斐波那契函数的语义到底是 “递归“ 还是 ”斐波那契“?这其实折射出了代码的某种特殊性,它同时具备字面义 ”斐波那契“ 和运行义 ”递归“(甚至”低效“、”算法“、“python”),而在一个代码仓库之中,代码还具备了上下游的属性:谁是我的caller,谁是我的callee?

这件事情为什么重要呢,因为传统的embedding向量相似度产出的是一个标量,它只能衡量一个维度的相似性!

  • 当你在说“查找与function A相关的代码”时,你想要的到底是什么?
    • function A的字面义相关的代码?
    • function A的运行效果相关的代码?
    • function A的caller/callee?
    • ...

然后你就发现从这个角度上来说,Code2Code的向量搜是很诡异的一件事情,至少传统的cos相似度无法干这件事——

  • 更悲伤(从学术研究的角度上来说或许是兴奋)的是,在现实需求中,我们真的不在意找到和一个function的字面意相关的function...假设你想要补全一段代码,你可能更需要关注谁会是它的caller,假设你需要优化一段代码,你可能搜索的方法是某种低效的pattern...而这些embedding相似度全部做不到
    • 据我所知,企业对这个接近摆烂了,只有MSRA有一个group还在研究,我之前溯源到的比较早的上下游建模技术是Order Embedding,感兴趣的或许可以试试做

直接结果是:我们只有NL2Code了

并且这个NL也只能关注一个方面...

什么样的NL才是真实会问会写的NL呢?Coding Agent的轨迹数据

没有轨迹数据怎么办?从大规模代码中挖掘注释作为NL

一个人写注释的方法和提问的方法不一样,这个语义空间的unmatch如何处理?各家自显神通

  • 例如2025年5月快手的OASIS: Order-Augmented Strategy for Improved Code Search,认为现有的code embedding往往关注的是代码的字面相似性,即只把代码认为是一种特殊的“语言”,而忽略了代码的非文字意义上的相似性

    • 对代码片段(结合其他静态分析信息),用LLM产生其作用的描述文本
    • 计算这个描述文本和其他代码片段的相似性,以此来挖掘难负样本
    • 因为这个文本描述的是相对High Level的函数作用,能够一定程度上避免变量名字等带来的影响,专注于实际作用
  • 24年12月的Nomic AI的cornstack

    • 强调<文档,代码>对的相关性重要性,并采用双重过滤:如果文档与代码间相似度低,或者并不在topK,只要有一个满足就筛掉
    • 动态硬负例挖掘策略:对于批内负样本挖掘,采用softmax概率采样,但是在训练过程中,逐步改变softmax的温度,前期温度高提高多样性,后期温度低,注重难负样本的区分
  • 25年5月的BAAI Towards A Generalist Code Embedding Model Based On Massive Data Synthesis

    • 强调退火训练,第一阶段纯文本,第二阶段全数据训练text-code能力,第三阶段纯代码

码仓Index - DeepWiki/CKG

这个说起来就比较简单了,deepwiki重要的始终是LLM的能力,而CKG则是静态分析的质量,开源的tree-sitter固然可用,企业也有一些统一各种语言的私有AST,静态分析已经日趋成熟,困难的是如何将这个信息给到LLM

很早在Google的博客中就有论述: 现在的vibe coding就像是把一个几千行的代码粘贴到记事本里面,然后让程序员来修改bug —— LLM看到的就是 “记事本”,而不是程序员的带有各种Lint和跳转的IDE界面

AST分析树如此庞大,除了摆烂式地提供一个获取caller/callee的mcp工具给agent之外,还可以做些什么?

先前在web领域有一个llms.txt的旨在LLM-friendly的格式,代码领域却暂时缺乏哪怕是新兴的统一处理格式

AI-friendly IDE可能对于Coding Agent的能力提升相当重要,这也是moonbit社区他们宣传的,不过我没有实际上手用过,也不是学PL的,就不瞎讲了,感兴趣的可以看张宏波的演讲

【AI时代下的基础软件 | 张宏波 刘子悦 蚂蚁&MoonBit Meetup杭州站回放】

https://www.bilibili.com/video/BV1wL8DzgEXZ/

除此之外,可能我们原先认为不能搜索或者没必要搜索的部分,现在也正在发挥着额外的作用

  • python的package,众所周知(可能并非),pip包只是一个特殊的压缩包,可以直接看到文本格式的原始代码,现代的claude-sonnet-4等模型在环境出错的时候会主动读/搜 pyproject.toml, .venv等特殊文件,遇到import的不了解api还会尝试进入package观看源码
  • 而某些binary风格或是一串神秘哈希引入的包可能对于LLM并不是很友好...

除此之外也有很多新兴的想法,例如注释本身可不可以作为一个天生的码仓Index? ...

CLI版本的Coding Agent好处是可以在各处方便的引用,尤其利于大规模并发采集数据

IDE版本的Coding Agent则会更加地“懂人”,原因是AI IDE在背后做了一大堆不仅仅是diff等格式渲染的工作,用户环境信息,用户系统信息,... 这些都被从后台塞入了Coding Agent的system prompt,使得你在Linux上运行Copilot的时候,模型不会让你执行 brew install命令

但文本本身依然有着局限,或许在某个未来,我们能看到真正的code native架构,不在绞尽脑汁地想把编译器的报错,AST的分析等等原本结构化的东西转成markdown再塞入永远不够的agent上下文...

Paper reading - Context Pruning and beyond hard pruning

· 17 min read
ayanami

引子

我们知道,在现在Agent需要处理的一大问题是长上下文下性能的开销问题,对此infra团队有非常多的优化,从attention架构的优化如各种windowed attention到kv的压缩重用如cacheblend和megicdec等都提出了一系列的解决方案,但有一个最本质的方法是:有没有可能直接减少上下文的长度(去掉不必要的上下文?) ,这就是Context Pruning的出发点。

而截止2025年8月,相关的方法已经发展了两三年了,大体上可以分成几个类别,本文会对此做一些简单的介绍和总结。

借用naver lab最新的相关论文里面的说法,现在的方法可以被一个四方格归纳:

其中,Hard和Soft代表裁剪方法是直接作用于token上(hard,相当于裁剪结束后,输入的是一个新的prompt),还是作用于token的embedding上(soft,相当于裁剪结束后,输入的是一个新的qkv和其他东西,无法还原出“原始”的token输入)

在线和离线一般代表着这个裁剪方案是否依赖于用户查询q,依赖q的方案是在线的,不依赖q的方案是离线的,可以提前做好。但传统上,如果你的裁剪方法也需要用到和原始模型一样大的LLM,也一般称之为离线,或许“是否会对在线推理造成明显时延影响”做划分更好一些

离线硬裁剪

在最早期的时候,就有相关的一些朴素方法,例如直接对查询文本段做一次总结摘要,再用总结后的文本段去做后续的任务,这种方法是离线硬裁剪的典型代表。如果用的是llm就是离线的,如果用轻量级模型做摘要或者总结就是在线的

而在后面的时候,出现了例如微软的llmlingua这样的工作,直接用一个小模型(gpt2 small,llama-7b, etc)去预测哪些token是重要的,哪些token是不重要的,然后把不重要的token直接裁剪掉,这种方法也是离线硬裁剪的典型代表。(llmlingua2 换成了微调的 BERT 来做这个事情,所以可以说在线的), 其出发点和常规的硬裁剪可能有部分地方不同,例如llmlingua认为,裁剪本身是可以得到一些人类不可读但是大模型可以理解的token序列的,所以可解释性上可能并没有想象的那么强。

离线软裁剪

和硬裁剪同时推进的是软裁剪相关的工作,其想法很简单: 如果我牺牲解释性,直接调整prompt的embedding这类,即使产生的是不对应任何token的"fake embedding",其在高维空间中也应该融合了多个token的语义,理应得到更高的压缩率(可以理解为,在训练过程中为llm 扩充为无限词表,然后定义了一些高效的"额外语言")

比较早期的工作是 xRAG, 其裁剪策略非常极端,将整个段落都压缩成1个embedding X,怎么训练呢?一个在此类论文中经常出现的是重建loss,即

压缩前Doc+queryxDoc + query \to x

压缩后Doc+queryxDoc' + query \to x', L=L(x,x)=DKL(x,x)L=L(x',x)=D_{KL}(x',x), 即自蒸馏,希望压缩后依然能重建原始的输出,论文实际中可能会用变体版本来实现指令遵循等

xRAG的做法是,使用一个通用编码器E,把这个编码器E视作一个新的模态,仿照CLIP的方法直接用MLP projector做通用编码器和实际使用的LLM token embedding的模态对齐

但[大家实测下来](笔记:RAG 的相关优化方法之六(xRAG/PISCO) - 刀刀宁的文章 - 知乎 https://zhuanlan.zhihu.com/p/29292925032),xRAG的效果并不好,而相对较好的是更新的Pisco方法

Refer to caption

Pisco将检索到的文档D和memory tokens一起送到LLM中,产生embeddings

再将embeddings +query送到相同的LLM中,产生输出,这个 q+E 和原始的 q+D 比较, 计算交叉熵损失

这里有一些复杂的地方:

  • 虽然叫解码和编码,但是Student LLM都是同一个LLM, 只是训练不同LoRA模块

  • 交叉熵是怎么得出的? teacher模型和student模型都是采用的最大长度128的贪婪解码,就可以直接令 L=1logp+0log(1p)=ilogP(aiq,e,a<i,θc,θd)L=-\sum 1logp + 0log(1-p) = - \sum_i log P(a_i|q,e,a_{<i},\theta_c,\theta_d) , 优化目标是 θc\theta_cθd\theta_d 还有 memory_tokens

  • 如何理解memory token? 我觉得文章是借用了之前的一些研究比如ICAE, 在这些文章之中,训练的压缩机制是,将上下文压缩成一个定长的memory slot, 这里的memory token实际上只是多个embedding向量而已,而更关键的是LoRA微调的θc\theta_c,我的理解是,memory tokens只是一个后置的、可以看到Documents的所有信息(假设它没有魔改注意力)的语义位置,叫tokens也可以理解为直接扩了词表加入了l个特殊token,类似BERT里面的[BOS] ,只是decoder llm需要后置。

    • 文章并没有细说这里的注意力是怎么设置的,但从后文中发现的memory tokens具有明显的位置特性(例如1位mem token主要注意最开头一段),感觉应该是没改过
  • 文章的另一个重要的实验结论是,微调student llm(θd\theta_d)是必要的,之前的研究中没有相关模块,会导致性能的大幅度下降。这细想其实是一个很有趣的事情,可以注意到,压缩的时候是没有接触到query信息的(这也是为什么称为离线的原因),可以理解为某种意义上的LLM as an embedder,而加入了query和embedding再训练的时候,θd\theta_d一边学会了如何理解自己产生的embedding,另一方面学会了如何根据query去选择embedding,整体上类似于ColBERT架构的Reranker(前面是multi-vec embed, 后面是maxsim)

在线硬裁剪

Provence

之前的裁剪方案只注重于“自然语言是有冗余的”,所以主要做的都是token-level的pruning,而provence则更注重实际一些,它发掘了一个问题是,其实现在RAG里面的 “Chunk” 是一个特别微妙的概念

如果chunk切得大了,那上下文自然就长了,甚至效果也会明显下降(详见ground truth在chunk中的不同位置的position bias相关的研究,现有embedder对这个bias耐受性不佳,会狠狠掉点);但如果chunk切得小了,语义信息的丢失、检索的困难又是很恼人的事情(先不论检索,检索到了多个小块之后信息不够怎么办?一种是合并,但策略怎么定?另一种是Anthropic的Contextual Retrieval,把上下文放进来,本质上还是变成大块(我说这个a一串真是炒作勾啊.jpg))。

而Provence给了一个折中的方案,既然我们有句子级别的语义,为什么不用呢?分几步走

  1. 训练一个接受q,d的BERT,给每一个token打0~1分,并根据用户指定的阈值进行二值化变为0/1, 表示删除/留下
  2. 进行句子级别的聚类,裁剪掉0的token数量大于1的token数量的句子

如何训练呢?选取有5~10个句子的段(可以多次选取来拓展到更长的上下文),标上句子序号,让LLM选择相关句子来产生label,从而训练模型

这里其实做了很有意思的工程设计,

  • 如果让LLM来打token-level的标,肯定是收集不到足够的样本的,并且真的无所谓多出来的几个token,更在意句意的完整性

  • BERT带来了相当多的好处:

    • 这样进行的句子裁剪,每个句子都可以和整个chunk里面的所有上下文交互,使得一个句子的保留与否不仅取决于这个句子和查询的相关性,还取决于其于其他(和查询相关性高的)句子的相关性,这就使得这个方法必然会优于按句子切分的朴素方法

    • 我们的 reranker 也就是个BERT啊,完全可以训裁剪和训rerank一起进行,推的时候也一样,相当于和rerank overlap了

      image/png

在线软裁剪

Oscar

Pisco为代表的离线软裁剪有一个问题是,它的压缩需要微调,并且受限于难以对齐encoder-only架构的预训练编码器模态和实际推理使用的decoder LLM的模态,难以把压缩这一步在线做

Oscar就提出了一种方法是,我的对齐既然难做,我直接不对齐了,使用LLM的前L层 + memory token(他们也做了用Llama硬对齐的版本),足以得到够好的embedding,文章最大的贡献其实是实验证明了这样表达能力已经足够,能训出来(太神奇了LLM)。当然,L越大效果越好

而还是复用Provence的工程技巧,把裁剪和rerank overlap起来,OSCAR的compressor留了一个RR头,在这个头和Teacher Reranker对齐,整体的Loss就是rerank loss + generation loss

而令LLM理解embedding这件事情还是通过LoRA adapter来做,这篇文章其实像是序列工作的延申,综合了PISCO的训练方法,把PISCO的压缩部分从LLM + LoRA换成目标模型的前N层transformer,然后压缩器全参微调、生成器LoRA微调,再使用和Provence相同的技巧进行rerank的overlap

image-20250907213839337

异曲同工

从HyDE到“投机解码”

另一个有趣的工作是广义上的“裁剪”,或者就是更好的搜索吧。我们知道HyDE的思想是原始query一般都比较短,而生成的假设文档可能会更好地与索引文档对齐,所以使用 q‘ = q + generated d 来进行搜索。

而智谱的memorag 则提出了这样一种场景,我们是否能以低成本训练一个小模型,来根据源文本生成这个假设答案呢?(例如,使用Llama3-8B在哈利波特上训练比用Deepseek-R1在哈利波特上训练成本要低廉的多,将HyDE的生成方从R1自己换成小模型)这就非常像是投机解码的思想了

外接模块: memory decoder, catridges

其实这种将memory训为embedding的方法确实不少,如果说前面的压缩器是在训一个meta network,能够从doc生成embedding的话,外接模块的工作就是在训练embedding本身 -> 我能否直接从一个大的文档库中训练出一个参数化的memory?

最近的memory decoder选择的是直接扭曲生成过程,将一个小模型在目标适配数据集上训练,在大模型生成token时,将小模型的概率和大模型的概率相加(再重归一化),认为这样会带来领域知识的纠正(比较暴力www)

而另一篇catridges 则是在使用类似P-tuning的方式训一个Prefix KVCache,在推理时实时加载,而希望这个KV中有相关的memory

包括一系列的kvcache evict的工作也是在做类似的东西,为了决定evict哪些甚至都把搜索又搬上来了,比如clusterKV的knn(笑)

总结

总体而言,我感觉相关工作已经进入了深水区了,硬裁剪可能在某些程度上到头了,现在主流在探索一些牺牲解释性的,更能scale out的方法来进行参数化memory来解决长上下文、领域适配等一系列问题

而大家方法逐渐趋向于无标签学习的统一也再次证明了scale out能力在广义embedding能力的训练上的重要性

另一个很有意思的是,可以看到搜索中的多向量和多memory token有一些很有趣的相似性,或许在后续的一些工作中,我们能看到一些多向量的方法被用到memory之中,希望会让memory这个很多时候靠prompt编故事的领域更多可验证性吧

而从另一个方面,正如这里列出的部分文章说的,自从eagle在投机解码中得到了确实很好的效果之后,大家都开始用 token + embedding的混合来捕捉更强的信息了,还有HyDE和投机解码这种很有趣的对应

ColBERT-后期交互方法

· 10 min read
ayanami

如果简单引入语义搜索,那么第一时间想到的肯定是向量搜索的方法

先不论小的优化,向量方法现在大体上就是两种架构,单塔和双塔,对应Cross-Encoder和普通的Encoder模型。

双塔模型如下,查询qq和文档dd分别通过两个独立的编码器,得到向量表示qvq_vdvd_v,然后计算相似度(内积,余弦,等等)。

overview traditional text embedding models

而单塔模型则是将查询和文档拼接在一起,输入到一个交叉编码器中,这个交叉编码器很多时候就直接输出相关性得分score了,即为我们所说的reranker

单塔虽然精度远高于双塔,但有无法离线计算的缺点

而双塔的一大精度困境在于,当编码的文档变长时,文档的大部分内容可能都和查询没什么关系,这会导致查询向量和文档向量的相似度计算不准确。实际上,在楼主之前的一些实验之中,一整个很大的文档集合内,和某个查询最无关和最相关的文档的余弦相似度相差也就0.2左右,这就是长文档带来的问题。

但客观地讲,长文档是无法避免的,如果把文档切成更细粒度的句子,在上下文补齐语义,后续合并等麻烦可能更多,并且会出现"长文档实际上是在让相似度检索考虑上下文"这样的情况,一个例子是,问题是"上海交大的用户论坛中,....",而文档可能是"...水源社区是上海交大的用户论坛。水源社区....." 如果仅在句子等短文本上面匹配,那缺少了上下文的情况下,"水源社区"当然和"上海交大"没什么关系。

那么,如何保证精度的同时又能离线计算呢?

ColBERT的思路是,使用双塔模型来计算相似度,但在编码文档时,使用了一个更细粒度的向量表示

ColBERT给每个token一个向量表示,而不是给每个文档一个向量表示。这样,查询和文档的相似度计算就可以在token级别进行。

如下图,ColBERT在拿到最后一层的输出之后(这一层有非常多的语义信息!),将每一个token对应的vector都存下来,这一部分是离线的。

而在计算相似度的时候,将query的tensor和文档的tensor进行一个MaxSimMaxSim算子

MaxSimMaxSim是一个最大池化操作,取出每个token的向量中与查询向量最相似的那个向量,然后计算相似度。

overview colbert

ColBERT的性能是逼近reranker的,这个也很好理解,毕竟交叉编码器的优势就是可以考虑q,dq,d之间的交互,而ColBERT除了保留语义嵌入之外,比起更暴力的加大embedding维度,更重要的是它保存了上下文次序的信息

而ColBERT的最后一层MaxSim,而没有采用神经网络的方案,让他带来了良好的可解释性

colbert snippet

那看了上面立刻就会想到,这每一个token保存一个768/1024/...维的向量,存储开销不会很大吗?

ColBERT也考虑到了这个问题,因此在ColBERTv2中,采用了这样质心编码的方法来降低存储开销,能降低8倍

  1. 对每个token的向量进行聚类,得到kk个质心(k是一个预定义的数字)

  2. 对每个token的向量,找到距离最近的质心,并将其索引存储下来,也就是从 (vd,)>(1,)(v_d, ) ->(1,)

  3. 将质心向量库构建ANN索引,例如FAISS, ScaNN

  4. 在计算相似度时,查询向量也进行同样的处理,找到距离查询最近的质心索引,然后从质心向量库中取出对应的质心向量进行相似度计算

在实际使用的时候,商业rag公司甚至对大规模检索做更狠的二值化向量压缩(说实话这也能检索出来真的有点现代模型神力了),让ColBERT的开销可以和单独的embedding媲美

colbert token

二值化的说法是这样的:

压缩方法通过将正维度表示为 1、负维度表示为 0 来简化文档标记向量。这种二进制表示有效地指示了文档标记向量中重要语义特征的存在与否正维度有助于增加点积,表明相关的语义相似性,而负维度则被忽略。

ColBERT的使用上,很多公司都有了支持,例如vespa, jina等等,开源方案则有早期的ragatouile和后来的上下游如milvus,llamaindex的支持

但是,文档ColBERT还不是它发挥全部潜能的时候,据说SPLADE算法就比他效果好不少(这个我没有实测过),它在图像又活出了第二世,即所谓的ColPali架构

ColPali是MRAG、MLLM那边的新论文和解决方案,几个月的时间砍了1.9k star,ColPali的想法是这样的

  • OCR的多个组件和分块带来误差传播,且预处理流程耗时也长,能不能直接端到端一次使用文档截图解决
  • 但是如果将整页的文档编码成一个向量,肯定精度不够
  • 我的ViT等视觉编码器会将整页文档变成一系列的patch(可以理解为子图),进而变成一系列视觉token,那我重用ColBERT,不就又有了多向量吗?并且这个存储和交互上比每个token存一个向量更合理! 子图本身就有很多的空间位置信息

image-20250529232012895

并且,你会发现ColBERT的强可解释性在图像上有更关键的作用!模型在文本中关注了什么可能是某个词,还需要人进行一点逻辑推理来判断关系是否合理,而图像中关注了什么,直接看图就知道了!

image-20250529232211333

作为一种新的RAG范式,ColPali从源头上解决了复杂的OCR和切块的问题

虽然其在重文字领域上的泛化性还留待验证,精度的提升也依赖于未来VLM的发展,但无疑社区已经认同了这个想法的价值

基于 OCR 的文本提取,以及随后的布局和边界框分析,仍然是重要文档 AI 模型(例如 LayoutLM)的核心。例如, LayoutLMv3 对文档文本进行编码,包括文本标记序列的顺序、标记或线段的 OCR 边界框坐标以及文档本身。这在关键的文档 AI 任务中取得了最佳成果,但前提是第一步——OCR 文本提取——能够顺利完成。

但通常情况并非如此。

根据我最近的经验,OCR 瓶颈导致现实世界生产文档档案中的命名实体识别 (NER) 任务的性能下降近 50%。

目前例如ColQwen2这种ColBERT + Qwen2.5-VL-3B-Instruct的方案也很火,很多榜上都刷到了SOTA,感兴趣的同学也可以自己试试