GFlowNets Need Automorphism Correction for Unbiased Graph Generation

Authors: Hohyun Kim, Seunggeun Lee, Min-hwan Oh
Primary area: generative models
Keywords: GFlowNet, graph generation, molecule optimization

Abstract:
Generative Flow Networks (GFlowNets) are generative models capable of producing graphs. While GFlowNet theory guarantees that a fully trained model samples from an unnormalized target distribution, computing state transition probabilities remains challenging due to the presence of equivalent actions that lead to the same state. In this paper, we analyze the properties of equivalent actions in the context of graph generation tasks and propose efficient solutions to address this problem. Our theoretical analysis reveals that naive implementations, which ignore equivalent actions, introduce systematic bias in the sampling distribution for both atom-based and fragment-based graph generation. This bias is directly related to the number of symmetries in a graph, a factor that is particularly critical in applications such as drug discovery, where symmetry plays a key role in molecular structure and function. Experimental results demonstrate that a simple reward-scaling technique not only enables the generation of graphs that closely match the target distribution but also facilitates the sampling of diverse and high-reward samples.

Decision: reject
Reviewer scores: 6, 5, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1729911285760

thread_type:
official_review

<thread_object>
creator: reviewer_mrgy

<content>
summary:
This paper focuses on Generative Flow Networks (GFlowNets), which are generative models used to produce graphs. While GFlowNet theory ensures that a fully trained model can sample from an unnormalized target distribution, the challenge lies in computing state transition probabilities, particularly due to equivalent actions that lead to the same state. The paper analyzes these equivalent actions in graph generation tasks and proposes efficient solutions to mitigate the associated challenges.

strengths:
1. The paper is well organized and the theories are well formulated.

2. The motivation is well introduced.

weaknesses:
1. How does the proposed method compare with other graph generative models, such as flow-based and discrete diffusion-based models?

questions:
N.A.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730426216897

thread_type:
official_review

<thread_object>
creator: reviewer_6mw9

<content>
summary:
This paper aims to address issues in graph generative GFlowNets that may fail to construct the target distribution due to equivalent actions. Specifically, it analyzes how discrepancies in the number of automorphism groups cause GFlowNets to incorrectly estimate the true reward. To address this, the paper incorporates the number of automorphism groups into the reward function and proves how this corrects reward underestimation. Notably, this paper also considers practical implementation for correction in fragment-based generation. Experimental results show that the proposed method better captures the target distribution.

strengths:
- The paper is well-written and easy to follow, and the proposed method is conceptually straightforward.
- This work is the first to address a significant pitfall in GFlowNets, i.e., errors due to equivalent actions, within primary setting of GFlowNets, i.e., graph generation.
- The authors provide a solid theorem for the corrected objectives showing how their global optima enable GFlowNets to construct the correct target distribution. Although the proofs consider TB and DB, these can also be easily extended to other objectives, such as subTB.
- The experiments are thorough and consider both important settings, namely atom-wise and fragment-wise graph generation.

weaknesses:
No weakness in the major flows. It seems that there are no errors in the proof.

questions:
- To better highlight the pitfalls, I wonder if the authors provide or illustrate a toy-example or toy-experiments where the conventional approaches induce an incorrect generative distribution, e.g., a distribution significantly biased towards the graphs with a low number of automorphism groups.
- Can authors provide the experimental computational costs for computing $|\text{Aut}(s)|$? I am curious about how much overhead the proposed method requires in practice, although the authors provide the time complexity in **Line 378**. Could this overhead be minor relative to time for reward computation or time for sampling trajectories?
- In DB-based implementation, I wonder if there might be improvement in convergence speed when we reparameterize the flow function $F(s)=\tilde{F}|\text{Aut}(s)|$ (like a prior flow reparameterization approach [1]), although this preserves the asymptotic optimality to induce the target distribution.

---

[1] Pan et al., Better Training of GFlowNets with Local Credit and Incomplete Trajectories, ICML 2023
</content>
</thread_object>

