Skip to main content

3 posts tagged with "agent"

View All Tags

迈向真实世界的软件智能体:如何将语义高亮功能融入智能体编程?

· 26 min read
ayanami

引言

我们训练了一个代码语义高亮模型SWE-Pruner并开源了配套的Agentic接入框架,在真实的Coding Agent多轮任务(SWE-bench, SWE-QA)上,在保持性能不变甚至提升的情况下提供多轮场景下30%甚至更多的token开销减少。该模型基于语义理解,自动识别并高亮检索到的文档中语义相关的句子。而框架作为agentic的接入层,可以轻松接入claude code,openhands等SOTA agent系统

模型发布

在本文中,我们将分享我们的技术方案。

问题:编码代理的上下文膨胀

在先前的RAG工作中,上下文的粗粒度膨胀已经是一个很严重的问题 https://huggingface.co/blog/zilliz/zilliz-semantic-highlight-model

在生产环境中的 RAG 系统中,一个典型的查询会检索 10 个文档,每个文档包含数千个词元,每次查询消耗数万个词元。问题在于: 只有几十个句子真正包含相关信息 ,其余的都是噪声,这会增加成本并降低答案质量。

然而,对于真实环境下的软件工程任务,问题甚至更加严重:一个典型的查询可能涉及上百个代码文件,而一些大型项目中,巨型文件的代码token包含的远不止数千个,现代的编程agent却仍然在广泛使用关键词匹配的方法,即grep,来对大型代码仓库进行探索,匹配项的过分庞大使得现代agent中都配备了相当多的上下文硬截断逻辑

image.png

在人工智能代理场景中,这个问题会变得更加严重,因为查询是经过推理和分解后的复杂指令。传统的高亮显示方式只是机械地标记匹配的词语,却会遗漏真正有价值的分析结论。

这就迫切需要一种有针对性的高亮模型,该模型仅保留上下文相关的代码,并高亮显示它们,同时剪掉所有无关的噪声内容——这种技术也被广泛称为上下文剪枝。类似的纯文本任务也被Provence为代表的一系列工作较为完善地解决,但对于究竟如何在代码任务上引入语义高亮模型,社区尚未给出一个答案。

现有模型的困境

我们调研了现有的解决方案,但发现它们并不完全符合我们的需求。

Provence/XProvence,OpenProvence,zilliz/semantic-highlight-bilingual-v1沿用Provence的思路继续优化Provence的上下文长度和多语言性能。然而,它们并没有指出在代码等非纯文本场景下,如何构造合理的粗细粒度,如何满足额外的AST约束等。

arxiv-provence-paper

另一个困境是,现有的模型沿用BGE的类BERT管道,但大部分现代代码嵌入模型都是LLM架构,decoder-only的模型如何设计裁剪头来处理特征的交互和输出的合理性仍是一个待解决的问题——BERT自然双向注意,但decoder模型需要某些“代偿”因果掩码单向注意的方案。

我们的选择

既然市面上没有这样的代码剪枝模型,那我们自己训练一个

第一个问题是在代码中,如何将纯文本里面的细粒度“句子”转换成代码的对应概念?

我们提出两套方案,基于代码行和基于AST statements,鉴于数据的生产难度,我们选择了以行取代句子作为细粒度

第二个问题是,选择什么样的模型?

BGE-m3是非常好的模型,但它没有针对代码数据进行大规模预训练;BAAI的另一个优秀模型是BGE-Code-v1,但社区对它的适配相对较少

我们注意到了社区的另一个工作modernbert,然而出于数据分布或是模型表达能力等原因,我们并没有在这个模型上训练成功

我们的猜想是——代码任务对基础模型的能力提出了可能超越大部分旧BERT模型的要求

最后我们采用的是Qwen的新模型,Qwen/Qwen3-Reranker-0.6B, 它的海量预训练数据和世界知识产生了相当好的代码嵌入,0.6B的大小又和传统bert类模型开销接近,LLM模型作为基底的另一个好处是自然完成了多语言query的适配。

类似Qwen这种LLM based reranker依然采用Casual Mask,这对token打分实际造成了一些影响,添加一个Full self-attention模块能让模型在测试集上的F1上升接近0.1

