StateTree: Enhancing Long-Term Dialogue Reasoning via Reinforcement Learning¶

会议: NeurIPS 2026
arXiv: 2609.38809
OpenReview: https://openreview.net/forum?id=Wtmbj6bI9C
代码: https://github.com/bluedream02/StateTree
领域: 强化学习
关键词: 长程对话推理, 强化学习, 课程学习, 状态树, 多会话记忆

一句话总结¶

针对长程多会话对话中证据跨会话分散、时序偏好动态演变及全长训练算力高昂等难题，StateTree 将多会话对话推理重塑为二叉状态树路径追踪伪任务，通过四阶段渐进式课程 GRPO 与混合判决奖励，仅凭 616 条样本在 10K 上下文下训练，即可外推至 128K 极端长上下文对话推理并显著超越更大体量基线。

研究背景与动机¶

大语言模型被广泛部署为写作、咨询与推荐等长程交互中的个性化助手，但在多会话（multi-session）长交互历史中，模型普遍难以有效调用并推理用户的沉淀知识。现有多会话对话评测揭示出模型存在严重的“迷失于长程对话”（lost-in-the-middle）和容错恢复失效现象；而现有的记忆增强外挂系统虽然能通过检索压缩上下文，其最终回答质量依然严重受制于底层基座大模型的内在推理上限。更为棘手的是，大多数长上下文研究将交互历史视作扁平的静态长文档，完全忽视了个性化对话特有的非平稳动态特性。

长程对话推理面临三重本质挑战：首先是极度的数据匮乏（\(\mathbb{C}1\)），真实多会话对话标注数据集极小（如 LoCoMo 仅包含 10 个对话约 1500 个问答对），而纯合成数据难以保证答案真实性；其次是高度复杂的非平稳推理要求（\(\mathbb{C}2\)），对话证据往往分散在不同时间戳的会话中，且用户偏好与个人状态会随时间发生动态修正，要求模型同时具备跨会话检索、时序比对、知识更新以及多跳组合推理能力；最后是全长上下文训练的高昂算力成本（\(\mathbb{C}3\)），在目标长文本（如 64K–128K）尺度直接开展强化学习训练开销巨大，并且极易破坏模型的短文本基础常识与通用数学推理能力。

解决上述矛盾的切入角度在于：多会话推理在本质上具有时序约束与层级依赖结构，完全可以被建模为受约束的路径搜索问题。既然真实对话数据稀缺且难以获取精细监督，便可在真实对话背景中植入带有确定真值标签的结构化辅助任务。核心 idea：将多会话对话推理重构成二叉树路径追踪伪任务（StateTree），把无语义的键值记录离散嵌入各会话并利用时间戳判别分支真伪，配合阶梯式加深的课程强化学习（Curriculum GRPO）与精确匹配/语义裁判混合奖励，以 10K 紧凑上下文训练激发模型泛化至 128K 的长程时序与多跳推理潜能。

方法详解¶

整体框架¶

StateTree 的核心流水线围绕“结构化伪任务构造”与“四阶段课程强化学习”展开。在数据构造端，系统在给定的真实多会话对话 \(\mathcal{S}=\{s_1, \ldots, s_S\}\) 基础上，围绕一个事实问答对 \((q_i, a_i)\) 构建深度为 \(D\) 的完全二叉树 \(\mathcal{T}\)，将根节点到目标叶子节点的各条边转化为 JSON 键值对嵌入不同会话中，并通过会话时间戳判定正确分支。在模型训练端，系统采用基于分组相对策略优化（GRPO）的多阶段课程训练策略，由浅入深逐步提高树深度与语义组合要求，并辅以混合验证奖励进行稳定优化。

整体系统的数据流向与训练架构如下所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["真实多会话对话数据<br/>带时间戳会话与问答对"] --> B["基础状态树生成<br/>无语义UUID二叉树路径嵌入"]
    B --> C["组合状态树生成<br/>三维语义轴分解与碎片化挂载"]
    C --> D["四阶段课程强化学习<br/>Warmup → D=2 → D=3 → 组合树"]
    D --> E["混合奖励判决机制<br/>字符串精确匹配与模型裁判结合"]
    E --> F["推理输出与长程外推<br/>提取答案并泛化至128K评测"]

关键设计¶

