Generalization and Distributed Learning of GFlowNets

Thank you for attentively reviewing and appreciating our work. If our answers below appropriately address your concerns, please consider raising your score. Otherwise, we will be more than happy to discuss our contributions in further detail.

The main idea of SAL (i.e., divide and conquer) is novel for GFlowNet. However, I wonder whether the aggregation accuracy will be compromised for large networks. In other words, is there a trade-off in the theoretical results about $m$?

Thank you for the interest. By Theorem C.1 and Proposition C.5, a GFlowNet learned via SAL is accurate when its corresponding leaf models are adequately learned. In the limit, when each learning objective is globally minimized, Theorem C.1 ensures that the resulting GFlowNet samples from the right distribution — irrespective of $m$.

In this context, for large $m$, two practical problems might arise. Statistically, a larger $m$ could increase the likelihood of catastrophic error propagation from an insufficiently trained leaf GFlowNet to the global GFlowNet. However, we did not observe this issue throughout our experiments. Computationally, enlarging $m$ would lead to a potential memory overhead when carrying out the aggregation phase. Nonetheless, most applications of GFlowNets are built upon small-scale models such as MLPs [1, 2] and GNNs [3, 4]. Hence, the memory cost of SAL is typically not a critical problem.

[1] Deleu et al. Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network. NeurIPS 2023.

[2] Pan et al. Generative Augmented Flow Networks. ICLR 2023.

[3] Bengio et al. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. NeurIPS 2021

[4] Zhang et al. Let the Flows Tell: Solving Graph Combinatorial Optimization Problems with GFlowNets. NeurIPS 2023.

The time cost of the Fixed-horizon partition does not appear to be addressed. Does this operation impact model efficiency? More precisely, how can a Fixed-horizon partition be found efficiently?

Thank you for the question. In practice, the FHP is only implicitly defined by an assignment function and need not be explicitly constructed during either training or inference. Also, the FHP is problem-dependent and fixed prior to training. As a consequence, there is no additional computational overhead related to the construction of the FHP.

We thank you again for the review. We hope that our answers elevate your appraisal of our work. Should any clarifications be needed, we would be glad to discuss our contributions in more detail.

5. Are there empirical limits which have been identified in the experiments regarding the effectiveness of SAL as the complexity of the underlying GFlowNet increases?

As mentioned, our experiments in Sections 6 and C.1 ran smoothly. As long as a FHP can be efficiently implemented, there does not seem to be clear computational constraints for SAL — independently of the complexity of the underlying GFlowNet. Currently, indeed, the main empirical limitation of SAL is its restriction to the problems described in Section C.1, namely, autoregressive, set-structured, and hypergrid navigation tasks. For specialized domains (e..g, drug discovery), devising effective FHPs is still an open challenge — worthy of a work of its own.

6. How do the authors envision adapting the partitioning scheme to be more flexible or dynamic as the policy evolves during training?

This is a tricky and compelling issue. Fundamentally, SAL is built upon a static FHP; it is unclear whether the theoretical results outlined in Section C.1 would be preserved for dynamically partitioned state graphs. We think, however, that a meta-learning approach would be ideal for addressing this issue by amortizing the learning of the partition-wise flow functions. Instead of learning a flow assignment for each component of a given FHP, we would learn a meta-model that receives a FHP as input and returns a flow assignment for each of the subgraphs. In doing so, we could harmlessly update the FHP as the policy evolves during training. Nonetheless, we are uncertain about the best way to define and train such a meta-model. A natural question is: how can such meta-learning be performed in a distributed or embarrassingly parallel fashion? A clear answer is currently out of reach, being an interesting venue to be investigated by future endeavors.

As discussed in Section E, a related and simpler problem consists of finding an optimal FHP for a given problem and utility function defining the goodness of a specific partition. In essence, this problem may be reduced to NAS: in both cases, we are searching for a discrete and complex structure that will enhance the training efficiency of the resulting model. Although there are standard strategies for approximately solving these problems, e.g., Bayesian optimization [1], they remain a largely unresolved issue in the machine learning literature.

