StateTree: Enhancing Long-Term Dialogue Reasoning via Reinforcement Learning
cs.CL, cs.AI
Submitted: 2026-09-30
Updated: 2026-09-30
Comments: NeurIPS 2026
Code: https://github.com/bluedream02/StateTree
License: http://creativecommons.org/licenses/by/4.0/
The gist: Large language models deployed as personalized assistants must reason over long, evolving interaction histories.
Terminology
Abstract
Large language models deployed as personalized assistants must reason over long, evolving interaction histories. However, in long-term dialogue reasoning, relevant evidence is scattered across sessions, preferences may be revised over time, and standard long-context training fails to address these challenges under data scarcity and prohibitive computational costs. We propose StateTree, a data-driven RL method that constructs a challenging auxiliary task from scarce dialogues with verifiable ground truth. StateTree augments multi-session dialogues with a tree-structured path-tracing task: key-value records are embedded across sessions to form a binary tree. Solving the task requires the model to traverse from root to leaf by retrieving records across sessions and comparing timestamps to resolve branches, then recover the hidden target question among distractor leaves. We apply curriculum RL training progressively increasing tree depth and introduce a compositional variant whose edges carry step-level reasoning fragments, training the model to compose partial cues into coherent queries. Trained on 10K-token contexts, StateTree generalizes to 128K tokens without full-length RL costs and exhibits capabilities including cross-session retrieval, temporal reasoning, knowledge update, and compositional multi-hop reasoning. StateTree outperforms both SFT and RL-based baselines while preserving short-context general reasoning. StateTree-7B achieves gains up to +23.60% on LongMemEval (128k), and StateTree-14B reaches 59.00% accuracy on LongMemEval, surpassing QwenLong-L1-32B (45.20%).
Sources
- LongReason: A Synthetic Long-Context Reasoning Benchmark via Context Expansion
- DeepSeek-V3 Technical Report
- MiniMax-01: Scaling Foundation Models with Lightning Attention
- YaRN: Efficient Context Window Extension of Large Language Models
- LongRoPE2: Near-Lossless LLM Context Window Scaling
- APIGen-MT: Agentic Pipeline for Multi-Turn Data Generation via Simulated Agent-Human Interplay
- ProRL: Prolonged Reinforcement Learning Expands Reasoning Boundaries in Large Language Models
- SFT Memorizes, RL Generalizes: A Comparative Study of Foundation Model Post-training
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
- Know Me, Respond to Me: Benchmarking LLMs for Dynamic User Profiling and Personalized Responses at Scale
- Large Language Models Can Self-Improve in Long-context Reasoning
- QwenLong-L1: Towards Long-Context Large Reasoning Models with Reinforcement Learning
- Qwen2 Technical Report
- Qwen3 Technical Report
- DAPO: An Open-Source LLM Reinforcement Learning System at Scale
- Measuring Massive Multitask Language Understanding
- Let's Verify Step by Step
- Instruction-Following Evaluation for Large Language Models
- GPT-4o System Card
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering