PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning GFlowNets from partial episodes for improved convergence and stability

Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, Nikolay Malkin

arXiv:2209.12782v3cs.LGstat.ML

TL;DR

Existing GFlowNet objectives trade off local updates against trajectory-level credit assignment and gradient variance. This paper introduces SubTB(λ), which learns from partial subtrajectories of varying lengths, and reports faster, more stable convergence plus applicability to sparse-reward and long-horizon environments. The analysis attributes these benefits to interpolation between detailed balance and trajectory balance.

  • Problem

    Existing GFlowNet objectives lie at opposite ends of a credit-assignment and gradient-variance tradeoff, motivating an objective that can combine their advantages.

  • Method

    SubTB(λ) trains GFlowNets from partial subtrajectories of varying lengths, interpolating between detailed balance and trajectory balance.

  • Results

    SubTB(λ) accelerates convergence, reduces sensitivity to hyperparameter choices, and enables training in environments with sparse rewards or long action sequences.

  • Takeaways & Limitations

    SubTB(λ) combines variance reduction from local objectives with the fast credit assignment of trajectory-level objectives.

  • Takeaways & Limitations

    The paper identifies empirical evaluation of connections between SubTB and related variational methods as future work.

Abstract

from arXiv · show

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.

1. Introduction

GFlowNets face a tradeoff between local objectives with limited credit assignment and trajectory-level objectives with higher gradient variance. SubTB(λ) addresses this by learning from partial action subsequences of varying lengths, improving convergence, robustness, and applicability to difficult environments.

  • 1. Introduction: GFlowNets train sequential samplers whose object probabilities are proportional to a specified nonnegative reward, supporting probabilistic modeling tasks.
  • 1. Introduction: SubTB(λ) learns from partial action subsequences of varying lengths, combining local and trajectory-level training signals.The objective is inspired by the bias-variance tradeoff associated with TD(λ)-style methods.
  • 1. Introduction: SubTB models approach the target distribution in fewer training steps and are less sensitive to hyperparameter choices than previous approaches.
  • 1. Introduction: SubTB enables GFlowNet training in environments with sparser rewards or longer action sequences, where past approaches perform poorly.
  • 1. Introduction: SubTB(λ) benefits are explained by lower stochastic-gradient variance, while λ interpolates between high-bias, low-variance DB and low-bias, high-variance TB.

2. Method

GFlowNets learn a forward policy whose terminal-state distribution is proportional to rewards, using objectives that enforce flow consistency locally or across complete trajectories. Subtrajectory balance generalizes these objectives to partial trajectories, with λ controlling the bias–variance tradeoff between detailed balance and trajectory balance.

  • Preliminaries: GFlowNets construct terminal objects by sampling actions from a forward policy on a directed acyclic state graph, targeting a terminal distribution proportional to reward.Complete trajectories begin at the unique initial state and end at terminal states; the terminal likelihood marginalizes over all trajectories reaching each object.
  • GFlowNet training objectives: Flow matching, detailed balance, and trajectory balance introduce auxiliary flow, policy, or normalization variables whose global optima imply reward-proportional sampling.Flow matching constrains incoming and outgoing flows, detailed balance matches state flows across actions, and trajectory balance constrains complete trajectories.
  • Subtrajectory balance: Subtrajectory balance enforces the GFlowNet consistency condition on partial trajectories, while detailed balance and trajectory balance appear as its one-action and complete-trajectory special cases.When the subtrajectory loss is zero for all partial trajectories, the resulting policy satisfies the desired reward-proportional sampling condition.
  • Subtrajectory balance: Subtrajectory losses can be computed from sampled complete episodes by optimizing a convex combination over their nontrivial partial subtrajectories.An episode of length n contains O(n^2) nontrivial subtrajectories; the paper notes that network-gradient computation still requires only one forward and one backward pass.
  • Subtrajectory balance: Subtrajectory balance is motivated as an interpolation between trajectory balance’s high variance and detailed balance’s lower variance but higher bias relative to the trajectory-balance gradient.The learned state-flow terms replace stochastic components of the trajectory-balance target with estimates of their expectations, introducing bias while reducing variance.

3. Related work