[1] Kandasamy et al. Neural architecture search with bayesian optimisation and optimal transport. NeurIPS 2018.

Again, many thanks for the thoughtful review and for the interesting questions. We would be delighted to engage in further discussion if needed.

Thank you for the thoughtful feedback and for appreciating our work. We hope our answers resolve your concerns satisfactorily. If not, we would be happy to continue the discussion.

Weaknesses

1. While the paper presents experimental results for both synthetic and real-world tasks, the scope of these experiments might be limited in terms of diversity of datasets, complexity of environments, and breadth of tasks.

Thank you for the suggestions. While we understand that these experiments could enhance the impact of our work, it unfortunately will not be possible to execute them during the short discussion period. Indeed, as combinatorial optimization was separately addressed by [1, 2] in the context of GFlowNets, we believe that both applications (ii) and (iii) are worthy of a work of their own. We focused our empirical analysis on common benchmark tasks in the literature.

For NAS, in particular, each leaf GFlowNet in SAL would only have to learn the final layers of the trained neural network. Nonetheless, defining an iterative generative process in the space of back-propagable neural network architectures is a challenging and unresolved issue in the literature. As a consequence, it is mostly unclear how to implement a GFlowNet for addressing this problem.

To strengthen our empirical evidence, however, we have included a large-scale experiment in Section E.1 (page 42) for generating $50$-sized subsets of $\\{1, … 100\\}$. In consistency with our prior results, SAL improves upon a centralized approach in terms of both mode discovery rate and distributional accuracy of the trained model.

[1] Zhang et al. Let the Flows Tell: Solving Graph Combinatorial Optimization Problems with GFlowNets. NeurIPS 2023.

[2] Zhang et al. Robust Scheduling with GFlowNets. ICLR 2023.

2. The SAL approach, while promising for distributed training, may have significant computational and memory overhead due to the need to train multiple GFlowNets in parallel and then aggregate their results. This could be a limiting factor in resource-constrained environments or for applications where computational efficiency is critical.

Regarding training time, we would like to highlight that SAL is always faster than centralized learning, as shown in Section C.1. In this case, the main potential bottleneck in SAL is the memory required for the aggregation step — which entails loading all leaf GFlowNets into a centralized server.

In spite of this, the leading applications of GFlowNets utilize small-scale models, such as MLPs [1, 2] and GNNs [3, 4], to parameterize the policy networks. As a consequence, the memory overhead will be negligible for most practitioners. However, for applications such as NLP with LLM-based policy networks, we acknowledge that lack of resources might be an issue.

[1] Deleu et al. Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network. NeurIPS 2023.

[2] Pan et al. Generative Augmented Flow Networks. ICLR 2023.

[3] Bengio et al. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. NeurIPS 2021.

[4] Zhang et al. Let the Flows Tell: Solving Graph Combinatorial Optimization Problems with GFlowNets. NeurIPS 2023.

Thank you for the detailed review and for the suggestions to enhance the presentation of our work. If you believe our answers satisfactorily address your concerns and clarify our contributions, please consider raising your score. Otherwise, we would be glad to discuss our results in further details.

Weaknesses

A full section is devoted to describing when GFN do not generalize.

We would like to clarify that the purpose of Section 4 is to argue that, depending on the exploratory distribution, learning a generalizable policy network is not possible even after observing an arbitrarily large portion of the state graph. In particular, Proposition 4.2 demonstrates that a small observed loss is not enough to ensure that the learned distribution is similar to the target. From these results, we conclude that a statistical guarantee for GFlowNets should be fine-tuned to the specific training data distribution, which is the approach we pursue in Section 5.1.

motivating corner cases are not discussed

As mentioned above, we regard Section 4 as a motivating corner case for the technical results presented in Section 5. However, if you have further suggestions, we would gladly consider additional edge cases and discuss them in our work.

The SAL algorithm does not need the bounds to be motivated, as divide and conquer is a very common strategy in CS in general.

Thank you for the chance to underline our motivation for SAL.

