SWAP：结构感知规划与准确世界模型

引言

SWAP（Structure-aware Planning with an Accurate World Model）[¹是由佐治亚理工学院和 Cisco Research 联合提出的推理框架，发表于 ACL 2025。该框架旨在解决大语言模型（LLM）在复杂多步决策任务中的推理能力不足问题。

与传统的 Chain-of-Thought（CoT）方法不同，SWAP 创新性地将结构化知识表示融入推理过程，通过蕴含图（Entailment Graph）编码前提条件与结论之间的依赖关系，实现中间步骤的符号化验证。这一方法标志着从纯自然语言推理向世界模型规划与推理融合的重要转变。

核心组件

SWAP 框架由四个核心组件构成，形成一个完整的结构感知规划系统：

蕴含图（Entailment Graph）

蕴含图是 SWAP 的核心数据结构，记作 $\mathcal{G}=(\mathcal{V},\mathcal{E})$，其中：

• 节点集 $\mathcal{V}$：每个节点 $v\in \mathcal{V}$ 代表一个陈述（Statement），如证据、假设、引理或规则
• 边集 $\mathcal{E}$：每条（有向）超边 $e=({\mathcal{V}}_{src},{\mathcal{V}}_{tgt})$ 表示从源节点集（前提）到目标节点集（结论）的蕴含关系

例如，语句”所有经常喝咖啡的人都对咖啡因依赖”和”Rina 是一个经常喝咖啡的学生”共同蕴含结论”Rina 对咖啡因依赖”。

蕴含图显式地捕获了前提如何导出中间结论，从而促进对claims的验证。与纯自然语言 CoT 相比，这种结构化表示提供了两个关键优势：

1. 策略模型可以利用结构信息做出更明智的决策
2. 世界模型可以预测更准确的后继状态

策略模型（Policy Model）

策略模型 ${\mathcal{P}}_{\pi }$ 负责提出候选扩展。给定当前状态 $(s_t,{\mathcal{G}}_t)$、目标 $G$ 和计划 $H$，策略模型从分布 ${\mathcal{P}}_{\pi }(a_t|G,H,s_t,{\mathcal{G}}_t)$ 中采样，生成下一步的行动建议。

策略模型的提议分布为：

$${\mathcal{P}}_{\pi }\left(a_t|G,H,s_t,{\mathcal{G}}_t\right)$$

该模型在训练阶段学习正样本轨迹的分布，并在推理阶段通过多样性建模（DM）机制增强生成多样性。

世界模型（World Model）

世界模型 ${\mathcal{P}}_{wm}$ 是 SWAP 的核心创新之一，负责预测结构更新。给定当前状态 $(s_{t-1},{\mathcal{G}}_{t-1})$ 和行动 $a_{t-1}$，世界模型预测后继状态：

$${\mathcal{P}}_{wm}\left(s_t,{\mathcal{G}}_t|s_{t-1},{\mathcal{G}}_{t-1},a_{t-1}\right)$$

具体流程：

1. 从世界模型分布中采样 $N_s$ 次，生成多个候选后继状态
2. 对候选状态进行符号验证（Symbolic Verification），过滤无效状态
3. 由判别器选择最优的后继状态

判别器（Discriminator）

判别器 ${\mathcal{P}}_d$ 负责重排多个候选的合理性。它采用对比排序（Contrastive Ranking）机制，通过直接比较候选选项来评估其质量，而非简单地分配单一数值分数。

判别器的评分函数整合了：

• 多个候选的对比信息
• 元知识 ${\mathcal{K}}_{meta}$（如常见陷阱和错误类型）
• 对每个候选的解释 $E$

$$E,a_t^{best}\sim {\mathcal{P}}_d\left(E,a_t^{best}|{\mathcal{K}}_{meta},G,H,s_t,{\mathcal{G}}_t,\{a_t^j,s_{t+1}^j,{\mathcal{G}}_{t+1}^j{\}}_{j=1}^K\right)$$

关键机制

多样性建模（Diversity-based Modeling, DM）

DM 旨在解决规划推理中的生成多样性问题。核心思想是从剩余概率质量中采样，确保每个新选项与之前的选项有所不同，从而缓解自偏差并促进探索。

对于第 $n$ 次生成，第 $l$ 个 token 的概率为：

$${\mathcal{P}}_{\pi }^{sem}\left(a_{t,l}^n|a_t^{\mathrm{1..}n-1},a_{t,\mathrm{1..}l-1}^n\right)=\frac{1}{n-1}\sum _{j=1}^{n-1}{\mathcal{P}}_{\pi }^{sem}\left(a_{t,l}^n|a_t^j,a_{t,\mathrm{1..}l-1}^n\right)$$

最终分布通过以下公式计算：

$${\mathcal{P}}_{\pi }\left(a_{t,l}^n|G,H,s_t,{\mathcal{G}}_t,a_t^{\mathrm{1..}n-1},a_{t,\mathrm{1..}l-1}^n\right)=\text{Norm}\left({\mathcal{P}}_{\pi }^{ori}\left(a_{t,l}^n|G,H,s_t,{\mathcal{G}}_t,a_{t,\mathrm{1..}l-1}^n\right)-{\gamma }_l{\mathcal{P}}_{\pi }^{sem}\left(a_{t,l}^n|a_t^{\mathrm{1..}n-1},a_{t,\mathrm{1..}l-1}^n\right)\right)$$

其中：

• ${\gamma }_l={\gamma }_0\cdot {\alpha }^l$ 为衰减因子（$\alpha \le 1$），在生成早期强调多样性
• $\text{Norm}(\mathcal{P})=\frac{\max (\mathcal{P},0)}{1^{\top }\max (\mathcal{P},0)}$ 确保输出为有效概率分布

对比排序（Contrastive Ranking, CR）

CR 旨在解决规划推理中的判别准确性问题。传统的过程奖励模型（Process Reward Model）将每个候选简化为单一数值分数，过于简化了复杂决策的方方面面。

CR 的自动标注流程：

1. 从正样本轨迹中随机选择中间步骤 $t$
2. 生成 $K$ 个替代推理轨迹
3. 通过符号验证过滤无效轨迹
4. 通过语义等价检查和树搜索结果验证，识别负样本轨迹中的首个错误步骤

判别器的训练使用正确和错误选项的配对数据，结合元知识和从 GPT-4o 引导的解释进行微调。

数学框架

MDP 形式化

SWAP 将规划任务形式化为马尔可夫决策过程（MDP）$(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R})$：

