Information-Geometric Forward Policy Training in GFlowNets¶
会议: NeurIPS 2026
arXiv: 2608.03967
代码: https://github.com/rodsveiga/infogeometric_gflows
领域: 强化学习
关键词: GFlowNets, 信息几何, 自然梯度, 费歇尔-拉奥度量, 轨迹平衡
一句话总结¶
将 GFlowNet 前向策略视为由其诱导的轨迹采样器统计流形，推导了轨迹费歇尔信息矩阵的单步鞅差精确时间解耦，并据此构建了精确、采样与结构化近似三类自然梯度优化方案，在不依赖辅助探索策略的前提下显著提升了多模态探索与分布拟合效率。
研究背景与动机¶
生成流网络（GFlowNets）作为一种用于离散与组合对象的摊销推断框架，核心在于通过学习局部前向构建策略，使得终止状态的边际概率与未归一化的奖励函数成正比。然而，GFlowNets 依赖于当前策略自主生成的采样轨迹进行学习，当面对多模态或稀疏高奖励空间时，当前策略赋予极低概率的区域极少被访问，导致关键奖励信号缺失，极易陷入局部最优或发生严重模式塌陷。
现有工作应对该探索瓶颈的主流方案多属于“外部辅助型”，例如引入自适应教师（Adaptive Teacher）训练额外的行为探索网络、设计同胞增强机制（SA-GFN）或在残差奖励上级联训练（Boosted GFlowNets）。这些方案改变了数据收集分布或增加了额外模型，却忽视了一个深层次的内在问题：在基准前向采样器与轨迹平衡（Trajectory Balance, TB）目标保持不变的前提下，参数更新本身的优化几何应当遵循怎样的准则？标准的欧氏梯度更新强依赖于特定的参数坐标表象，在面临状态访问频率极不均衡或局部决策饱和时无法反映概率测度的本质变化。
信息几何通过费歇尔-拉奥（Fisher-Rao）黎曼度量刻画概率分布空间中的局部 KL 散度位移。然而，由于 GFlowNet 构建轨迹具有序列性、状态相关动作掩码以及神经网络参数全局共享等特性，直接引入黎曼优化面临严重的计算与结构耦合瓶颈。核心 idea：将 GFlowNet 前向策略视为轨迹空间上的统计模型，利用单步分数的鞅差中心化消除跨时序交叉项，将轨迹费歇尔信息解耦为每步条件二阶矩之和，并根据目标与模型因子化结构建立精确、采样与结构化近似三类黎曼预条件优化体系。
方法详解¶
整体框架¶
该方法的核心思想是在诱导的轨迹测度空间建立内在的一阶自然梯度更新规则。GFlowNet 前向策略 \(\pi_\theta\) 在状态有向无环图（DAG）上依序采样转移动作 \(a_t \sim \pi_\theta(\cdot \mid s_t)\)，构成完整构建轨迹 \(\tau = (s_0, a_0, \dots, s_T)\)，诱导出的轨迹概率为 \(q_\theta(\tau) = \prod_{t=0}^{T-1} \pi_\theta(a_t \mid s_t)\)。信息几何通过轨迹费歇尔信息矩阵（Fisher Information Matrix, FIM）定义可识别切空间上的黎曼度量。
针对欧氏优化忽略分布流形曲率的问题，该框架通过时间解耦消除跨时刻交叉项，再结合具体问题结构选择闭式动态规划（DP）、蒙特卡洛轨迹估计或图模型因子化近似构造费歇尔预条件算子，最终通过预条件共轭梯度（PCG）或阻尼线性求解器对梯度进行自然梯度预条件投影，端到端地更新前向策略参数。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["前向策略轨迹采样<br/>q_θ(τ) = ∏ π_θ(a_t | s_t)"] --> B["鞅差条件中心化与时序解耦<br/>E[u_t | F_t] = 0 消除跨时刻交叉项"]
B --> C["三类费歇尔计算与近似范式<br/>精确 DP / 蒙特卡洛采样 / 目标因子化代理"]
C --> D["相对误差与阻尼扰动界保障<br/>验证 ||F_λ^(-1/2) (F_tilde - F) F_λ^(-1/2)|| < 1"]
D --> E["阻尼自然梯度预条件求解<br/>θ+ = θ - η (F_hat + λ I)^(-1) h_hat"]关键设计¶
1. 鞅差条件中心化与时序解耦：消除轨迹费歇尔的跨时序交叉耦合
针对序列决策过程中多步动作联合分布导致轨迹费歇尔矩阵极其稠密、无法直接计算的痛点，该设计利用马尔可夫决策过程的时序因果滤过结构实现精准降维。令 \(u_t(\theta; \tau) \coloneqq \nabla_\theta \log \pi_\theta(a_t \mid s_t)\) 为第 \(t\) 步的局部评分增量，\(\mathcal{F}_t \coloneqq \sigma(s_0, a_0, \dots, s_t)\) 为做出动作 \(a_t\) 前的历史 \(\sigma\)-代数。由于动作依策略分布采样且局部概率和为 1，局部评分在给定历史条件下的期望恒为零： $\(\mathbb{E}[u_t(\theta; \tau) \mid \mathcal{F}_t] = \sum_{a \in \mathcal{A}(s_t)} \nabla_\theta \pi_\theta(a \mid s_t) = \nabla_\theta \sum_{a \in \mathcal{A}(s_t)} \pi_\theta(a \mid s_t) = 0\)$ 由此，\(u_t\) 构成鞅差序列。当计算轨迹费歇尔矩阵 \(F(\theta) \coloneqq \mathbb{E}_{\tau \sim q_\theta}\left[\nabla_\theta \log q_\theta(\tau) \nabla_\theta \log q_\theta(\tau)^\top\right]\) 时，对于任意不同时刻 \(t < t'\)，利用条件期望的重叠塔式性质可知 \(\mathbb{E}[u_t u_{t'}^\top] = \mathbb{E}[u_t \mathbb{E}[u_{t'}^\top \mid \mathcal{F}_{t'}]] = 0\)。所有跨时间的非对角交互项全部精确归零，从而将全局轨迹费歇尔严格化简为逐时刻条件二阶矩的代数和： $\(F(\theta) = \sum_{t \ge 0} \mathbb{E}_{\tau \sim q_\theta}\left[u_t(\theta; \tau) u_t(\theta; \tau)^\top\right]\)$ 该性质不仅严格证明了基于单步梯度外积的累加具备精确流形度量意义，而且为后续不同参数化结构下的轻量化求解奠定了理论基石。
2. 三类费歇尔计算与近似范式：覆盖从精确制导到高维神经网络的统一计算谱系
在消除时序交叉项后，参数空间的耦合程度完全取决于前向策略的参数化表象。本文系统地将计算途径划分为三种范式： - 精确范式（Exact Regime）：在表格化 Softmax 策略下，各状态拥有独立的对数几率向量 \(\theta_s\)。单步局部协方差表现为分类分布协方差矩阵 \(C(\pi_s) = \operatorname{Diag}(\pi_s) - \pi_s \pi_s^\top\)。此时全局费歇尔矩阵严格退化为状态访问占用率 \(d_\theta(s) \coloneqq \mathbb{E}_{\tau}\left[\sum_{t} \mathbb{I}\{s_t = s\}\right]\) 加权的分块对角矩阵 \(F(\theta) = \bigoplus_{s \in \mathcal{S}} d_\theta(s) C(\pi_s)\)，可通过有向无环图上的前向动态规划（DP）精确闭式求出。 - 采样范式（Sampled Regime）：当策略采用共享参数神经网络时，状态间的解耦被破坏，对数几率雅可比 \(J_\theta(s) = \nabla_\theta \ell_\theta(s)\) 引入稠密曲率： $\(F(\theta) = \mathbb{E}_{\tau}\left[\sum_{t \ge 0} J_\theta(s_t)^\top C(\pi_\theta(\cdot \mid s_t)) J_\theta(s_t)\right]\)$ 该范式基于当前批次采样的 \(N\) 条轨迹，通过重要性校正权值 \(\omega_i = q_\theta(\tau^{(i)})/\mu(\tau^{(i)})\) 构造无偏蒙特卡洛估计算子 \(\widehat{F}_N(\theta) = \frac{1}{N} \sum_{i=1}^N \sum_{t} \omega_i u_t(\tau^{(i)}) u_t(\tau^{(i)})^\top\)，并借助无矩阵向量积与共轭梯度（CG）实现隐式高维求解。 - 结构化近似范式（Structure-Exploitable Regime）：当环境奖励具备局部因子化分解 \(R(x) = \prod_\alpha \psi_\alpha(x_\alpha)\)（如图模型结构学习或低树宽马尔可夫网）时，利用目标图模型的条件独立性、分离子集或置信传播构建稀疏块代理矩阵 \(\widetilde{F}\)，以极低的计算代价逼近主导曲率。
3. 相对误差与阻尼扰动界保障：量化代理算子对自然下降方向的几何保真度
在实际任务中使用采样费歇尔或结构化代理矩阵 \(\widetilde{F}\) 替代真实费歇尔 \(F\) 时，传统的加性误差界不足以保证预条件梯度的稳定。本文给出了在带阻尼算子 \(F_\lambda = F + \lambda I\) 与 \(\widetilde{F}_\lambda = \widetilde{F} + \lambda I\) 下严格的相对误差控制理论。定义相对算子误差为： $\(\rho_\lambda \coloneqq \left\| F_\lambda^{-1/2} (\widetilde{F} - F) F_\lambda^{-1/2} \right\|_{\mathrm{op}}\)$ 当 \(\rho_\lambda < 1\) 且梯度估计为 \(\widetilde{h}\) 时，所得阻尼自然梯度更新方向 \(\widetilde{v}_\lambda = \widetilde{F}_\lambda^{-1} \widetilde{h}\) 相对于理想方向 \(v_\lambda = F_\lambda^{-1} h\) 在诱导模长下的相对偏离满足严格上界： $\(\|\widetilde{v}_\lambda - v_\lambda\|_{F_\lambda} \le \frac{\rho_\lambda}{1 - \rho_\lambda} \|v_\lambda\|_{F_\lambda} + \frac{1}{1 - \rho_\lambda} \|\widetilde{h} - h\|_{F_\lambda^{-1}}\)$ 特别地，当使用精确梯度时，代理方向与梯度的内积被严格界定在 \([(1+\rho_\lambda)^{-1}, (1-\rho_\lambda)^{-1}]\) 之间，从而在理论上彻底杜绝了因曲率代理估计不准而产生反向或发散更新的隐患，给出了选择阻尼因子 \(\lambda\) 的可计算充分条件。
损失函数 / 训练策略¶
前向策略参数 \(\theta\) 基于轨迹平衡损失函数 \(\mathcal{J}_{\mathrm{TB}}\) 的评分形式梯度进行优化。设目标轨迹平衡残差为： $\(\delta_{\theta,\phi,\xi}(\tau) = \log Z_\xi + \sum_{t=0}^{T-1} \log \pi_\theta(a_t \mid s_t) - \log R(x) - \sum_{t=0}^{T-1} \log P_\phi(a_t \mid s_{t+1})\)$ 其关于前向策略参数的梯度 \(h(\theta) = \mathbb{E}_{\tau \sim \mu}[\ell'(\delta(\tau)) \nabla_\theta \log q_\theta(\tau)]\) 严格位于费歇尔矩阵的值域可识别切空间 \(\mathcal{U}_\theta\) 内（\(h \in \operatorname{range}(F(\theta))\)），满足 \(h^\top F^\dagger h \le \mathbb{E}[(\ell')^2]\)。
参数更新结合了类似 AdaFisher 的自适应动量架构： 1. 维护前向梯度的一阶动量偏差校正估计 \(\widehat{m}_k\)； 2. 保持后向策略参数 \(\phi\) 和流标量 \(\log Z_\xi\) 采用标准的欧氏 SGD/Adam 更新； 3. 对前向策略的一阶动量应用阻尼费歇尔逆求解更新量 \((\widehat{F}(\theta) + \lambda I) \Delta \theta = \widehat{m}_k\)。在神经网络参数下，利用无矩阵向量积执行 20 次 PCG 迭代完成求解，更新参数 \(\theta \leftarrow \theta - \alpha_\theta \Delta \theta\)。
实验关键数据¶
主实验¶
论文在随机游走基准（8-Gaussians、Rings）和欺骗性网格世界（Deceptive Hypergrid，\(H \in \{128, 256\}\)）上，全面对比了欧氏基线（Adam、SGD、带探索噪声版本）、代表性探索方法（ACE、Adaptive Teacher）以及自然梯度算法（Natural Exact、Natural Sample、Natural Approx/CG/K-FAC）。
在多模态随机游走任务中，8-Gaussians 包含 29 个模式，Rings 包含 36 个同心环狭窄模式（15,000 步训练）：
| 数据集 | 指标 | 本文方法 (Natural Exact / CG) | 传统基线 (Euclidean Adam / TB) | 前沿基线 (ACE / Adaptive Teacher) | 相对增益 / 差异分析 |
|---|---|---|---|---|---|
| 8-Gaussians (29 modes) | TV 距离 \(\downarrow\) | 0.1706 (Exact) | 0.1902 | 0.2078 (ACE) | 相对欧氏降低 10.3%，较 ACE 降低 17.9% |
| 8-Gaussians (29 modes) | 奖励间隙 (Gap) \(\downarrow\) | 0.0259 (Exact) | 0.3501 | 0.3634 (ACE) | 奖励间隙缩小至欧氏基线的 1/13 |
| 8-Gaussians (29 modes) | 发现模式数 \(\uparrow\) | 29 / 29 | 29 / 29 | 29 / 29 | 全部覆盖 |
| Rings (36 modes) | 发现模式数 \(\uparrow\) | 25 / 36 (Exact) | 16 / 36 | 16 / 36 (ACE) | 成功突破同心环低奖励屏障，多覆盖 9 个模式 |
| Rings (36 modes) | ELBO \(\uparrow\) | -0.3788 (Exact) | -0.4835 | -0.4616 (ACE) | 显著提升下界质量 |
| Deceptive Grid (\(H=128\)) | 发现模式数 \(\uparrow\) | 666 / 676 (CG) | 676 / 676 | 676 / 676 (Adaptive Teacher) | 纯几何预条件接近双模型教师探索水平 |
| Deceptive Grid (\(H=256\), 15k步) | 发现模式数 \(\uparrow\) | 2405 / 2601 (CG) | 2031 / 2601 | 2572 / 2601 (Adaptive Teacher) | 在大网格半程预算下显著快于普通 Adam |
消融实验与结构学习评估¶
在真实生物信号网络 Sachs 蛋白质因果 DAG 结构学习任务（\(d=11\) 节点，\(N=853\) 观测样本）中，采用共享 Linear Transformer 参数化策略，在 3 个成对随机种子上对比了局部优化与结构推断指标（500 步持续训练）：
| 配置 / 方法 | 损失曲线 AUC \(\downarrow\) | P90 残差 AUC \(\downarrow\) | 有向图 AUPRC \(\uparrow\) | 期望结构汉明距离 (E-SHD) \(\downarrow\) | 期望图边数 \(\mathbb{E}\|G\|\) |
|---|---|---|---|---|---|
| Euclidean (Adam) | 0.751 \(\pm\) 0.005 | 0.881 \(\pm\) 0.009 | 0.359 \(\pm\) 0.021 | 28.83 \(\pm\) 0.02 | 26.56 \(\pm\) 0.22 |
| KL-controlled Adam | 0.734 \(\pm\) 0.006 | 0.876 \(\pm\) 0.009 | 0.361 \(\pm\) 0.018 | 28.98 \(\pm\) 0.11 | 26.66 \(\pm\) 0.14 |
| Damped Fisher (PCG，本文) | 0.734 \(\pm\) 0.004 | 0.870 \(\pm\) 0.007 | 0.349 \(\pm\) 0.047 | 28.80 \(\pm\) 0.19 | 26.28 \(\pm\) 0.45 |
在 2D 超网格（Hypergrid \(H=16\)）与 8-Gaussians 任务上，消融不同费歇尔计算范式的计算与性能权衡： - Exact vs. Sample vs. Approx：在 8-Gaussians 上，采样近似（Natural sample, TV=0.1699）和因子化边际代理（Natural approx, TV=0.1708）几乎完美复现了精确费歇尔（Natural exact, TV=0.1706）的收敛质量； - 计算开销对比：因子化代理方法完全避免了全局反向 DP 扫表与显式高维逆矩阵求解，步耗时较精确 DP 降低一个数量级以上，但其收敛速度和分布对齐指标（JSD=0.0269）与精确版（JSD=0.0268）处于同一水平。
关键发现¶
- 在几何流形结构复杂的连续环形模式（Rings）中，欧氏梯度、Wasserstein 优化以及基于发散排斥的 ACE 全部陷入 16 个局部模式的死锁，而自然梯度方法凭借费歇尔流形曲率预条件，能够顺着环形流形切向稳定探索，成功解锁 24~25 个模式。
- 在高维共享参数的欺骗性网格实验中，自然梯度更新（Natural CG 与 K-FAC）展现出了与引入外生辅助策略网络（Adaptive Teacher）相当的模式覆盖加速能力，证实了调整参数流形优化几何本身即可直接赋能探索。
- 在 Sachs 因果图发现任务中，Damped Fisher 相比原始 Adam 改善了局部优化（Loss AUC 从 0.751 降至 0.734），但与 KL 约束 Adam 相当；这说明在复杂的参数共享与无环约束下，曲率预条件的主要收益源于防止策略在概率空间剧烈漂移的有效步长缩放。
亮点与洞察¶
- 将前向生成策略内化为统计采样流形：不同于以往仅仅将 GFlowNets 视作强化学习策略或离散能量匹配模型，该工作首次从信息几何的 Fisher-Rao 流形切入，建立了一套内生于分布空间而非参数坐标轴的原则性更新体系。
- 鞅差正交性瓦解时序交叉耦合：证明局部评分 \(u_t\) 是严格的鞅差序列，利用概率公理在理论上消除了时序间的全部非对角交叉协方差，将看似极其复杂的序列长程曲率优雅化简为单步二阶矩累加。
- 目标结构直接反哺优化度量：提出将无向图因子化（如马尔可夫网的势函数分解）作为优化几何中构建稀疏阻尼代理算子的指导原则，开辟了“物理/目标先验 \(\to\) 算法优化度量曲率”的通用跨界范式。
局限与展望¶
- 高维共享神经网络的逆矩阵开销：对于超大规模深度网络，即便采用 CG 或 K-FAC，逐批次采样轨迹的 Fisher-Vector 积计算仍然带来不容忽视的训练壁垒，亟需发展适应离散序列生成的低秩流形更新算法。
- 非规范构建（Non-canonical）中的终点方差遗漏：当前框架理论基于轨迹级费歇尔，在存在多条轨迹合成同一终止状态的非规范 DAG 中，轨迹度量与真实的终止状态度量相差一个终点条件协方差矩阵，在高度重联图上的理论偏差尚待弥合。
- 缺乏端到端自适应阻尼调节机制：当前的阻尼参数 \(\lambda\) 仍依赖超参搜索；若能结合局部曲率谱估计实现动态置信区间阻尼，有望进一步提升在稀疏回报极限场景下的鲁棒性。
相关工作与启发¶
- vs Adaptive Teacher / ACE: 传统方法通过引入额外的辅助网络或修改经验回放探索策略来缓解高回报模式塌陷；本文完全在单模型原生框架内，通过重构前向策略的流形几何曲率实现了同等级别的多模态覆盖。
- vs Natural Policy Gradient (NPG / TRPO): 传统 RL 中的自然策略梯度旨在最大化累积回报期望，目标梯度并不自动位于费歇尔值域内；而 GFlowNet 轨迹平衡的评分形式梯度被严格证明恒在可识别切空间内，与费歇尔度量具备天然的几何相容性。
- vs Fisher-Flow / 连续松弛: 连续松弛（如 Gumbel-Softmax）破坏了离散图构建的拓扑合法性约束；本文在完全保留离散有效 DAG 转移特性的同时，利用参数空间的 Fisher-Rao 流形实现了严谨的黎曼连续优化。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立 GFlowNet 前向策略的轨迹信息几何理论，鞅差时序解耦与相对误差界分析非常优美严密。
- 实验充分度: ⭐⭐⭐⭐☆ 涵盖典型代数图、多尺度连续网格、高维欺骗网格与真实蛋白质网络，评测指标全面，但大模型上的评估仍显不足。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导层层递进，概念区分透彻清晰，对理论成立条件与局限性阐述坦诚客观。
- 价值: ⭐⭐⭐⭐☆ 为离散生成流网络的优化器设计与内在探索机理研究奠定了坚实的信息几何基石。