Reasoning · ICML 2025 · arXiv 2501.18858

BRiTE：把「想要的那种思考」写成一个后验分布，用 RL 去逼近它

论文把 LLM 的生成写成概率图模型 X → Z → Y，再加一个评估信号 O。真正想要的是「已知答案正确时，思考过程长什么样」——一个难算的后验。 BRiTE 的 ψ 步用 RL 去近似这个后验，θ 步再拿近似出来的数据做 SFT 或 DPO。代价是每组数据要多跑一遍 RL，好处是同一套推导把 PPO、STaR、拒绝采样微调都收编了。

Han Zhong, Yutong Yin, Shenao Zhang, Xiaojun Xu, Zhaoran Wang 等 · 北大 / 西北 / 字节 / 明大 arXiv 2501.18858 ↗ ICML 2025 ↗

TL;DR · 30 秒抓住核心

拒绝采样类方法的 E 步是「采样 → 筛」，本质是用经验分布近似后验 P(z | x, y*, O*)。 BRiTE 指出这一步可以换成 RL：构造一个 MDP，使它的最优策略正好等于那个后验。

关键不是「用 PPO」这三个字，而是 reward shaping：把整条序列的总奖励取成 log P(z, y, O* | x, θ)，再按自回归拆成逐 token 奖励。β=1 时熵正则最优策略的概率正好正比于总奖励，于是「答案越可能正确的思考」被抽得越多——这是它和普通 outcome-reward RL 的根本分野。

理论贡献是同一件事的另一面：BRiTE 就是把 X, Z, Y, O 四个空间取不同子集，就分别还原出 SFT、 PPO、latent DPO、STaR 与 ReSTEM，收敛率统一为 1/T。实验上，Qwen2.5-7B 在 MATH500 上从基线的 44.1 提到 ψ-update 的 79.1，而同样数据量的拒绝采样只到 54.3。

79.1 vs 54.3

MATH500（Qwen2.5-7B，40K 混合数据）
BRiTE ψ-update vs 拒绝采样

50.5 vs 41.5

Gemma-2-9B-it 上 MATH 准确率
BRiTE vs 人工标注推理链 SFT

1/T

迭代 T 次后的收敛率
与 KL 初态−最优之差成正比

推理链坏掉的地方不在答案，在「怎么想到的」这一步

明白为什么「只看答案对不对」不足以解释推理质量，以及为什么思考过程天然是一个潜变量。

CoT 提示让模型把题目拆成步骤，确实提升了准确率。但论文指出的问题是：生成的 rationale 经常缺少逻辑完整性或有效性，质量高度依赖任务特定提示。同样的提示换个模型，或同一模型换个温度，正确的思考路径就会消失。推理时扩展（test-time scaling）也只是在缓解症状，没有解决「生成高质量推理过程」这个核心。

BRiTE 的做法是把它形式化成一张概率图：X 是问题，Z 是内部思考过程（潜在变量），Y 是答案，O 是评估信号。令 P(o | x, z, y) 与模型参数 θ 无关——它代表一个真值裁判，比如标准答案或单元测试。于是

P(z, y, o | x, θ) = P(z | x, θ) · P(y | x, z, θ) · P(o | x, z, y)

式 (3.1)：引入 Z 把 P(Y | X) 拆成 P(Z | X) 与 P(Y | X, Z)；引入 O 给「正确的思考」提供判别信号

学习目标是最大化 L(θ) = log P(z ∈ 𝒵, y ∈ 𝒴, o ∈ 𝒪 | x, θ)——生成出来的思考、答案和评估同时落在「想要的那个集合」里的概率。问题在于：这个目标对 z 求和，直接算不可行，必须引入某个辅助分布去近似。

核心矛盾

答案正确 ≠ 推理正确，而「正确推理」是一个不好采样的分布

如果一个模型 10 条里有 1 条答对，那这 1 条的推理也未必可信；如果 10 条里都答对，那根据答案筛出来的推理又长得差不多、没有多样性。真正想要的是「给定 (x, y*)，Z 应该如何分布」——这是对思考过程做的条件后验，不能靠暴力枚举得到。

两条旧路：人工标推理链烧钱，拒绝采样只会「筛」不会「学」

看清 EM 视角下 E 步的两种实现方式，才知道 BRiTE 的新意落在哪里。

Human SFT

人写推理链

数据集里给的是 (x, z*, y*)，直接把人类写的 z* 做 SFT。这是公认的金标准。

代价：标注贵、慢、不一致；前沿问题上，人类专家本身也未必有唯一的「标准思考」。

Rejection Sampling

生成后筛选

每题采 N = 30 条，留下答案正确的那条 z_rs 做 SFT（STaR / RS 微调 / ReSTEM 都属于这一类）。

代价：本质是「采样后的经验分布」近似后验。它不会主动提高正确推理的生成概率，只会从已有的生成里挑；样本一多就被稀释。

从 EM 的视角看，两者都是「E 步估计潜变量后验 → M 步最大化期望对数似然」，只是 E 步的实现分别是「用人工标注当后验」和「用拒绝采样近似后验」。 BRiTE 的切入点很明确：E 步为什么不能直接是一个可优化的策略？

θ_t+1 = argmax_θ Σ_(z,y) log P(z, y| x, θ) · P(1 | x, z, y) · P(1 | x, z, y) = 1{y 正确}

式 (3.7)：拒绝采样 EM 的形式——权重是 0/1 的硬指示函数，E 步只是在做筛选

BRiTE 换掉的正是这个 0/1 权重：不再「过就留下、不过就扔掉」，而是让一个 RL 策略去学习怎样的 z 更可能通向正确，输出的样本带着概率质量，而不是被硬切一刀。

把「E 步」写成 RL：最优策略正比于目标后验，奖励就是逐 token 的对数概率

理解 Lemma 3.1 的变分下界、Reward shaping 的选取方式，以及为什么它和普通的 outcome reward 不同。

第一步：变分下界（与 EM / VAE 同一套思想）

引入一个辅助分布 Q(z, y, o | x, ψ)，用 Lemma 3.1（即 KL 非负）把对数边缘似然写成对 Q 的下界最大化：

L(θ) = max_Q Σ_(z,y,o) log P(z, y, o | x, θ) · Q(z, y, o | x, ψ) − Σ_(z,y,o) log Q(z, y, o | x, ψ) · Q(z, y, o | x, ψ)

式 (3.3)：Lemma 3.1（KL 非负）给出 log ΣP ≥ E_Q[log P − log Q]，取 Q = P / ΣP 时等号成立

最优 Q 有闭式解，正是我们要的后验：

Q^*(z, y, o | x, ψ) = P(z, y, o | x, θ_t) / Σ P(z′, y′, o′ | x, θ_t) ∝ P(z, y, o | x, θ_t)

式 (3.4) · ψ-update：当 Y = {y*}、O = {1} 时，这就是 P(z | x, y*, 1)——「答案已知正确时，思考该是什么样」

接下来 θ 的更新是一个加权对数似然，对应 M 步：

θ_t+1 = argmax_θ Σ_(z,y,o) log P(z, y | x, θ) · Q(z, y, o | x, ψ_t+1)

式 (3.5) · θ-update：log P(o | x, z, y) 与 θ 无关，在优化中可丢弃

第二步：Q 难采样，那就交给 RL——关键是奖励怎么设

当 Y = {y*}、Z = 𝒵 时，Q* 就是后验 P(z | x, y*, 1)。LLM 是自回归的，要「在答案已知为 y* 的条件下」生成 z，天然不可直接采样。论文构造一个确定性熵正则 MDP，用 Proposition 3.7 说明：最优策略满足 π*(a_h…a_H | s_h) ∝ exp(1/β · Σ r)。

所以只要令 β = 1，并把整条序列的总奖励设为 log P(z, y, o | x, θ_t)，最优策略就恢复成 Q*。又因为生成是自回归的，这条总奖励可以按 token 拆开：

r_total(z, y, o) = Σ_j=1..τ log P(a_j | a_1:j−1, x)

Section 3.4：第 j 个 token 的奖励就是它在参考策略下的对数概率——不是结果奖励，是「这个 token 本来就有多可能出现」

与普通 RL 的分野

BRiTE 的 reward 不是「答对 +1」，而是「这个回答本来有多大概率」

通常的 outcome-reward RL 在所有正确回答上给同样的奖励（0/1）。BRiTE 对每个 token 给它的对数概率： 参考模型本来就更倾向生成的 token，奖励更高；被模型视为「意外」的 token，奖励更低。 加上 β=1 的温和正则，最终学到的策略不是「在所有正确回答上均匀分布」，而是「正确回答的概率像参考模型本身一样被加权」——这正是后验 ∝ 联合概率的形式。

论文的实现允许两种 reward 形式： (i) log P(z, y* | x)（无外部 verifier，训练集就是正确答案的集合）与 (ii) log P(z, y | x) + R(x, z, y)/β（有外部 verifier / 奖励函数）。前者对应「只用正确的题目」，后者对应「用一份额外的打分」。

自己动手：看到后验采样和拒绝采样的密度差异

