本文是一个超短文，不过后面具体方案倒是不短。
正文Post-training阶段的RL直接从reward学习这点大家都已经知道，OpenAI o1和DeepSeek R1都为我们打了样，这里不展开。
从API往应用层来看，workflow也可以纳入到RL过程中，直接针对一个多轮之后的业务reward进行学习。当然这里如果能把多个节点都使用同一个LLM的话，RFT时的显存占用会更小一些，不必维持多个模型。
考虑到目前LLM模型context大多还不是特别长，那么tool调用和RAG过程似乎也应该进行优化。例如ODR产品访问了大量的网页，但这些网页无论有用没用，真的都放入到context中了吗？o3的context应该有200k以上，但真的这么干了么？感觉不像。这里是否有一种单独的网页筛选组件，从搜索到的网页中召回相关的内容放入context？这个组件该如何实现呢？它可以当作一个tool，但这个tool的实现仍然需要被优化，而似乎它应该也可以被纳入到整个RL过程中进行优化。
反过来考虑RL是否能够深入到模型内部，而不只是整体优化LLM的参数。
在模型架构上，对于各种稀疏注意力、线性注意力、混合注意力方案等，也在进行各种尝试，DeepSeek、MiniMax、Moonshot分别交了一份不同的答卷。但这些attention结构的设计仍然有拍脑袋的感觉。这个attention结构本身不能被作为RL优化的一部分么？似乎无法否定。不过RL仍然是在post-training阶段的，在pretrain阶段目前并不做RL，而attention的结构在pretrain阶段就要固定。似乎也可能设计一个attention结构，可以在pretrain阶段先进行训练，并在post-training阶段仍然可以被训练并显著地改变其召回策略。这样就可以通过RL过程对其进行直接针对最终reward的优化，并纳入到Full-Stack RL的过程中。
本文的新想法是：把attention的结构优化也纳入到针对业务reward的RL优化过程中，以此来解决玄学的attention结构炼丹问题。
模型层方案的抛砖引玉
为了说明思路确实还有某种可行性，这里给一个更具体的想法。当然现在LLM模型层的结构设计已经是相当专业的领域，需要同时考虑模型效果和硬件特性，已经不是我能随便想想就完全可行的。所以这只能算抛砖引玉，给大家一些思路启发。
这个设计针对以下场景：Long Context（>200k）；MoE架构。
一般认为，LLM在decode过程中，依赖的Context应该具有局部性，即相邻的token大概会使用相近的context。而粗粒度的Context block召回可以利用这特性，DeepSeek的NSA（Native Sparse Attention）方案就包括对于Context的粗粒度block召回。
然后再考虑MoE架构：目前MoE的E并不是针对“完整语义单元”，而是基于token的。也就是说一小段完整语义/能力可能被以更细粒度拆分到多个Expert上，在不同的token位置需要召回不同的Expert。但从人认知的角度上，同样的能力被聚合到少数Expert上似乎是更加合适的，而且这可以增强召回的Expert候选集在token生成过程中的局部性，减少Expert的切换。而这个Expert的局部性应该可以通过在训练时对其进行激励（优化）来实现。
那么我们就在Context block和Expert层面都具有一定的decode过程中的局部性，有希望减少decode过程中的context重新召回和Expert的重新召回均摊次数。当然正如NSA的方案中所体现的，Context的召回方式需要不止一种来满足各种场景的需求。而本文讨论的只是从超长Context中如何召回相关部分这一块。
如果单纯的只考虑RL，在post-training阶段可以有很多种设计思路。但我们很难脱离pretrain阶段来实现，而Pretrain阶段巨大的训练成本基础并不能让我们轻易的采用很多暴力的RL方式。
基于以上考虑，可以在context block和Expert的召回中设计一个两阶段的召回范围，第一层召回与现在一般的习惯相同，直接用于当前decode token。然后引入第二层召回范围，训练目标是召回未来decode过程中W个token需要的context block和Expert。在推理时，两层的召回范围都参与计算。如果发现第二层的候选元素的score超过了阈值，应该进入第一层范围了，就说明当前序列的生成正在脱离当前的局部性环境，应该再下一个token计算时重新触发更大范围的召回计算。在只有两层召回的情况下，这个召回计算就是全局重新计算。
第一层召回范围的选择方式可以使用目前NSA、MoE等现有方案的思路。关键在于新引入的第二层召回范围如何动态选择，选择组件肯定是参数化的，而学习目标就可以选择为未来W个token需要召回的元素的并集。不过计算每个token时，并不清楚别的位置上的token的召回情况，所以这需要在每次整个序列计算后增加一个新的环节，汇总别的位置的token召回情况，作为当前位置上第二层召回的拟合目标。
这样虽然增加了单独的过程，但在pretrain阶段应该仍然可以保持一个高效训练的方式。具体的开销增加预估我就不做了，这应该找模型层的团队进行专业测算。
在以上pretrain阶段的过程中，虽然有点RL的感觉，但其实仍然是一种监督学习方式，不涉及针对延迟反馈reward的优化。但在post-training阶段和RFT阶段，可以对第一层召回选择组件、第二层召回选择组件、判定何时该触发下一轮完整召回的组件等针对目标reward进行RL学习。
通过这种方式，就实现了在LLM模型内部，除了单纯的固定参数外，增加了可以被RL直接优化的部分，并实现计算成本的优化。这种方式的效果收益靠单纯的LLM参数直接优化无法完成。
希望能够给读者在模型结构设计上提供一些启发。
相关材料
Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
https://arxiv.org/abs/2502.11089
