Beyond Markovian: Reflective Exploration via Bayes-Adaptive RL for LLM Reasoning

Paper: arXiv:2505.20561 Code: shenao-zhang/BARL Code reference: main @ 1a6dc27 (2025-05-27)

1. Motivation (研究动机)

这篇论文要解决的核心问题不是“RL 能不能让 LLM 变强”,而是更细的一个机制问题:为什么 RL 之后的推理模型会在测试时出现“重新想一遍、修正错误、换一种策略”的自反行为,以及这种行为到底是否由传统 Markovian RL 目标本身解释。作者指出,常规 RL 把策略写成只依赖当前状态的 ,即使状态在推理轨迹中被“回到”了一次,策略也没有理由因为历史里多观察到的失败或局部奖励而改变动作分布。这样一来,训练期探索只是为了找到高回报动作序列;部署期策略则更接近固定的 exploitation,无法自然解释 LLM 在同一语义状态下因为“刚才尝试失败了”而切换思路。

论文把这一点和 LLM CoT 中的反思现象联系起来:一次“Let me reconsider”或“这个思路不对”的文字片段,本质上不是多写几个 token,而是在轨迹中携带了关于当前策略是否有效的新证据。若策略只看当前文本状态的表面表示,它可能把这种回退当成冗余;若策略维护“哪个 MDP/哪种解法假设更可能正确”的 belief,它就会把失败尝试当成信息,进而调整之后的动作。作者因此把 reflective exploration 定义为:轨迹在某种抽象状态上重复出现,但策略分布因为历史奖励和 belief 改变而不同。

上图中的二叉树例子给了最小直觉: 中, 是一次探索, 是“回到先前状态”的反思,随后 表示根据新证据切换策略。传统 Markovian 策略在两次到达 时应给出同一动作分布;Bayes-adaptive 策略则允许“状态相同但 belief 不同”,因此可以解释为什么第二次访问 时要选择另一条边。论文的动机就是把这种直觉从经验现象变成一个可优化的 RL 目标,并给出一个可在 LLM reasoning 训练中近似实现的算法 BARL。

2. Idea (核心思想)

核心新意可以概括为一句话:把 LLM 推理时的多个候选解法看成 posterior over MDPs 中的假设,把中间步骤带来的 answer-probability / reward 变化看成证据,再用 Bayes 更新后的候选权重给每个 token 或 step 分配训练信号。这样,模型不是被奖励“更频繁地反思”,而是被奖励“在当前策略被证据否定后切换到更合理的候选策略”。

这和普通 outcome-reward GRPO 的区别在于,GRPO 主要比较同一 prompt 下完整回答的最终正确性;BARL 试图回答“在第 步时,继续当前候选解法还是换到另一个候选解法更有信息价值”。它也不同于简单的 process reward:process reward 通常给每一步的局部质量打分,而 BARL 的 step value 来自多个候选 MDP 的 posterior-weighted value,权重会随着已经观察到的 reward/progress 与候选预测是否一致而变化。

论文的玩具任务说明了这个思路:训练 prompt 只覆盖 token 0/1,部署时 prompt 变成 token 2;普通 RL 可能记住训练三连串 000/111,却不知道在新 token 上要枚举和排除候选。BARL 把 000、111、222 等候选视作不同 MDP 假设,失败尝试会剔除或降低某个假设权重,最终促使策略探索 222。这也是标题里 “Beyond Markovian” 的含义:部署期推理不是无记忆状态机,而是在用历史观察改写后续策略。

3. Method (方法)

3.1 Bayesian RL formulation

论文先把 LLM reasoning 形式化为有限时域 MDP:初始状态 是 prompt,动作 是第 个 reasoning step,轨迹状态 是已经生成的 CoT 前缀。数据集 诱导出一个关于 MDP 的后验 ;部署时最大化的不是单一真实 MDP 下的回报,而是 Bayesian objective: 为了把 belief 引入策略,作者把问题等价转成 epistemic POMDP:隐藏状态包含真实 MDP ,观察到的是语言轨迹状态 ,历史 会更新 belief 。Theorem 4.2 说明 Bayes-optimal policy 可以表示成 ;Theorem 4.3 进一步构造出严格 uncertainty-adaptive policy 比任何 Markovian policy 都好任意多的例子。在二叉树构造中,Markovian policy 最多以 (深度扩展后 )命中奖励叶子,而 adaptive policy 通过失败反馈逐步排除假设,最终回报可达 1。

3.2 From posterior-weighted value to BARL

