PaperlayerView exact revision on arXiv

Source-linked AI summary

StateTree: Enhancing Long-Term Dialogue Reasoning via Reinforcement Learning

Naen Xu, Wanqing Cui, Yibo Hu, Shixin Hong, Hengyu An, Meiguang Jin, Junfeng Ma, Tianyu Du

arXiv:2609.38809v1cs.CLcs.AI

TL;DR

Long-term dialogue reasoning must recover evidence scattered across sessions as user preferences evolve, yet training is constrained by scarce verifiable data and the cost of long-context RL. StateTree embeds a tree-structured path-tracing task in dialogues and trains with curriculum RL at 10K-token contexts. It improves over SFT and standard RL baselines, while generalizing to longer contexts and preserving short-context abilities.

  • Problem

    Long-term dialogue reasoning requires using scattered evidence and revised preferences across sessions, but available training data are scarce and full-context RL is costly.

    Source: §1 · p. 1

  • Method

    StateTree embeds a timestamped binary-tree path-search task in authentic multi-session dialogues and uses curriculum RL, including a compositional variant with step-level reasoning fragments.

    Source: §1 · p. 2

  • Results

    StateTree improves over SFT and standard RL baselines across scales on LoCoMo, with StateTree-14B reaching 60.91% average accuracy.

    Source: §4.2 · p. 7

  • Takeaways & Limitations

    Training on 10K-token contexts generalizes to 128K while preserving short-context reasoning abilities.

    Source: §1 · pp. 2–3

    All training data come from one English dialogue dataset, so generalization to other languages or dialogue domains remains unvalidated.

    Source: §A.1 · p. 14

Abstract

from arXiv · show

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%).

1 Introduction

Long-term dialogue reasoning is difficult because evidence is scattered across evolving sessions, while scarce verified data and costly long-context RL constrain training. StateTree turns dialogues into a structured path-search task and reports gains in dialogue reasoning with training at 10K-token contexts.

1 Introduction

  • Training data are scarce: LoCoMo has only 10 dialogues and 1,540 QA pairs, while LLM-synthesized dialogues cannot guarantee answer correctness.Long-context RL also incurs prohibitive computational costs and risks degrading short-context and general reasoning abilities.Source: §1 · pp. 1–2
  • StateTree-trained models exhibit cross-session retrieval, temporal reasoning, knowledge updates, and multi-hop reasoning in the illustrated dialogue trajectories.The figure contrasts these behaviors with QwenLong-L1-32B getting lost across sessions and missing information.Source: Figure 1 · p. 2
  • Long-term dialogues challenge models to retrieve sparse evidence across sessions, compare timestamps, track revised preferences, and combine multiple evidence pieces.Memory-augmented systems still rely on the backbone model’s reasoning, while much long-context research treats histories as flat static text.Source: §1 · pp. 1–2
  • StateTree embeds key-value records across dialogue sessions as a binary tree whose timestamped branches lead to a target question, with curriculum RL and a compositional stage adding step-level reasoning fragments.Training on 10K-token contexts avoids the prohibitive cost of full-length RL.Source: §1 · pp. 2–3
  • 60.91% average accuracy on LoCoMo is achieved by StateTree-14B, which surpasses SFT and standard RL baselines across parameter scales and base architectures.The introduction also reports generalization to 128K benchmarks and preservation of short-context reasoning.Source: §1 · p. 2

2 Related Work

Prior work evaluates long-term dialogue over extended interaction histories and develops memory-augmented and long-context reasoning approaches. These approaches address retrieval or long-context reasoning, but remain limited for non-stationary multi-session dialogue.

Evaluating Long-term Multi-Session Dialogues

  • Dialogue evaluation has shifted toward realistic interactions spanning 10K to over 100K tokens, with LoCoMo, LongMemEval, and PersonaMem among representative benchmarks.The cited evaluations find that even strong models struggle with long-term dialogue reasoning.Source: §2 · p. 3

Long-Context Reasoning

  • Memory-augmented systems externalize memory, while synthetic-data SFT and RL approaches target long-context reasoning; these methods are designed for static documents rather than non-stationary multi-session dialogue.The passage describes SEALONG, QwenLong-L1, and LoongRL as examples of these approaches.Source: §2 · p. 3

3 StateTree