</thread>
<thread>
creation date:
1730664734391

thread_type:
official_review

<thread_object>
creator: reviewer_rf27

<content>
summary:
This work first studies the properties of equivalent actions when applying GFlowNets for graph generation. Equivalent actions denote the set of actions that lead to isomorphic graphs at each step of the autoregressive generation process. This work provides a theoretical analysis on the impact of ignoring equivalent actions and points out that it would introduce bias in the sampling distribution. With this insight, this work further proposes a simple correction on the GFlowNet objectives by using the order of the automorphism group to account for equivalent actions. This can correct the reward for highly symmetric graphs.

Experiments on small graph generation and small molecule generation are conducted to show the performance of the proposed correction.

strengths:
(1) It is really interesting and valuable to the community to identify the impact of ignoring equivalent actions in GFlowNets for graph generation. The theoretical analysis is quite sound from my reading and I think it is valuable to other readers.

(2) The theoretical results are quite elegant, thus leading to a simple correction to the original GFlowNets objectives. It is quite enjoyable to see that the correction term is the order of the automorphism group.

(3) The experiments can show that with such a simple correlation, the sampling bias and resulting performance are notably improved, which can support the theoretical analysis and the proposed corrected objectives straightforwardly.

(4) The paper is well written.

weaknesses:
(1) It looks really computationally expensive to evaluate the order of the automorphism group and the complexity could increase exponentially with the size of the graph. I understand that the paper provides some analysis on the computation. However, the experimental study on the complexity is missing, while it is very important to assess the practical usefulness of the proposed idea.

(2) I am a bit concerned about the practicality of the method. The experiments are mainly on small graph and small molecule generation. It is unclear if this method can be scalable to generate large molecules.

questions:
See the weaknesses section
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691169418

thread_type:
official_review

<thread_object>
creator: reviewer_6mnv

<content>
summary:
This work points out that automorphic actions may cause GFlowNets to artificially over/undersample terminal states compared to the target distribution. They also propose a fix by re-scaling the reward to account for the size of the automorphism group of terminal states. They illustrate the pathology and their fix in a toy example and a molecule generation task.

strengths:
* This work is overall well-written and easy to follow;
* It shows that a common practice of treating graphs as if they were the GFlowNet states leads to incorrect sampling --- implying that, to some extent, there is a series of incorrect experiments in the GFlowNet literature;
* Authors provide a quick fix to the issue.

weaknesses:
* It appears Figure 3 uses Equation 3 to compute the final state probabilities. I am not sure this is a fair evaluation. I suggest the authors use the empirical approximations of the distributions over terminal states (based on GFlowNet samples) for comparison. For instance, measuring L1 between the empirical sampling distribution and the target.

* The metrics in Table 1 have no direct relationship to goodness-of-fit. I understand enumerating the terminal states is impossible for extensive supports, making computing the L1 distance to the target unfeasible. Nonetheless, authors could use the FCS [1] as a proxy. Otherwise, we cannot draw conclusions about sampling correctness in large environments.

* Authors said the additional cost of running BLISS in the final states is negligible. I reckon this should be task-specific. This shouldn't intuitively be negligible if all intermediate states are also final. Please elaborate on the discussion and provide numbers/experimental evaluations. 

* The experimental campaign is relatively short compared to recent works on GFlowNets. 

* While I value the authors' contribution, I believe their contributions and derivations are somewhat straightforward and the work's novelty is limited.

[1] https://openreview.net/forum?id=B8KXmXFiFj

questions:
* It would be nice to see an illustration of the bias authors point to using a uniform target. Then, plotting the marginal over the size of automorphism relations for each sample should highlight this bias.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693492520

thread_type:
official_review

<thread_object>
creator: reviewer_ynfw

<content>
summary:
This paper shows that the so-called problem of equivalent actions (appearing from graph symmetries) biases graph generation processes in GFlowNets. To tackle this, the paper proposes a simple correction procedure that scales the reward function by the number of symmetries of the associated graph (terminal state).  Experiments on artificial data and molecule generation tasks aim to show the effectiveness of the proposed approach.