我们实验了不同的架构设计和特征trick,发现几个训练技巧对这个任务比较好用:

  • 冻结大部分层:我们发现只解冻最后两层和全参数微调差距极小,而冻结之后更能避免过拟合,训练速度也更快
  • 特征融合:对于LLM不同层数的特征作用,在先前已经有了相当多的工作进行讨论,一般认为浅层会有更多的语法信息,而深层则更重语义,将浅层、中间层和深层做简单的concat接入下游即可显著提升token打分头性能(我们使用的是7, 14和28层)
  • Sample Loss Norm: 代码片段的长短有着很大的差别,更关键的是,与文本领域大部分的信息都是稀疏不同,保留的代码行既有稀疏的片段(大部分行被删除)、又有稠密的片段(大部分行被保留),且长度上也有比较大的区别。如果简单在token level计算,模型训练相对不稳定,将每一个样本的loss使用样本的长度归一化能更好的平衡长短样本
  • CRF Layer:我们特殊设计了Token打分头的模型结构,使用CRF层进行序列建模,取代BERT风格的mlp层,实验证明对于改进模型的稳定性有一定帮助

模型架构&训练

我们设计了一个如下图的模型架构,它在思路上借鉴了Provence,但细节做了大量改动

image.png

此外,我们使用了和Provence相同的做法保留了基础模型的原始rerank头, 即图中的Path B,这种Rerank-While-Prune的方式进一步降低了在RAG管道等场景接入模型的成本。

最终,我们使用8 * H100 训练了3个epoch,用时大约是4个小时。

我们的具体训练超参数可以在论文中找到。

推理过程:

推理过程非常直接:

  1. 将输入拼接为 指令 + 查询 + 代码
  2. 对上下文中的每个 token 进行打分(范围在 0 到 1 之间)
  3. 对每行内所有 token 的得分取平均,得到该行的得分
  4. 根据阈值保留得分高的行,同时移除得分低的行

一个有趣的细节是阈值的选择,在Provence的开源模型之中,他们推荐将threshold设置为0.1来保守裁剪;而我们阈值调优实验发现,和训练保持一致的0.5阈值就能在下游任务中达到最佳性能,而0.1似乎过于保守——0.3已经会保留非常多的部分,这也可能是代码数据的分布特点和文本不同所导致的,在下文中会详细说明。

训练数据:如何构造“部分相关”的查询-代码对?

我们使用Qwen-Coder-30B-A3B-Instruct,从github 2025年的新仓库上自行构建数据(再次感谢qwen团队的优秀工作!),在来自6k个仓库、19w个文件的20w条代码片段上,使用了9个种子任务构建数据,并使用Qwen3-Next-80B-A3B-Thinking进行质量筛选,最后得到了约5万条高质量“部分关联query code对”

为什么是“部分相关”?

与文本领域不同的是,代码领域的QA数据集相对较少,缺少类似ms macro这样规模的query,code对;另一个问题是,由于缺少搜索引擎商开源的数据,现有的code query数据集的query往往分布偏窄,常是“这个函数的功能是什么?”“这段代码的运行结果是什么?”这种与整段代码相关的查询

但我们想要在代码之中引入semantic highlight功能,需要的是对“部分代码”的查询,例如,“这一条if分支会产出什么结果?” “这个类的属性A的作用是什么?”,这样我们才能划分出highlight的部分

也就是说,我们最后的数据集格式会是类似 <查询,原始代码,highlight部分> 的三元组

如何保证数据质量

数据质量直接决定了最后模型的性能,乃至训练是否能够成功。我们采用了以下几种方法来提升数据质量和解决实际遇到的问题。

  • 问题:模型产出的query同质化比较严重
    • 解法:定义代码重构、代码debug、代码解释等9个种子任务,构建时随机抽取,让模型根据种子任务写短/中/长不同长度,不同相关程度的查询
  • 问题:代码片段太短,较难从中捕捉不同的执行流分支;代码片段太长,行数会变得过多(相比于Provence大部分段落在5-10个句子的情况,我们的代码段取的大约是100-150行)
    • 解法:让模型使用range格式表示最终的“相关片段”,即1,4-10,16-28这样的区间,来避免模型陷入不断输出下一个数字的坏情况
  • 问题:模型会认为大部分代码行都和query “有点关系”,相关度难以量化
    • 解法:借鉴Provence的prompt设计,让模型仅根据提供的代码片段回答query,并要求引用,仅将回答中引用到的片段视为相关部分
  • 问题:由于行数过多还是会出现一些bad case
    • 解法:将构建的数据集使用Qwen3-Next-80B-A3B-Thinking筛选,判断其在代码结构保留程度,query质量和相关性质量上的表现。代码结构的筛选也是我们的方法在裁剪后AST正确率上领先其他方法的重要基础,在编程语言这种高度结构化的文本中,缺少一个 :} 都会导致整个语法树崩溃。

      压缩方法AST 语法正确率 (%)
      无压缩 (Full Context)98.5
      Token 级压缩 (如 LLMLingua2)0.29
      Function RAG92.3
      Pruner + Function RAG87.3