SubTB(λ) is related to eligibility-trace methods in reinforcement learning and to maximum-entropy or path-consistency approaches. In tree-structured environments, its subtrajectory loss is equivalent to path consistency learning, while prior applications without intermediate rewards used more restricted subtrajectories.

  • Eligibility traces: SubTB(λ) draws on TD(λ), where λ combines multi-step returns to control a bias–variance tradeoff.Larger λ approaches a Monte Carlo estimate, yielding lower bias and higher variance in the reinforcement-learning analogy.
  • MaxEnt RL: In directed-tree environments, the SubTB loss on individual subtrajectories is equivalent to path consistency learning.Prior PCL applications without intermediate rewards computed losses only on length-one subtrajectories or subtrajectories containing a terminal state.
  • MaxEnt RL: GFlowNets are connected to maximum-entropy reinforcement-learning methods, especially when the environment’s Markov decision process has a tree structure.Related methods also maximize entropy over policy or state-visitation distributions to achieve objectives similar to GFlowNets.

4. Experiments

Across hypergrid, gradient, molecule, sequence, peptide, and protein experiments, SubTB(λ) converges faster or more reliably than prior GFlowNet objectives while extending training to sparse-reward and long-sequence settings.

  • 4.1. Hypergrid: Robustness to sparse rewards: SubTB(λ) converges faster and with less variation across random seeds than TB on all tested hypergrid sizes.The comparison uses L1 distance between empirical and target distributions over training.
  • 4.1. Hypergrid: Robustness to sparse rewards: With background reward 10^-4, SubTB(λ) remains strong, whereas TB fails to discover all target modes on grids larger than 8×8.
  • 4.1.1. A CLOSER LOOK AT GRADIENT VARIANCE: At batch size 64, DB has the highest gradient self-consistency, TB the lowest, and SubTB(λ=0.8) lies between them.
  • 4.1.1. A CLOSER LOOK AT GRADIENT VARIANCE: SubTB(λ) interpolates between DB's high-bias, low-variance updates and TB's low-bias, high-variance updates through λ.
  • 4.3.1. BIT SEQUENCES: SubTB(λ), especially λ=1, outperforms DB and TB at optimal hyperparameters and is more robust to hyperparameter choices on the molecule task.
  • 4.3.1. BIT SEQUENCES: SubTB(λ) achieves higher final Spearman correlation and discovers modes faster than other GFlowNet objectives and non-GFlowNet baselines.
  • 4.3.2. ANTIMICROBIAL PEPTIDE GENERATION: On AMP generation, SubTB(λ) significantly improves both reward and diversity over TB, FM, DB, and other baselines.
  • 4.3.3. FLUORESCENT PROTEIN GENERATION: On GFP generation, SubTB(λ) significantly outperforms TB in reward while achieving similar diversity, with a larger advantage than on AMP.

5. Discussion and conclusion

SubTB(λ) combines local and trajectory-level learning to reduce gradient variance while preserving fast credit assignment. The authors identify adaptive trajectory selection and broader applications to incomplete episodes as future directions.

  • 5. Discussion and conclusion: SubTB(λ) harnesses local objectives’ variance reduction while retaining trajectory-level objectives’ fast credit assignment.The paper frames this as a bias-variance tradeoff between DB and TB.
  • 5. Discussion and conclusion: The authors identify dynamic λ selection and active-learning-based trajectory sampling as promising future work.
  • 5. Discussion and conclusion: Incomplete-episode learning motivates applying subtrajectory objectives to RL environments and settings with rewards for incomplete states.

A. Experiment details: Hypergrid

The Hypergrid experiments use MLP-based forward and backward policies with shared layers and train models using Adam across extensive trajectory batches. The harder 8 × 8 grid is also used for gradient bias and variance experiments.

  • Model architecture: MLPs model P_F(-|s), P_B(-|s), and log F(s) from one-hot state coordinates, sharing all layers except the last.The initial-state flow log Z is represented independently.
  • Optimization: All models use Adam with batch size 16 for 10^6 trajectories, while learning rates are selected from seven candidate values.The candidate learning rates range from 0.0005 to 0.01.
  • Optimization: The optimal SubTB parameter is λ = 0.9, selected from {0.8, 0.9, 0.99}.
  • Gradient experiments: Gradient bias and variance experiments use the harder 8 × 8 grid with a tabular GFlowNet trained by Adam at learning rate 0.007 and SubTB(λ = 0.8).

A.1. Additional experiments

Additional experiments examine SubTB(λ) across harder grids, restricted subtrajectory lengths, λ values, exploration settings, and alternative state-flow computations. These results consistently support strong convergence and variance-reduction behavior, while identifying computational limits for exact flow calculations.

  • Additional environments: Additional results include experiments on more difficult grid environments.
  • Restricted subtrajectories: SubTB(λ) remains strong when training uses only subtrajectories of length up to 4.The restriction truncates the objective by excluding pairs with j−i > 4.
  • Effect of λ: SubTB(λ) converges fastest for λ values slightly below 1 while interpolating gradually between DB and TB.
  • Exploration behavior: TB can benefit from off-policy training and higher policy-logit temperature, but SubTB(λ) still learns faster and finds more modes within the training budget.
  • State-flow analysis: The modified gradient analysis replaces learned state flows with exact true forward or backward flows computed from the current policies, but only for similarity evaluation.Exact computation is tractable in the small environment but not generally possible.
  • State-flow analysis: Using true forward state flows produces gradient-similarity behavior similar to learned forward flows, suggesting learned flows remain close to optimal and retain SubTB(λ)'s variance-reduction benefits.