下面是一个高度简化的示意：一个题目的思考过程空间被记为 20 个格子（代表 20 种候选 z），每格颜色深浅是它的「质量」 q(z) = p_θ(y*|x,z) · p_θ(z|x)。拖动滑块改变采样条数，对比两种方法最终拿去训练的数据分布。参数说明：T = 0.5 是采样温度，β = 1 是 BRiTE 的熵正则系数。

采样条数 N 30

采样温度 T 0.50

当前判断 RS

表中数字为解析计算的期望，不是随机模拟。

这个模拟假设了什么？点开看与真实算法的对应

构造：设 q(z) 在 z = 1, 4, 7, 12, 18 这五格上非零，取值的原始权重为 p_θ(z|x) 与「答对率」，二者相乘并归一化： q = [0.50, 0.20, 0.15, 0.10, 0.05]——一条典型的「少数几类推理容易被采到，其他稀薄」的分布。

RS 在温度 T 下的采样分布为 q^1/T 再归一化，被接受的样本（答对）又按「答对率」过滤一次，最终「被保留的比例 ∝ q(z)」。但这里衡量的不是保留数，而是我们期望能拿到多少条不同的高质量 z——RS 保留的条数 ≈ N·Σq，其中大部分集中在前几格；BRiTE 的最优策略则正比于 q(z)，每条采样都直接反映质量。

真实模型里，思考过程是长文本，q(z) 是高度复杂、无法枚举的分布；这个 20 格示意只保留了「分布很尖、且采样分布与目标分布不一致」这一层因果。

工程上就是「PPO 求 Q，再接 SFT / DPO」——开销集中在 ψ 步

明白三套实现（BRiTE-SFT、BRiTE-DPO、扩展数据）的超参和算力，以便判断复现成本。

ψ 步和 θ 步的实现开销差异很大：前者要跑 RL，后者是标准训练，可以接进任何已有的后训练管线。

ψ-update（RL）

贵：要跑 PPO / GRPO

论文用 PPO 或 GRPO 优化熵正则 MDP。训练配置（Appendix D）：

数学（4 个 instruct 模型）：4×H100，PPO，lr 5×10⁻⁷，batch 1
扩展实验（Qwen2.5-7B）：8×H100，GRPO/verl，lr 1×10⁻⁶，batch 512×8

θ-update（SFT / DPO）

便宜：标准训练

拿 Q 采出来的数据做常规训练，与已有的后训练管线兼容：

BRiTE-SFT：用 (x, z_Q, y*) 做 SFT，lr 5×10⁻⁵，batch 2
BRiTE-DPO：用 Q 生成 30 条/题，构造 (最优, 最差) 偏好对，跑 3 轮 DPO

两组实验都统一采用 LoRA（r = 32，α = 128）。三套实现的差异：

实现	ψ 步	θ 步	前提
BRiTE-SFT	PPO 求 Q（奖励形式 i 或 ii）	用 `(x, z_Q, y*)` 做 SFT	训练集全部为正确答案，或有 verifier
BRiTE-DPO	每轮用 Q 采 30 条	按答案正确性构造最佳/最差对，做 DPO（3 轮）	偏好数据是自动生成的
BRiTE 扩展（Qwen2.5）	40K 混合数据（STILL-3-Preview + MATH + 历史 AIME），GRPO	θ 步用 Q 采出的数据做 SFT	只是「训练数据更多 + 基座更强」

论文坦言，与小数据下（Table 1）BRITE 和 RS 的差距只有 1–10 分；扩展到 40K 数据 + 更强的 Qwen2.5-7B 之后，优势才拉开（4.3 节）。也就是说， ψ 步的 RL 开销换来的是「同样的数据量下更高质量的 z」，数据量本身也仍是瓶颈。

一条主循环：X → Z → Y → O → 更好的 Z

用这张流程图看清 BRiTE 在推理时和训练时各自做了什么，以及它对应 EM 的两步。

图 1（本页自绘，对应论文 Fig. 1 + 3.2 / 3.4 节流程）：BRiTE 的 ψ 步接受含 O 的联合信息，输出采样自后验的思考过程；θ 步把它当作训练数据。

这张图解释了为什么作者说「BRiTE 在推理时用到 RL」：ψ 步本身就是一次推理/生成——生成 z 时 Q 已经知道自己想要哪种推理，而不是等生成完再回来看答案。反过来说，θ 步可以落到任何标准后训练算法上，所以框架本身并不绑定 PPO。

效果：小数据 + 或强或弱的基座能追平人类标注，大数据 + 强基座拉开 25 分

Table 1、Table 2 与 Figure 2 里的数字，可以直接告诉你它能赢在哪，赢不了在哪。