StateTree frames long-term dialogue reasoning as traversing a structured path through evidence scattered across sessions. It pairs this task with progressively harder RL training, including a compositional variant that requires combining reasoning fragments.

3 StateTree

  • StateTree embeds a hidden question in a tree of records scattered across dialogue sessions, requiring retrieval, temporal discrimination, and multi-hop traversal to recover and answer it.The design targets cross-session retrieval, temporal reasoning and knowledge update, and multi-hop reasoning.Source: §3 · p. 3, §3.1 · p. 4

3.1 Data Construction Pipeline

  • Basic StateTree uses UUID-keyed binary-tree records and plausible distractor questions, with the correct edge placed in a more recent session at each on-path fork.The model must follow the unique root-to-target-leaf path rather than choose a question based on its apparent relevance.Source: §3.1 · p. 5
  • Compositional StateTree replaces bare pointers with edge-level reasoning fragments that the model accumulates along the path and aggregates into a question before answering.Distractor leaves differ from the target along at least one semantic axis.Source: Figure 3 · p. 4, §3.1 · pp. 5–6

3.2 Long-Context Multi-turn Dialogue Reinforcement Learning

  • Training progresses from dialogue QA warmup through Basic StateTree at depths 2 and 3, then to depth-3 Compositional StateTree.The stages add cross-session retrieval and temporal discrimination, another traversal hop, and then fragment composition.Source: §3.2 · p. 6
  • The combined reward pairs exact-match verification with an LLM judge's binary semantic-equivalence score for the extracted answer.The model is prompted to place its final answer in a boxed format for unambiguous extraction.Source: §3.2 · p. 6

4 Experiments

StateTree improves LoCoMo performance while retaining short-context reasoning, and the experiments analyze training dynamics, emergent behaviors, and the contributions of curriculum, task structure, and reward design. The setup evaluates dialogue reasoning and short-context capabilities against frontier, long-context, and memory-augmented baselines.

4.1 Experimental Setup

  • Experiments train Qwen2.5-7B, Qwen2.5-14B, and Qwen3-8B with four curriculum stages, augmenting 616 training samples with StateTree in Stages 1–3.Evaluations cover LoCoMo, LongMemEval, PersonaMem, and short-context MMLU, MATH-500, and IFEval; baselines include frontier, long-context, and memory-augmented models.Source: §4.1 · pp. 6–7

4.2 Main Results

  • 60.91% average accuracy on LoCoMo for StateTree-14B is +11.82% over its base and surpasses GPT-4o by 8.18% and QwenLong-L1-32B by 14.55%.The largest reported lift is Temporal reasoning (F1 +26.19 for 14B), attributed to explicit timestamp comparison at each branching step.Source: §4.2 · p. 7, Table 2 · p. 7
  • +0.1% average gain over base Qwen3-8B; 7B and 14B drops are -0.2% and -0.6%, respectively, with competitive MMLU, MATH, and IFEval results.The passage reports the 8B variant has zero degradation and describes the other models’ drops as minimal.Source: §4.2 · p. 8, Table 5 · p. 8

4.3 Analysis

  • Validation reasoning F1 rises monotonically across training stages as average response length increases, with no reported sign of reward hacking.The authors interpret the trends as evidence of a stable combined-reward signal and curriculum scaling of reasoning.Source: §4.3 · p. 8, Figure 4 · p. 8
  • Qualitative analysis identifies cross-session retrieval, multi-hop synthesis, timestamp comparison, and privileging newer information as emergent behaviors.The passages describe these behaviors as aligned with the pseudo-task design and as patterns that transfer beyond the training distribution.Source: §4.3 · pp. 8–9

4.4 Ablation Study

  • Removing Stage 1 (D=2) drops performance by -4.93, the largest curriculum ablation; removing warm-up, Stage 2 (D=3), and Stage 3 drops it by -4.42, -3.38, and -2.86.The authors associate warm-up with basic conversation QA, Stage 2 with multi-hop temporal reasoning, and Stage 3 with compositional reasoning.Source: §4.4 · p. 9, Table 6 · p. 9
  • Structural ablations indicate that shared-key forks, cross-session scattering, and opaque identifiers contribute distinct signals, while generic CoT prompting yields only marginal gains.The passage reports that entity keys capture 80% of the full gain and attributes the remaining deficit to lexical shortcuts.Source: §4.4 · p. 9, Table 7 · p. 9
  • The combined reward balances exact correctness with semantic flexibility; EM, standalone LLM-Judge, and dense lexical metrics each have reported drawbacks.The max-selection mechanism is described as providing a dense, reliable training signal aligned with answer correctness.Source: §4.4 · pp. 9–10, Table 8 · p. 9