为什么选择这些模型/数据集?

我们认为,github的原始代码最能够保证和真实用户数据的分布接近,同时具备足够的、真实场景的多样性

使用Qwen-Coder-30B-A3B-Instruct有以下几个原因:

  • 性能稳定,我们抽取了部分样例人工校验模型认为的“相关代码”是否符合人类的观点,它的符合程度相当高
  • 速度非常快——我们使用vllm的offline推理模式在24h左右的时间,8 * H100机器上完成了20万条数据的构建和筛选;而对应的GLM-Z1-32B等dense模型则需要一周

而使用Qwen3-Next-80B-A3B-Thinking则是在含有思维链情况下平衡速度和质量的一个优解

评估1:单轮任务

我们在多个数据集上对比了不同模型的性能,包括:

  • LongCodeQA
  • LongCodeCompletion

评估的方法包括 我们的模型, RAG, LLMlingua2, selective context, LongCodeZip 等, 实验中swe-pruner的压缩率向下取整,其他方法的压缩率向上取整(实在难调平)

image.png

主要发现:

  • 上下文噪声的去除不会影响QA的性能,甚至可能因为噪声的去除有所提升
  • 我们的模型在压缩率上要远远超出其他方法 —— 事实上为了调整不同方法的压缩率到近似可比,我们花了非常大的力气调整超参数
  • 我们的模型取得了所有压缩率下的SOTA

还有一个有趣的发现是,以Qwen3-Reranker-0.6B这种LLM based reranker作为基底,它的自定义instruction能力让下游任务的泛化更加轻松了。例如代码补全任务的query是上一行代码行而不是自然语言文字,这对于大部分模型其实是难以理解的,但我们可以在instruction中描述补全任务,使得模型在训练相对少的补全任务上性能也相当不错。

智能体服务:不止于 RAG

在构建了这样的模型之后,我们回到了我们的初始想法,如何将其融合到现代的Code Agent呢?虽然CodeRAG也是很广泛的范式,但也有相当多Agent是使用Grep作为搜索的工具,比如我们用得最多的Claude Code

更进一步的,会令上下文爆炸的,其实很多时候是ReadFile而不是Grep,我们想要削减上下文的开销的话,应该也对ReadFile处理好冗余上下文的问题

所以,我们设计了一个最小侵入的highlight模型接入方式,它不仅能服务于CodeRAG,还能服务于Grep, ReadFile乃至任何阅读工具

具体的做法是,假设有一个工具A,会产生一些输出,现在在工具A的原始参数基础上,增加一个参数context focus question,SWE-Pruner就可以根据这个question对A的输出进行裁剪 —— 而A可以是cat, grep, rag search, 甚至像是我们论文之中更暴力的bash

image.png

image.png

Agent模型负责根据当前运行的需要产生question, 甚至不产生question(留空则不进行裁剪),让整个系统留下了更多的灵活性和随基础模型成长的可能——在我们的实验之中,我们确实看到了顶尖的模型会评估自己当前的需求、预估命令输出的长度,从而灵活地选择简单或是复杂的 question,又或是不使用question来保证关键命令的输出完整性

这种仅对工具做包装的方法保留了最大的兼容性,我们实现了claude code sdk和openhands sdk等框架的兼容示例,可以在 https://github.com/Ayanami1314/swe-pruner/blob/public/examples/README.md 中找到

评估2:多轮任务

我们使用SWE Bench Verified和SWE-QA两个benchmark来评估这套方案的agentic能力,并在claude-sonnet-4.5和glm-4.6上做了实现,接入mini-swe-agent进行测试,结果也是比较喜人的——不仅明显地节约了token的开销,并且在正确率上能打平甚至微升(sweqa, sonnet)

我们也针对不同的方法再次进行了消融对比,即使不考虑速率的因素,我们的方法也比简单的大模型总结更加优秀——代码是一个相当难总结的形式。

真实场景案例研究:精准识别核心语句