While we agree that divide-and-conquer algorithms are common in computer science, this fact alone is not enough to justify the significant superior performance of SAL relatively to a centralized GFlowNet. By taking a PAC-Bayesian stance, we directly address an issue at the core of machine learning algorithms, namely, generalization. From a scientific standpoint, Sections 5.2 and 6 respectively ask (i) what makes generalization difficult? and (ii) how to facilitate the learning of a generalizable policy network?.

In (i), we demonstrated via Theorems 5.2 and 5.4 that the trajectory length (represented by $\log t_{m}$) and the sample size ($n$) play an important role in the generalization of GFlowNets. In this sense, we introduced SAL as an actionable answer to question (ii) in view of (i): both the reduced $t_{m}$ (through the FHP) and the increased $n$ (by employing a distributed scheme) aim at facilitating the learning of a generalizable policy network.

We feel the first four pages are too wordy.

The recursive SAL algorithm (which we find of interest) is only introduced in the appendix.

We appreciate your feedback. We integrated the discussion of recursive SAL into the main text to better emphasize our work’s contributions (page 10). Also, we shortened Section 3 to improve the paper’s readability.

Section 4 is unclear; the set generation problem is barely understandable.

Some mathematical results are strangely placed. For instance, the purpose of Lemma 4.1 is unclear in the section context

Thank you for providing the chance to elaborate on Section 4.

In short, the set generation task consists of sampling fixed-size subsets of a given set in proportion to a specific reward function; it is a standard benchmark in the GFlowNet literature (see, e.g., Bengio et al. 2023, Section 5.1).

To generate them, we start from an initially empty set ($s_{o}$) and iteratively add elements to it until it achieves a prescribed size. We have included both an illustration of the corresponding state graph in Figure 5 on page 20 of the document and an updated (informal) description of the problem in the main text (Section 4).

In this sense, Lemma 4.1 and Proposition 4.2 show that minimizing the empirical loss on an arbitrarily large portion of the state graph might not be enough to ensure that the model learns a generalizable policy network for the set generation task. As mentioned earlier, this demonstrates that the construction of non-vacuous statistical guarantees for GFlowNets should take into account the specific trajectories observed during training rather than merely their quantity, which motivates our empirical analysis in Section 5.1.

Lemma D.1 is both trivial and stated after it is used.

We agree that Lemma D.1 represents a simple and well-known fact. We have included it only to ensure that our proof is both comprehensive and complete; that is the reason it was deferred to the appendix.

Questions

1. Could you provide insight and/or examples illustrating the meaning of the bounds you introduce?

We are happy to elaborate further on the insights from our bounds.

Insight 1: Small empirical loss does not imply generalization. Proposition 4.2 shows that, depending on the trajectory distribution, a small empirical loss does not ensure that the GFlowNet accurately approximates its target. Contrapositively, this entails that a meaningful statistical guarantee can only be obtained through data-dependent bounds, as demonstrated in Section 5.1.

Insight 2: Generalization poses a trade-off between trajectory length and model size. Theorems 5.2 and 5.4 present a trade-off between the description length (measured by the state graph’s diameter, $t_{m}$) and the dimensionality of the policy network’s parameter space. For instance, when generating sequences of size $2$, we may either start from an empty object and add one element at a time ($t_{m} = 2$) or simultaneously add both elements ($t_{m} = 1$). For $t_{m}=1$, we would need bigger neural networks representing larger action spaces in their output, which would negatively affect our generalization bounds through the KL term. On the other hand, setting $t_{m} = 2$ allows using smaller models ($\downarrow \mathrm{KL}$) at the cost of an increased description length ($\uparrow \log t_{m}$). However, we acknowledge that striking a perfect balance between these two aspects remains open, and we believe it to be a fruitful venue for future works.

2. I am confused about the link between the FCS and TB losses. Could the author clarify this link and how they differ? Is FCS loss used for training? Is it computable as a metric during training?

Thank you for the question. From a fundamental viewpoint, when the distributional approximation of a GFlowNet is perfect, both the FCS and TB vanish.

However, we note that FCS is not an appropriate learning objective for GFlowNets. First, computing FCS requires evaluating the marginal probability $p_{\top}(x)$ of a terminal state $x$, which makes it costlier than TB. Second, in contrast to TB, FCS is not in the log-domain; its optimization could be numerically unstable.

