最近我在折腾检索增强生成，底座试了一圈之后，盯上了 RWKV 的状态。确切地说，我在尝试“EmbeddingRWKV State-Centric Retrieval with Reusable States”这条技术路线：用 RWKV 的循环状态而不是普通 embedding 向量来组织检索。这篇文章适合正在做 RAG、想尝试状态化记忆方案，或者对 RWKV 推理优化感兴趣的读者。我会把“为什么状态适合检索”“状态索引怎么建”“状态怎么复用”“训练怎么对齐”这四件事讲透，中间穿插我在实际项目里踩过的坑和实测数据。
1. 为什么我要把检索的载体从 Embedding 换成 State
1.1 传统切块向量检索在长文档上的真实痛点
先从一个很常见的场景说起。假设知识库里放的是产品排障手册，一个故障的处理过程分为“现象确认—原因定位—参数修改—结果验证”四个阶段，每个阶段彼此依赖。传统做法是把这个手册按段落切块，每块单独过 embedding 模型，存进向量库。问题很快就来了。
切块之后的单块往往失去了语境。第一阶段的文本里会写着“记录当前设备的错误码”，如果不看后文，根本不知道这个错误码是用来跟第二阶段定位结果做比对的。向量库检索时，query 是“设备重启后错误码仍存在”，它能召回“错误码定义”那段，却很难召回“参数修改后如何验证错误码消失”那段，因为后者表面上看起来和 query 没有词面重叠。
我试过各种切块策略：重叠窗口、按标题切、按语义段落切。重叠窗口能缓解一点，但本质问题没有解决：embedding 表达的是“这一小段文本单独看起来在说什么”，而不是“这一小段文本在整个上下文流程中承担什么角色”。FAQ 类短文本无所谓，因为答案本身是自包含的；但流程型、排障型、技术手册型内容，恰恰是语境依赖最重的场景，传统方案在最需要它的地方反而最弱。
1.2 RWKV 的 State 是“演化的记忆”而向量是“静止的照片”
RWKV 是线性复杂度的递归式架构，它跟普通 RNN 最大的不同在于：每读入一个 token，内部会更新一组固定大小的循环状态，这个状态既包含了对历史信息的压缩，又是预测下一个 token 时的重要依据。你可以把状态理解成模型读完整段文本后的“动态摘要”，它的维度通常是 hidden size 的倍数，比如 768 hidden 对应 3072 维状态。
这和 embedding 有本质区别。embedding 是对一段文本做一次前向得到的一个静态向量，文本内部的顺序关系被压缩成单一的语义点；RWKV 的状态则是一个时间演化过程，每读一个 token 就变一次，前一步的状态参与后一步的计算。换句话说，状态天然携带了“这段文本接下来会往哪个方向走”的预测性信息。
而检索的本质恰恰是判断“哪一段内容接在 query 后面最合理”。这里有个信息论层面的动机：RWKV 的状态是为了预测下一步 token 而优化的，它包含了条件分布意义上的充分统计量；你要检索的东西，往往正是在这个条件分布里概率最高的区域。所以拿状态做检索，理论和实际都对得上。
1.3 State-Centric 比 Embedding-Centric 多出来的本质差异
“State-Centric”这个词，我想强调的不只是换一个向量那么简单。普通向量检索是 query-embedding 和 doc-embedding 之间算一次静态相似度，query 和文档是两座孤岛。状态中心检索让 query 也走一遍 RWKV 的递归编码，得到一个 query-state，再与文档的状态快照序列做匹配。
这两者的交互方式完全不同：
| 对比维度 | Embedding 检索 | State 中心检索 |
|---|---|---|
| query 编码 | 一次静态前向，得到固定向量 | 递归编码，状态沿序列演化 |
| 文档编码 | 每块独立编码，块间无关联 | 同一文档共享前序状态，块间上下文连续 |
| 匹配对象 | 空间中的两个点 | 两个演化轨迹的状态点 |
| 上下文感知 | 弱，单块语义 | 强，蕴含前文积累 |
| 对长流程文档 | 表现常受切块影响 | 结构信息保留较完整 |
这里有一个容易被忽略的差异：状态中心检索的“方向性”。embedding 相似度是对称的，query 和 doc 换个位置结果一样；但状态匹配天然有方向，因为 query-state 是“我要找一段能接下去的内容”的状态，doc-state 是“我这段内容自身演化的终点”，两者在语义角色上就不同。这也是后来我坚持给 query 和 doc 分别做投影的原因之一，后面会细讲。
2. 状态快照、投影对齐与匹配打分：三段式设计
2.1 文档侧：按 token 序列生成状态快照
文档侧要做的事，是把一篇文本编码成一条状态序列。实现上很简单：把文档 token 化之后逐 token 前向 RWKV，每一步都拿到一个当前状态。关键问题是采样策略——状态快照不能每个 token 都存，否则索引体积完全失控。
我用的模型 hidden size 是 768，state 维度 3072。如果每个 token 存一个 float32 状态，那就是每个 token 12KB；一篇 10 万 token 的文档就要 1.2GB 索引，这还没算向量索引本身的开销。我实测的结论是：按 32 token 间隔采样，索引体积缩到原来的 1/32，检索效果却几乎不掉；间隔压到 16，收益也仅仅提升 1-2 个点，成本却翻倍。所以默认 32 的间隔是性价比比较高的选择。
采样时还有两个细节要注意。第一，状态要取在 token 读入之后、更新完毕的状态，不要取读入之前的状态，否则当前 token 的信息没进去。第二，每个状态快照必须附带 token 偏移和文本边界信息。做重排阶段要回源拿到对应的原始文本片段，没有这两个字段就得重算一遍，非常浪费。
2.2 Query 侧：状态投影对齐与 EmbeddingRWKV 的含义
query 侧的编码方式和文档一样，也是走 RWKV 得到 query-state。但直接拿 query-state 和 doc-state 做相似度，效果会比较差。原因有两个：其一，query 通常很短，状态演化的步数很少，和长文档状态处于不同的分布区间；其二，RWKV 的状态是为“预测下一个 token”服务的，它的坐标系对语义判别并不友好。
所以我在状态之上加了一个投影层。具体做法是 query-state 和 doc-state 分别过一个 MLP，把它们映射到一个公共检索空间。投影层做两件事：把状态从“语言建模特征”筛选成“检索判别特征”，同时拉齐 query 和 doc 之间的分布差异。投影之后再算余弦相似度，才有稳定的区分度。
这就解释了标题里 EmbeddingRWKV 的含义：它不是一个输入侧的静态 embedding，而是从 RWKV 循环状态中派生出来的、动态生成的检索 embedding。同一段文本在不同上下文位置会得到不同的状态，投影后也会得到不同的检索向量——这是它和静态 embedding 最大的区别。
2.3 匹配打分：先余弦，后打分器
匹配阶段我推荐分两步走。第一步用归一化余弦相似度加温度缩放。高维状态向量经过投影之后，如果不做归一化直接算内积，个别维度的方差会主导结果；归一化之后配合一个温度系数把分数拉开，Top-K 才稳定。
第二步是当余弦不够用的时候再上浅层打分器。我遇到过一种情况：故障排查类文档里，两个片段的前因描述高度相似，但结论方向截然相反，余弦相似度很难区分。这种场景可以在投影之后接一个两层 MLP，把 query 特征和 doc 特征拼起来预测匹配分数。代价是索引和检索阶段多一次推理，但对精度要求高的场景值得。如果预算有限，我建议先扛住余弦方案，把评测口径搞清楚了再决定要不要上打分器。
3. Reusable States：一次状态计算在三个层级上的复用
3.1 文档内部的前缀复用：增量索引的核心收益
RWKV 的状态是增量更新的，这带来一个直接的好处：同一篇文档内，后面的块可以复用前面的状态，不用重新编码整篇文档。
我在做切块索引时，把同一文档的状态快照按 token 偏移组织成一个链表结构。每次新增或追加文档，只需要从最后一条已保存的状态继续前向计算，不需要重头跑一遍全文。实测下来，一篇 50 万 token 的长文档追加 1 万 token 的新内容，索引耗时几乎只等于那 1 万 token 的编码时间，而不是整篇文档的编码时间。这对于知识库高频更新的项目来说是一个很大的工程收益。
这里有个隐含的约束：RWKV 的状态只对同一段连续文本有效。如果文档中间被插入或改写了内容，那么插入点之后的状态全部失效，必须从最近的失效点重新编码。好在知识库更新通常是追加为主，真正随机改写的情况不多，所以前缀复用策略在大多数场景下都成立。
3.2 跨检索任务的复用：状态是记忆基底，不是一次性钥匙
传统向量索引里，一条 doc-embedding 被检索一次之后，它的价值就消耗完了。状态不一样：同一份状态保存下来之后，不管未来有多少个 query 都可以复用，而且还能作为生成侧的上下文缓存一起参与解码。
我把这个能力总结成“状态是记忆基底”。具体做法是：文档编码只做一次，得到的状态序列同时写入检索索引和生成状态缓存。检索命中后，直接把对应位置的连续状态载入生成器，而不需要重新把检索到的文本再编码一遍。一次计算、两处复用，在长文本场景下省下的推理时间非常可观。
要注意的是跨模型复用和跨任务复用的边界。不同规模的模型状态维度不同，不能混用；同一个模型在同一个部署场景下做任务迁移，状态才能直接复用。如果换了 prompt 模板或者改了采样参数，建议重新评估一下状态分布是否偏移，再决定是否复用旧缓存。
3.3 把检索结果写回状态：真正的 State-Centric 生成注入
RAG 的常规做法是检索结果拼进 prompt 里，作为文本前缀喂给生成模型。我在做状态方案时的另一个核心改动是：把命中片段重新编码成状态后，直接写入生成模型的循环状态通道，而不是写进文本前缀。
这个改动解决了两个实际痛点。第一，长文档片段拼进 prompt 之后，模型注意力容易被超长前缀带偏，出现位置分心；状态注入不污染 token 序列的位置编码，模型读到的是“我已经理解了参考内容”，而不是“我又要被迫读一段长前缀”。第二，状态通道和文本通道共享同一套 RWKV 内部计算，命中片段的状态可以直接参与下一个 token 的预测，信息利用效率更高。
我在这块做过一个对比实验：同样命中率下，状态注入比文本拼接的方式流畅度更高，事实一致性也更好。原因其实不难理解：状态在递归更新过程中已经做了信息融合，而文本拼接相当于把所有参考内容重新交给模型去读一遍，后者浪费了额外的注意力开销。
4. 训练目标与状态对齐的工程细节
4.1 双目标训练：保留语言能力，叠加检索能力
把状态直接丢进对比学习损失，一个直觉的方案是让正样本对（query-state 和对应 doc-state）距离拉近、负样本拉远。但实际跑一版就发现问题：梯度穿过整个循环状态链向前传播，RWKV 的语言建模能力迅速劣化，生成质量掉得很明显。
所以我采用了双目标训练：主任务保留原始的语言建模损失，辅助任务叠加检索损失。这一步的关键是权重。检索损失权重太小不收敛，太大伤语言建模。在 768 hidden 的模型上，我把检索损失权重压在 0.05 到 0.15 之间，语言建模指标基本不掉，检索测试的召回率则稳步上升。
一个更稳妥的简化方案是：直接冻结 RWKV 主干，只训练投影层和评分模块，检索损失只作用于投影层。这个方案在任何规模的模型上都不会伤语言能力，适合不想大动训练流程的团队。代价是检索指标的上限会略低于主干参与微调的版本。
4.2 负样本、温度系数与 Batch 构造的实践值
检索损失我用的是 InfoNCE 式的对比损失，公式不展开了，但三个实践参数值得记录。
温度系数 τ 我落在 0.05 到 0.1 之间，最后常用 0.07。状态向量经过投影后分布相对聚拢，如果 τ 取 1，所有相似度都会被压平，损失几乎没有梯度；取太小又容易让训练发散，0.07 是实测比较稳的点。
负样本不能只靠 in-batch negatives。同一个 batch 里如果放的多是同一类文档，负样本之间的区分度很差，模型很快就分不清“有点相关的干扰项”和“真正不相关的噪声”。我用了一个简单有效的补救：训练过程中临时用当前模型当检索器，从索引里打捞出相似度最高的非正样本作为 hard negatives 加入损失。batch size 64 配 32 个 hard negatives，收敛速度和最终指标都明显好于纯 in-batch。
还有一个细节：同一个 batch 内尽量混入不同来源的文档类型。如果 batch 里全是同一套流程手册的不同片段，对比损失会引导模型关注“流水账式”的相似度，反而损害真正需要的语义判别能力。
4.3 状态主干冻结还是微调：一个踩坑经验
这一节全是我实际踩坑换来的经验。最初做全参数微调，检索测试集指标上得非常快，我一度以为方案已经成功了；结果一跑端到端生成，困惑度涨了快十个点，回答变得前言不搭后语。
后来的解决方式是分层学习率。RWKV 主干的学习率降为投影层的十分之一，投影层正常学习。这样跑下来，检索指标只比全参数微调低两三个点，但语言建模能力完全保住了。从那次之后我形成一个判断：状态主干的语言建模能力是检索能力的底座，丢了底座，线下检索指标再好看，端到端一跑生成就会露馅。
另外要注意数值稳定性。RWKV 的状态是逐步累积出来的，长文本下状态量级可能涨得很大。不同 chunk 的状态如果量级差异悬殊，对比损失的梯度会失衡。我在做状态对齐前会给状态做一次层归一化，这个简单操作对训练稳定性的提升非常显著。
5. 评测与消融：什么时候这个方案会输给纯向量检索
5.1 评测口径：检索指标和生成指标会出现背离
评测不能只看检索命中率。我同时看三组指标：检索侧 Recall@k 和 MRR、生成侧困惑度变化、以及人工评级的引用相关性。这三组经常出现背离。
最典型的一次：检索召回率提了七个点，但生成问答的评分反而掉了两个点。排查后发现，召回片段虽然准确，但状态注入时把检索单元后面的多余状态也带进了生成器，干扰了推理。这让我意识到，状态检索单元和分析窗口需要配套设计。我最后把检索单元限制在 256 token 左右，采样时只保留单元结尾的状态，生成侧的效果明显改善。
所以建议所有做这条路线的人：先建好端到端评测，再优化检索指标，否则很容易被线下指标误导。
5.2 消融实验结论：投影层比想象中重要得多
我做了三组消融：去掉投影层、冻结主干只训投影、放开主干全参微调。结论很明确。
去掉投影层，只用原始状态做余弦相似度，检索性能下降非常大。原始状态的信息虽然丰富，但它处于“为预测下一个 token 而组织”的坐标系里，并不直接以语义可比的形式排列，必须经过判别式的投影变换，这个环节省不得。
冻结主干只训投影层，是性价比最高的一档：检索指标提升超过十个点，生成指标基本不掉。放开主干还能再涨两三个点，但语言建模代价不小。除非你的推理任务完全不要求通用生成能力，否则我不建议直接全参微调，分层学习率是更好的折中。
5.3 双路召回：工程上最推荐的混合方案
最后给出我的最终选型：状态检索和纯向量检索不互相替代，而是互补。
离散知识点型的内容，比如 FAQ 单条答案、名词解释、API 参数说明，纯向量检索效果已经很好，成本又低，没必要换状态方案。长流程文档、强上下文依赖的数据，比如排障手册、操作流程、设计方案文档，状态检索明显胜出。
我最终用双路召回：状态检索做主通道，向量检索做兜底，两侧的结果合并后送交互式重排。工程实现并不复杂——两侧共用同一套文档解析流水线，只是落到不同的索引格式。在这个配置下，生成质量比单用任何一侧都高，而且状态索引的增量维护能力也能完整保留下来。
最后说点个人体会。做这条路线踩过的最大坑，是前期低估了投影层的作用，总以为 RWKV 的状态已经够强，拿来直接比就行，结果被真实数据教育了。如果重新做一次，我会把预算的七成花在状态对齐和评测口径上，而不是花在前端的索引工程。还有一个值得留的小技巧：做状态快照采样时，同一批次里一定缓存 token 偏移和文本边界，重排环节会反复用到这两个字段，省得到时候再回源查一次文本。
