技术摘要：GFlowNets 中的信息几何前向策略训练

问题陈述

生成式流网络（GFlowNets）提供了一个框架，用于对离散和混合离散-连续对象进行摊销推理（amortised inference），其核心是学习一个随机构建策略（前向策略），该策略的终端分布与奖励函数 $R(x)$ 成正比。训练 GFlowNets 的一个核心挑战是“摊销采样”（amortised sampling）问题：策略从其生成的轨迹中学习，这意味着被当前策略分配低概率的区域即使拥有高奖励，也难以被采样。因此，在高奖励区域的训练信号是稀疏的。

虽然现有方法通过改变数据收集方式（例如 Adaptive Teachers）或结合多种采样器来解决覆盖率不足的问题，但本研究探讨了另一个问题：在保持规范的前向采样器和平衡目标（如 Trajectory Balance）不变的情况下，什么样的几何结构应该支配参数更新？标准的欧几里得梯度更新依赖于坐标，且无法考虑到概率分布空间的内在几何结构。本文提出，将前向策略视为统计流形上的诱导轨迹采样器，可以应用信息几何来提高训练的稳定性、收敛性和模态发现能力。

方法论

信息几何公式化

作者将前向策略训练公式化在诱导轨迹分布族 $\{q_\theta\}$ 的统计流形上。他们确定了 Fisher-Rao 度量 是该流形的内在几何。Fisher 信息矩阵 $F(\theta)$ 定义为轨迹得分外积的期望：
$$F(\theta) := \mathbb{E}_{\tau \sim q_\theta} [\nabla_\theta \log q_\theta(\tau) \nabla_\theta \log q_\theta(\tau)^\top]$$
自然梯度定义为 $\tilde{\nabla}_\theta J(\theta) = F(\theta)^\dagger \nabla_\theta J(\theta)$（其中 $\dagger$ 表示伪逆），它代表了单位局部轨迹分布变化下的最陡下降方向，而非单位欧几里得参数变化下的方向。该方向在平滑重参数化下具有不变性。

轨迹 Fisher 分解

一个关键的理论贡献是推导出了轨迹 Fisher 的精确分解。通过定义每步得分增量 $u_t(\theta; \tau) = \nabla_\theta \log \pi_\theta(a_t | s_t)$，作者表明轨迹 Fisher 可以分解为每步条件二阶矩之和：
$$F(\theta) = \sum_{t \ge 0} \mathbb{E}_{\tau \sim q_\theta} [u_t(\theta; \tau) u_t(\theta; \tau)^\top]$$
这种分解依赖于得分的鞅性质（条件中心化），从而消除了时间交叉项。

• 表格策略（Tabular Policies）： 在表格设置中，这产生了一个由状态占用度 $d_\theta(s)$ 加权的精确块对角 Fisher 矩阵。
• 共享参数策略（Shared Parameter Policies）： 在具有共享权重的神经网络设置中，由于跨状态和跨时间步的参数耦合，Fisher 变得稠密，尽管时间分解仍然是精确的。

计算范式

论文根据访问 Fisher 算子的方式，将自然梯度训练组织为三个计算范式：

1. 精确范式（Exact Regime）： 适用于精确边际化（例如通过动态规划）是可行的场景，允许计算精确的状态占用度和精确的 Fisher。
2. 采样范式（Sampled Regime）： 使用来自采样轨迹的蒙特卡洛估计量来近似 Fisher。论文为这些估计量提供了相合性（consistency）保证。
3. 结构可利用范式（Structure-Exploitable Regime）： 利用目标分解或构建局部性来构建代理（surrogates）。例如，在贝叶斯 DAG 学习中，虽然奖励可以分解，但无环约束耦合了 Fisher。作者建议使用图形模型工具（精确边际化、分离器方法、置信传播）来构建准确的 Fisher 期望代理。

近似保证

作者推导了阻尼自然梯度更新的相对误差界限（定理 2）。他们证明，如果代理 Fisher $\tilde{F}$ 与精确 Fisher $F$ 之间的相对误差是有界的，则产生的更新方向将接近真实的自然梯度。这为使用不完美的 Fisher 而不丢失几何益处提供了一种原则性的方法。

核心贡献

1. 公式化： 本文将 GFlowNet 前向策略训练公式化在诱导轨迹分布族的统计流形上，并确定了轨迹 Fisher 为相关的度量。
2. 分解： 推导出了轨迹 Fisher 到每步条件二阶矩的精确分解，阐明了表格策略中的占用度加权块结构与共享参数策略中的稠密耦合之间的区别。
3. 训练框架： 建立了三种计算范式（精确、采样、结构可利用），用于自然梯度训练，并提供了基于代理更新的相对误差保证。
4. 实验验证： 通过表格和神经基准测试验证了该框架，证明了 Fisher-Rao 预处理与欧几里得训练相比，在收敛性和分布匹配方面具有改进。

结果

作者在以下基准测试中评估了所提方法：

• Triangle-Count ERM G 模型： 在具有精确占用度的表格设置中，自然梯度更新比欧几里得训练更快地加速了收敛并降低了 KL 散度。
• Hypergrid： 在 2D 网格基准上，自然梯度比欧几里得下降和 ACE 基线更有效地降低了 Trajectory Balance (TB) 损失和全变分 (TV) 距离。在 $H=16$ 的情况下，结构化 Fisher 近似成功复现了目标的多层奖励结构。
• Deceptive Grid： 使用共享神经策略时，Fisher 预处理更新（通过共轭梯度或 K-FAC）相比于欧几里得 Adam 和 SGD 显著改善了模态发现，实现了与 Adaptive Teacher 相当的覆盖率，且无需辅助探索策略。
• 慢速随机游走（8-Gaussians & Rings）： 自然梯度变体比欧几里得、Wasserstein 和 ACE 基线恢复了更多的模态并实现了更高的 ELBO。值得注意的是，在 Rings 基准上，自然梯度恢复了 24–25 个模件，而基线则坍缩到 16 个。
• 贝叶斯 DAG 学习（Sachs 数据）： 在真实的蛋白质信号数据上，阻尼 Fisher（通过使用非对角采样 Fisher 的 PCG）相对于普通 Adam 改善了局部优化（Loss AUC），但实际上与受 KL 控制的 Adam 基本持平。结果表明，在这种特定的真实数据设置中，大部分收益可能来自于共享步长范数和 KL 控制，而非 Fisher 结构本身。

意义与主张

本文声称，将前向策略视为统计采样器，可以将目标结构转化为优化几何。主要意义在于提供了一条通往具备结构感知能力的、可处理的前向策略训练的路径。

作者对自己的主张保持谦逊，特别是在关于 Sachs 蛋白质信号数据方面。他们指出，虽然 Fisher-PCG 相比于普通 Adam 改善了局部优化，但并未显示出与受 KL 控制的 Adam 有可靠的区分，这表明“大部分收益来自于共享步长范数和 KL 控制”。这项工作被呈现为一种概念验证，即在全局约束图空间中使用采样、非对角神经 Fisher 的可行性。

论文强调，其贡献并非一种通用的基于 Fisher 的优化器，而是一种将自然梯度应用于 GFlowNet 平衡目标的特定应用，保留了由前向策略诱导的轨迹律来定义局部度量。它通过保持保持有效性的离散构建过程，与连续松弛方法（如 Concrete, Gumbel-Softmax）区分开来。

最终，这项工作表明，目标结构不仅可以用于奖励分解或信用分配，还可以用于构建采样器本身的优化几何，为替代启发式探索机制提供了一种原则性的方案。