跳转至

TSLM: Tree-Structured Language Modeling for Divergent Thinking

会议: ICLR 2026
OpenReview: https://openreview.net/forum?id=PV5Dy4lW3t
领域: LLM推理
关键词: 树形语言建模, 发散思维, 搜索树监督, 测试时扩展, 系统化探索

一句话总结

TSLM 用几个特殊 token 把一棵完整搜索树(含成功路径与失败死路)线性序列化,让普通自回归语言模型在一次生成里原生地产出多分支探索结构,从而以监督学习的方式内化系统性搜索能力——在 Game of 24 上 pass@1 达到 100%(基线 17%),在更大的 Gridworld 外推任务上 91.5% 远超 Tree-of-Thought 的 42.7%,且推理速度显著更快。

研究背景与动机

领域现状:当前语言模型按 \(p(y\mid x)=\prod_t p(y_t\mid x,y_{<t})\) 顺序逐 token 生成,推理也是一条线性轨迹。要做"多路探索",主流做法要么靠 o1 / DeepSeek-R1 这类拉长推理链的方法,要么靠 Tree-of-Thought(ToT)这种外部搜索脚手架——在每一步采样多个候选、用外部算法(beam search 等)做事后编排。

现有痛点:顺序生成天然有四个毛病:对单一路径的线性承诺、错误沿链传播、需要多个解时重复计算、无法并行系统地探索备选。而外部搜索方法 ToT 是独立采样 k 条轨迹,每条都要把共享前缀重新算一遍,计算量随深度指数爆炸;更糟的是这些样本之间没有协调,导致某些区域冗余重叠、某些关键分支却被漏掉。

核心矛盾:模型本身没有任何机制去连贯地重建整个搜索空间。顺序模型把所有推理坍缩成一条路;外部搜索把探索拆成互不知情的并行采样,得到的是一个"碎片化"分布,而非一棵结构完整的树。

本文目标:让一个标准 Transformer 在单次生成里就吐出完整的分支结构,把"系统性探索"这件事从外部脚手架搬进模型参数内部。

切入角度:作者的观察是——只要能把一棵搜索树无损地线性化成 token 序列,就能用最普通的语言建模损失去训练它;模型学到的不仅是答案,还有"在每个节点该展开哪些分支、哪条是死路"的完整探索过程。

核心 idea:用特殊 token([SEP]/[FAIL]/[GOAL])把搜索树编码进线性序列,让顺序语言模型通过监督学习内化树形探索,推理时再把分支"拼接"回独立上下文并行展开。

方法详解

整体框架

TSLM 要解决的是"如何让一个只会顺序生成的 Transformer 产出一棵带分支的搜索树"。它的做法是:先把搜索树(结构化任务直接由规则生成;开放域任务用自举算法合成)通过特殊 token 序列化成线性序列,再用标准语言建模损失训练——但训练信号被树拓扑约束(每个节点只看祖先和左侧兄弟,不看无关子树),从而学会"上下文解耦";推理时模型一次性生成分支结构,解析特殊 token 找出可行分支,把它们拼接进各自独立的上下文里做 BFS 并行展开,直到命中 [GOAL] 或穷尽。

整条管线如下:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入任务<br/>(状态 + 目标)"] --> B["搜索树监督自举<br/>规则生成 或 ToT 合成伪树"]
    B --> C["Token 序列化<br/>[SEP]/[FAIL]/[GOAL] 编码分支"]
    C --> D["上下文解耦的树形训练<br/>标准 LM 损失 + 拓扑约束"]
    D --> E["连贯树推理<br/>生成树→解析→拼接 BFS 展开"]
    E -->|命中 [GOAL]| F["输出解"]
    E -->|穷尽无解| G["判定不可解"]

关键设计

1. 搜索树监督自举:开放域任务没有现成树,就合成一棵

结构化任务(Game of 24、Gridworld)有明确的转移函数 \(T(s,a)\) 和有限动作空间 \(A(s)\),可以直接枚举出完整搜索树让 TSLM 模仿。但大多数真实任务(ProntoQA、GSM8K)只有正确答案或金标轨迹,没有显式树结构。作者用一个简单的自举算法(Algorithm 1)合成训练树:对每个实例,用一个监督语言模型按 beam search 在每个状态采样 \(k\) 个候选动作并向前传播;把已知金标轨迹当作高优先级分支强制并入,保证每棵树至少含一条有效解;剩余分支用 RAP 的奖励函数 \(R(s,a)\) 排序以优先扩展有希望的动作;最后对冗余路径去重但保留树结构。这样得到的伪树既含成功路径、也含模型探索出的失败死路,正是 TSLM 需要的"完整探索过程"监督信号。

2. Token 序列化:用特殊 token 把树拓扑无损塞进线性序列