On the other hand, FCS is bounded and perfectly fits in as a risk functional in a PAC-Bayesian analysis. Also, it was recently shown to be an appropriate tool for diagnosing the distributional accuracy of GFlowNets [1]. We have clarified this in the updated PDF (Section 3).

[1] Silva et al. Analyzing GFlowNets: Stability, Expressiveness, and Assessment. SPIGM workshop @ ICML 2024.

3. Are you able to provide deeper insight into the structural properties of the graphs or into the reward leading to lower generalization, deeper than the dimensionality curse?

The “curse of dimensionality” (CoD) typically refers to the difficulty of solving a problem with a prescribed error depending exponentially on the dimension of the problem’s state space [1, 2]. In this sense, we are unaware of a deeper mathematical connection between the CoD and our results. In a concurrent study, for instance, Boussif et al. [3] showed that a reduction on the trajectory length of the state graph through the introduction of abstract actions improved the learning convergence of GFlowNets, in agreement with our generalization bounds. In doing so, however, the state space’s “dimension” remains unaffected — implying no direct association between the trajectory’s length and the problem’s “dimension”.

[1] Donoho et al. High-Dimensional Data Analysis: The Curses and Blessings of Dimensionality. AMS Math Challenges Lecture, 2000.

[2] John Rust. Using Randomization to Break the Curse of Dimensionality. Econometrica, 1997.

[3] Boussif et al. Action abstractions for amortized sampling. https://arxiv.org/abs/2410.15184

4. How could one make your bounds actionable or act as a guardrail ?

As demonstrated in Proposition 5.1, our empirical bounds could be used as a guardrail for provably controlling the distributional accuracy of a trained GFlowNet with high probability. As a practical safeguard, if a practitioner deems the bound to be above his tolerance threshold, he could consider further training the GFlowNet before deployment.

We appreciate the reviewer’s thoughtful questions. We would be happy to engage further, should additional clarifications be needed.

1." We would like to clarify that the purpose of Section 4[..]" Understood, however, I fail to see how this is specific to GFNs: the reward may be located on a arbitrarily small fraction of the states and any algorithm would have similar difficulties. I agree with the conclusion of the need to fall-back to data-dependent bounds. 2. "As mentioned above, we regard Section 4 as a motivating corner". Let me clarify: Section 4 motivates data-dependent bounds, ok. My point is however different: when I have a theoretical bound, I want to know what it means. Examples are important to this end as they show when the bound is sharp, when some term dominate the other, when some part are big or small. This is essential to explaining such bounds. These examples should be provided at the very least in appendix, and its not my role to give them.

3. "In (i), we demonstrated via Theorems 5.2 and 5.4 that the trajectory length (represented by ) and the sample size () play an important role in the generalization of GFlowNets". My point is that, such a statement falls back to your section 4 example: say that the graph is a binary tree of depth D with reward on the leaves. Then, obviously, the length of the paths (which is D) have a significant impact on the ability to generalize (or even learn anything) as the number of relevant states (the leaves) grows as 2^D. The bound then says nothing more than "long paths imply high combinatorics hence hard problem".

Such an information does not help to understand generalization of GFN and is not specific to GFN

4. "it is a standard benchmark in the GFlowNet literature" I know this is a common benchmark, my point was on the understandability of the way it is written in the present work.

Can the authors provide more details on the way the Fixed-horizon DAG partition (definition 6.1) should be constructed?

Thank you for your question. In practice, a FHP is only implicitly defined via an assignment function (AF) that separates the states into (possibly overlapping) subgraphs. In doing so, we do not need to explicitly construct a FHP. This AF is defined via either a hash- or ranking-based indexation of the state space. The former is standard in, e.g., PostgreSQL-based databases, and ensures that the state graph is approximately homogeneously divided between each component of the partition. For a lengthy discussion on how to construct a FHP, we refer the reviewer to Section C.1 in the appendix.

May be there are some heuristics or theoretically justified approached to it?