除了基准分数之外,让我们来看几个更有趣的例子,以直观地展示我们的模型在实际应用中的性能。

模型在查看文件时,使用了 Focus on the MRO resolution logic in Inherit Docstring 作为context focus question,而SWE-Pruner在大量代码中精准定位到了高亮逻辑 —— Docstring的类定义,mro的部分代码逻辑

image.png

另一个例子来自SWE-bench的 django__django-10554 问题,Baseline 智能体和搭载了 Pruner 的智能体表现出了截然不同的行为模式:

配置方案运行步数读取操作搜索操作编辑操作Token 总消耗
Baseline1645939257,001,934
Pruner562010111,170,160
  • 传统的 Baseline 智能体倾向于采用“广度优先”的搜索策略。它会盲目执行类似 find . -name "*.py" | grep union 的命令,并不断使用 sed 命令在大堆文件中反复读取片段。这种方式虽然覆盖面广,但会导致大量无关紧要的代码进入上下文,最终造成上下文溢出,让智能体在冗余信息中迷失方向。
  • 相比之下,搭载了 Pruner 的智能体表现得像个经验丰富的架构师。它直接定位到核心文件 django/db/models/sql/query.py,通过带行号的 cat -n 进行读取,并迅速锁定了关键代码分支(if self.combinator: ...),避免了 Baseline 的无效试错

站在巨人的肩膀上

该模型的开发建立在大量前期工作的基础上,我们谨此感谢所有为我们的这项工作做出贡献的人:


  1. Provence的理论基础: 提出了一种有效的训练Rerank-While-Prune模型的方法,即句级打标与双头损失
  2. Qwen团队的优秀模型:Qwen3-Reranker, Qwen3-Coder-30B-A3B-Instruct, Qwen3-Next-80B-A3B-Thinking

在这些基础之上,我们提出了若干创新

  1. 将Provence双头设计首次迁移到Qwen3-Reranker这样的decoder-only LLM上
  2. 从github代码仓库中构建部分相关查询-代码对的数据集的完整管道与数据
  3. 模型结构和训练方法的优化
  4. 实际Agentic系统的接入方法创新,提出了context focus question这一工具包装机制并实验验证

我们诚挚感谢 Provence 团队和 Qwen 团队的奠基性工作。

Conclusion

在本文中,我们分享了我们从识别生产环境中 Coding Agent 中的上下文成本问题到构建最先进的代码语义高亮模型SWE-Pruner的历程。

SWE-Pruner 填补了粗粒度文件检索与细粒度 Token 压缩之间的空白, 并为代码、多轮的场景引入了一套新式的上下文压缩解决方案,比起对Agent的轨迹进行“事后”压缩,SWE-Pruner强调在“事前”过滤掉无关的代码,从根源上降低Agent在生产环境的token开销和注意力噪声。通过引入任务感知的行级剪枝,我们不仅显著降低了 AI 编码的经济成本,还提高了智能体在复杂代码库中的推理质量。

该框架具有极强的即插即用性,可作为标准中间件集成至 OpenHands,Claude Code 等主流 Agent 框架中,也可以与任意轨迹压缩方案集成,为构建更高效、更具扩展性的 AI 软件工程师提供了新的路径

Future Works & Limitations

  • 由于我们的算力有限,我们只进行了20万条数据、单语言环境下的迭代,多语言和更大规模数据量的scale或许是将来采用这项技术的团队的命题
  • 在Agent方案之中,模型产出的question质量对实验效果的影响较大,在将来社区如果能对基础模型进行后训练以专门适应包装形式的工具,预计效果仍会出现显著的提升
  • 生产环境的“屎山”代码仓库benchmark较为稀少,swebench等开源代码仓库还是不够“脏”,对于semantic highlight模型的发挥有一定劣势,希望在将来这个方法能被用于更冗长的代码库之中,并对这种极端情况下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和投机解码这种很有趣的对应

结构化输出与AI工具与Agent

· 17 min read
ayanami

假如大伙接到一个需求,需要把claude code接入jupyter前端(例如,在jupyter前端直接输入魔法指令和claude code交互,而后台claude code展示claude code的一些关键节点,工具调用,费用开销,输出结果等),会怎么做?

一种想法是,将claude code的输出塞到一个文件里面去,起一个后台线程读取这个文件,尝试解析之中的某些部分,再以插件的形式加载到jupyter前端