5 Conclusion

StateTree reframes long-term dialogue reasoning as temporally constrained path search and trains with embedded tree tasks and a structured curriculum. The paper reports gains across in-domain and out-of-domain benchmarks, generalization from 10K to 128K contexts, and preserved short-context capabilities.

5 Conclusion

  • StateTree uses 616 examples at 10K context to elicit cross-session retrieval, temporal reasoning, knowledge update, and compositional multi-hop reasoning, with gains across in-domain and out-of-domain benchmarks.The conclusion reports that models surpass larger baselines, generalize to 128K contexts without degrading short-context capabilities, and exhibit reasoning behaviors aligned with the task design.Source: §5 · p. 10

A.1 Limitations

The appendix identifies limitations in StateTree’s data scope, compositional-task generation, and fixed tree depths, and notes privacy risks associated with more capable dialogue reasoning.

A.1 Limitations

  • Training uses one English dialogue dataset, leaving generalization to other languages and dialogue domains unvalidated; compositional generation may inherit GPT-4o semantic biases, and only fixed depths D=2 and D=3 are explored.The authors note that adaptive depths or dynamic tree structures are not explored.Source: §A.1 · p. 14

A.2 Broader Impacts

  • More capable dialogue reasoning may also enable sensitive-information extraction or unauthorized surveillance, motivating privacy safeguards, user consent, and data-retention policies.Source: §A.2 · p. 14

B.1 Existing Assets; B.2 LLM Usage Declaration

  • The appendix lists LoCoMo, LongMemEval, PersonaMem, Qwen models, and GPT-4o among the assets used, and declares GPT-4o for decomposition and Qwen2.5-1.5B-Instruct for reward judging.Source: §B.1 · p. 14, §B.2 · p. 14

C.1 Source Data; C.2 Dataset Split

  • All curriculum stages reuse 616 LoCoMo conversation–question pairs, while the held-out test set uses the original direct-QA format without tree augmentation or inserted records.The remaining pairs are split between training/validation and held-out test data.Source: §C.1 · p. 14, §C.2 · p. 14

C.3 Dataset Statistics and Construction Parameters

  • Table 9 summarizes dataset statistics and construction parameters by curriculum stage, including UUID format and record insertion mode.Source: Table 9 · p. 15, §C.3 · p. 15

C.4 Question Decomposition

  • Stage 3 decomposes each target question into a 2 × 2 × 2 tree of person, event, and question-detail fragments, yielding eight paths per conversation.The example decomposition’s three fragments combine into the target question.Source: §C.4 · p. 15, Table 10 · p. 15

C.4.1 Decomposition LLM Prompt

  • GPT-4o generates decompositions from the target question and conversation question pool, and structural checks validate tree shape, placeholders, target-leaf count, and sibling diversity before acceptance.The prompt requires concrete factual questions and exactly one target leaf.Source: §C.4.1 · pp. 15–16

C.5 Distractor Question Pool; C.6 Record Insertion Strategy

  • Distractors come from the same conversation or differ from the target along at least one semantic axis, while inserted records are spread across sessions and temporal forks are checked so correct-path records are newer.Records are placed at sentence boundaries inside speaker quotations; repeated keys are separated across sessions.Source: §C.5 · p. 16, §C.6 · p. 17

C.7 Entity Key Pool for Ablation

The entity-key ablation replaces nonsemantic UUIDs with familiar concrete nouns while preserving the key-pairing mechanism.

C.7 Entity Key Pool for Ablation

  • In the entity-key ablation, UUIDs are replaced by high-frequency concrete nouns sampled without replacement from a 14-word pool.Source: §C.7 · p. 17, Table 11 · p. 17
  • The ablation is designed to distinguish effects of semantic key content from topical relevance.Source: §C.7 · p. 17
  • The key-pairing mechanism remains unchanged, allowing the model to use familiar word meanings rather than rely on pure string matching.Source: §C.7 · p. 17

