Memory-T1: Reinforcement Learning for Temporal Reasoning in Multi-session Agents
Memory-T1 employs reinforcement learning for time-aware memory retrieval, boosting multi-session dialogue reasoning to 67.0% accuracy, outperforming baselines by over 10%.
Key Findings
Methodology
Memory-T1 adopts a coarse-to-fine retrieval framework, first predicting the query’s time window with a large language model (LLM) to filter dialogue sessions, then using a reinforcement learning (RL) policy to select precise evidence sessions. The multi-level reward function combines answer accuracy, evidence grounding, and a novel temporal consistency reward that evaluates session-level proximity and utterance-level temporal fidelity. The RL training employs Group Relative Policy Optimization (GRPO) to ensure stability over long contexts up to 128k tokens. This approach effectively filters noisy long histories, enabling robust temporal reasoning in multi-session dialogues. Experimental results on the Time-Dialog benchmark show a 7B model reaching 67.0%, surpassing a 14B baseline by 10.2%, demonstrating the method’s scalability and robustness.
Key Results
- Memory-T1 achieves 67.0% overall score on Time-Dialog, outperforming existing open-source models and significantly improving complex reasoning tasks like event order and range inference. It maintains performance even with 128k tokens of noisy history, validating robustness. Ablation studies reveal that combining temporal consistency and evidence grounding rewards yields a 15% performance boost. Compared to models trained solely on answer accuracy, Memory-T1’s multi-level reward structure enables better temporal alignment and evidence selection.
- In comparison with baselines like GPT-4, Time-R1, and MemAgent, Memory-T1 demonstrates superior performance on multi-session temporal reasoning benchmarks, especially in event sequencing and time span inference. The model generalizes well across datasets like Time-Dialog and LoCoMo, confirming its robustness and transferability.
- Ablation experiments show that relying only on answer correctness causes performance degradation in long contexts, while integrating temporal and evidence rewards stabilizes results. This highlights the importance of multi-level reward design for temporal reasoning in noisy, extended dialogues.
Significance
This work advances the state-of-the-art in multi-session dialogue understanding by integrating reinforcement learning with temporal-aware memory retrieval. It addresses longstanding challenges of noisy long histories and ambiguous time expressions, enabling models to reason more accurately about temporal relationships. The approach enhances factual consistency and temporal coherence, crucial for applications like virtual assistants, customer support, and long-term conversational agents. By demonstrating robustness up to 128k tokens, it paves the way for scalable, real-world deployment of temporally grounded dialogue systems, contributing significantly to AI’s ability to handle complex, long-term interactions.
Technical Contribution
The core innovation lies in the design of a multi-stage, RL-based memory retrieval framework that incorporates a novel temporal consistency reward at both session and utterance levels. This reward explicitly guides the model to select evidence that aligns temporally with the query, overcoming the limitations of sparse supervision. The use of a coarse-to-fine retrieval process—initial temporal filtering, relevance ranking, followed by RL fine selection—reduces noise and improves efficiency. The employment of GRPO ensures stable training over large-scale, long-context data. These technical advances collectively enable robust, scalable temporal reasoning without relying on explicit supervision or handcrafted structures.
Novelty
This study is the first to integrate reinforcement learning with multi-stage, time-aware memory retrieval specifically for multi-session dialogue reasoning. The introduction of a dense, multi-level temporal consistency reward distinguishes it from prior work that relied solely on explicit annotations or static knowledge bases. Unlike existing models that treat dialogue history as flat text, Memory-T1 dynamically learns to select evidence based on implicit temporal cues, significantly improving robustness and generalization in noisy, extended conversations.
Limitations
- Despite robustness, the model may still struggle with extremely ambiguous or contradictory time expressions, especially when the dialogue contains conflicting temporal cues. Its performance depends on the accuracy of initial time window prediction, which can be affected by noisy data.
- Training with large models and reinforcement learning incurs high computational costs, limiting real-time deployment without further optimization.
- The current reward design, while effective, may not fully capture complex temporal nuances such as time zone shifts or implicit temporal inferences, suggesting future work in more sophisticated temporal modeling.
Future Work
Future directions include integrating multi-modal signals (visual, auditory) to enhance temporal understanding, developing more fine-grained temporal models, and reducing computational costs through model compression or distillation. Additionally, exploring unsupervised or semi-supervised approaches to further reduce reliance on annotated data, and extending the framework to handle more complex temporal phenomena like causality and probabilistic time, are promising avenues.
AI Executive Summary
Deep Dive
Plain Language Accessible to non-experts
Imagine you’re trying to piece together a story from a huge collection of photo albums. Each album is like a different session of a conversation, with photos taken at different times. To understand what happened, you need to find the right photos that match the story’s timeline. But with so many albums and photos, it’s easy to get confused—some photos look similar but are from different days. Memory-T1 is like having a smart assistant who first narrows down the albums based on the approximate time you’re interested in, then carefully picks out the most relevant photos that truly fit the story’s timeline. The assistant also learns over time which clues help him pick the right photos faster, so he gets better at understanding the sequence of events. This way, he helps you see the real story without getting mixed up by irrelevant or confusing pictures.
ELI14 Explained like you're 14
Imagine you’re trying to remember what you did last weekend. You have a huge photo album with pictures from many days. Some photos look alike but are from different days, so it’s tricky to tell exactly when something happened. Now, think of a super-smart friend who helps you find the right photos. First, he guesses roughly which days you’re talking about, so he only looks at photos from those days. Then, he carefully picks out the best photos that match your story, making sure they’re from the right time. Over time, he learns which clues—like the date written on the photo or what you were wearing—help him pick the right pictures faster. Thanks to him, you can tell your story accurately, without mixing up the days or getting confused by similar photos. That’s what Memory-T1 does for computers: it helps them understand the timeline of long conversations by smartly choosing the right pieces of past dialogue.
Abstract
Temporal reasoning over long, multi-session dialogues is a critical capability for conversational agents. However, existing works and our pilot study have shown that as dialogue histories grow in length and accumulate noise, current long-context models struggle to accurately identify temporally pertinent information, significantly impairing reasoning performance. To address this, we introduce Memory-T1, a framework that learns a time-aware memory selection policy using reinforcement learning (RL). It employs a coarse-to-fine strategy, first pruning the dialogue history into a candidate set using temporal and relevance filters, followed by an RL agent that selects the precise evidence sessions. The RL training is guided by a multi-level reward function optimizing (i) answer accuracy, (ii) evidence grounding, and (iii) temporal consistency. In particular, the temporal consistency reward provides a dense signal by evaluating alignment with the query time scope at both the session-level (chronological proximity) and the utterance-level (chronological fidelity), enabling the agent to resolve subtle chronological ambiguities. On the Time-Dialog benchmark, Memory-T1 boosts a 7B model to an overall score of 67.0\%, establishing a new state-of-the-art performance for open-source models and outperforming a 14B baseline by 10.2\%. Ablation studies show temporal consistency and evidence grounding rewards jointly contribute to a 15.0\% performance gain. Moreover, Memory-T1 maintains robustness up to 128k tokens, where baseline models collapse, proving effectiveness against noise in extensive dialogue histories. The code and datasets are publicly available at https://github.com/Elvin-Yiming-Du/Memory-T1/