但带来了一个问题是,效果(尤其是工具数量upup,上下文长度upup后的效果)不稳定,纯prompt的形式约束claude code及时向这个文件中写入以向前端通信,在经过长的交互过程后,claude经常会把这个文件忘掉

那claude code直接全塞前端呢?

在claude code里面问一个问题,可能就是几千上万token的交互,全塞前端,那用户体验就烂掉了。

另一个很容易想到的方案是,那我们不要让他输出文件了,直接当场处理把,定义一些特殊块叫 display 之类的东西,在prompt里面指定这个块里面是什么格式,让他如果想要和前端输出的话,放到这个块里面

这样看起来比文件好一些,但带来了新的问题没解决,长上下文下,display块的结构偶尔会有不稳定,会有不少特殊的渲染格式如html等由于几个字符的差异退化成了纯文本

如何修复这个呢?一个简单的方法,也是你能在任意一个现在的agent中看到的,是及时判错,再把把错误的部分发给模型让他修复一下,但又带来了额外的开销,并且前端的呈现也收到影响

有没有更优雅的办法呢?

如果你做AI应用比较多的话,肯定注意到了这实际上是一个结构化输出(约束解码)的场景,但现在的问题是,输出不止是一个json,而是正常文本块和display块的交错

(对于不了解约束解码的简单介绍一下,就是把上层的json等约束编译成状态机之后,用于动态建立llm output logits的mask,从而杜绝输出非法输出的技术)

看起来似乎不能约束解码?但display块本身是可以约束解码的,好恶心。

让我们打开vllm文档,翻到 Structured Outputs,你会发现,除了常见的regex约束解码之外,还有两种更强语义的解决方案,救赎之道就在其中,ebnf解码和structure tags解码

实际上,json解码只不过是ebnf解码的特殊情况罢了,毕竟实际都是状态机 (不知道ebnf是什么的同学,可以搜索一下编译前端,BNF范式,就能看懂下面的示例啦)

官方给的一个ebnf解码的例子如下, 用于执行一个简化sql的约束解码以提升sql正确率

simplified_sql_grammar = """
root ::= select_statement

select_statement ::= "SELECT " column " from " table " where " condition

column ::= "col_1 " | "col_2 "

table ::= "table_1 " | "table_2 "

condition ::= column "= " number

number ::= "1 " | "2 "
"""

completion = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": "Generate an SQL query to show the 'username' and 'email' from the 'users' table.",
}
],
extra_body={"guided_grammar": simplified_sql_grammar},
)
print(completion.choices[0].message.content)

如果放到这个问题,我们可以快乐地写出类似这样的定义

output := (display | normal text) *
display := (```display json ```)
json = ...
normal text = others

其中,display, json都是容易得到的,但恶心的地方在于什么是“others” 未拓展的ebnf是没有“非”定义的,从实操上虽然感觉可行(mask token取反),但这下已经没有支持了

(但ebnf解码肯定是有大用的,还是以Text2SQL举例,任何一个数据库都会给你他们的解析引擎的ebnf定义,都不需要你写)

怎么办呢,就带来了最后一个冷门工具,structured tags, 我先上代码,

def get_structural_tag_params(
tags: list[StructuralTag], triggers: list[str]
) -> dict:
return {
"type": "structural_tag",
"structures": [model.model_dump() for model in tags],
"triggers": triggers,
}

model_v2 = ChatOpenAI(
base_url=base_url,
model=model_name,
api_key=api_key,
temperature=0.15,
top_p=0.9,
extra_body={
"response_format": get_structural_tag_params(
tags=[
StructuralTag(
begin="<block=text>",
end="</block>",
schema=TextMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=image>",
end="</block>",
schema=ImageMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=tool_use>",
end="</block>",
schema=ToolUseMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=todo_list>",
end="</block>",
schema=TodoListMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=html>",
end="</block>",
schema=HTMLMsgSchema.model_json_schema(),
),
],
triggers=["<block="],
)
},
)

这个tags + triggers, 就是structured output的关键之处,它允许我们在trigger触发的时候才开始约束解码,在end结束的时候停止约束解码

至此,这个工作已经做完了


那约束解码和不约束带来的效果差距有多大呢,我在24B的Mistral-Small上做了个实验 最后的结果直接尝试解析后渲染到前端

Prompt如下,

