Symmetry-Aware GFlowNets

TL;DR

This work identifies a fundamental bias in generative flow networks (GFlowNets) arising from graph symmetries, where equivalent actions due to automorphisms lead to systematic mis-sampling. It introduces Symmetry-Aware GFlowNets (SA-GFN), a principled framework that incorporates symmetry corrections by scaling rewards with the automorphism-group order $|Aut(G)|$, enabling unbiased sampling without explicit one-step transition computations. The theory shows that orbit-equivalence suffices to correct action biases, and the approach extends to both node-by-node and fragment-based graph generation, including an unbiased model-likelihood estimator and considerations for GNN expressiveness. Empirically, SA-GFN removes symmetry-induced biases, improves diversity and high-reward molecule generation, and offers computational advantages over per-transition equivalence corrections, with exact and approximate variants demonstrated across illustrative, synthetic, and molecular tasks.

Abstract

Generative Flow Networks (GFlowNets) offer a powerful framework for sampling graphs in proportion to their rewards. However, existing approaches suffer from systematic biases due to inaccuracies in state transition probability computations. These biases, rooted in the inherent symmetries of graphs, impact both atom-based and fragment-based generation schemes. To address this challenge, we introduce Symmetry-Aware GFlowNets (SA-GFN), a method that incorporates symmetry corrections into the learning process through reward scaling. By integrating bias correction directly into the reward structure, SA-GFN eliminates the need for explicit state transition computations. Empirical results show that SA-GFN enables unbiased sampling while enhancing diversity and consistently generating high-reward graphs that closely match the target distribution.

Figures (15)

• Figure 1: Illustration of graph transitions from $G_1$ to various successor graphs. The blue oval highlights graphs $G_2$ and $G_3$ are isomorphic.
• Figure 2: Graphs representing transitions $(G_1, G_2, G_3)$, where the first transition is performed by AddNode and the second by AddEdge. The number of forward/backward actions are represented as the number of arrows. Symmetries in each graph is related to orbit-equivalent actions, as seen in the ratio $|{\mathrm{Aut}}(G_1)|/|{\mathrm{Aut}}(G_2)| = |{\mathrm{Orb}}(G_1,$$)|/|{\mathrm{Orb}}(G_2,$$)|$. Nodes in the same orbit are given the same color.
• Figure 3: (a) Terminating probabilities of trained models in the uniform-reward environment. States are sorted according to the number of graphs in the state, $|x|$. (b), (c)$L_1$ errors between the target probabilities and the model’s terminating probabilities during training in the synthetic environment. (d) Errors in the estimated model log-likelihood, defined as the difference between estimated and exact log-likelihood. "Random" denotes the errors of an initial random model, while "Trained" refers that of a trained model. Numbers in brackets indicate the number of edges in the terminal states used for estimation.
• Figure 4: Correlations between log rewards and model log-likelihoods during training in fragment experiment.
• Figure 5: Two actions induce isomorphic graphs, making them transition-equivalent. However, they are not orbit-equivalent. This example was originally presented by ma2024baking.

Theorems & Definitions (46)

• Definition 3.1: Isomorphism
• Definition 3.2: Automorphism
• Definition 3.3: Orbit
• Definition 4.1: Transition equivalence
• Definition 4.2: Orbit equivalence
• Theorem 4.3
• Theorem 4.4: Sufficiency of orbit equivalence
• Lemma 4.5
• Theorem 4.6: Automorphism correction
• Corollary 5.1: TB correction