AI知识库

53AI知识库

学习大模型的前沿技术与行业应用场景


LLM模型层+应用层 Full-Stack RL的一种设想

发布日期:2025-03-05 14:41:32 浏览次数: 1569 来源:孔某人的低维认知
推荐语

探索全栈强化学习在大型语言模型中的应用前景。

核心内容:
1. 应用层中将工作流纳入RL过程的设想
2. LLM模型context长度限制与优化策略
3. 注意力结构设计和RL优化的可能性与挑战

杨芳贤
53A创始人/腾讯云(TVP)最具价值专家

本文是一个超短文,不过后面具体方案倒是不短。

正文

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


53AI,企业落地大模型首选服务商

产品:场景落地咨询+大模型应用平台+行业解决方案

承诺:免费场景POC验证,效果验证后签署服务协议。零风险落地应用大模型,已交付160+中大型企业

联系我们

售前咨询
186 6662 7370
预约演示
185 8882 0121

微信扫码

添加专属顾问

回到顶部

加载中...

扫码咨询