TR-ICRL: Test-Time Rethinking for In-Context Reinforcement Learning

📄 arXiv: 2604.00438v1 📥 PDF

作者: Wenxuan Jiang, Yuxin Zuo, Zijian Zhang, Xuecheng Wu, Zining Fan, Wenxuan Liu, Li Chen, Xiaoyu Li, Xuezhi Cao, Xiaolong Jin, Ninghao Liu

分类: cs.CL

发布日期: 2026-04-01

备注: 14 pages, 7 figures

🔗 代码/项目: GITHUB

💡 一句话要点

TR-ICRL：面向上下文强化学习的测试时重思考框架，提升推理和知识密集型任务性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 上下文强化学习 大型语言模型 伪标签 测试时学习 迭代优化

📋 核心要点

  1. ICRL面临推理时缺乏真实奖励信号的挑战，导致奖励估计不准确。
  2. TR-ICRL通过检索相关实例，生成伪标签作为奖励信号，迭代优化LLM的答案。
  3. 实验表明，TR-ICRL在MedQA和AIME2024等任务上显著提升了LLM的性能。

📝 摘要（中文）

本文提出了一种名为TR-ICRL（Test-Time Rethinking for In-Context Reinforcement Learning）的上下文强化学习（ICRL）新框架，旨在解决ICRL中奖励估计这一核心挑战，尤其是在推理和知识密集型任务中。TR-ICRL首先从无标签评估集中检索与给定查询最相关的实例。在每个ICRL迭代中，大型语言模型（LLM）为每个检索到的实例生成一组候选答案。然后，通过多数投票从这组答案中导出一个伪标签。该伪标签作为代理，提供奖励信息和形成性反馈，引导LLM进行迭代优化。最后，将合成的上下文信息与原始查询集成，形成一个全面的提示，并通过最后一轮多数投票确定答案。在主流推理和知识密集型任务上的评估表明，TR-ICRL显著提高了性能。例如，在MedQA上，Qwen2.5-7B的性能平均提高了21.23%，在AIME2024上甚至提高了137.59%。大量的消融研究和分析进一步验证了该方法的有效性和鲁棒性。

🔬 方法详解

问题定义：上下文强化学习（ICRL）旨在使大型语言模型（LLM）能够直接在上下文窗口中从外部奖励进行在线学习。然而，在推理阶段，模型通常无法访问真实标签，导致奖励估计成为一个关键瓶颈，严重影响ICRL的性能，尤其是在需要复杂推理和知识的任务中。现有方法难以有效利用上下文信息进行准确的奖励预测。

核心思路：TR-ICRL的核心思路是利用测试时检索到的相关实例，通过生成伪标签来模拟真实的奖励信号。通过迭代地让LLM基于这些伪标签进行“重思考”，逐步优化其答案。这种方法的核心在于利用无标签数据来弥补推理时缺乏真实奖励的不足，从而提高ICRL的性能。

技术框架：TR-ICRL框架主要包含以下几个阶段：1) 实例检索：对于给定的查询，从无标签的评估集中检索最相关的实例。2) 候选答案生成：LLM为每个检索到的实例生成一组候选答案。3) 伪标签生成：通过多数投票从候选答案中生成伪标签。4) 迭代优化：使用伪标签作为奖励信号，引导LLM进行迭代优化，生成形成性反馈。5) 最终答案确定：将合成的上下文信息与原始查询集成，并通过最后一轮多数投票确定最终答案。

关键创新：TR-ICRL的关键创新在于利用测试时检索到的无标签数据，通过生成伪标签来模拟真实的奖励信号，从而解决了ICRL在推理时缺乏真实奖励的问题。与传统的ICRL方法相比，TR-ICRL不需要预先训练的奖励模型，而是直接利用LLM的生成能力和多数投票机制来生成奖励信号，更加灵活和高效。

关键设计：在实例检索阶段，可以使用余弦相似度等方法来衡量查询与实例之间的相关性。在候选答案生成阶段，可以调整LLM的生成参数（如温度）来控制答案的多样性。在伪标签生成阶段，需要选择合适的多数投票策略来平衡准确性和鲁棒性。迭代优化的次数可以作为一个超参数进行调整，以平衡性能和计算成本。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TR-ICRL在MedQA和AIME2024等主流推理和知识密集型任务上取得了显著的性能提升。具体而言，在MedQA上，使用Qwen2.5-7B作为基础模型，TR-ICRL的性能平均提高了21.23%。更令人瞩目的是，在AIME2024上，TR-ICRL的性能提升高达137.59%。这些结果表明，TR-ICRL能够有效提高LLM在复杂任务中的推理能力和知识利用效率。

🎯 应用场景

TR-ICRL具有广泛的应用前景，例如在医疗诊断、法律咨询、金融分析等需要复杂推理和知识的任务中，可以利用该框架提升LLM的性能，辅助专业人员进行决策。此外，该方法还可以应用于智能客服、教育辅导等领域，提高LLM的智能化水平和服务质量。

📄 摘要（原文）

In-Context Reinforcement Learning (ICRL) enables Large Language Models (LLMs) to learn online from external rewards directly within the context window. However, a central challenge in ICRL is reward estimation, as models typically lack access to ground-truths during inference. To address this limitation, we propose Test-Time Rethinking for In-Context Reinforcement Learning (TR-ICRL), a novel ICRL framework designed for both reasoning and knowledge-intensive tasks. TR-ICRL operates by first retrieving the most relevant instances from an unlabeled evaluation set for a given query. During each ICRL iteration, LLM generates a set of candidate answers for every retrieved instance. Next, a pseudo-label is derived from this set through majority voting. This label then serves as a proxy to give reward messages and generate formative feedbacks, guiding LLM through iterative refinement. In the end, this synthesized contextual information is integrated with the original query to form a comprehensive prompt, with the answer determining through a final round of majority voting. TR-ICRL is evaluated on mainstream reasoning and knowledge-intensive tasks, where it demonstrates significant performance gains. Remarkably, TR-ICRL improves Qwen2.5-7B by 21.23% on average on MedQA and even 137.59% on AIME2024. Extensive ablation studies and analyses further validate the effectiveness and robustness of our approach. Our code is available at https://github.com/pangpang-xuan/TR_ICRL.