B. Experiment details: Molecules

The molecule experiments reuse the published Malkin et al. codebase and the prior molecular setup, including its proxy reward, held-out correlation set, graph neural network, and tuned exploration settings.

  • Experimental setup: Molecule experiments use the same proxy reward, held-out molecule set, graph neural network architecture, exploration rate, and early-stopping likelihood as prior work.The implementation is based on published Malkin et al. (2022) code extending Bengio et al. (2021a).

C. Experiment details: Bit sequences

The bit-sequence experiments use a common Transformer policy and fixed training protocol, while matching previously reported FM and baseline implementations.

  • C. Experiment details: Bit sequences: All methods use a 3-layer, dimension-64 Transformer with 8 attention heads as the policy architecture.Training uses 50,000 iterations, minibatches of 16, and Adam.
  • C. Experiment details: Bit sequences: The modes and test sequences are selected following Malkin et al. (2022).
  • C. Experiment details: Bit sequences: FM and baseline methods use the exact implementation and hyperparameters reported by Malkin et al. (2022).

D. Experiment details: Antimicrobial peptide generation

The antimicrobial-peptide experiments use DBAASP sequences to train a proxy reward model and compare GFlowNet policies under a shared Transformer-based setup.

  • D. Experiment details: Antimicrobial peptide generation: The dataset contains 6438 known AMP sequences and 9522 non-AMP sequences from DBAASP.
  • D. Experiment details: Antimicrobial peptide generation: A classifier trained on these sequences provides the proxy reward function, using 20 PERCENT of the data for validation.
  • D. Experiment details: Antimicrobial peptide generation: The reward model is a 4-layer Transformer, while all policy methods use a 3-layer Transformer with hidden dimension 64 and 8 attention heads.The reward model is trained with minibatches of 256, learning rate 10^-4, and validation-based early stopping.

E. Experiment details: Fluorescent protein generation

The fluorescent-protein experiments evaluate sequence generation using a fluorescence regressor and tune exploration, reward shaping, and objective-specific optimization settings.

  • E. Experiment details: Fluorescent protein generation: The GFP variant uses 56,086 proteins with fluorescence scores, retaining 20 PERCENT for validation and early stopping.
  • E. Experiment details: Fluorescent protein generation: A 4-layer Transformer regressor with hidden dimension 64 and 8 attention heads supplies the learned fluorescence reward.It is trained with minibatches of 256, learning rate 10^-4, and validation-based early stopping.
  • E. Experiment details: Fluorescent protein generation: Exploration uses random actions with probability 0.01, and reward exponent beta is set to 3.
  • E. Experiment details: Fluorescent protein generation: SubTB(lambda) selects lambda=0.99 as the best value among {0.7, 0.8, 0.9, 0.99}.Trajectory balance, flow matching, A2C, and SAC use separately selected optimization settings.

F. Inverse protein folding: Non-autoregressive sequence generation

The inverse-folding task samples fixed-length amino-acid sequences through non-autoregressive replacements and evaluates how well learned sampling matches the target distribution. Intermediate lambda values provide the best fit on held-out terminal states.

  • F. Inverse protein folding: Non-autoregressive sequence generation: The task samples amino-acid sequences of fixed length L=40 from a Boltzmann distribution defined by physics-model energies.
  • F. Inverse protein folding: Non-autoregressive sequence generation: Actions modify one amino-acid position at a time, beginning from a uniformly random sequence and terminating after N=40 replacements.
  • F. Inverse protein folding: Non-autoregressive sequence generation: Intermediate values of lambda provide the best fit to the target distribution according to held-out correlation results in Fig. F.1.The metric correlates log R(x) with marginal sampling likelihood log p_theta(x).
  • F. Inverse protein folding: Non-autoregressive sequence generation: The appendix evaluates empirical L1 curves across lambda values on an 8 x 8 grid and compares harder grid variants under exploratory policies.
  • F. Inverse protein folding: Non-autoregressive sequence generation: Gradient similarity with analytically computed state flows is examined in a separate comparison with Fig. 4.
Loading 2209.12782v3…