D Full Prompt Templates

The appendix supplies prompts for direct dialogue QA, tree navigation, compositional question reconstruction, and answer evaluation, alongside examples of records embedded in conversations.

D.1 System Prompt

  • The shared system prompt requests a step-by-step reasoning trace in <think> tags and a final answer in \boxed{}.Source: §D.1 · p. 17

D.2 Warm-up: Direct Dialogue QA

  • The Warm-up prompt follows the LoCoMo evaluation format, asking for a short answer grounded in the conversations and favoring exact wording when possible.Source: §D.2 · p. 18

D.3 Stage 1 & 2: Basic StateTree (D=2 / D=3)

  • Basic StateTree prompts instruct the model to follow UUID key-value chains, selecting the record from the most recent session when a key appears more than once, until it reaches and answers a question.Source: §D.3 · p. 18

D.4 Stage 3: Compositional StateTree (D=3)

  • The compositional prompt directs the model to follow a tree, accumulate each step fragment, combine the fragments into a question at the leaf, and answer it from the dialogue.Source: §D.4 · p. 18

D.5 CoT Prompt Baseline

  • The CoT baseline adds explicit step-by-step reasoning instructions but uses neither StateTree structure nor RL training.It tests whether generic reasoning prompts can replicate gains from structural navigation.Source: §D.5 · pp. 18–19

D.6 Evaluation Prompts

  • Evaluation prompts judge answers against gold responses, with LongMemEval variants addressing standard tasks, temporal reasoning, knowledge updates, and personalized preferences.The LoCoMo prompt labels an answer correct or wrong; LongMemEval prompts use yes/no judgments with task-specific criteria.Source: §D.6.1 · p. 19, §D.6.2 · pp. 19–20

E.1.1 Case Study

  • The Basic StateTree case study embeds UUID records in dialogue sessions and illustrates following the newest-session forks to a question leaf before answering from the conversation.The example distinguishes the target path from distractor paths.Source: §E.1.1 · pp. 20–21

E.2 Compositional StateTree Record (Stage 3)

  • In the compositional example, person, event, and question-detail fragments occupy successive records; aggregating the path reconstructs the full question.Source: §E.2 · p. 21

F Hyperparameters and Training Configuration

Training uses GRPO to optimize against groups of rollout trajectories with a combined reward, and reports compute resources and repeated-run averaging.

F.1 Group Relative Policy Optimization (GRPO)

  • For each question, augmented dialogue context, and gold answer, GRPO samples rollout trajectories from the old policy and optimizes the policy using a clipped objective with a KL penalty.The clipping range and KL-penalty weight are controlled by ε and β.Source: §F.1 · pp. 21–22
  • Trajectory advantages are computed from group rewards, and the combined reward function is used to mitigate reward hacking.Experiments are run with three random seeds and report averages.Source: §F.1 · p. 22

F.2 Compute Resources

  • Training used 32 NVIDIA H20 GPUs; reported runtimes were approximately 90 hours for Qwen2.5-7B-Instruct, 100 hours for Qwen3-8B, and 160 hours for Qwen2.5-14B-Instruct.Source: §F.2 · p. 22

G.1 Benchmark Descriptions

The evaluation covers in-domain and out-of-distribution dialogue reasoning, plus short-context general reasoning. On PersonaMem-32k, StateTree-14B reports 64.52% average accuracy, including 86.87% on recalling update reasons.

G.1 Benchmark Descriptions

  • LoCoMo is the in-domain dialogue benchmark, while LongMemEval and PersonaMem test out-of-distribution long-term memory and dialogue reasoning.LoCoMo includes Multi Hop, Temporal, Open Domain, and Single Hop categories; LongMemEval covers multi-session, temporal, knowledge-update, and single-session dimensions, and PersonaMem evaluates evolving profiles and personalized responses.Source: §G.1 · p. 22
  • MMLU tests general knowledge and reasoning across 57 academic subjects, MATH-500 is a curated subset of 500 competition-level mathematics problems, and IFEval measures compliance with specific formatting and content constraints.Source: §G.1 · p. 22

G.2 Inference Configuration

  • Evaluations use temperature 0.6, a maximum 128K-token input, and a maximum 4096-token output.Source: §G.2 · p. 22

