[论文解读] Trajectory balance: Improved credit assignment in GFlowNets

Nikolay Malkin, Moksh Jain+3 more|arXiv (Cornell University)|Jan 31, 2022

Topic Modeling被引用 12

一句话总结

本文提出轨迹平衡（trajectory balance），一种用于GFlowNets的新训练目标，通过在完整轨迹上进行优化而非局部流动约束，改善了长动作序列中的信用分配。该文证明，轨迹平衡的全局最小化可确保从目标分布中精确采样，且在实验中验证了其在长时序和大动作空间设置下具有更快的收敛速度、更高的样本多样性以及更强的鲁棒性。

ABSTRACT

Generative flow networks (GFlowNets) are a method for learning a stochastic policy for generating compositional objects, such as graphs or strings, from a given unnormalized density by sequences of actions, where many possible action sequences may lead to the same object. We find previously proposed learning objectives for GFlowNets, flow matching and detailed balance, which are analogous to temporal difference learning, to be prone to inefficient credit propagation across long action sequences. We thus propose a new learning objective for GFlowNets, trajectory balance, as a more efficient alternative to previously used objectives. We prove that any global minimizer of the trajectory balance objective can define a policy that samples exactly from the target distribution. In experiments on four distinct domains, we empirically demonstrate the benefits of the trajectory balance objective for GFlowNet convergence, diversity of generated samples, and robustness to long action sequences and large action spaces.

研究动机与目标

• 为解决现有GFlowNet训练目标（如流匹配和细致平衡）在长动作序列中信用传播效率低下的问题，这些方法在时间维度上存在缓慢的信用分配问题。
• 开发一种新型学习目标，以实现在组合物体生成中复杂且长轨迹上的更高效、更精确的信用分配。
• 证明轨迹平衡的全局最小化可产生一个策略，该策略能从目标未归一化密度 R(x) 中精确采样。
• 通过实证验证轨迹平衡在收敛速度、样本多样性以及在长序列和大动作空间设置下的鲁棒性方面优于基线方法，涵盖四个不同领域。
• 首次对GFlowNets中的细致平衡目标进行实证评估，支持与轨迹平衡和流匹配的对比分析。

提出的方法

• 提出轨迹平衡作为新型训练目标，其作用对象为从初始状态到终止状态的完整动作序列（轨迹），而非局部边流或成对状态转移。
• 将轨迹平衡目标定义为对采样轨迹的总和，其中每条轨迹的贡献基于路径上流动与目标奖励 R(x) 之间的差异。
• 使用神经网络参数化边流 F(u→v)，并通过 F(u→v)/∑v′F(u→v′) 推导出策略，确保动作选择的随机性。
• 通过梯度下降优化轨迹平衡目标，梯度使用路径导数（pathwise derivatives）和重参数化技术计算，适用于可微模型。
• 证明轨迹平衡目标的任意全局最小值均可产生一个策略，该策略能从目标分布 p(x) ∝ R(x) 中精确采样。
• 在多个基准环境上，使用标准GFlowNet训练流程，实现并比较轨迹平衡与流匹配及细致平衡目标。

实验结果

研究问题

• RQ1 与基于局部流的客观（如流匹配和细致平衡）相比，基于轨迹的客观是否能提升GFlowNets中信用分配的效率？
• RQ2 轨迹平衡是否能带来在长时序生成任务中更快的收敛速度和更高的样本多样性？
• RQ3 在组合物体生成中，轨迹平衡对大动作空间和高维状态空间是否具有鲁棒性？
• RQ4 在样本质量与训练稳定性方面，轨迹平衡与流匹配和细致平衡相比，实证表现如何？
• RQ5 轨迹平衡是否能实现比以往更长序列的GFlowNet训练？

主要发现

• 在所有评估领域中，包括长序列和大动作空间设置，轨迹平衡的训练收敛速度均快于流匹配和细致平衡。
• 与基线目标相比，该方法在熵和覆盖率等指标上显著提升了样本多样性。
• 轨迹平衡成功实现了对长达100步的动作序列的GFlowNet训练，而此前的训练目标在该设置下无法收敛。
• 实证结果表明，轨迹平衡降低了生成样本与目标分布 R(x) 之间的差异，表明其与期望奖励函数的对齐性更优。
• 首次对细致平衡目标进行实证验证，确认其理论合理性，但发现其收敛速度更慢且样本方差更高，相较轨迹平衡。
• 全局最小化轨迹平衡目标可严格保证从目标分布 p(x) ∝ R(x) 中精确采样，从而确立了其理论正确性。