Universal Reasoner: 冻结 LLM 的可组合即插即用推理器¶

会议: ICML 2026
arXiv: 2505.19075
代码: https://github.com/hangeol/UniR
领域: LLM 推理
关键词: 推理增强, 模块化推理, 可组合推理, 冻结 LLM, 可验证奖励

一句话总结¶

提出通用推理器（UniR）——通过训练独立的轻量推理模块来捕获奖励导向的推理行为，在推理时通过逻辑叠加与冻结 LLM 组合，实现无需微调冻结模型、跨模型大小转移和多任务可组合的推理增强。

研究背景与动机¶

领域现状：当前通过 RL 微调（RFT）增强 LLM 推理能力，但需大量计算和内存资源。PEFT 如 LoRA 试图降低成本，但仍有两大缺陷——（1）强烈依赖模型架构，跨不同尺度模型（3B→14B）转移性差；（2）多个 LoRA 适配器线性组合缺乏理论支撑。

现有痛点：无法在不访问 LLM 内部参数前提下灵活高效地增强推理能力；无法跨模型尺度复用训练好的推理能力；无法组合多个不同任务的推理模块。

核心矛盾：推理增强需要参数更新（传统微调）但模型通常冻结；多任务学习需要端到端重训（多目标冲突）。

本文目标：设计模块化、可转移、可组合的推理增强方法。

切入角度：观察到可验证奖励（如数学问题正确性）可通过轨迹级信号转换为令牌级指导。将轨迹级奖励建模为推理模块的对数概率之和，使得逻辑叠加成为自然的组合机制。

核心 idea：分离奖励模型训练与策略更新，训练专用推理模块 \(\pi_r\) 以最大化可验证奖励，推理时通过添加其逻辑 \(\log\pi_r\) 到冻结骨干 \(\pi_b\) 的逻辑来指导令牌生成。

方法详解¶

整体框架¶

UniR 将 LLM 推理增强分解为两阶段——训练阶段在较小骨干模型上训练推理模块 \(\pi_r\)，使用可验证奖励和 GRPO 算法；推理阶段将训练好的 \(\pi_r\) 与任意冻结 LLM 组合，通过逻辑叠加进行令牌级指导。

关键设计¶

轨迹奖励转令牌指导的理论映射:
- 功能：将全局轨迹奖励分解为每个令牌的指导信号。
- 核心思路：假设轨迹奖励可表示为推理模块对数概率之和 \(\frac{1}{\beta}r(x,y)=\sum_{t=1}^{|y|}\log\pi_r(y_t|x,y_{<t};\phi)\)。通过替代 KL 正则目标中的奖励，导出最优策略 \(\log\pi_\theta(y_t|x,y_{<t})=\log\pi_b(y_t|x,y_{<t})+\log\pi_r(y_t|x,y_{<t})-\log Z'(x,y_{<t})\)。定理 4.1 证明收敛时 \(\log\pi_r(y_t|x,y_{<t})=\frac{1}{\beta}Q^*(y_t|x,y_{<t})\)。
- 设计动机：传统轨迹级奖励无法直接指导每一步生成；通过结构化假设将其转换为令牌级指导。
推理模块的 GRPO 训练:
- 功能：在不修改冻结骨干前提下训练推理模块最大化可验证奖励。
- 核心思路：采用 GRPO，从 \(\pi_b\) 采样 G 个候选响应，计算每个响应的外部奖励 \(r_i\)，归一化优势 \(A_i=\frac{r_i-\text{mean}(\{r_1,...,r_G\})}{\text{std}(...)}\)，然后在 GRPO 目标上优化 \(\phi\)。梯度中的比率项自动消除骨干项只影响 \(\pi_r\)。
- 设计动机：GRPO 无需显式价值函数且对奖励稀疏友好，适合可验证奖励场景。
多模块逻辑叠加的可组合性:
- 功能：支持多个不同任务的推理模块在推理时无缝组合。
- 核心思路：对 \(N\) 个不同奖励函数 \(\{r_1,...,r_N\}\)，分别训练 \(N\) 个推理模块 \(\{\pi_r^1,...,\pi_r^N\}\)。推理时组合为 \(\log\pi_\theta(y_t|x,y_{<t})\propto\log\pi_b(y_t|x,y_{<t})+\sum_{i=1}^{N}\alpha_i\log\pi_r^i(y_t|x,y_{<t})\)，权重 \(\alpha_i\) 可动态调整。
- 设计动机：推理往往涉及多个约束；逻辑叠加既是原则性解又允许零代价组合。

实验关键数据¶

主实验¶

模型配置	GSM8K	MATH-500	AIME24	平均提升
Llama3.2-3B 骨干	65.6	33.0	3.7	基线
+ GRPO LoRA (3B)	65.8	32.1	6.0	-0.3%
UniR (1B + 3B)	77.5	48.8	7.3	+35.2%
Qwen2.5-3B 骨干	74.4	44.2	6.3	基线
UniR (1.5B + 3B)	75.6	48.3	8.1	+6.8%

消融实验¶

实验项	描述	性能
推理模块单独	1B \(\pi_r\) 独立生成	远低于组合
无推理模块	仅冻结 3B 骨干	65.6 (GSM8K)
+推理模块（有调优）	结合后	77.5
多模块（α₁=1, α₂=0.5）	数学+翻译加权	72.3

关键发现¶

弱强转移——1B 推理模块能有效指导更大模型（14B），无需重训。
可组合性验证——多模块加权组合在不同权重下表现稳健。
奖励内化——训练后推理模块对正确响应（r=1）的对数概率显著高于错误响应（r=0）。

亮点与洞察¶

优雅的理论基础：从 KL 正则多目标优化出发严格推导逻辑叠加。
零重训转移：推理模块跨模型尺度转移无需微调。
推理时组合灵活性：多模块加权融合在推理阶段即时可调。
实证 Q 函数学习：通过对数概率分离验证了推理模块确实在学习令牌级最优决策信号。

局限与展望¶

假设局限——方法假设轨迹奖励可分解为令牌对数概率之和。
推理模块能力瓶颈——受其初始大小和训练数据限制。
权重选择问题——多模块组合时权重 \(\alpha_i\) 需手调。
改进：探索非分离式奖励的令牌级分解；设计自适应权重学习。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 将轨迹级奖励系统地映射到推理模块的令牌级对数概率。
实验充分度: ⭐⭐⭐⭐⭐ 多个数学基准对比 + 消融验证弱强转移和多模块组合。
写作质量: ⭐⭐⭐⭐⭐ 动机清晰、推导严格、实验有据。
价值: ⭐⭐⭐⭐⭐ 解决冻结 LLM 推理增强的关键问题，通用、低成本、易部署。