G.3 Additional OOD Results

  • 64.52% average accuracy (+5.61% over base) on PersonaMem 32k, with 86.87% on recalling update reasons, is reported for StateTree-14B.The results are presented in the OOD PersonaMem-32k comparison.Source: §G.3 · p. 22, Table 12 · p. 23

H Category-Level Curriculum Ablation

The curriculum ablation finds that shallow tree navigation and QA warm-up contribute most, with deeper and compositional stages adding smaller but appreciable gains. Baseline comparisons use released checkpoints where available and retraining otherwise.

H Category-Level Curriculum Ablation

  • LoCoMo gains rank Stage 1, warm-up, Stage 2, and Stage 3 at +4.93, +4.42, +3.38, and +2.86; LongMemEval ranks them at +4.80, +4.40, +3.20, and +2.60.Stage 1 is the largest contributor on both benchmarks, while deeper compositional stages have diminishing but appreciable marginal returns.Source: H Category-Level Curriculum Ablation · p. 23, Table 13 · p. 23
  • Stage 1 contributes most to Temporal and Multi-Session reasoning, Stage 2 to LoCoMo Multi-Hop, and the compositional stage contributes more uniformly across categories.Warm-up contributes most on Open Domain, Single-Session User/Assistant/Preference, and foundational QA.Source: H Category-Level Curriculum Ablation · p. 23

I Baseline Implementation Details

  • SEALONG and RL-MemAgent are evaluated with their released checkpoints, while LoongRL is retrained on Qwen2.5 base models using its public dataset and published configuration.Source: I Baseline Implementation Details · p. 23

J Complete Training Examples

The examples use a LoCoMo conversation to show both basic UUID-chain traversal and a compositional tree whose semantic fragments combine into a question. In both examples, the recovered question is answered from the dialogue.

J.1 Basic StateTree Training

  • Basic traversal searches all records for the current UUID and follows the value from its most recent session, repeating until it reaches a question to answer.The example distributes six edge records across six distinct sessions, and its correct path follows two edges.Source: §J.1 · pp. 24–25
  • The Basic StateTree example hides “What is Caroline’s identity?” on a two-edge path and gives “Transgender woman” as the answer.The example introduces the training input for conversation conv-26; only turns containing tree records are shown.Source: J Complete Training Examples · p. 23, §J.1 · p. 24, Table 15 · p. 25

J.2 Compositional StateTree Training (Stage 3)

  • The compositional example uses nested records whose step field carries a semantic fragment and whose next field points to a child UUID.Its correct three-edge path combines “[A] Caroline,” “[B] discussing her identity,” and “What is [A]’s identity while [B]?” to recover the original question.Source: §J.2 · pp. 25–26
  • The compositional example distributes 14 edge records across eight sessions and answers the recovered identity question with “Transgender woman.”Source: §J.2 · p. 26

K.1 Cross-Session Retrieval

The examples contrast StateTree’s handling of scattered evidence with base-model errors in retrieval, composition, chronology, and preference updates. Across these cases, StateTree retrieves and combines evidence across sessions or privileges newer information.

K.1 Cross-Session Retrieval

  • StateTree enumerates relevant records across sessions before aggregating an answer, while the base model stops at the first matching session and returns a partial answer.In one example, StateTree combines fish counts from two aquariums to reach the ground-truth total of 17.Source: §K.1 · pp. 26–27

K.2 Multi-Hop Reasoning

  • StateTree chains scattered facts into coherent answers where the base model fails to compose them, including inferring that the user was 11 when Alex was born.In a LoCoMo example, the model also uses a March 2022 conversation and a one-month duration to infer that John resumed drumming in February 2022.Source: §K.2 · pp. 27–28

K.3 Temporal Reasoning

  • StateTree compares session timestamps to order events, while the base model mis-orders them using superficial lexical cues.For the workshop-versus-webinar example, the cited comparison identifies the “Data Analysis using Python” webinar as earlier.Source: §K.3 · pp. 28–29

K.4 Knowledge Update

  • StateTree tracks the latest state by privileging newer sessions, whereas the base model conflates old and new mentions and retrieves outdated values.In the postcard example, the base model reports 17 added postcards, while the later conversation gives the updated total as 25.Source: §K.4 · pp. 29–30
Loading 2609.38809v1…