strengths:
- **Clarity**: Overall, the text is well-written and easy to follow (despite some overloaded notation);
- **Motivation and Relevance**: The motivation is clear, and the problem is relevant as graph (molecule) generation is one of the main applications of GFlowNets;
- **Flexibility**: The proposed correction procedure is flexible as it applies to different training schemes (balance conditions).

weaknesses:
- **Computational cost**: While the paper mentions the additional cost didn't lead to "significant delays in computation", it is not clear why. I believe the paper deserves a more comprehensive discussion about the computational complexity of the proposal. Also, I wonder if the proposed approach becomes prohibitive in some settings.

- **Experiments**: The theoretical analysis does not seem to support the claimed gains on real-world datasets. What are the implications of correctness to top-k diversity/reward? Also, although the paper cites ZINC250K in the Introduction, the experiments only include the QM9 dataset.

- **Technical novelty**: The theoretical contributions of the paper are straightforward. I wonder if the GFlowNet community already knows about the equivalent action problem. 

- **Notation**: I found the notation overloaded, which may confuse readers unfamiliar with GFlowNets. For instance, the paper uses the same $P_F$ to refer to the graph-level, state-level policies, and the marginal distribution over terminal states (i.e., $P_F(x)$). 

- **Limitations**: The paper does not discuss limitations.

questions:
1. Could the authors provide a detailed analysis of the computational complexity of the proposal? Are there environments where the proposed method becomes prohibitive?

2. Could you provide time comparisons for the real-world experiments?

3. The paper says "A reward exponent of 1 is used for the atom-based task, and a value of 16 is used for the fragment-based task". Was this choice based on prior works? If not, could you elaborate on this choice?

4. Is this the first paper to bring attention to the "action equivalent problem"? Could you elaborate on the impact of your findings on previous works that use GFlowNets for graph generation? 

5. I suggest turning Theorem 2 into a Corollary of Theorem 1.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732076144787

thread_type:
public_comment

<thread_object>
creator: ~emmanuel_bengio1

<content>
comment:
Dear authors, dear reviewers,