You correctly noted that a poorly designed FHP would not be beneficial for the training of GFlowNets. In Section E.1, we have included additional experiments illustrating this point in extreme and fabricated cases. However, similar observations could be made regarding an inappropriately chosen exploratory policy (as examined in Section 4) or loss function in conventional GFlowNets. In this sense, the efficacy of a particular FHP is ultimately an empirical problem that should be experimentally investigated.

In our work, our partitioning strategy described in Section C.1 was based on the heuristic of allocating similarly sized subgraphs to each of FHP’s components, which was demonstrably effective in practice (Section 6.2). In particular, we did not observe either mode collapse or insufficient sample diversity throughout our experimental campaign. Intuitively, the main reason for this is that SAL greatly increases the exploratory potential of a standard GFlowNet by distributing the task of traversing the state graph. As a consequence, issues such as mode collapse are far less likely under SAL than under a centralized GFlowNet.

That said, although our heuristics for constructing a FHP significantly enhanced GFlowNet training, we agree that the design of theoretically optimal FHPs is an interesting direction for future works — as stated in Section E.

We again thank the reviewer for their feedback. We would be happy to further discuss our contributions if needed.

From the experimental point of view, I would also say that the proposed experimental investigation of the Subgraph Asynchronous Learning algorithm is insufficient. All considered environments can be considered as rather toyish, with actual number of terminal states less than $10^{10}$.

To the best of our knowledge, our experimental campaign is on par with the wider GFlowNet literature in terms of diversity of the generative tasks [1, 2, 3, 4, 5], encompassing six different problems that are often used to benchmark GFlowNets. As we demonstrated in the main text and in the supplement, SAL significantly improved upon a centralized approach for all the considered experiments. We believe this provides strong evidence for the effectiveness of our distributed algorithm.

As we acknowledged in Sections 7 and E, however, experiments on specialized domains such as drug discovery (molecule generation) comprise a promising venue for the evaluation of SAL. Nonetheless, designing a fixed-horizon partition for these environments would require the development of expert techniques that deserve a work of their own. In our paper, we developed strategies for devising a FHP in the context of set generation, sequence design, and hypergrid navigation (please refer to Section C.1).

To ensure that SAL scales beyond the size of the considered generative tasks, we executed additional experiments on the generation of $50$-sized subsets of a $100$-sized fixed set (~$10^{30}$ elements) in Section E.1 (page 42). As we can see, a SAL-trained GFlowNet also achieves a better distributional approximation and coverage of the state graph than its centralized counterpart in this relatively large domain.

[1] Malkin et al. GFlowNets and Variational Inference. ICLR 2023.

[2] Malkin et al. Trajectory balance: Improved credit assignment in GFlowNets. NeurIPS 2022.

[3] Pan et al. Better Training of GFlowNets with Local Credit and Incomplete Trajectories. ICML 2023.

[4] Silva et al. On Divergence Measures for Training GFlowNets. NeurIPS 2024.

[5] Shen et al. Toward Understanding and Improving GFlowNet training. ICML 2023.

Although minor, I would also highlight some stylistic nuances in the exposition. For instance, the authors describe their own results in Theorem 5.2 as 'insightful' (lines 204 and 2130). It is uncommon in the literature to use such language when referring to one's own findings.

Thank you for your suggestion. By “insightful”, we intended to mean that the derived bound carried useful information regarding the generalization properties of GFlowNets. To avoid a potential overstatement, we have replaced “insightful” by “informative”. We are open to further feedback.

Similarly, the statistical guarantees presented in the current submission, are often referred to as "non-vacuous" bounds.

We are thankful for your detailed feedback. However, we note that “non-vacuous” is a standard terminology in the PAC-Bayesian literature [1, 2]. In a nutshell, a statistical guarantee is vacuous if it is trivially true. In the context of Proposition 5.1, for example, we would call a high-probability bound of the form $L_{FCS} \le 1.8$ vacuous, as $L_{FCS} \le 1$ by design. By stating the verifiable non-vacuousness of our results in Section 5.1, we mean that the chosen risk functional ($L_{FCS}$) is non-trivially upper bounded with high probability; this is shown in Figure 2. Having said that, we understand that the term may appear an overstatement. We added a clarification on its context and meaning in Section 2.