这是 TSLM 的核心表示。作者引入几个结构 token,让标准 Transformer 能在线性序列里读写树:[SEP] 标记一个可继续展开的可行动作;[FAIL] 标记死路(不可行动作);[GOAL] 标记目标状态;[BOS]/[EOS] 标记每个节点子展开序列的边界。序列化采用深度优先遍历把整棵树拍平。每个节点 \(s_i\) 生成的 token 序列是 \(y_{s_i}=[a_i, m_i]\),其中 \(a_i\) 是动作描述、\(m_i\in\{[SEP],[FAIL],[GOAL]\}\) 是结构标记。关键在于:这种序列化同时记录了成功路径和失败探索,教会模型完整的搜索过程,而不只是最终答案——这与"只克隆一条金标 CoT"的顺序基线(SC/PC)有本质区别。

3. 上下文解耦的树形训练目标:每个节点只看祖先和左侧兄弟

虽然损失形式就是普通的交叉熵 \(L=-\sum_t \log p(y_t\mid y_{<t},x)\),但它带着一个"非标准"的训练信号。在顺序建模里每个 token 只依赖它的前缀;而 TSLM 里 token 依赖的是树结构。形式化地,节点 \(s_i\) 的上下文被定义为 $\(\text{ctx}(s_i)=\bigcup_{s_j\in\pi(s_i)} y_{s_j}\ \cup\ \bigcup_{s_k\in\text{siblings}(s_i),\,k<i} y_{s_k}\)$ 即只包含它的祖先路径 \(\pi(s_i)\)此前已生成的左侧兄弟\(k<i\) 强制每个分支因子内部从左到右有序),而不包含无关子树。训练目标变成对所有节点求和的 \(L=-\frac{1}{|N|}\sum_{s_i\in N}\sum_t \log p(y_{s_i,t}\mid y_{s_i,<t},\text{ctx}(s_i))\)。这个"选择性条件"正是 TSLM 区别于"把所有探索路径简单首尾拼接"的关键:虽然输入序列是线性的,但条件依赖由树拓扑决定,于是 [SEP]/[FAIL] 这类标记会被训练成"取决于该动作是否可展开、是否与之前兄弟不同",模型因此学会在每个决策点生成多个互异动作并打上正确的可行性标签。

4. 连贯树推理与测试时扩展:在一棵树里探 k 个分支,而不是采 k 条独立轨迹

推理时 TSLM 重建树结构:先生成带多个候选动作的下一步,解析结构 token 找出可行分支([SEP]),把不同分支 fork 成独立序列"拼接"(stitching)进各自上下文做并行展开,再递归扩展每个可行状态直到 [GOAL] 或穷尽。默认用广度优先(BFS)遍历。这带来一种全新的测试时扩展范式:传统方法把模型当固定采样器,增大 \(k\) 就是从同一分布多采 \(k\)独立冗余的轨迹;TSLM 则在一次树构造里取前 \(k\) 个终端状态来验证,候选来自同一棵连贯的树而非 \(k\) 次独立采样。因为共享前缀只算一次、且推理时能忽略无关子树,TSLM 既比 ToT(重复采样、无计算共享)快很多,也比 Procedure Cloning(必须处理整条线性化搜索轨迹)快。更进一步,显式树结构让搜索策略可控——可以选 BFS(重解的最优性)或 DFS(重模型置信度),甚至自适应地更激进展开高奖励分支,这是独立采样根本做不到的。

损失函数 / 训练策略

训练就是对整条序列化序列施加标准语言建模交叉熵损失(公式见上),损失同时覆盖推理内容 token 和结构 token。基座模型默认用 Llama-3-8B,每个任务用不到 10K 实例做后训练(刻意控制数据量以比较架构差异而非规模效应)。推理时 fork 出的分支选取前 \(k=5\) 个生成动作并用精确匹配去重。

实验关键数据

主实验

TSLM 用 pass@1(单次尝试)评测,而 ToT 用 pass@100(100 个终端节点任一正确即算成功,相当于采样脚手架的收敛上界),即便如此 TSLM 仍多数任务持平或超越 ToT。

任务 SC (pass@1) PC (pass@1) GRPO (pass@1) TSLM (pass@1) ToT (pass@100)
Game of 24 17.0% 47.0% 15.0% 100% 32.0%
Gridworld (i.d, 10×10) 78.2% 99.7% 24.0% 100% 95.0%
Gridworld (o.o.d, 20×20) 33.0% 81.1% 6.0% 91.5% 42.7%
ProntoQA 99.7% 97.5% 99.8% 100% 100%
GSM8K 55.8% 55.9% 60.8% 61.6% 62.3%

最醒目的是 Gridworld 外推:ToT 在域内 10×10 拿到 95.0%,但放大到 20×20 灾难性崩到 42.7%;TSLM 则从 100% 只掉到 91.5%,说明内化的搜索过程比外部搜索算法泛化得更好。GSM8K 上 ToT 略高(62.3% vs 61.6%),但一旦需要系统性探索(Game of 24),ToT 就退化到与纯顺序方法同级(17%)。

