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

arXiv: https://arxiv.org/abs/2603.21663 · HTML: https://mlsanity.com/paper/2603.21663/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 27 of its 31 references that are in the map.
- First posted on arXiv 23 Mar 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)
- Efficient LLM inference and serving / [Prompt and context compression](https://mlsanity.com/inference/prompt-and-context-compression/)
- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

## Builds on (most cited first, 27 of 27)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Games](https://mlsanity.com/paper/2103.01955/) (2021)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [The NarrativeQA Reading Comprehension Challenge](https://mlsanity.com/paper/1712.07040/) (2017)
- [A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers](https://mlsanity.com/paper/2105.03011/) (2021)
- [LLM-Planner: Few-Shot Grounded Planning for Embodied Agents with Large Language Models](https://mlsanity.com/paper/2212.04088/) (2022)
- [Prometheus 2: An Open Source Language Model Specialized in Evaluating Other Language Models](https://mlsanity.com/paper/2405.01535/) (2024)
- [The Lessons of Developing Process Reward Models in Mathematical Reasoning](https://mlsanity.com/paper/2501.07301/) (2025)
- [Benchmarking Large Language Models in Retrieval-Augmented Generation](https://mlsanity.com/paper/2309.01431/) (2023)
- [Teaching Small Language Models to Reason](https://mlsanity.com/paper/2212.08410/) (2022)
- [Chain of Agents: Large Language Models Collaborating on Long-Context Tasks](https://mlsanity.com/paper/2406.02818/) (2024)
- [Compact Language Models via Pruning and Knowledge Distillation](https://mlsanity.com/paper/2407.14679/) (2024)
- [An Emphatic Approach to the Problem of Off-policy Temporal-Difference Learning](https://mlsanity.com/paper/1503.04269/) (2015)
- [PlanGenLLMs: A Modern Survey of LLM Planning Capabilities](https://mlsanity.com/paper/2502.11221/) (2025)
- [From Calculation to Adjudication: Examining LLM judges on Mathematical Reasoning Tasks](https://mlsanity.com/paper/2409.04168/) (2024)
- [KG-Augmented Executable CoT for Mathematical Coding](https://mlsanity.com/paper/2508.04072/) (2025)

## Built on by (most cited first, 1 of 1)

- [Learning to Adapt SFT Data for Better Reasoning Generalization](https://mlsanity.com/paper/2605.26924/) (2026)

## Abstract (arXiv)

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.