要素	含义
状态 $s_t\in \mathcal{S}$	当前图结构 ${\mathcal{G}}_t$，捕获所有已知或推断的信息及蕴含关系
行动 $a_t\in \mathcal{A}$	扩展步骤，从当前状态推导或推断新信息
转移概率 $\mathcal{P}(s_{t+1}	s_t, a_t)$
奖励 $\mathcal{R}(s_t,a_t)$	衡量行动质量，SWAP 用判别器直接比较不同行动

蕴含图更新规则

初始图 ${\mathcal{G}}_0$ 由世界模型根据问题 $Q$ 生成：

$${\mathcal{P}}_{wm}(G,s_0,{\mathcal{G}}_0|Q)$$

推理过程中，世界模型通过添加新节点和边来逐步扩展图：

$$(s_t^{\prime },{\mathcal{G}}_t^{\prime })\sim {\mathcal{P}}_{wm}(s_t,{\mathcal{G}}_t|s_{t-1},{\mathcal{G}}_{t-1},a_{t-1})$$

最终图 ${\mathcal{G}}_T$ 包含最终答案，验证过程确保图的正确性。

策略模型提议分布

计划提议（第 0 步）：

$${\mathcal{P}}_{\pi }(H|G,s_0,{\mathcal{G}}_0)$$

行动提议（第 $t$ 步）：

$${\mathcal{P}}_{\pi }(a_t|G,H,s_{t-1},{\mathcal{G}}_{t-1})$$

判别器评分函数

判别器采用对比学习的评分机制，对候选动作进行排序：

• 动作排序：使用模拟的即时后继状态 $(s_{t+1},{\mathcal{G}}_{t+1})$
• 计划排序：使用模拟的终端状态 $(s_T,{\mathcal{G}}_T)$

$${\mathcal{P}}_d(E,H^{best}|{\mathcal{K}}_{meta},G,s_0,{\mathcal{G}}_0,\{H^j,s_T^j,{\mathcal{G}}_T^j{\}}_{j=1}^K)$$

实验结果

SWAP 在多个推理密集型基准上进行了广泛评估：

数学推理基准

数据集	SWAP (LLaMA3-8B)	基线模型
MATH-500	显著提升	+15.2%
GSM8K	大幅改进	+12.8%

逻辑推理基准

数据集	SWAP 表现
FOLIO	一阶逻辑推理显著提升
LogiQA	逻辑推理能力明显增强

编码任务基准

数据集	性能提升
HumanEval	代码生成准确性提升
MBPP	编程问题解决能力增强

消融实验

消融实验验证了各组件的贡献：

• 多样性建模（DM）：对扩展搜索空间至关重要，使正确解的发现概率显著提高
• 对比排序（CR）：提升判别准确率，有效区分正确与错误候选
• 世界模型：准确的状态预测是策略模型有效性的关键

扩展性分析

实验表明，随着总 rollout 次数和宽度限制（breadth limit）的增加，SWAP 在各基准和基础模型上的准确率持续提升，证明了框架的可扩展性和鲁棒性。

总结

SWAP 框架代表了 LLM 推理能力增强的重要进展，其核心贡献包括：

1. 结构化推理表示：通过蕴含图编码前提与结论的依赖关系，实现符号化验证
2. 准确世界模型：增强策略模型的决策质量，通过多候选生成和判别重排
3. 多样性探索：DM 机制从剩余概率质量采样，缓解自偏差
4. 精确判别：CR 机制通过对比学习提升判别准确率

SWAP 的成功表明，将结构化知识表示与学习规划相结合，是突破纯自然语言 CoT 局限性的有效途径，为构建更可靠、可验证的推理系统奠定了基础。

Footnotes

1. Xiong, S., Payani, A., Yang, Y., & Fekri, F. (2025). Deliberate Reasoning in Language Models as Structure-Aware Planning with an Accurate World Model. ACL 2025. ↩