I hope this comment is taken in the constructive spirit it is intended. This paper is a close analog to our NeurIPS 2023 AI for Science Workshop paper, "[Baking Symmetry into GFlowNets](https://arxiv.org/abs/2406.05426)", by George Ma, myself, Yoshua Bengio, & Dinghuai Zhang. In addition, the code provided in the supplementary is, by extrapolating from the appendix and simple inspection, a derivative of an open source `gflownet` repository; of which I am the main contributor and maintainer; and in which features to correct for so-called equivalent actions (idempotent actions in `gflownet`) were [merged into trunk](https://github.com/recursionpharma/gflownet/pull/42) on Jan 19, 2023. 

To be clear, the core methodological contribution of this paper is to correct flows by counting automorphisms exactly. In contrast, our paper proposes exactly that, as well as the use of positional encoding matching as an efficient alternative to auto/isomorphism testing.

We understand that this could have happened unintentionally, and we appreciate the effort put into this research, which formalizes and tests this issue more thoroughly than our prior work. That being said, we believe our work provides valuable context and prior art for this submission, and we find this situation disappointing. A simple search using relevant keywords would have easily revealed our paper and the associated code.

Thank you for your understanding.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732288461465

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

We provide a detailed comparison with the “Baking Symmetry into GFlowNets” paper.

First, we were pleased to discover that the original authors of GFlowNets had previously recognized this issue and proposed a method to address it—further evidence that this problem is both significant and well-motivated.

---
|  | **“Baking Symmetry into GFlowNets”** | **Ours** |
| --- | --- | --- |
| **Motivation** | Incorporate internal symmetries within the generation process | Train with exact target distribution |
| **Method** | Approximately compute equivalent actions at each transition using positional encodings of a graph. Then, sum their probabilities. | Scale rewards by the order of automorphisms. |
| **Types (Generality)** | Node types  | Node types, edge types, and fragments |
| **Theory** | No theoretical guarantees | Theoretical guarantees on exact learning is provided. The bias without corrections amounts to $\|\text{Aut}(s)\|$. |
| **Experiment** | Synthetic graphs | Both molecules from real-world data and synthetic graphs |
| **Computation** | Multiple positional encoding computations are required at each transition. | Computation of $\|\text{Aut}(s)\|$ is required once for each trajectory. For the approximate method, a summation operation is required over the number of fragments.|
---

The paper “Baking Symmetry into GFlowNets” does not provide sufficient context for full reproduction. However, the open-source `gflownet` code includes the function `get_idempotent_actions` , which we believe serves as an implementation of the paper’s method for exact isomorphism tests. We present the computational cost of performing isomorphism tests using `get_idempotent_actions`, measuring only the compute time of the function call and excluding any preparation code. The QM9 test data (size: 13,389) were used, with trajectories sampled using a uniform backward policy, resulting in an average of 12.72 transitions per trajectory. Although the function must be called for each transition, we report the total cost per trajectory for comparison. Compute times were measured separately for forward and backward actions, and these must be summed to allow a direct comparison with our method.

---
|  | `get_idempotent_actions` (“Baking Symmetry into GFlowNets”) | Our method using *bliss* |
| --- | --- | --- |
| Forward actions | 26.69 ms ± 19.93 | - |
| Backward actions | 4.56 ms ± 5.38 | - |
| Compute time per trajectory | **31.24 ms ± 21.16** | **0.010 ms ± 0.008** |
---

The analysis we provided in our paper for exact isomorphism tests assumes the use of a graph hashing algorithm. The large computational cost of `get_idempotent_actions` arises from its computationally expensive pairwise comparison of actions.

We also observed that the code selectively applies corrections, skipping those for backward equivalent actions when uniform backward policy is used (i.e., when `do_parameterize_p_b` is set to `False`). However, as demonstrated in our paper, equivalent actions should be accounted for regardless of the type of backward policy. We believe that misconceptions on this topic may stem from it being relatively underexplored.

We hope this provides sufficient context for comparisons with prior work. Our method offers an efficient solution to the automorphism problem.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732563800576

thread_type:
public_comment

<thread_object>
creator: ~tiago_silva4

<content>
title:
A question to the authors

comment:
Dear authors, 

I recently stumbled upon your paper and, despite being nicely written, I am struggling to understand why Equation (1) is the right way of correcting the automorphism-induced bias in a GFlowNet. To clarify my point, I will assume that the (state-level) reward function is constant, i.e., the target distribution is uniform and that the graphs are unlabeled. 

In this case, a naive implementation of a GFlowNet (on the graph-level state graph, $(\mathcal{G}, \mathcal{A})$, with unmodified reward function) would sample each state $s$ in proportion to the size of its equivalence class when balanced, namely, $p_{T}(s) \propto |s|$, instead of uniformly at random. Under the flow network epistemology, the reason for this is that the terminal flow associated to the equivalence class $s$ equals the sum of the terminal flows associated to each of its members. Hence, as the authors noticed, the sampling distribution is inherently biased. 

We both disagree, however, in how to eliminate this bias. On the one hand, I would _reduce_ the reward associated to a graph $G$ by a factor of $1/|[G]|$. On the other hand, the authors _increase_ this same reward by a factor of $|\text{Aut}(G)|$ in Equation (1). Nonetheless, in doing so, each state $s$ would be sampled proportionally to $|\text{Aut}(G)| \cdot |s|$, as explained, and the bias would persist. In this regard, it is mostly unclear to me how plugging the equation in Theorem 1 into the TB loss solves the biased distribution problem. Also, it appears to me that the empirical analysis would yield the same conclusions if the scaling factor $|\text{Aut}(G)|$ was replaced by any discrete-valued isomorphism-invariant function. 

Additionally, the text contains seemingly conflicting statements, e.g., "If we allow individual graphs to represent states, the equivalence class of a larger graph will be sampled exponentially more often." and "If we do not scale the reward, we are effectively reducing the rewards for highly symmetric graphs". Does the unmodified GFlowNet tends to sample highly symmetric graphs more or less frequently? 

Importantly, I may have failed to properly understand this work, and I hope my comments above do not disrupt the review process. Nevertheless, I would be happy to understand the issue with my reasoning. :)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732802550177

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear reviewers,