BARL 的 policy gradient 把常规 policy gradient 中“真实 MDP 的 value”替换为后验加权 value: 这个式子里的关键是 。作者把它拆成两部分:一部分是根据 CoT 前缀判断某个候选答案/候选 MDP 是否仍然 plausible;另一部分是历史 reward 与该候选 MDP 预测 reward 的一致性。若候选 预测某一步应提升正确答案概率,但实际观察到的 progress 不匹配,则其权重被 压低。最终得到的近似量可以理解为:

  1. 采样 条 CoT,抽取候选答案 ,并把 ground-truth answer 视为
  2. 在每个位置 ,计算当前前缀后模型生成候选答案的概率,作为“该候选当前仍可信”的 plausibility;
  3. 用历史 observed reward / progress 与候选预测 reward 的差异更新候选权重;
  4. 用 posterior-weighted value 给当前位置的动作分配 advantage,使模型学会在候选被证据否定时切换。

3.3 Algorithm and intuition

论文 Algorithm 1 很短:对每个训练 prompt,先从当前 LLM 采样 条 CoT;从每条 CoT 中抽取 candidate answer;对每条 CoT 的每个 timestep 计算公式中的 posterior-weighted value;再用 BARL policy gradient 更新 。直觉上,这相当于把 best-of-N 的多个候选路径“线性化”进一条训练轨迹:不是先完整生成 N 条答案再选最好,而是在单条 CoT 过程中持续问“如果当前假设看起来不对,哪些候选还可能解释已观察到的证据”。

图中对比了 Markovian RL 与 BARL 在玩具任务中的行为。Markovian RL 容易把训练 prompt 对应的三连串当作固定解法;BARL 会维护候选 triplet 集合,失败一次就减少一个候选的 posterior mass,因此更像在做“带记忆的枚举与排除”。这里的反思不是装饰性文本,而是 belief 更新后的策略改变。

def barl_training_loop(prompts_with_answers, actor, args):
    # openrlhf/trainer/barl_trainer.py::fit
    for prompt_batch in prompts_dataloader:
        experiences = experience_maker.make_experience_list(
            prompt_batch,
            max_new_tokens=args.generate_max_len,
            max_length=args.max_len,
        )
        for exp in experiences:
            replay_buffer.append(exp)
        replay_buffer.normalize("advantages", strategy)
        status = barl_train(global_step)  # PPO-style actor update
        replay_buffer.clear()
 
def make_experience_list(prompt_batch, args):
    # openrlhf/trainer/ppo_utils/experience_maker.py::make_experience_list
    samples = generate_samples(prompt_batch, n=args.n_samples_per_prompt)
    candidates = extract_boxed_answers(samples.completions)
    for each_prompt:
        candidates = deduplicate(candidates)
        candidates = [c for c in candidates if len(c) <= 20]
        if no_candidate_matches_ground_truth:
            candidates.append(ground_truth)
    experiences = [make_experience(sample, prompt_candidates) for sample in samples]
    posterior_qs, outcome_rewards = process_experiences(experiences)
    for exp, q in zip(experiences, posterior_qs):
        exp.advantages = q.detach()
    return experiences
def posterior_q_for_one_sequence(seq, candidates, gt_answer, step_size=128):
    # openrlhf/trainer/ppo_utils/experience_maker.py::posterior_value
    weight = uniform_over(candidates)
    for position in chunked_positions(seq, step_size):
        # candidate_probs: probability of inserting each candidate answer after current context
        candidate_probs = actor_answer_probs(context=seq[:position], answers=candidates)
        answer_prob = actor_answer_prob(context=seq[:position], answer=gt_answer)
        if position == 0:
            temp_weight = normalize(weight * candidate_probs)
        else:
            r_m = candidate_probs - last_candidate_probs
            r_gt = answer_prob - last_answer_prob
            value_error = abs(r_m - r_gt)
            weight = weight * exp(-1.0 * value_error)  # beta fixed to 1 in released code
            temp_weight = normalize(weight * candidate_probs)
        weighted_q[position_chunk] = sum((candidate_probs - candidate_probs_at_position) * temp_weight)
    return weighted_q + outcome_reward

3.4 Source code mapping

Code reference: main @ 1a6dc27 (2025-05-27) — pseudocode and mapping based on this commit.

