DOCSEARCHbyAlgolia
Submit Docs

torchgfn: PyTorch Implementation of GFlowNets

repository·master·Indexed 43 days ago

https://github.com/gfnorg/torchgfn

A Python library for fast prototyping of GFlowNet (GFN) algorithms. It provides a decoupled framework for defining environments, sampling processes, and function approximators. Version 2.4.1 supports discrete, continuous, and graph-based states, and is optimized for scaling with trajectory length and large batch sizes.

Tokens
41.4K
Snippets
105
Records
158
Agent score
53%

What's inside torchgfn

  1. Compare torchgfn performance with other GFlowNet libraries

    master

    This benchmark analysis compares the per-iteration training performance of torchgfn against gflownet (PyTorch/Hydra) and gfnx (JAX/Equinox).

    Supported Environments

    • Hypergrid (2D, height=8): Discrete grid navigation. Supported by all three libraries.
    • Box/CCube (2D, delta=0.1): Continuous environment with Beta mixture policies. Supported by torchgfn and gflownet only.
    • Ising (6x6, 10x10): Discrete Ising model. Supported by all three libraries.
    • BitSequence: Non-autoregressive binary sequence generation. Supported by torchgfn and gfnx only.

    Benchmark Configuration

    All benchmarks use Trajectory Balance (TB) loss with the following standard settings:

    • hidden_dim: 256
    • layers: 2-layer MLPs
    • lr: 1e-3
    • lr_logz: 0.1
    • batch sizes: 32, 128, and 512

  2. Performance comparison: torchgfn vs gflownet vs gfnx

    master

    Benchmark results for GPU (CUDA) environments show significant performance differences across different scenarios:

    Hypergrid Scenarios

    • gfnx is the fastest (5-100x faster than others) due to JAX JIT compilation and jax.lax.scan.
    • torchgfn scales efficiently, increasing ~1.5x for a 16x batch increase.
    • gflownet scales poorly (~10x increase for 16x batch increase) and suffers from super-linear scaling as environment complexity grows.

    Ising Scenarios

    • torchgfn is significantly faster than gflownet (12–200x faster). The performance gap widens as grid size and batch size increase.
    • gflownet's bottleneck is a cubic relationship: Total work ∝ batch_size × action_space_size × trajectory_length.

    Box 2D Scenarios

    • Expected to scale similarly to hypergrid_small or better, as continuous actions avoid the actions2indices bottleneck found in discrete environments.

  3. Overview of GFlowNet Loss Functions

    master

    In torchgfn, a GFlowNet is a parametrization that includes one or multiple Estimators. To train a GFlowNet, the chosen loss function must be compatible with the provided estimators and the data format. The loss function takes either States, Transitions, or Trajectories Container instances as input, depending on the specific loss implementation.

    Commonly used loss classes include:

    • TBGFlowNet (Trajectory Balance)
    • DBGFlowNet (Detailed Balance)
    • SubTBGFlowNet (Sub-Trajectory Balance)
    • FMGFlowNet (Flow Matching)
    • LogPartitionVarianceGFlowNet (ZVar)
    • RelativeTrajectoryBalanceGFlowNet (RTB)
    • PolicyGradientGFlowNet (VPG)
    • EntPPOGFlowNet (Ent-PPO)

  4. Use Containers to manage transitions and trajectories

    master

    Containers are collections of States and associated metadata (like rewards or densities). They are used to organize data for training.

    Available Container Types:

    • Transitions: Represents a batch of transitions $s \rightarrow s'$.
    • Trajectories: Represents a batch of complete trajectories $\tau = s_0 \rightarrow s_1 \rightarrow \dots \rightarrow s_n \rightarrow s_f$.
    • StatesContainer: Represents a batch of states (optimized for flow-matching GFlowNets).

    Data Transformation for Training

    Containers provide methods to convert complete trajectories into the specific formats required by different loss functions. These methods automatically filter out dummy states and meaningless transitions used for padding:

    Target FormatMethodUse Case
    TransitionsTrajectories.to_transitions()Edge-decomposable losses
    StatesTrajectories.to_states()State-decomposable losses
    State PairsTrajectories.to_state_pairs()Flow matching losses

    Convenience Method: Use GFlowNet.loss_from_trajectories() to handle these conversions internally.

  5. What are Conditional GFlowNets?

    master

    A Conditional GFlowNet learns to sample from a family of distributions $p(x | c) \propto R(x, c)$ parameterized by an external condition variable $c$, rather than a single distribution $p(x) \propto R(x)$.

    This approach is used when:

    • The reward function depends on external parameters (e.g., temperature, task ID).
    • You need a single model that can interpolate between different target distributions.
    • The condition represents side information that modifies the desired outputs.

  6. Key differences between torchgfn and other GFlowNet libraries

    master

    When comparing torchgfn to gflownet or gfnx, note the following architectural and implementation differences:

    Aspecttorchgfn
    State representationStates objects
    Backward policy (box)Configurable (uniform or learned)
    PF log-prob caching (box)Yes (via save_logprobs=True)
    LR schedulerNone
    Input normalization (box)2*x - 1

    Comparison with other libraries:

    • gflownet: Uses raw tensor batches instead of States objects, uses a fixed uniform backward policy for Box environments, and applies a StepLR scheduler every iteration.
    • gfnx: Uses JAX arrays and performs extra computation during trajectory rollout (entropy computation, redundant softmax, and double forward passes).

  7. Understand torchgfn performance scaling and trajectory length

    master

    When choosing between torchgfn and gflownet for your workloads, consider the expected trajectory length of your environment:

    • Short trajectories (e.g., Box environment, ~3 steps): gflownet is generally faster (approx. 1.13x) because torchgfn has a higher fixed overhead per iteration due to States/Actions object construction.
    • Long trajectories (e.g., Hypergrid, ~14+ steps): torchgfn scales better and is faster (approx. 1.3x) because its simpler backward pass becomes more efficient as the per-step cost dominates the fixed overhead.

    In summary, torchgfn is optimized for scaling with trajectory length and larger batch sizes.

  8. Implement a recurrent policy with RecurrentPolicyMixin

    master

    RecurrentPolicyMixin is used for per-step, non-vectorized evaluation (is_vectorized=False). It threads a carry through the rollout steps.

    Requirements:

    • The estimator must implement init_carry(batch_size, device) to initialize the state stored in ctx.carry.
    • compute_dist must call the estimator using the signature (states_active, ctx.carry) -> (est_out, new_carry), then update ctx.carry with the new state.
    • log_probs follows the non-vectorized path, which includes padding to (N,) and strict checks for ±inf after masking.

    Use this for sequence models (like RNNs or GRUs) where the next action depends on the history of previous actions/states.

  9. Use RolloutContext for per-rollout state

    master

    The RolloutContext is a lightweight container created once per rollout to manage state and buffers. It contains:

    • batch_size, device, and optional conditions.
    • carry: Optional state for recurrent policies.
    • trajectory_log_probs: Per-step buffer for log-probabilities.
    • trajectory_estimator_outputs: Per-step buffer for estimator outputs.
    • current_estimator_output: Cached raw model output for immediate retrieval.
    • extras: dict: A dictionary for storing arbitrary policy-specific data or diagnostics.

  10. Understand the PolicyMixin rollout API

    master

    Estimators in torchgfn become policy-capable by mixing in a rollout API. This allows a single Sampler and probability utilities to work with different estimator families (discrete, graph, conditional, recurrent) without custom glue code.

    A policy-capable estimator must expose the following interface:

    • is_vectorized: bool: Indicates if the estimator can be evaluated in a single vectorized call (no per-step carry).
    • init_context(batch_size, device, conditions): Allocates a per-rollout RolloutContext.
    • compute_dist(states_active, ctx, step_mask, ...) -> (Distribution, ctx): Runs the model and builds a torch.distributions.Distribution.
    • log_probs(actions_active, dist, ctx, step_mask, vectorized, ...) -> (Tensor, ctx): Evaluates log-probabilities, optionally padded to the batch.
    • get_current_estimator_output(ctx): Accesses the last raw model output.

    All per-step artifacts (log-probs, raw outputs, recurrent state) are owned and recorded by the RolloutContext.

  11. On-Policy vs Off-Policy Training in GFlowNets

    master

    GFlowNet training can be configured as either on-policy or off-policy. The primary distinction is how log-probabilities are handled during loss computation.

    On-Policy Training

    Used when trajectories are sampled from the current policy and used immediately. To avoid redundant forward passes, you should cache log-probabilities during sampling and reuse them during loss calculation.

    • Sampling: Set save_logprobs=True.
    • Loss: Set recalculate_all_logprobs=False.

    Off-Policy Training

    Used when you decouple sampling from learning (e.g., using a replay buffer, epsilon-greedy exploration, temperature scaling, or noisy layers). Because the policy has changed since the trajectories were collected, cached log-probabilities are stale and must be recomputed.

    • Sampling: Set save_logprobs=False.
    • Loss: Set recalculate_all_logprobs=True.
    # On-policy example
trajectories = gflownet.sample_trajectories(env, n=batch_size, save_logprobs=True)
training_samples = gflownet.to_training_samples(trajectories)
loss = gflownet.loss(env, training_samples, recalculate_all_logprobs=False)

# Off-policy example
trajectories = gflownet.sample_trajectories(env, n=batch_size, save_logprobs=False)
# ... store in buffer, retrieve later ...
training_samples = gflownet.to_training_samples(trajectories)
loss = gflownet.loss(env, training_samples, recalculate_all_logprobs=True)

  12. How carry management works in recurrent policies

    master

    In recurrent policies, the hidden state (the "carry") is managed automatically via the RecurrentPolicyMixin:

    1. Initialization: The sampler calls init_carry(batch_size) at the start of each rollout to allocate the initial hidden state.
    2. Updates: The carry is updated at every step and stored within the RolloutContext.
    3. Log-probability Computation: When computing log-probabilities for existing trajectories (e.g., during off-policy training), the carry is re-threaded step-by-step through the sequence to reconstruct the state at each point.