Thank you for your thoughtful and valuable feedback. We have carefully revised our paper to address your comments, aiming to improve clarity and provide a more comprehensive presentation. Notably, we have included thorough comparisons to the workshop paper *“Baking Symmetry into GFlowNets”* to offer readers additional context and to appropriately credit prior work, highlighting how our contributions build upon and extend the existing literature.

We remain confident that our work provides significant contributions to the community by addressing the fundamental issue of equivalent actions in GFlowNets in a rigorous and comprehensive manner and proposing a much more efficient and practical solution.

Below, we summarize the key updates:

- **Comparison with Ma et al. (2024)**: To further position our work in context, we made a detailed comparison to  *“Baking Symmetry into GFlowNets”* paper (Section 2 and Appendix B).
- **Evaluation:** In response to **Reviewer 6mnv**, we included FCS metric (Silva et al., 2024) in our evaluation (Section 6.2).
- **Discussion:** In response to **Reviewer ynfw**, we updated “Discussion and Conclusion” section, which includes limitations as well as implications our work has on prior work in the presence of bias (Section 7).
- **Notation:** We have clarified notations to distinguish between state-level and graph-level policies, eliminating overloaded terms, as suggested by **Reviwer ynfw** (Section 3.2, Section 4.1):
    - $p_\mathcal{S}$, $p_\mathcal{G}$, and $p_\mathcal{S}^{\top}$ now denote the state-level policy, graph-level policy, and marginal state probability, respectively.
    - For backward policy, $q_\mathcal{S}$ and $q_\mathcal{G}$ are used for state-level and graph-level policies, respectively.
    - Forward and backward graph-level actions are now denoted by $\overrightarrow e$ and $\overleftarrow e$, respectively.
- **Experiments:** In response to **Reviewer 6mnv**, we have included additional results using a uniform target distribution to further validate and illustrate our approach (Section 6).
- **Computation:** To **Reviewer ynfw, Rf27**, a new section on computation has been added to the Appendix for completeness (Appendix G).

Overall, we have made several revisions to enhance the paper's readability.

### **Relation to “Baking Symmetry into GFlowNets”**

While the issue of equivalent actions in GFlowNets was identified and partially addressed in the "Baking Symmetry into GFlowNets" paper, the discussion was limited to experimental validation, indicating that equivalent actions could lead to “potentially” incorrect flow probabilities. In contrast, our work provides the first rigorous theoretical foundation for automorphism correction, demonstrating that this issue is not just experimental but a fundamental and systematic challenge tied to graph symmetries, both for atom-based and fragment-based generation. 

Now, given the well-supported motivation for addressing this equivalent actions problem, the key question then becomes: **Is there an efficient solution to resolve this fundamental issue** **of equivalent actions in GFlowNets?**

### **Novel Contributions**

In addition to establishing a theoretical foundation, we propose an efficient solution to this problem. Our method applies the correction only once at the end of trajectories, rather than at every transition within a trajectory. This correction involves computing the order of automorphisms, which we found to be computationally efficient even for the largest molecules in the PubChem dataset. The solution is both simple and easy to implement. Any alternative, more "sophisticated" approaches to this problem would, in essence, amount to approximating the order of automorphisms.