sys_prompt = f"""

你是一个agent模型,你负责处理用户的问题,发起工具调用, 绘制图片、html、获取文本等。

由于你的token交互量很大,不是所有信息都需要展示给前端。

你可以正常思考和输出,但你需要将你认为需要展示给用户的有效信息包裹在 `<block={{tag}}> {{schema}} </block>` 中。

前端会将这部分内容进行渲染,交给用户。

你现在可用的tag有:

tags: "text", "image", "tool_use", "todo_list", "html"

对应的schema(pydantic格式)如下:

- {schemas_str}

例如,你可以先产生一个todo list,然后不断执行子任务,并更新todo list,直到所有任务完成。

由于你现在没有接入工具调用,所以对于所有工具调用交互,你只需要“假装”执行了工具调用并得到一个合理的响应就行,这是一个debug环境,

你需要根据用户的问题尽可能多的展示不同的block,并给出一个合理的响应。

"""

这个prompt下,<block=text>111</block> 这种就取代了上文所述的display块的效果

只定义了五种特殊的前端展示格式,文本,图片,TODO list,工具调用和HTML块

效果对比如下:

用户:帮我完成编写一个论坛帖子,打开浏览器的水源社区论坛,登录之后在discourse发帖的流程。

alt text alt text alt text alt text

可以看到,左侧没有约束解码的模型,在这样的任务负载下,json 参数就已经频频出现失误了,而右边的即使是24B模型的fp8量化非思考版本,却跑出了几百B agent的气势,并且token开销是来回倒腾的几分之一


一点感想: 我们常说一个子领域的知识对于另一个子领域是用处寥寥的,然而,这不是拒绝新领域知识的理由啊,vllm和xgrammer、outlines这种框架都把几种更强大的结构化解码方法摆到人们的脸上了,还是能在知乎看到“ebnf好像是编译原理的内容,(作为后端程序员)跳过”,或者是在各种开源仓库中还在广泛使用的拿prompt指导llm输出,完全不考虑(甚至不知道)结构化输出这样的东西

现在的后端、infra、算法,又有多少更深的优化方案是独立的呢?今天在看snowflakes优化方案,真是把上层算法和底层infra相辅相成,只是缺乏探索性的人们,会拿"这不是我的工作,这是专攻模型/infra/算法的人的工作"搪塞,最后又堆起来一个prompt史山罢了

在现在的agent框架中,充斥的也是prompt的兜底方案,带来的是qwen3-coder几个问题爆掉用户百万token,带来的是claude code问个“你是谁”都要花一角钱,但有没有一种可能,我们本可以用更确定的东西呢?LLM是一种万能的模糊推理,但好钢也要用在刀刃上啊。

参考完整代码如下

# ruff: noqa: E501
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import argparse
import asyncio
import enum
import json
import os
import re
from pathlib import Path
from typing import Any

import colorlog
import openai
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from dotenv import load_dotenv

load_dotenv()


class StructuralTag(BaseModel):
begin: str
end: str
schema: dict[
str, Any
] # JSON schema for validation, model_dump by pydantic model


class TextMsgSchema(BaseModel):
text: str = Field(..., description="Text message")

def to_html(self) -> str:
"""Render text message as HTML"""
return f'<div class="text-message">{self.text}</div>'


class HTMLMsgSchema(BaseModel):
raw_html: str = Field(..., description="raw html str, like <div></div>")

def to_html(self) -> str:
"""Render HTML message as HTML"""
return f'<div class="html-message">{self.raw_html}</div>'


class ImageMsgSchema(BaseModel):
image_url: str = Field(..., description="Image URL")
image_name: str = Field(..., description="Image name")

def to_html(self) -> str:
"""Render image message as HTML"""
return f"""<div class="image-message">
<img src="{self.image_url}" alt="{self.image_name}" style="max-width: 100%; height: auto;">
<p class="image-caption">{self.image_name}</p>
</div>"""


class ToolUseMsgSchema(BaseModel):
tool_name: str = Field(..., description="Tool name")
args: dict[str, Any] = Field(..., description="Tool args")
tool_output: dict[str, Any] = Field(..., description="Tool output")

def to_html(self) -> str:
"""Render tool use message as HTML"""
args_html = json.dumps(self.args, indent=2, ensure_ascii=False)
output_html = json.dumps(self.tool_output, indent=2, ensure_ascii=False)
return f"""<div class="tool-use-message">
<h4>Tool: {self.tool_name}</h4>
<div class="tool-args">
<strong>Arguments:</strong>
<pre>{args_html}</pre>
</div>
<div class="tool-output">
<strong>Output:</strong>
<pre>{output_html}</pre>
</div>
</div>"""


