[论文解读] Learning GFlowNets from partial episodes for improved convergence and stability

Kanika Madan, Jarrid Rector-Brooks+7 more|arXiv (Cornell University)|Sep 26, 2022

Reinforcement Learning in Robotics被引用 5

一句话总结

本文提出Subtrajectory Balance（SubTB），一种用于GFlowNets的新颖训练目标，可从不同长度的部分动作子序列中进行学习，平衡梯度偏差与方差。SubTB加速了收敛速度，在稀疏或长时域环境中提升了稳定性，并在合成任务与真实世界任务（包括肽和蛋白质生成）中，于奖励与多样性指标上均优于现有的TB与FM等目标。

ABSTRACT

Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks. Existing training objectives for GFlowNets are either local to states or transitions, or propagate a reward signal over an entire sampling trajectory. We argue that these alternatives represent opposite ends of a gradient bias-variance tradeoff and propose a way to exploit this tradeoff to mitigate its harmful effects. Inspired by the TD($λ$) algorithm in reinforcement learning, we introduce subtrajectory balance or SubTB($λ$), a GFlowNet training objective that can learn from partial action subsequences of varying lengths. We show that SubTB($λ$) accelerates sampler convergence in previously studied and new environments and enables training GFlowNets in environments with longer action sequences and sparser reward landscapes than what was possible before. We also perform a comparative analysis of stochastic gradient dynamics, shedding light on the bias-variance tradeoff in GFlowNet training and the advantages of subtrajectory balance.

研究动机与目标

• 为解决GFlowNet训练中的偏差-方差权衡问题，其中局部目标（如详细平衡）具有低方差但高偏差，而轨迹级目标（如轨迹平衡）具有低偏差但高方差。
• 开发一种灵活的训练目标，能够从不完整或部分轨迹中学习，从而在具有长动作序列或稀疏奖励的环境中实现更高效且稳定的训练。
• 提升GFlowNet训练的收敛速度与对超参数选择的鲁棒性，特别是在具有挑战性的概率建模场景中。
• 提供一个统一框架，通过可学习的$\lambda$参数，在局部一致性与完整轨迹信用分配之间实现插值。

提出的方法

• 提出SubTB($\lambda$)，一种受强化学习中TD($\lambda$)启发的训练目标，可在不同长度的子轨迹间平衡信用分配。
• 使用子轨迹上时序差分目标的加权和，其中权重由资格迹参数$\lambda$决定，从而在局部（类似DB）与全局（类似TB）学习之间实现插值。
• 将该目标推导为轨迹平衡的推广形式，通过基于未来奖励为中间状态分配信用，实现从不完整轨迹中学习。
• 采用基于子轨迹段的随机梯度更新规则，相比完整轨迹反向传播，显著降低了方差。
• 引入一种可微分机制，动态采样并加权子轨迹，提升数据效率与收敛速度。
• 将该目标应用于通过策略梯度方法训练GFlowNets，其中策略由神经网络参数化，可泛化至未见状态。

实验结果

研究问题

• RQ1 是否一种从部分轨迹中学习的GFlowNet训练目标，相较于完整轨迹或局部目标，能实现更快的收敛速度与更高的稳定性？
• RQ2 SubTB($\lambda$)中的超参数$\lambda$如何影响GFlowNet训练中的偏差-方差权衡？
• RQ3 SubTB($\lambda$)是否能在长动作序列或稀疏奖励函数环境中实现有效训练，而此前的方法在此类场景中失败？
• RQ4 与轨迹平衡相比，SubTB($\lambda$)在多大程度上降低了梯度方差，以及这对训练动态有何影响？
• RQ5 SubTB($\lambda$)是否能同时提升序列生成任务中的样本质量（奖励）与多样性？

主要发现

• 在合成环境与真实世界环境中，SubTB($\lambda$)相比TB、FM与DB目标，实现了更快的收敛速度，并对超参数选择的敏感度更低。
• 在抗菌肽生成任务中，SubTB($\lambda$)实现了0.96 ± 0.02的平均top-100奖励，显著优于TB（0.90 ± 0.03）与FM/DB（0.78 ± 0.05）。
• 在237个残基长度的荧光蛋白生成任务中，SubTB($\lambda$)实现了1.18 ± 0.10的top-100奖励与204.44 ± 0.45的多样性，优于TB（0.76 ± 0.19与204.31 ± 0.44）及其他基线方法。
• 采用SubTB($\lambda$)训练的模型更早发现更多模式，且与目标分布的Spearman等级相关系数更高，表明其探索更快速且更高效。
• 与TB相比，SubTB($\lambda$)下的随机梯度方差显著降低，验证了SubTB在保持低偏差的同时有效降低梯度方差的假设。
• SubTB($\lambda$)使在稀疏奖励与长轨迹环境中实现有效训练成为可能——这些正是以往GFlowNet目标难以收敛或完全失败的场景。