Another notable contribution is the introduction of a novel model likelihood estimator that accounts for equivalent actions. Since the estimation of model likelihood is a fundamental measure for all generative models—essential for evaluating model bias, performance, generalization, and more—this contribution has the potential to significantly influence future research.

### **General Implications**

Our findings carry significant implications, especially given that GFlowNets were initially popularized for their reward-matching capabilities. We emphasize that correcting for automorphisms is a fundamental requirement for unbiased sampling, which is critical for applications like molecule discovery. In this regard, our results highlight the need for future work to explicitly detail the methods used to address equivalent actions, ensuring reproducibility and rigorous evaluation. We believe the problem was identified previously, but its significance and seriousness have been emphasized solely in our paper.

We hope this provides additional clarity regarding the contributions of our work, and we would sincerely appreciate any reconsideration of the assigned score in light of these clarifications.

Regards, 

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
**Summary**: This work brings attention to the equivalent action problem with GFlowNets (an important class of generative models, inspired by reinforcement learning, for discrete structured data/graphs). Specifically, the authors analyse this problem theoretically, demonstrating that failing to account for equivalent actions may introduce systematic bias in the sampling distribution for atom-based as well as fragment-based graph generation.  They further relate this bias to the number of symmetries associated with the graph and propose a corrective automorphism-based reward-scaling approach for unbiased sampling, providing empirical validation of its effectiveness.      

**Strengths**:  The reviewers appreciated several aspects of this work, notably, the motivation and relevance of the problem, clarity of the presentation, sound theoretical analysis, flexibility of the corrective procedure to accommodate different training schemes  (trajectory balance, detailed balance, and their extensions), and the empirical substantiation.  

**Weaknesses**:  The reviewers also raised several concerns and questions pertaining to the additional computational overhead due to the corrective procedure, overloaded notation, lack of discussion on the limitations of the proposed approach, insufficiency of the experiments e.g. in terms of dataset size, fairness of evaluation, metrics not being aligned with goodness-of-fit hindering validation of sampling correctness, and technical novelty.  They also provided several constructive suggestions.   

**Recommendation**
Most reviewers provided detailed, insightful reviews. However, one of the reviewers (mRgy) asked only rather generic questions about generative models and did not participate subsequently in the discussions. Therefore, I decided to not consider their evaluation in my recommendation. 

I commend the authors for their thoughtful response and additional experiments, which were generally appreciated by the reviewers, prompting many of them to upgrade their assessment of this work.  In contrast to the point raised by one of the reviewers, I do not think the technical machinery needs to be sophisticated/straightforward so long as the theory provides meaningful/actionable/clear insights or interpretations, which this work does.  

During the response period, a public comment was posted pointing to a prior work ``Ma, Bengio, Bengio, and Zhang, Baking Symmetry into GFlowNets, NeurIPS 2023 AI for Science Workshop" drawing everyone's attention. This workshop paper had clearly identified the issue of equivalent/idempotent actions in GFlowNet and proposed a closely related method for correcting the flows to address the problem, although without theoretical justification. A part of the code was also leveraged by the authors of the current work. 

To their credit, the authors of this work acknowledged their oversight in failing to appropriately position the multiple contributions of that workshop paper (including it being the first to identify the equivalent action issue and providing valuable context for the current work). 

At least one reviewer felt that despite the authors' response on this particular issue, the current work loses its novelty - as presented in the original submission - significantly. While I think that the theoretical formalism and the insights provided by the authors here form an important contribution in itself, I cannot disregard this concern about novelty.  

Given all the facts and discussion, I believe it would only be fair to the authors of the workshop paper if this paper goes through another review cycle so that the new set of reviewers can make a more informed assessment of the merits of this paper, especially its novelty and repositioning of the contributions with respect to prior work.  I'm therefore not able to recommend a positive decision for this paper at this time. However, I'd like to state that should the program chairs decide to overrule this recommendation, I won't have any (strong) objections. 

I hope the authors take this recommendation in the right spirit (though I understand it might be disappointing for them), and use all the feedback and discussion to make a stronger submission.