class TodoListMsgSchema(BaseModel):
todo_list: list[tuple[bool, str]] = Field(..., description="Todo list")

def to_html(self) -> str:
"""Render todo list message as HTML"""
items = []
for done, item in self.todo_list:
checked = "checked" if done else ""
item_class = "completed" if done else "pending"
items.append(
f'<li class="{item_class}"><input type="checkbox" {checked} disabled> {item}</li>'
)
items_html = "\n".join(items)
return f"""<div class="todo-list-message">
<h4>Todo List</h4>
<ul class="todo-list">
{items_html}
</ul>
</div>"""


def get_structural_tag_params(
tags: list[StructuralTag], triggers: list[str]
) -> dict:
return {
"type": "structural_tag",
"structures": [model.model_dump() for model in tags],
"triggers": triggers,
}


def parse_structured_response(response: str) -> str:
"""Parse structured response and convert blocks to HTML"""
# Schema mapping
schema_classes = {
"text": TextMsgSchema,
"image": ImageMsgSchema,
"tool_use": ToolUseMsgSchema,
"todo_list": TodoListMsgSchema,
"html": HTMLMsgSchema,
}

def replace_block(match):
tag_type = match.group(1)
content = match.group(2).strip()

if tag_type not in schema_classes:
return match.group(0) # Return original if unknown tag

try:
# Parse JSON content
data = json.loads(content)
# Create schema instance
schema_instance = schema_classes[tag_type](**data)
# Return HTML
return schema_instance.to_html()
except (json.JSONDecodeError, ValueError) as e:
return (
f'<div class="error">Error parsing {tag_type} block: {e}</div>'
)

# Replace all <block=type>content</block> with HTML
pattern = r"<block=(\w+)>\s*(.*?)\s*</block>"
return re.sub(pattern, replace_block, response, flags=re.DOTALL)


def create_comparison_html(response1: str, response2: str) -> str:
"""Create a comparison HTML page with both responses"""
parsed_response2 = parse_structured_response(response2)

css = """
<style>
body {
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
margin: 0;
padding: 20px;
background-color: #f5f5f5;
}
.container {
max-width: 1200px;
margin: 0 auto;
background: white;
border-radius: 8px;
box-shadow: 0 2px 10px rgba(0,0,0,0.1);
overflow: hidden;
}
.header {
background: #2563eb;
color: white;
padding: 20px;
text-align: center;
}
.comparison {
display: flex;
min-height: 600px;
}
.column {
flex: 1;
padding: 20px;
border-right: 1px solid #e5e5e5;
}
.column:last-child {
border-right: none;
}
.column h3 {
margin-top: 0;
color: #1f2937;
border-bottom: 2px solid #e5e5e5;
padding-bottom: 10px;
}
.content {
line-height: 1.6;
color: #374151;
}

/* Schema-specific styles */
.text-message {
background: #f8fafc;
padding: 15px;
border-radius: 8px;
margin: 10px 0;
border-left: 4px solid #3b82f6;
}
.image-message {
background: #f0fdf4;
padding: 15px;
border-radius: 8px;
margin: 10px 0;
border-left: 4px solid #10b981;
text-align: center;
}
.image-caption {
margin: 10px 0 0 0;
font-style: italic;
color: #6b7280;
}
.tool-use-message {
background: #fefce8;
padding: 15px;
border-radius: 8px;
margin: 10px 0;
border-left: 4px solid #eab308;
}
.tool-use-message h4 {
margin: 0 0 10px 0;
color: #92400e;
}
.tool-args, .tool-output {
margin: 10px 0;
}
.tool-args pre, .tool-output pre {
background: #1f2937;
color: #f9fafb;
padding: 10px;
border-radius: 4px;
overflow-x: auto;
}
.todo-list-message {
background: #fdf2f8;
padding: 15px;
border-radius: 8px;
margin: 10px 0;
border-left: 4px solid #ec4899;
}
.todo-list-message h4 {
margin: 0 0 10px 0;
color: #be185d;
}
.todo-list {
list-style: none;
padding: 0;
}
.todo-list li {
margin: 5px 0;
padding: 5px 0;
}
.todo-list li.completed {
text-decoration: line-through;
opacity: 0.7;
}
.html-message {
background: #f5f3ff;
padding: 15px;
border-radius: 8px;
margin: 10px 0;
border-left: 4px solid #8b5cf6;
}
.error {
background: #fef2f2;
color: #dc2626;
padding: 15px;
border-radius: 8px;
margin: 10px 0;
border-left: 4px solid #dc2626;
}
pre {
white-space: pre-wrap;
word-wrap: break-word;
}
</style>
"""