1. 基础状态树路径追踪任务：以时间戳判别与无语义键打破检索捷径 为了解决模型在跨会话多跳检索中缺乏明确监督信号的痛点，基础状态树（Basic StateTree）构建了一棵深度为 \(D\) 的满二叉树。树中的每个内部节点均被赋予一个随机生成的无语义 UUID 字符串（如 f391e945-...），彻底切断模型利用词面语义相似度走捷径检索的可能性。内部节点的两条出边编码为 JSON 记录（{kv: kuj}），而叶子节点则挂载自然语言问题（其中唯一的正确叶子节点挂载目标问题 \(q_i\)，其余叶子节点由对话中的其他合法问题作为干扰项）。整棵树的 \(2(2^D-1)\) 条边记录被均匀打散插入各会话的句子边界内。更关键的是，在正确路径上的每个分叉处，正确边记录所在的会话时间戳被严格约束为晚于干扰边所在会话的时间戳（\(\text{timestamp}(s_{\text{corr}}) > \text{timestamp}(s_{\text{dist}})\)）。模型必须从 Prompt 提供的根节点 UUID 出发，逐层比对各会话时间戳以选取最新边，最终遍历到达目标叶子并回答问题，同时强化了跨会话检索、时序比对与知识更新能力。

2. 组合状态树：语义维度解构与多跳推理碎片聚合 基础状态树虽然锻炼了多跳指针寻径与时序判别，但边之间传递的仅是裸指针，模型无须在推理链路中对中间线索进行语义整合。针对复杂多跳推理需要综合不同阶段证据的痛点，组合状态树（Compositional StateTree）引入了层级语义分解机制。针对目标问题 \(q_i\)，利用大模型沿人物主体（Level 1）、事件类别（Level 2）以及问题细节（Level 3）三个正交语义轴解构成 \(2 \times 2 \times 2\) 的三层二叉决策空间。树上的每条边不再只是简单的下一跳地址，而是携带着对应的推理碎片字段（step），如根节点边携带 [A] Jon，二级边携带 [B] expanding his studio's social media presence，三级边携带 When did [A] start [B]?。模型必须沿正确时间线一路搜集这些局部推理片段，将它们拼装组合成完整的复合问题，再在全局对话历史中定位答案。这种机制迫使模型在长程记忆中习得组合式多跳推理（Compositional Multi-Hop Reasoning）的原语。

3. 渐进式课程强化学习：从直接对话到高深树的阶梯进阶 直接在高深度树或复杂组合树上进行强化学习会面临极大的奖励稀疏与策略震荡问题。为此，论文设计了平滑的四阶段课程训练流水线（Curriculum RL）： - 预热阶段（Warm-up，40 steps）：直接在无树状结构的原始多会话 QA 数据上进行微调，让模型熟悉基本对话长上下文理解； - 第一阶段（Stage 1，100 steps）：引入深度 \(D=2\) 的基础状态树（4 个叶子节点，6 条插入记录），建立浅层跨会话检索与时序二选一判别能力； - 第二阶段（Stage 2，100 steps）：推进至深度 \(D=3\) 的基础状态树（8 个叶子节点，14 条插入记录），强化多步跳跃链路的记忆维持； - 第三阶段（Stage 3，60 steps）：切换为深度 \(D=3\) 的组合状态树，引入语义碎片聚合任务，并在该阶段混入 2,500 条 DAPO-Math 数学推理样本，以防止强化学习在长上下文专业任务微调中发生通用逻辑与数学能力的灾难性遗忘。

4. 规则匹配与语义裁判结合的双轨混合奖励 在长对话 QA 的强化学习优化中，仅用精确匹配（Exact Match, EM）会导致表达不同但语义正确的同义回答被判定为 0 奖励，抑制有效探索；而若完全依赖大模型判决（LLM-as-a-Judge），则会引入评判噪声、计算延迟与评分漂移。StateTree 提出了规则与模型判决取最大值的混合奖励函数：

\[r_\phi(x, y) = \max\Big(r_{\text{EM}}(y_{\text{ans}} = y_{\text{gold}}), \; r_{\text{LLM}}(x, y_{\text{ans}}, y_{\text{gold}})\Big)\]

Prompt 要求模型将最终答案统一置于 \boxed{...} 框内以便精确提取 \(y_{\text{ans}}\)。当模型输出与真实答案 \(y_{\text{gold}}\) 字面完全一致时，\(r_{\text{EM}}\) 直接赋予 1 奖励；若字面不匹配，则由部署在本地的轻量化开源判决模型 Qwen2.5-1.5B-Instruct（在温度 0 下运行严格模板）输出二元语义等价判断 \(r_{\text{LLM}} \in \{0, 1\}\)。通过两者取上限，既保证了刚性正确答案的确定性，又赋予了开放表达的语义容错能力。

损失函数 / 训练策略¶

