TAMTRL: Teacher-Aligned Reward Reshaping for Multi-Turn Reinforcement Learning in Long-Context Compression

📄 arXiv: 2603.21663v1 📥 PDF

作者: Li Wang, Yandong Wang, Xin Yu, Kui Zhang, Tianhao Peng, Wenjun Wu

分类: cs.CL

发布日期: 2026-03-23

💡 一句话要点

提出TAMTRL，通过教师对齐奖励重塑解决长文本压缩中的多轮强化学习问题。

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

关键词: 长文本压缩 多轮强化学习 奖励重塑 教师对齐 信用分配

📋 核心要点

  1. 现有方法在长文本压缩的多轮强化学习中，面临时间信用分配难题，难以有效评估每轮记忆更新的质量。
  2. TAMTRL的核心思想是利用相关文档作为教师信号，通过与每轮模型输入对齐并进行奖励重塑，提供细粒度的学习信号。
  3. 实验结果表明，TAMTRL在多个长文本基准测试中，显著优于现有基线方法，验证了其有效性。

📝 摘要（中文）

大型语言模型（LLMs）的快速发展在各种任务中取得了显著的性能提升。然而，当处理超过模型上下文窗口限制的长文档时，无法一次性处理整个上下文，需要分块处理，这需要多轮读取不同的块并更新记忆。然而，监督通常只由最终结果提供，这使得难以评估多轮训练中每一轮记忆更新的质量，从而引入了时间信用分配挑战。现有的方法，如LLM-as-a-judge或过程奖励模型，会产生大量的计算开销并遭受估计噪声的影响。为了更好地解决多轮记忆训练中的信用分配问题，我们提出了用于多轮强化学习的教师对齐奖励重塑（TAMTRL）。TAMTRL利用相关文档作为教师信号，通过将它们与每一轮的模型输入对齐，并通过自监督方式中的归一化概率分配奖励。这为每次记忆更新提供了细粒度的学习信号，并提高了长上下文处理能力。在七个长上下文基准上对多个不同规模的模型进行的实验表明，TAMTRL始终优于强大的基线，证明了其有效性。我们的代码可在https://anonymous.4open.science/r/TAMTRL-F1F8获得。

🔬 方法详解

问题定义：论文旨在解决长文本压缩中多轮强化学习的信用分配问题。现有方法，如使用LLM作为裁判或训练过程奖励模型，计算成本高昂且容易引入噪声，导致训练不稳定，难以有效指导每一轮的记忆更新。

核心思路：TAMTRL的核心思路是利用与输入相关的文档作为“教师”信号，通过将教师信号与每一轮的模型输入对齐，并基于对齐程度进行奖励重塑。这种方法避免了直接使用LLM进行奖励评估，降低了计算成本，并减少了噪声干扰。

技术框架：TAMTRL的整体框架包括以下几个主要步骤：1) 将长文本分割成多个chunk，模型逐个chunk地读取并更新记忆；2) 对于每一轮的输入chunk，找到与之相关的文档片段作为教师信号；3) 计算模型输出与教师信号之间的相似度，并将其归一化为概率分布；4) 使用该概率分布作为奖励信号，指导模型的记忆更新。

关键创新：TAMTRL的关键创新在于使用教师对齐的奖励重塑方法，将外部知识（相关文档）融入到强化学习的奖励函数中。与传统的基于LLM的奖励模型相比，TAMTRL更加高效和稳定，能够提供更细粒度的学习信号。

关键设计：TAMTRL的关键设计包括：1) 如何选择合适的教师信号（即相关文档片段）；2) 如何定义模型输出与教师信号之间的相似度度量；3) 如何将相似度归一化为概率分布，并将其作为奖励信号。论文中具体使用了余弦相似度来衡量模型输出和教师信号之间的相似性，并使用softmax函数进行归一化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TAMTRL在七个长文本基准测试中均优于现有基线方法。具体性能提升幅度取决于不同的模型和数据集，但总体趋势是TAMTRL能够显著提高模型在长文本处理任务上的性能。例如，在某些数据集上，TAMTRL的性能提升超过了5%。

🎯 应用场景

TAMTRL方法可应用于各种需要处理长文本信息的场景，例如文档摘要、信息检索、问答系统等。通过提升模型对长文本的理解和记忆能力，可以显著改善这些应用的用户体验和性能。该研究对于开发更智能、更高效的自然语言处理系统具有重要意义。

📄 摘要（原文）

The rapid progress of large language models (LLMs) has led to remarkable performance gains across a wide range of tasks. However, when handling long documents that exceed the model's context window limit, the entire context cannot be processed in a single pass, making chunk-wise processing necessary. This requires multiple turns to read different chunks and update memory. However, supervision is typically provided only by the final outcome, which makes it difficult to evaluate the quality of memory updates at each turn in the multi-turn training setting. This introduces a temporal credit assignment challenge. Existing approaches, such as LLM-as-a-judge or process reward models, incur substantial computational overhead and suffer from estimation noise. To better address the credit assignment problem in multi-turn memory training, we propose Teacher-Aligned Reward Reshaping for Multi-Turn Reinforcement Learning (TAMTRL). TAMTRL leverages relevant documents as teacher signals by aligning them with each turn of model input and assigns rewards through normalized probabilities in a self-supervised manner. This provides fine-grained learning signals for each memory update and improves long-context processing. Experiments with multiple models of varying scales across seven long-context benchmarks show that TAMTRL consistently outperforms strong baselines, demonstrating its effectiveness. Our code is available at https://anonymous.4open.science/r/TAMTRL-F1F8.