return f"""
<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>模型响应对比</title>
{css}
</head>
<body>
<div class="container">
<div class="header">
<h1>模型响应对比</h1>
<p>左侧:无结构化标签 | 右侧:带结构化标签(已渲染)</p>
</div>
<div class="comparison">
<div class="column">
<h3>无 Structure Tag</h3>
<div class="content">
<pre>{response1}</pre>
</div>
</div>
<div class="column">
<h3>Structure Tag(已渲染)</h3>
<div class="content">
{parsed_response2}
</div>
</div>
</div>
</div>
</body>
</html>
"""


if __name__ == "__main__":
base_url = "localhost:8000/v1"
model = openai.OpenAI(base_url=base_url, api_key="sk-")
schemas = [
TextMsgSchema.model_json_schema(),
ImageMsgSchema.model_json_schema(),
ToolUseMsgSchema.model_json_schema(),
TodoListMsgSchema.model_json_schema(),
HTMLMsgSchema.model_json_schema(),
]
schemas_str = "\n- ".join([json.dumps(s, indent=4) for s in schemas])
sys_prompt = f"""
你是一个agent模型,你负责处理用户的问题,发起工具调用, 绘制图片、html、获取文本等。
由于你的token交互量很大,不是所有信息都需要展示给前端。
你可以正常思考和输出,但你需要将你认为需要展示给用户的有效信息包裹在 `<block={{tag}}> {{schema}} </block>` 中。
前端会将这部分内容进行渲染,交给用户。

你现在可用的tag有:
tags: "text", "image", "tool_use", "todo_list", "html"
对应的schema(pydantic格式)如下:
- {schemas_str}

例如,你可以先产生一个todo list,然后不断执行子任务,并更新todo list,直到所有任务完成。
由于你现在没有接入工具调用,所以对于所有工具调用交互,你只需要“假装”执行了工具调用并得到一个合理的响应就行,这是一个debug环境,
你需要根据用户的问题尽可能多的展示不同的block,并给出一个合理的响应。

"""
base_url = "http://localhost:8000/v1"
model_name = "stelterlab/Mistral-Small-3.2-24B-Instruct-2506-FP8"
api_key = "sk-"
model = ChatOpenAI(
base_url=base_url,
model=model_name,
api_key=api_key,
temperature=0.15,
top_p=0.9,
)
print("-" * 50)
logger = colorlog.getLogger("Agent")
msgs = [
SystemMessage(content=sys_prompt),
HumanMessage(
content="帮我完成编写一个论坛帖子,打开浏览器的水源社区论坛,登录之后在discourse发帖的流程。"
),
]

model_v2 = ChatOpenAI(
base_url=base_url,
model=model_name,
api_key=api_key,
temperature=0.15,
top_p=0.9,
extra_body={
"response_format": get_structural_tag_params(
tags=[
StructuralTag(
begin="<block=text>",
end="</block>",
schema=TextMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=image>",
end="</block>",
schema=ImageMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=tool_use>",
end="</block>",
schema=ToolUseMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=todo_list>",
end="</block>",
schema=TodoListMsgSchema.model_json_schema(),
),
StructuralTag(
begin="<block=html>",
end="</block>",
schema=HTMLMsgSchema.model_json_schema(),
),
],
triggers=["<block="],
)
},
)

logger.info("=== 测试开始 ===")
response1 = model.invoke(msgs).content
logger.info(f"=== 测试结束 ===\n{response1}")

logger.info("=== 测试开始 ===")
response2 = model_v2.invoke(msgs).content
logger.info(f"=== 测试结束 ===\n{response2}")

# 生成对比HTML文件
comparison_html = create_comparison_html(response1, response2)
Path("tmp/test_comparison.html").write_text(
comparison_html, encoding="utf-8"
)
logger.info("已生成对比HTML文件: tmp/test_comparison.html")

# 保留原有的Markdown文件
with Path("tmp/test_diff.md").open("w", encoding="utf-8") as f:
f.write("无structure tag: \n")
f.write(response1)
f.write("\n\nstructure tag: \n")
f.write(response2)