Local Gram Flow Loss

• Local Gram Flow Loss is an optimization criterion in GFlowNets that uses additive energy functions for localized credit assignment across both complete and incomplete trajectories.
• It overcomes the limitations of terminal-reward based methods like detailed balance and trajectory balance, thereby enhancing convergence and reducing variance.
• Empirical results in domains such as molecule discovery and bit-sequence generation show faster convergence, increased mode diversity, and improved performance with local flow updates.

Local Gram Flow Loss, also referred to as the local flow loss or the forward-looking flow objective, is an optimization criterion within the Generative Flow Network (GFlowNet) framework that enables effective parameter updates using both complete and incomplete trajectories. This objective leverages additive energy functions to provide "dense" local credit assignments, improving both the efficiency and flexibility of training by incorporating intermediate state information.

1. Background and Motivation

GFlowNets are designed to sample compositional structured objects via sequential decision processes that construct an object $x$ through a trajectory of states, with the probability of generating $x$ proportional to a specified reward function $R(x)$ or equivalently $\exp(-E(x))$, where $E(x)$ is an energy function. Traditional training objectives for GFlowNets, such as detailed balance (DB) and trajectory balance (TB), require evaluation of rewards at terminal states, limiting the ability to update parameters based on partial or incomplete trajectories. In situations where only the terminal reward is available, training can suffer from slow credit assignment and increased variance, particularly for long trajectories. The local flow objective addresses these limitations by exploiting additive decompositions of the energy, facilitating localized gradients and updates without necessitating knowledge of the final reward (Pan et al., 2023).

2. Formal Definitions and Notation

Let $S$ denote a directed acyclic graph (DAG) representing the state space, with unique initial state $s_0$ and set of terminal states $X \subset S$. For a nonterminal state $s$, $Ch(s)$ and $Pa(s)$ denote its children and parents, respectively. The reward $R(x): X \to \mathbb{R}_+$ is specified at terminal $x$, commonly written as $R(x) = e^{-E(x)}$. Under the additive energy extension (Assumption 4.1), $E$ can be defined over all $s \in S$ with per-transition energy $E(s \to s') = E(s') - E(s)$. A policy comprises a forward Markov kernel $P_F(s'|s;\theta)$, a backward kernel $P_B(s|s';\theta)$, and a scalar flow function $F(s;\theta)$ satisfying $F(x) = R(x)$ for $x$ terminal.

The forward-looking (local) flow reparameterization is defined as

$$\widetilde F(s) \equiv e^{E(s)} F(s) = \sum_{x \succeq s} P_B(s|x) e^{-E(s \to x)}$$

where $x \succeq s$ denotes all terminal descendants of $s$.

3. From Global to Local Flow-Matching Objectives

The classic global DB constraint, applied for each edge $(s \to s')$, is

$F(s)\,P_F(s'|s) = F(s')\,P_B(s|s').$

TB and SubTB objectives extend this to products over full trajectories or subtrajectories but require access to terminal rewards, limiting practical applicability for incomplete data.

By leveraging additive energies, the local (forward-looking) detailed-balance (FL-DB) constraint is formulated as

$$\widetilde F(s)\,P_F(s'|s) = \widetilde F(s')\,P_B(s|s')\,e^{-E(s \to s')}.$$

This local constraint integrates immediate per-transition energy increments and supports parameter updates based solely on local (edge) information, irrespective of trajectory completion.

4. Local Flow Loss Definition and Computation

For each observed transition $(s \to s')$ in a dataset comprising possibly incomplete trajectories, the per-edge residual is defined by

$$\delta(s, s'; \theta) = \log \widetilde F(s;\theta) + \log P_F(s'|s;\theta) - \log \widetilde F(s';\theta) - \log P_B(s|s';\theta) + E(s \to s').$$

The local flow loss aggregates squared residuals over all sampled transitions: $$L_\text{local}(\theta) = \sum_{(s \to s') \in \mathcal{D}} \left[ \delta(s, s';\theta) \right]^2$$ where $\mathcal{D}$ is the set of all transitions, from both complete and incomplete trajectories.

A subtrajectory (SubTB-style) variant may impose analogous constraints over finite subpaths, but still leverages only local increments and does not depend on full terminal reward evaluation.

5. Training Algorithm and Implementation

The operational training procedure for the local flow objective is as follows:

• Parameters $\theta$ for $P_F$, $P_B$, and the flow-to-go network $\widetilde F$ are maintained.
• For each training iteration:
  • A trajectory prefix or batch of prefixes of arbitrary length is sampled or replayed.
  • The per-transition energy $E(s \to s')$ is computed.
  • The FL-DB residual $\delta(s, s';\theta)$ is calculated for each edge using the definitions above.
  • Parameters are updated via gradient descent: $$\theta \leftarrow \theta - \eta \nabla_\theta [\delta(s, s';\theta)^2]$$.

In practice, batches of transitions are used to build the empirical loss and enable efficient backpropagation. Since the loss terms rely only on pairs of consecutive states, updates can be performed even in the absence of terminal states or rewards (Pan et al., 2023).

6. Empirical Outcomes and Comparative Analysis

Across domains including set generation, bit-sequence generation, and molecular graph generation, the local flow loss (FL-DB objective) has demonstrated:

• Significantly faster convergence than previous DB, TB, or SubTB objectives, especially with long trajectories or large action/state spaces.
• The ability to identify more high-reward (and thereby diverse) modes compared to baselines.
• Robustness to training exclusively on incomplete trajectories; performance matches or exceeds full-trajectory training in various settings.
• In molecule discovery, both the mean top-100 reward and result diversity (lower average Tanimoto similarity) improved under FL-DB compared to global objectives, with better correspondence observed between log-sampled probability and log-reward for held-out molecules.

These benefits arise from the denser and more direct local credit signal provided by the per-edge FL-DB objective, mitigating the delay and variance inherent in global, terminal-reward-based estimators (Pan et al., 2023).

7. Implications and Future Directions

The local flow loss expands the applicability of GFlowNets to domains where complete trajectories or terminal rewards are expensive or impractical to obtain, enabling learning from arbitrary trajectory fragments. This suggests significant potential in scientific discovery tasks, probabilistic modeling over structured spaces, and other compositional generation contexts. A plausible implication is that future research may employ localized flow-based objectives for other classes of sequential decision processes to further reduce sample and computation inefficiencies. Ongoing theoretical investigations into the statistical properties and possible generalizations of FL-DB loss are likely to broaden the reach of GFlowNet methodologies.