Paper ConceptSource FileKey Class/Function
BARL CLI / training entrypointopenrlhf/cli/train_barl.pytrain(), parser args, BARLTrainer(...) construction
Released launch configtrain_barl.shQwen2.5-Math-7B, Big-Math-RL-Verified, batch and rollout settings
Rollout → replay buffer → PPO update loopopenrlhf/trainer/barl_trainer.pyBARLTrainer.fit(), barl_train(), training_step_actor()
Candidate answer extraction and experience creationopenrlhf/trainer/ppo_utils/experience_maker.pymake_experience_list(), generate_samples(), extract_boxed_content()
Posterior-weighted step valueopenrlhf/trainer/ppo_utils/experience_maker.pyposterior_value()
Correctness verifier for math answersopenrlhf/utils/verifier.pyverifier() used to test candidate vs ground truth
Dataset keying for prompt/answer pairsopenrlhf/datasets/prompts_answers_dataset.pyprompt-answer dataset consumed by trainer

论文公式与 released code 实现差异:论文把 写成超参数,并在实验设置中说明 BARL 使用 ;released code 的 posterior_value() 直接写成 torch.exp(-1.0 * value_error),等价于把 固定为 1,没有独立 CLI 参数。论文的候选 MDP 数量写作 ;released script 用 --n_samples_per_prompt 5 采样五条 CoT,但代码会对 boxed answer 去重、过滤长度超过 20 的候选,并在候选中没有 ground-truth 时追加正确答案,所以实际候选集合大小可能小于或大于 5。论文还描述了 Qwen-1.5B、Qwen-7B、Llama-8B 多模型实验;released train_barl.sh 只给出 Qwen/Qwen2.5-Math-7B 的 BARL 训练脚本,没有同步公开 GRPO、Progress baseline 或 1.5B/Llama 的完整 launch scripts。

4. Experimental Setup (实验设置)

论文包含一个 synthetic task 和一组数学推理实验。synthetic task 的动作空间是 ,目标是在 29 步内生成与 prompt token 相同的三连串;训练 prompt 覆盖 0/1,评估 prompt 是 2。这个设计故意制造 train-deploy shift:记住训练解 000/111 的策略无法泛化到 222,只有能维护候选并排除失败假设的策略才会在部署时成功。

数学推理实验使用 Big-Math 数据训练,评估 GSM8K、MATH、CollegeMath、OlympiadBench、AIME、AMC。模型包括 Qwen2.5-Math-1.5B、Qwen2.5-Math-7B 和 DeepSeek-R1-Distill-Llama-8B;论文正文还在附录报告 Llama-3.2-3B-Instruct。训练最大 prompt 长度 512,最大 response 长度 1024;BARL 设 ;online sampling temperature 为 1.0,evaluation temperature 为 0.0;BARL 和 Progress baseline 的 reasoning step size 为 128 tokens。Qwen 系列训练 110 个 rollout-batch iterations,Llama 系列训练 60 个 iterations。对照方法覆盖 outcome-reward GRPO 与带 progress reward 的 process-reward baseline;GRPO group size 也设为 5,并在 KL 系数 中选择 0.005。

released code 中可直接锚定的训练配置来自 train_barl.sh,而不是 parser 默认值:--pretrain Qwen/Qwen2.5-Math-7B--prompt_data SynthLabsAI/Big-Math-RL-Verified--input_key problem--output_key answer--train_batch_size 128--rollout_batch_size 1024--micro_train_batch_size 8--micro_rollout_batch_size 16--prompt_max_len 512--generate_max_len 1024--actor_learning_rate 5e-7--zero_stage 2--bf16--flash_attn--gradient_checkpointing--apply_chat_template--n_samples_per_prompt 5。这些数值与论文正文的大部分设置一致,但代码仓库的复现范围明显更窄:它是 OpenRLHF 上的 BARL training fork,而不是完整实验矩阵。

5. Experimental Results (实验结果)

5.1 Main accuracy

主表报告三次独立训练的 accuracy mean ± standard error。最关键的现象是 BARL 在多数任务和平均分上优于 GRPO / Progress,尤其在更需要搜索和排除错误思路的难题上收益更明显。

ModelMethodGSM8KMATHCollegeOlympiadAIMEAMCAvg
Qwen-1.5BBase40.034.16.621.816.732.525.3
Qwen-1.5BGRPO83.971.545.131.617.855.851.0
Qwen-1.5BProgress84.872.145.935.514.455.851.4
Qwen-1.5BBARL85.872.746.835.817.860.853.3
Qwen-7BBase59.153.721.919.013.342.534.9
Qwen-7BGRPO90.377.647.038.524.565.057.1
Qwen-7BProgress91.178.846.741.124.565.457.9
Qwen-7BBARL91.779.247.542.029.066.759.4
Llama-8BBase82.065.735.826.510.042.543.7
Llama-8BGRPO85.774.339.636.016.760.452.1
Llama-8BProgress85.973.939.835.315.055.851.0
Llama-8BBARL85.473.940.437.217.861.752.7

