Multi-Turn Reasoning When Context Arrives in Pieces: Scalable Sharding and Memory-Augmented RL

2026-06-11

AI Summary

This work addresses the significant performance degradation of large language models in multi-turn dialogues due to contextual fragmentation, a phenomenon often described as “getting lost in conversation.” To mitigate this, the authors propose a memory-augmented reinforcement learning framework that integrates a rolling memory mechanism with a low-cost data sharding strategy. This approach enables the model to reason using only a compact, dynamically updated memory representation instead of the full dialogue history. Remarkably, when trained solely on a sharded version of GSM8K, the method substantially improves multi-turn accuracy and demonstrates strong zero-shot generalization to more challenging mathematical reasoning and long-context question-answering tasks, outperforming baselines that rely on the complete conversational history.

Abstract

When a user reveals task-critical information across several conversation turns, LLM accuracy drops by up to 65% despite full context availability. We show that this Lost in Conversation degradation can be substantially mitigated by training models to maintain a compact rolling memory instead of attending to a growing history. To make such training scalable, we introduce a low-cost sharding pipeline that converts single-turn QA datasets into multi-turn fragmented-information episodes, eliminating the need for hours of manual annotation. Training only on sharded GSM8K, our memory-augmented policy significantly improves multi-turn accuracy and generalises zero-shot to harder math and out-of-domain long-context QA. Moreover, memory-trained models outperform full-history baselines even when given the full history at test time, suggesting that learning to compress induces more robust incremental reasoning than full-context exposure alone.

Problem

• multi-turn reasoning
• lost in conversation
• context fragmentation
• long-context QA
• incremental reasoning

Innovation

• memory-augmented RL
• scalable sharding
• multi-turn reasoning
• rolling memory
• fragmented context