[1] Pierre Alquier. User-friendly introduction to PAC-Bayes bounds. Foundations and Trends® in Machine Learning, 2024.

[2] Lotfi et al. Non-Vacuous Generalization Bounds for Large Language Models. ICML 2024.

Minor typos.

Thank you for catching these typos. We corrected them in the updated manuscript.

1. Proposition 5.1 - $\tau_{\alpha}$ is undefined, only $\tau_{1−\alpha}$ previously appeared;
2. line 979-980 - [ref] undefined;

Briefly, $\mathcal{T}\_{\alpha}$ is the complement of $\mathcal{T}\_{1 - \alpha}$ and [ref] refers to Boucheron et al. (2013).

Questions

Is it possible to provide larger-scale experiments?

As we mentioned, it will not be possible to run large-scale drug discovery experiments during the discussion period. Nonetheless, our results on the larger set generation tasks (see above) indicate that SAL can be scaled to bigger state graphs. We also stress that the considered benchmark tasks are of potentially practical interest and standard in the evaluation of novel GFlowNet learning algorithms.

Thank you for carefully reviewing our paper and for acknowledging its novelty. We hope our answers below elevate your appraisal of our work. If so, please consider raising your score. Otherwise, we will be happy to engage further.

Weaknesses

From the theoretical point of view, theoretical contribution of the paper is rather incremental, as the main theoretical findings of the paper are mostly corollaries from known results in PAC-Bayes literature

Thank you for the chance to highlight the novelty of our contributions. From a technical standpoint, Theorem 5.2 introduced a strategy for bounding the GFlowNet’s policy network away from zero and a corresponding upper-bound on the KL divergence learning objective. On the other hand, in Theorem 5.4, we developed an Azuma-Hoeffding-style inequality for independently distributed martingales that, to the best of our knowledge, was not previously established in the literature.

Although these results may not represent a major advance in the PAC-Bayesian realm, we would like to emphasize that our focus was on shedding light on the generalization properties of GFlowNets through the lens of PAC-Bayesian analysis. For this, we provided both empirically tight guarantees and a clearer understanding of GFlowNet generalization through oracle bounds. This is our primary theoretical contribution, which touches on an issue that has been of large interest in the GFlowNet literature [1, 2].

On a side note, we also underline that similar studies were carried out for, e.g., diffusion probabilistic models [3], GANs [4], and LLMs [5, 6], with the sole objective of better understanding the generalization behavior of the corresponding models. Although these works might be seen as incremental from the perspective of the PAC-Bayesian literature, they are relevant and important for their respective communities.

[1] Bengio et al. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. NeurIPS 2021.

[2] Shen et al. Towards Understanding and Improving GFlowNet training.ICML 2023.

[3] Li et al. On the Generalization Properties of Diffusion Models. NeurIPS 2023.

[4] Mbacke et al. PAC-Bayesian Generalization Bounds for Adversarial Generative Models. ICML 2023.

[5] Lotfi et al. Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models. NeurIPS 2024.

[6] Lotfi et al. Non-Vacuous Generalization Bounds for Large Language Models. ICML 2024.

Moreover, the supervised setting for training GFlowNets is questionable and is of limited interest, since its implementation requires a sampler from the target distribution.

Thank you for the opportunity to clarify our assumptions. We start by noticing that none of our experiments assume that samples from the target distribution are available. Also, except for Theorem 5.2, our theoretical results are not restricted to any data distribution in particular.

In Section 5.1, we (perhaps unfortunately) used the term “supervised learning” to illustrate the shift from the traditional GFlowNet training paradigm into a supervised learning problem, which was also considered by [1], for a better fit in the conventional PAC-Bayesian ontology. This was only a didactic artifact. In fact, the only supervision comes from the reward function. We integrated this discussion into the updated manuscript.

[1] Atanackovic et al. Investigating Generalization Behaviours of Generative Flow Networks. SPIGM workshop at ICML 2024.