从平均分看,BARL 对 Qwen-1.5B 的提升最明显:相对 GRPO +2.3、相对 Progress +1.9;对 Qwen-7B 也有 +2.3/+1.5 的平均提升。Llama-8B 上 BARL 的平均分仍最高,但 GSM8K 和 MATH 两列不是单项最优,说明 BARL 不是无条件支配所有任务;它更适合需要在多个候选思路之间有效切换的场景。

synthetic results 支持了机制解释:REINFORCE 很快记住训练 prompt 上的解,但部署到未见 prompt token 时泛化差;BARL 因为显式维护候选并根据失败证据切换,评估准确率显著更高。当候选集合更有先验结构(只含 repeated triplets)时,收敛更快,这也暴露出 BARL 的一个边界条件:候选集合必须既多样又 plausible,候选太宽会浪费 posterior mass,候选太窄则覆盖不了真实解。

5.2 Token efficiency and reflective behavior

token efficiency 结果是论文最有说服力的部分之一。作者把 pass@k accuracy 与平均 token 消耗一起看,发现 BARL 以更少 token 达到更高准确率:最多比 Progress 少 39% 平均 token、比 GRPO 少 50%、比 base model 少 90% 以上。这说明 BARL 并不是简单鼓励模型写更长 CoT;相反,它让“思考 token”更有效,把探索集中在能够改变 belief 的位置。

自反频率分析进一步区分了“表面反思”和“有效探索”。作者用关键词检测 self-reflection,并按每题六个 response 中错误数定义难度。结果显示,base model 可能有更高反思频率,但准确率更低;BARL 的优势不是反思词更多,而是反思更能带来策略切换和信息增益。这一点对理解 RL 后 reasoning model 很重要:如果只统计 “wait / reconsider / alternatively” 之类 token,会把预训练中学到的风格模式误判为探索能力;BARL 关注的是反思是否改变后续候选权重。

5.3 Bayesian value and Markovian optimality ablations

Bayesian value ablation 显示,BARL 生成的 actions 在 GSM8K、MATH、CollegeMath、OlympiadBench 上有更高 posterior-weighted value,特别是在 OlympiadBench 这类较难任务中,探索收益会在 CoT 中段达到峰值:早期先快速降低不确定性,中段通过候选切换获得更多信息,后期则逐渐 exploitation。这和方法部分的设计一致,也解释了为什么“频率高但价值低”的反思不一定有帮助。

length-controlled GRPO ablation 用最大 32 response length 训练,发现训练准确率上升、response length 下降,最终训练准确率可接近常规 GRPO。这支持论文的 Remark:Markovian RL 的最优性可以通过记忆训练解、跳过 CoT、直接输出答案来达成;它在训练集上看似最优,却不保证部署期 reflective exploration。BARL 的贡献因此不是单纯提高表格数值,而是把“为什么要在测试时继续探索”放进了目标函数。

5.4 Limitations and reading notes

这篇论文的理论和算法之间存在近似落差。理论部分讨论 Bayes-optimal policy over MDP posterior;实际代码无法枚举真实 MDP,也没有为每个候选做完整 branched Monte Carlo Q 估计,而是用候选答案概率、答案概率增量和 outcome reward 构造近似 posterior value。这是合理的工程折中,因为 LLM reasoning 中对每个 step 分叉 rollout 代价极高;但也意味着 BARL 的效果依赖答案抽取、verifier、候选长度过滤和模型对候选答案概率的校准。

另一个边界是候选生成质量。若当前 policy 采样不到有用候选,BARL 只能在错误候选之间重分配权重;若候选中总是包含 ground truth 但缺少能解释中间过程的策略多样性,训练信号可能退化成更复杂的 outcome reward。论文的 synthetic task 清楚展示了 candidate-set prior 的重要性;真实数学任务中,这个 prior 由当前模型 rollout 和 boxed answer 抽取共同决定,因此对格式、温度、采样数和 verifier 都敏感。

总体上,BARL 给 RL for reasoning 提供了一个很好的分析框架:自反不是“长 CoT 的副产品”,而是 Bayesian belief update 在语言轨迹中的可见行为。对后续工作来说,最值得沿用的是这个评价角度:不要只问模型是否会反思,而要问反思后是否降低了不确定性、是否切换到更高 posterior value 的策略、是否以更少 token 达到同样或更高正确率。