Table 1：四个 instruct 模型，GSM8K / MATH

方法	Mistral-7B-v0.2		Gemma-1.1-7B		Gemma-2-9B		Llama-3-8B
	GSM8K	MATH	GSM8K	MATH	GSM8K	MATH	GSM8K	MATH
Base	41.8	9.8	49.0	18.8	81.3	37.3	79.2	28.3
SFT（人工标注）	52.8	13.6	57.5	19.6	80.1	41.5	72.6	27.1
RS	47.7	10.3	58.4	18.7	87.6	47.5	79.5	28.9
BRiTE	52.2	11.2	59.2	23.7	89.7	50.5	81.0	30.0

值得注意的三点：

BRiTE 全面优于 RS，但不是每个格子都赢 SFT：Mistral-7B 的 GSM8K（52.2 vs 52.8）和 MATH（11.2 vs 13.6）都输给人工标注 SFT；Gemma-1.1-7B 的 GSM8K 赢 1.7 分（59.2 vs 57.5），MATH 赢 4.1 分（23.7 vs 19.6）。所以论文的措辞是「可匹敌，部分场景超越」，而不是全面超越。
Gemma-1.1-7B 的 MATH 提升最大：从 18.8（base）到 23.7，而 SFT 只有 19.6；RS 甚至比 base 还低（18.7），说明它在这个模型上筛出来的推理链没有帮上忙。
Gemma-2-9B 的 MATH：BRiTE 50.5，比人工标注 SFT 的 41.5 高 9.0 分，也比 RS 的 47.5 高 3.0 分——这是表里差距最显眼的一格。

Table 2：Qwen2.5-7B + 40K 混合数据

方法	MATH500	Minerva	OlympiadBench	AIME24	AMC23	GPQA Diamond
Base	44.1	12.9	16.1	0.9	10.1	25.9
RS	54.3	21.0	23.1	5.6	31.6	26.9
BRiTE（ψ-update）	79.1	35.0	35.7	14.3	57.7	28.5
BRiTE（θ-update）	76.9	40.6	37.0	14.4	57.1	29.8
BRiTE-iter-2（ψ）	80.6	41.3	37.3	14.3	57.9	29.9
BRiTE-iter-2（θ）	78.2	39.8	37.9	15.3	56.4	30.1

几个反直觉的发现：

ψ-update 本身就比 θ-update 更好，MATH500 上还高 2.2 分。这说明「生成高质量思考」本身就是一次能力提升，SFT 只是把它固化下来。
迭代几乎停滞：BRiTE-iter-2 在多数指标上只比第一轮高 1–2 分，作者自己解释是「第一轮就到达平台，限制了进一步收益」。
代码任务也不同——但幅度有限：deepseek-coder-6.7b-instruct 上 HumanEval Basic BRiTE 81.7 vs RS 79.3，BCB Hard 15.5 vs 11.5（Figure 2），但 BCB Full 只 36.3，还被 SFT 的 37.2 压过一头。

什么时候别用 BRiTE：验证器不可靠、基座已经很强、算力紧

从论文自身的实验设置和附录细节里，读出四条明确的边界。

边界一

必须有可信的 O，否则整个 ψ 步就是错的

MDP 的最优策略是「按参考分布加权后的正确推理」，如果 O 判错，错误推理也会被大量采出。论文全部实验都在数学（可自动验算）和代码（单元测试）上，没有覆盖需要人类主观判断的开放域任务。

边界二

基座本身很强时优势变小

Table 1 里 Gemma-2-9B 和 Llama-3-8B 的 base 已经很接近饱和，BRiTE 相对 RS 的改进只有 1-2 分。作者自己指出：instruct 模型已经做过后训练、数据规模有限时，差距不明显；要放大到 40K + Qwen2.5 才拉开。

边界三

RL 的开销不是免费的，而且迭代收益快速衰减

ψ 步要跑 PPO / GRPO，数学实验 4×H100，扩展实验 8×H100；而且 second iteration 在 Table 2 上只比第一轮高 1-2 分——这篇文章里看不到「越迭代越强」的长期增长趋势。

边界四

1/T 收敛依赖较强的假设

Theorem 3.3 要求目标 L 是凹的、且 logits 落在某个 RKHS 里。论文在附录 C.2 承认，去掉凹性只能得到驻点级别的弱保证。这给「BRiTE 一定收敛」的直觉加了一层需要留意的条件。

一句话：BRiTE 解决的是「在有可靠裁判的前提下，如何更聪明地生产高质量的思考过程」，而不是「如何判断什么样的思考是好的」。