模型基于分组相对策略优化（GRPO）算法展开训练。针对输入上下文 \(\mathcal{L}\)、输入问题 \(q\) 及真实答案 \(a\)，策略模型 \(\pi_{\theta_{\text{old}}}\) 采样生成一组由 \(G\) 个候选轨迹构成的候选组 \(\{o_1, \ldots, o_G\}\)。策略网络 \(\pi_\theta\) 通过最大化带截断重要性采样与 KL 惩罚的目标函数进行优化：

\[J_{\text{GRPO}}(\theta) = \mathbb{E} \left[ \frac{1}{G} \sum_{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \left( \min\left[ \rho_{i,t}(\theta) A_i, \text{clip}(\rho_{i,t}(\theta), 1-\varepsilon, 1+\varepsilon) A_i \right] - \beta D_{\text{KL}}(\pi_\theta(\cdot \mid q, o_{i,<t}) \parallel \pi_{\text{ref}}(\cdot \mid q, o_{i,<t})) \right) \right]\]

其中重要性权重为 \(\rho_{i,t}(\theta) = \frac{\pi_\theta(o_{i,t} \mid q, o_{i,<t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,<t})}\)，优势项 \(A_i\) 由组内各轨迹所得奖励 \(r_i\) 进行标准化计算：

\[A_i = \frac{r_i - \text{mean}(\{r_1, \dots, r_G\})}{\text{std}(\{r_1, \dots, r_G\})}\]

训练超参数方面，组大小 \(G=8\)，学习率为 \(1 \times 10^{-6}\)，梯度截断阈值 1.0，KL 散度系数 \(\beta=0.001\)。采样温度 0.6，top-\(p=0.95\)。7B/8B 模型的 Prompt 批大小为 64，14B 模型为 32。模型在 32 卡 NVIDIA H20 集群上完成训练，7B、8B 和 14B 模型训练耗时分别为 90、100 和 160 小时。测试与验证集完全保留无树状结构的原始对话问答格式，确保评测反映真实的下游推理能力。

实验关键数据¶

主实验¶

论文在域内基准 LoCoMo（约 10K 上下文）以及两个极端长上下文域外（OOD）基准 LongMemEval（128K）和 PersonaMem（128K）上进行了全面评测，主要对比结果如下表所示：

模型	LoCoMo Avg ACC (%)	LoCoMo Temporal ACC (%)	LongMemEval 128k Avg ACC (%)	PersonaMem 128k Avg ACC (%)
GPT-4o	52.73	14.47	-	-
QwenLong-L1-32B	46.36	34.59	51.00	52.40
R1-Distill-Qwen-32B	51.43	33.33	29.00	37.62
Qwen2.5-7B-Instruct (Base)	44.29	23.90	23.80	40.48
SEALONG-7B	45.19	20.75	45.40	41.66
LoongRL-7B	46.10	25.16	26.60	31.39
RL-MemAgent-7B	48.03	31.45	41.80	43.78
StateTree-7B (本文)	49.61 (+5.32)	32.70 (+8.80)	47.40 (+23.60)	47.30 (+6.82)
Qwen2.5-14B-Instruct (Base)	49.09	34.59	45.20	45.40
SEALONG-14B	48.96	32.70	47.80	46.02
LoongRL-14B	52.99	37.74	54.20	48.07
RL-MemAgent-14B	51.43	37.11	52.00	48.15
StateTree-14B (本文)	60.91 (+11.82)	41.51 (+6.92)	59.00 (+13.80)	52.59 (+7.19)
Qwen3-8B (Base)	43.64	29.56	49.20	45.36
StateTree-8B (本文)	52.08 (+8.44)	37.74 (+8.18)	58.66 (+9.46)	50.31 (+4.95)

消融实验¶

论文通过两组严密的消融实验，分别验证了课程阶段设计的必要性以及任务拓扑结构的关键作用（均基于 7B 模型）：

课程阶段消融配置	LoCoMo Avg ACC (%)	LongMemEval 128k Avg (%)	PersonaMem 128k Avg (%)	说明
StateTree-7B (完整模型)	49.61	47.40	47.30	完整四阶段课程
w/o Warm-up 阶段	45.19 (-4.42)	43.00 (-4.40)	44.00 (-3.30)	缺少对话基础问答对齐，初期优化不稳定
w/o Stage 1 (\(D=2\))	44.68 (-4.93)	42.60 (-4.80)	43.78 (-3.52)	降幅最大，证明浅层跨会话判别为基石能力
w/o Stage 2 (\(D=3\))	46.23 (-3.38)	44.20 (-3.20)	45.10 (-2.20)	削弱了深层多跳链路的维持能力
w/o Stage 3 (组合树)	46.75 (-2.86)	44.80 (-2.60)	45.51 (-1.79)	削弱了局部推理片段的语义拼装能力

结构机制与奖励消融配置	LoCoMo Avg ACC (%)	LongMemEval 128k Avg (%)	说明
基座模型 Qwen2.5-7B	44.29	23.80	无结构训练的基座基线
+ 思维链提示词 (w/ CoT prompt)	45.19 (+0.90)	24.60 (+0.80)	仅靠单纯 Prompt 引导思考提升极微
独立随机键 (w/ Independent Random Keys)	45.45 (-4.16)	25.20 (-22.20)	分叉失去同名键对比，无法建立时序决策
单会话放置 (w/ Single Session)	47.79 (-1.82)	42.80 (-4.60)	消除跨会话检索，失去长程分散检索信号
具名实体键 (w/ Entity Keys)	48.44 (-1.17)	45.60 (-1.80)	单词语义引发浅层匹配捷径，弱化逻辑严谨度
纯精确匹配奖励 (Exact Match only)	47.79 (-1.82)	46.80 (-0.60)	错杀同义表述，探索信号过刚性
纯模型裁判奖励 (LLM-Judge only)	48.70 (-0.91)	46.60 (-0.80)	评分存在噪声与浮动，训练收敛较慢

关键发现¶

以小博大的高样本效率：仅依靠 616 条训练样本在 10K 上下文上微调，StateTree-14B 的 LoCoMo 平均准确率达到 60.91%，超越了参数量大一倍以上的 QwenLong-L1-32B（46.36%）以及专有模型 GPT-4o（52.73%）。
极端上下文长度外推能力：在完全未见过的 128K 极端长上下文域外评测 LongMemEval 上，StateTree-7B 取得了惊人的 +23.60% 准确率飞跃（从 23.80% 提升至 47.40%），证明通过树状寻径强化的时序比对与多跳定位原语具有跨长度通用性。
浅层时序判别是核心基石：消融表明 Stage 1（\(D=2\)）对整体系统贡献最大（移除后 LoCoMo 掉点 4.93%，LongMemEval 掉点 4.80%）；若破坏键值配对机制（Independent Random Keys），LongMemEval 准确率直接雪崩 22.20%，证明收益并非来自模型对长文本的虚假拟合，而是源自真实学习到的分支决策算法。
零通用能力退化：在通用短文本基准测试中，StateTree-8B 在 MMLU、MATH-500 和 IFEval 上的平均表现反而微增 +0.10%，7B 与 14B 模型的性能浮动亦均在 0.6% 以内，彻底避免了长文本 RL 常伴随的通用能力坍塌。

亮点与洞察¶

将非平稳对话重塑为树搜索：传统长文本研究将对话视作平铺文章，而 StateTree 深刻抓住了多会话交互中“时间演进”与“偏好修订”的本质，将其形式化为二叉决策树上的时间戳偏序判定，设计极为巧妙。
无语义 UUID 打破表面词汇关联：使用随机 UUID 作为内部节点键，强迫模型在各会话中执行真实且精准的精确字符串寻径与跨会话指针追踪，杜绝了模型依赖软注意力模糊匹配走捷径的可能。
短训长用打破算力壁垒：完全避开了在 64K/128K 尺度进行高昂全量 RL 的工程瓶颈，通过在 10K 紧凑窗口中沉淀核心推理原语，实现了零显存溢出风险的超长距离推理外推。

局限与展望¶

数据源局限：所有训练构造均基于单一英语多会话数据集 LoCoMo，在跨语言（如中文长对话）及复杂跨领域场景下的泛化能力尚待验证。
组合树构造对商业模型的依赖：组合状态树中的三轴语义分解高度依赖 GPT-4o 的离线拆解与干扰项生成，可能引入专有模型的先验偏置。
固定深度的二叉拓扑限制：当前研究仅探索了静态固定深度（\(D=2, 3\)）的满二叉树，未来若能引入动态深度或多叉树结构，有望进一步提升复杂决策的表征能力。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [将多会话对话推理创造性地转化为时间戳偏序的二叉树路径追踪伪任务，设计独树一帜]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 7B/8B/14B 多种架构，跨 LoCoMo、LongMemEval、PersonaMem 及短文本基准，消融极其详实]
写作质量: ⭐⭐⭐⭐⭐ [逻辑链路严谨清晰，问题定义抽象精炼，数学与算法表述极为规范]
价值: ⭐⭐⭐⭐⭐ [为学术界和工业界在有限算力下开展长程个性化对话强化学习提供了极具启发性的高能效范式]