Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

Abstract

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.

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

Abstract

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.
Paper Structure (38 sections, 2 theorems, 19 equations, 11 figures, 5 tables)

This paper contains 38 sections, 2 theorems, 19 equations, 11 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Reinforcement Learning for LLMs
  4. Long Context Compression
  5. Temporal Credit Assignment in Reinforcement Learning
  6. Preliminary: DAPO
  7. Method
  8. TAMTRL Workflow
  9. Problem Formulation
  10. Teacher-Aligned Reward Reshaping
  11. Multi-Turn Reinforcement Learning Training
  12. Theoretical Analysis
  13. Experiment
  14. Experiment Setup
  15. Training
  16. ...and 23 more sections

Key Result

Theorem 1

Consider an optimization step at state $S_t = (q, D_t, M_t)$ within a multi-step reasoning process. Let $\pi_\theta(\cdot \mid S_t)$ denote the policy generating the next memory $M_{t+1}$, and let $r_i \in \{0,1\}$ be the binary indicator of final task success, whose distribution depends on $M_{t+1} This objective $\mathcal{J}(\theta)$ can be exactly decomposed into a weighted sum comprising a suc

Figures (11)

  • Figure 1: In multi-turn RLVR, the lack of temporal credit assignment may wrongly penalize turns with good answers or reward turns with poor answers, introducing noisy supervision that complicates training and degrades performance. Existing solutions either incur significant computational overhead or rely on external models.
  • Figure 2: Framework of the TAMTRL Method. Solving long-context problems through chunk-based processing, utilizing the probability scores from a teacher with local and global views for turn-level credit assignment to enable multi-turn reinforcement learning.
  • Figure 3: Training dynamics comparison of Qwen3-0.6B model.
  • Figure 4: Training dynamics comparison of Qwen3-1.7B model.
  • Figure 5: Analysis of the impact of different document quantities on the performance of TAMTRL on the HotpotQA dataset. (Left) Test results with varying training document quantities at a fixed test document quantity of 100; (Right) Test results with varying test document quantities at a fixed training document quantity of 100.
  • ...and 6 more figures

Theorems & Definitions (3)

  • Theorem 1: Information-Theoretic Decomposition of the TAMTRL Objective
  • Theorem A1: Information-Theoretic Decomposition of the TAMTRL Objective
  • proof