消融 / 分析实验

配置 / 分析 关键指标 说明
TSLM, k=1 (GSM8K) 61.3% 单候选就几乎追平 ToT 收敛值 62.3%
TSLM, 饱和 (GSM8K) 67.2% 扩展候选后显著超过 ToT 上界
BFS vs DFS (top-1) 61.3% / 63.1% DFS 跟随高置信动作,top-1 更高;BFS 覆盖更广
训练 k=5 → 推理 k=10 67.2% → 71.1% 探索能力外推到更大分支因子
不可解检测 (Game of 24, 100 例) 97/100 TSLM 正确终止;SC/PC/GRPO 几乎全部幻觉出无效解

关键发现

  • 内化搜索 > 外部脚手架:TSLM 用 pass@1 打 ToT 的 pass@100 仍占优,且在分布外(更大网格、更大分支因子)保持鲁棒,证明把搜索学进参数比外挂搜索算法更能泛化。
  • 涌现的不可解判定:在 100 个无解的 Game of 24 实例上,TSLM 靠"探尽全树后无路可走"正确拒答 97 个;顺序基线因被训练成"总要给个答案"而几乎全部幻觉。这把"系统性探索"和"少幻觉"直接挂上了钩。
  • 新的扩展范式更省:TSLM 在一棵树里探 \(k\) 个分支共享前缀,墙钟时间远低于 ToT 的独立重复采样;扩展即"更聪明地探索"而非"更多地采样"。

亮点与洞察

  • 把搜索树当语言来建模:核心 trick 是用 5 个特殊 token + DFS 序列化,把"树"无损压进"线性序列",于是不用改架构、不用 RL、不用外部编排,纯监督学习就能教会模型分支探索——这个"表示换框架"的思路非常巧妙且可迁移到任何带结构的生成任务。
  • 上下文解耦是隐藏的关键:序列虽线性,但条件依赖由树拓扑(祖先 + 左兄弟)决定而非整段前缀,这让模型推理时能"无视无关子树",既是效率来源也是泛化来源。值得迁移到其他"长序列里只有局部依赖"的场景。
  • 训练失败路径反而有用:让模型见过 [FAIL] 死路,使它学会判定"此题无解",把幻觉问题转成了搜索完整性问题——这是顺序 CoT 训练拿不到的信号。

局限与展望

  • 作者承认:开放域自举用的监督来自 ToT 采样 + RAP 奖励,并非训练 TSLM 的最优数据源;可以换更强的推理模型(如 ReJump)或自训练(如 SoS)来造更好的树,本文刻意只聚焦 TSLM 概念本身。
  • 任务范围偏窄:主战场是 Game of 24、Gridworld、ProntoQA、GSM8K 这类有明确转移/动作空间或可枚举的任务;对真正开放、动作空间无界的自然语言任务,"如何定义并合成搜索树"仍是开放问题。
  • 推理仍需多次 fork:虽然树构造连贯,但把分支 fork 成独立序列展开仍涉及多次调用,序列化带来的额外结构 token 也会拉长序列;在超大分支因子 / 超深树上的开销与上限有待进一步验证。
  • 公平比较的口径:ToT@100 vs TSLM@1 是作者声称的"保守估计",但不同任务难度、不同候选预算下的比较仍需读者带着 caveat 看待。

相关工作与启发

  • vs Tree-of-Thought (ToT):ToT 在每步独立采样多候选、靠外部 beam search 编排,样本间不协调、共享前缀重复计算、随复杂度指数爆炸且外推差;TSLM 在一次前向里连贯生成整棵树、共享前缀、可忽略无关子树,外推显著更稳、推理更快。
  • vs o1 / DeepSeek-R1 等 RL 推理模型:它们靠拉长顺序推理链 + RL(GRPO 等)激励推理,但本质仍是顺序生成器,无法在生成过程里显式解耦并行分支;TSLM 直接用监督学习内化分支结构,挑战了"鲁棒推理必须靠 RL 或推理时搜索"的默认假设。
  • vs Sequence Cloning / Procedure Cloning:SC 只克隆单条金标 CoT,PC 把整条搜索轨迹线性克隆但仍是一条序列、推理时必须处理整段;TSLM 显式保留树拓扑,既学失败路径又能跳过无关子树,性能与效率双赢。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ "用 token 序列化把搜索树变成可监督学习的语言建模"是一个干净且少见的框架级创新。
  • 实验充分度: ⭐⭐⭐⭐ 覆盖结构化 + 开放域、外推、测试时扩展、不可解检测多维度,但任务偏受控、开放域案例有限。
  • 写作质量: ⭐⭐⭐⭐ 动机与形式化清晰,框架图直观;个别表述(如结论里 76.5% vs 26% 与正文数字口径)需对照原文。
  • 价值: ⭐⭐⭐⭐⭐ 给"推理时扩展"提供了不靠 RL/外部搜索的监督学习替代路径,思路可迁移性强。