Reproduce “Learning GFlowNets from partial episodes for improved convergence and stability” on OpenResearch

About this paper

Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio and 1 more

This research introduces Subtrajectory Balance (SubTB(λ)), a new training objective for Generative Flow Networks (GFlowNets) designed to improve convergence and stability. SubTB(λ) addresses the inherent bias-variance tradeoff in GFlowNet training by interpolating between local and global objectives, leading to faster learning and enhanced robustness across diverse generative modeling tasks, including molecule and protein design.

Key insights

• Existing GFlowNet training objectives (Flow Matching, Detailed Balance, Trajectory Balance) occupy different points on a gradient bias-variance spectrum, with local objectives having high bias/low variance and global objectives having low bias/high variance.
• Subtrajectory Balance (SubTB) effectively interpolates between these extremes, reducing gradient variance by incorporating learned state flow estimates, similar to how value functions reduce variance in RL.
• Learning from partial episodes (subtrajectories) via SubTB(λ) facilitates faster credit assignment and improved generalization of learned state flows, especially in sparse reward or long-sequence environments.

Summary by alphaXiv.