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

• URL: http://arxiv.org/abs/2603.21663v1
• Date: Mon, 23 Mar 2026 07:42:04 GMT
• Title: TAMTRL: Teacher-Aligned Reward Reshaping for Multi-Turn Reinforcement Learning in Long-Context Compression
• Abstract summary: We propose Teacher-Aligned Reward Reshaping for Multi-Turn Reinforcement Learning (TAMTRL)<n>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.<n>Experiments with multiple models of varying scales across seven long-context benchmarks show that TAMTRL consistently outperforms strong baselines.
• Score: 24.10525956976745
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• From Player to Master: Enhancing Test-Time Learning of LLM Agents via Reinforcement Learning over Memory [27.810603317768997]
  Large language model (LLM) agents are increasingly deployed in long-running settings where improving through experience at test time is important.<n>We propose MemoPilot, a plug-in memory copilot that explicitly trains the memory update process to improve a frozen LLM's performance across sequential interactions.<n>We evaluate MemoPilot on two testbeds: multi-round Rock-Paper-Scissors (RPS) and Limit Texas Hold'em (LHE)
  arXiv  Detail & Related papers  (2026-06-07T14:53:19Z)
• Fast Unlearning at Scale via Margin Self-Correction [52.46927918952516]
  Language-model unlearning updates a trained model to behave as if it had not seen selected training examples.<n>MASC achieves a competitive forget-retain trade-off at a fraction of the computational cost of existing baselines.
  arXiv  Detail & Related papers  (2026-06-01T21:49:54Z)
• MemSifter: Offloading LLM Memory Retrieval via Outcome-Driven Proxy Reasoning [78.46301394559903]
  Large Language Models (LLMs) are increasingly used for long-duration tasks.<n>Current methods face a trade-off between cost and accuracy.<n>MemSifter is a novel framework that offloads the memory retrieval process to a small-scale proxy model.
  arXiv  Detail & Related papers  (2026-03-03T02:57:38Z)
• Document Reconstruction Unlocks Scalable Long-Context RLVR [60.74632963522131]
  Reinforcement Learning with Verifiable Rewards(RLVR) has become a prominent paradigm to enhance the capabilities (i.e. long-context) of Large Language Models(LLMs)<n>We investigate unsupervised approaches to enhance the long-context capabilities of LLMs, eliminating the need for heavy human annotations or teacher models' supervision.<n>We validate the effectiveness of our method on two widely used benchmarks, RULER and LongBenchv2.
  arXiv  Detail & Related papers  (2026-02-09T03:23:23Z)
• Reinforcement Learning for Machine Learning Engineering Agents [52.03168614623642]
  We show that agents backed by weaker models that improve via reinforcement learning can outperform agents backed by much larger, but static models.<n>We propose duration- aware gradient updates in a distributed asynchronous RL framework to amplify high-cost but high-reward actions.<n>We also propose environment instrumentation to offer partial credit, distinguishing almost-correct programs from those that fail early.
  arXiv  Detail & Related papers  (2025-09-01T18:04:10Z)
• Intention-Conditioned Flow Occupancy Models [80.42634994902858]
  Large-scale pre-training has fundamentally changed how machine learning research is done today.<n>Applying this same framework to reinforcement learning is appealing because it offers compelling avenues for addressing core challenges in RL.<n>Recent advances in generative AI have provided new tools for modeling highly complex distributions.
  arXiv  Detail & Related papers  (2025-06-10T15:27:46Z)
• DeepInsert: Early Layer Bypass for Efficient and Performant Multimodal Understanding [26.39397960987363]
  We propose a simple modification to pretrained transformer models.<n>Instead of concatenation with the language prompt at the start, we insert multimodal tokens directly into the middle.<n>Our results indicate that our method reduces computational costs during both training and inference.
  arXiv  Detail & Related papers  (2025-04-27T18:56:26Z)
• Efficient Transfer Learning for Video-language Foundation Models [13.166348605993292]
  We propose a parameter-efficient Multi-modalpatio Ssupervised-Temporal Adapter (MSTA) to enhance alignment between textual and visual representations.<n>We evaluate the effectiveness of our approach across four tasks: zero-shot transfer, few-shot learning, base-to-novel generalization, and fully-Temporal learning.
  arXiv  Detail & Related papers  (2024-11-18T01:25:58Z)
• Regressing the Relative Future: Efficient Policy Optimization for Multi-turn RLHF [47.19899163378932]
  We introduce REgressing the RELative FUture (REFUEL) to address multi-turn RLHF in Large Language Models (LLMs)<n>REFUEL frames the multi-turn RLHF problem as a sequence of regression tasks on iteratively collected datasets.<n>We evaluate our algorithm by using Llama-3.1-70B-it to simulate a user in conversation with our model.
  arXiv  Detail & Related papers  (2024-10-06T20:20:22Z)
• DataStates-LLM: Lazy Asynchronous Checkpointing for Large Language Models [3.3484462092188005]
  We introduce a lazy asynchronous multi-level approach that takes advantage of the fact that the tensors making up the model and state shards remain immutable for extended periods of time. The results show up to 48$times$ faster checkpointing and 2.2$times$ faster end-to-end training compared with state-of-art checkpointing approaches.
  arXiv  Detail & Related papers  (2024-06-15T18:30:40Z)
• Text Generation with Efficient (Soft) Q-Learning [91.47743595382758]
  Reinforcement learning (RL) offers a more flexible solution by allowing users to plug in arbitrary task metrics as reward. We introduce a new RL formulation for text generation from the soft Q-learning perspective. We apply the approach to a wide range of tasks, including learning from noisy/negative examples, adversarial attacks, and prompt generation.
  arXiv  Detail & Related papers  (2021-06-14T18:48:40Z)
• Efficient Retrieval Optimized Multi-task Learning [16.189136169520424]
  We propose a novel Retrieval Optimized Multi-task (ROM) framework for jointly training self-supervised tasks, knowledge retrieval, and extractive question answering. Our ROM approach presents a unified and generalizable framework that enables scaling efficiently to multiple tasks. Using our framework, we achieve comparable or better performance than recent methods on QA, while drastically reducing the number of parameters.
  arXiv  Detail & Related papers  (2021-04-20T17:16:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.