On Divergence Measures for Training GFlowNets

Summary

This paper investigates alternative training methods for Generative Flow Networks (GFlowNets) by evaluating various divergence measures, including Renyi-α, Tsallis-α, reverse, and forward Kullback-Leibler (KL) divergences. Traditional methods focusing on minimizing log-squared differences are shown to lead to biased and high-variance estimators. The authors propose efficient estimators for the stochastic gradients of these divergences and introduce control variates to reduce gradient variance. Empirical results across diverse tasks demonstrate that minimizing these divergence measures significantly accelerates training convergence and enhances stability compared to traditional methods. The paper establishes theoretical connections between GFlowNets and variational inference, extending these insights to arbitrary topological spaces, and highlights the practical effectiveness of control variates in improving training efficiency.

Strengths

- **Comprehensive Evaluation**: The paper thoroughly evaluates multiple divergence measures and their impact on GFlowNet training. - **Theoretical Insights**: Establishing theoretical connections between GFlowNets and VI broadens the understanding of these models. - **Practical Contributions**: The design of control variates to reduce gradient variance is a significant practical contribution that can be applied in various optimization scenarios. - **Variance Reduction Techniques**: The introduction of variance reduction techniques, including control variates and leave-one-out estimators, effectively addresses the high variance issue in gradient estimates, enhancing the learning stability and efficiency.

Weaknesses

- **Theory** The theoretical contribution builds upon previous works [4,5]. However, I do not think the theory itself has enough contribution. Because the Measurable pointed DAG is from previous work [5] and the main theoretical claim, Proposition 1, has limited novelty compared with Proposition 1 in [4]. Therefore, unlike previous works [4,5], they have their novel theoretical contributions and synthetic experiments would suffice. Thus, for this paper, stronger experiments are required, as will be discussed in the next. - **Experiments - Part 1** This paper does not contain enough real-world tasks, such as fragment-based molecule generation [1], graph combinatorial optimization [2], and RNA sequence generation [3] to illustrate its main contribution. These are standard tasks in evaluating GFN performances and are necessary to include. For the only real-world task in the paper, i.e., the BPI task, this paper admits that "not statistically significant" in line 334. Therefore, it is unclear whether the proposed methods to use other divergence measures will be meaningful in real-world scenarios. - **Experiments - Part 2** For the synthetic tasks, the plots in Figure 3 are also missing some lines. For example, why are there only two lines in the **Sets** plots? Also, there is no single best loss function that is uniformly better than the KL baseline. Therefore, additional divergence might seem unnecessary. There should be stronger evidence to support the use of other measures. Or, the authors can provide a guideline on how to choose the best measures given prior information about the tasks. - I would raise my scores if additional experiments are included and the proposed methods are indeed beneficial in more complex tasks. [1] Jain, Moksh, et al. "Biological sequence design with gflownets." International Conference on Machine Learning. PMLR, 2022. [2] Zhang, Dinghuai, et al. "Let the flows tell: Solving graph combinatorial problems with GFlowNets." Advances in Neural Information Processing Systems 36 (2024). [3] Kim, Minsu, et al. "Local search gflownets." arXiv preprint arXiv:2310.02710 (2023). [4] Malkin, Nikolay, et al. "GFlowNets and variational inference." arXiv preprint arXiv:2210.00580 (2022). [5] Lahlou, Salem, et al. "A theory of continuous generative flow networks." International Conference on Machine Learning. PMLR, 2023.

Questions

See **Weaknesses**.

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

See **Weaknesses**.

Summary

GFlowNets are a probabilistic framework for training amortized samplers for high-dimensional compositional spaces. The samplers are typically trained using local consistency objectives which are squared log losses. This paper examines alternatives to these obejctives in the form of general statistical f-divergence measures. The authors consider forward and reverse KL, Renyi-$\alpha$ and Tsallis-$\alpha$ divergences. The authors derive gradients for the divergences in the case of a fixed $p_B$ and on-policy samples. These gradient computations rely on REINFORCE-style estimators and consequently can suffer from high-variance gradient estimates which affect the optimization procedure. For variance reduction, the authors derive control variates for their gradient estimators. Through a series of experiments on a variety of tasks, the authors demonstrate the effectiveness of these divergences for training GFlowNets.

Strengths

* The paper is quite well written and clear. The authors are thorough and clear in introducing the central ideas and provide sufficient details making it easy to follow. * The paper studies the important problem of finding the "right" learning objective in the context of training GFlowNets. Following a long line of work in VI, the authors leverage statistical divergences and propose optimizing them directly to learn the GFN sampler. * Additionally, the authors improve upon prior work drawing a connection between GFNs and VI by proposing principled control variates to reduce the variance in the gradient estimates for the divergences from the REINFORCE estimators. The CVs seem to have a significant effect on the performance. * The empirical analysis covers a variety of problems including continuous and discrete spaces as wells as general DAGs and tree spaces in the case of discrete spaces. * The authors also include code with the submission aiding reproducibility.

Weaknesses

* The paper considers the on-policy setting for training GFlowNets and the proposed learning objectives based on the divergences assume on-policy smaples. However, existing flow-based objectives are all off-policy, and the advantage of GFlowNets on a lot of tasks (specifically challenging tasks with multi-modal target distributions) comes from the ability to train on off-policy trajectories (e.g. replay buffer[1] to local search [2]). So while the proposed learning objectives empirically perform better than TB on tasks where on-policy sampling is enough, it lacks the flexibility of accomodating off-policy training. * The authors also assume that the backward policy $P_B$ is fixed (L153). While this is true in some cases, learning $P_B$ results in significant improvements to the learned sampler[3,4]. As far as I can tell, modifying the proposed objectives to accomodate learning the P_B is non-trivial. * I appreciate the diversity of problems studied by the authors in their experiments, but there are some gaps in the empirical analysis. Specifically, the authors only include a TB baseline and not other objectives such as SubTB, DB which can perform better than TB (and have better training stability) in some cases. Additionally, the paper also does not include the VarGrad-style objective [5] which does away the need for estimating $Z$ in TB. * Moreover, the tasks consider relatively small tasks so it is not clear how scalable the proposed objectives are. [1] Towards Understanding and Improving GFlowNet Training. Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, Tommaso Biancalani. ICML 2023. [2] Local Search GFlowNets. Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, Jinkyoo Park. ICLR 2024. [3] Trajectory balance: Improved credit assignment in GFlowNets. Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, Yoshua Bengio. NeurIPS 2022. [4] A theory of continuous generative flow networks. Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, Nikolay Malkin. ICML 2023. [5] Robust Scheduling with GFlowNets. David W. Zhang, Corrado Rainone, Markus Peschl, Roberto Bondesan. ICLR 2023.

Questions

In addition to the points in Weaknesses: * L135 says TB requires estimating $Z$ but KL doesn't but I am not sure that is correct? Since even in the KL you need the normalizing constant in the $P_B$ term. * L146: [1] would be a more appropriate reference here I think? * L157: Missing reference to [2] * What is the computational performance (in terms of runtime) of the proposed objectives relative to TB? [1] Rubinstein, R. Y. (1981). Simulation and the Monte Carlo Method. In Wiley Series in Probability and Statistics. Wiley. https://doi.org/10.1002/9780470316511 [2] f-Divergence Variational Inference. Neng Wan, Dapeng Li, Naira Hovakimyan. https://arxiv.org/abs/2009.13093.

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors do not explicitly address the limitations of their approach (discussed in the weaknesses) section, though the assumptions are mentioned briefly in Section 2 and 3. The authors only mention the choice of $\alpha$ as a limitation. There is also no discussion of broader impacts (the reference in the checklist is broken too).

Summary

This paper investigates the potential of using a variety of divergence measures directly as training losses for GFlowNets, relying on many of the connections made between GFlowNets and variational inference. Training GFlowNets essentially consists in enforcing balance/flow-matching conditions between a proposal and a target distribution given some data. Given the links between GFlowNets and variational inference, training the GFlowNet to directly minimize the KL divergence was observed to result in biased and high-variance estimators in general. This paper aims to verify the latter claim, for different divergence measures (KL, reverse KL, Tsallis-\$\alpha\$, Renyi-\$\alpha\$). The paper alleviates the afore-mentioned limits by proposing low-variance gradient estimates to that of the latter quantities through control variates. Finally, the paper verifies the proposed methods experimentally through an extensive set of experiments.

Strengths

- The paper widens the scope of the theoretical links between GFlowNets and variational inference beyond the assumption of finitely supported measures. - The authors alleviate the high variance of divergence measures' gradients in practice, relying on control variates. - The paper is very well written and presented and is easy to follow. - The experimental setup is exhaustive, and shows the effect of each of the proposed design choices.

Weaknesses

See questions.

Questions

- How is the diversity of samples impacted? - There should be experiments that show that the different behaviors of each method with respected to the number of discovered modes throughout training (that should be compared to TB too!). For instance, similarly to Figure 4 in (Malkin et al., https://arxiv.org/pdf/2201.13259).

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

4

Limitations

Limitations adequately addressed throughout the paper.

Summary

This paper investigates divergence measures as learning objectives for Generative Flow Networks, which are amortized inference models designed for sampling from unnormalized distributions over composable objects. The authors review four divergence measures - Renyi-\alpha, Tsallis-\alpha, reverse and forward Kullback-Leibler divergences, and design estimators for their stochastic gradients in the context of training enerative Flow Networks. The authors verify that minimizing these divergences yields correct and empirically effective training schemes on several toy environment, and show that it often lead to faster convergence than previously proposed optimization methods. The authors also design control variates based on REINFORCE and score-matching estimators to reduce gradient variance.

Strengths

- The paper provides evaluation of different divergence measures as learning objectives for GFlowNets, showcasing their effectiveness in improving training convergence in several toy environments. - The authors develop control variates for reducing the variance of gradient estimates.

Weaknesses

1. While the paper provides empirical evidence for the effectiveness of divergence-based objectives, a more extensive comparison with traditional GFlowNet training methods across a wider range of datasets and applications would strengthen the claims. 2. The choice of the \alpha parameter in Renyi-\alpha and Tsallis-\alpha divergences is not extensively explored. More insights into the impact of \alpha on the learning dynamics and guidance on selecting an appropriate value would be beneficial. 3. The computational overhead introduced by the control variates and their impact on training time is not explicitly discussed.

Questions

1. How do the proposed divergence-based objectives perform on more complex and high-dimensional datasets commonly used in fields such as drug discovery and natural language processing? 2. Can the control variate techniques be extended to other GFlowNet training objectives beyond the divergence measures considered in this paper? 3. How does the choice of the \alpha parameter affect the learned GFlowNet's ability to capture multi-modal target distributions or discover diverse high-quality samples in combinatorial optimization tasks? 4. Are there any theoretical guarantees or bounds on the sample complexity or convergence rates when using the proposed divergence-based objectives for training GFlowNets?

Rating

4

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes.

Author-Reviewer Discussion Phase (ends Aug 13 midnight AOE)

Dear reviewers, Thank you for your efforts so far! The authors have performed a number of additional experimental evaluations in response to your reviews. Could you follow up with them to ask any clarifying questions that may further inform your opinion on the paper? With thanks, AC

Response to rebuttal

Thanks for the comments. I have a few follow-ups > amenable to a degree of off-policy learning by implementing an importance sampling estimator Indeed one can always use importance sampling to do off-policy training with an on-policy objective (e.g. as noted in [1]) but IS also introduces other challenges (e.g. high variance). So as you note there is certainly a trade-off. However recent work has also illustrated the importance of off-policy samples for training GFlowNets in challenging domains [2,3] which might make the objectives like TB better suited. > we will run additional experiments for divergence-based training of GFlowNets with a learnable $p_B$ Looking forward to the results. > computation overhead incurred by the proposed gradient estimation techniques is negligible The compute overhead is indeed not high, but my comment was more about the scale of the problems used in the paper. Prior work on learning objectives has considered much larger and challenging problems. I see the authors have added the hypergrid and causal DAG tasks but do not provide details about the size of the problems? What size was the hypergrid and the number of variables considered for the causal DAG task? > mathematical definition of TB and KL depends on the constant Z I agree that the Z does not play a role in the optimization but I would note that in a similar fashion TB can also be optimized without the Z (i.e. VarGrad) In the common response the authors claim "broadest experimental evaluation of GFlowNet objectives in the literature" but I think this loses a lot of nuance. While the paper does consider 6 tasks, they are all much smaller (and potentially simpler) than tasks considered in prior work (e.g. much larger molecular optimization and sequence design). This nuance is critical and I hope the authors avoid making sweeping claims. [1] GFlowNets and variational inference, ICLR 2023. [2] Amortizing intractable inference in large language models, ICLR 2024. [3] Improved off-policy training of diffusion samplers, arXiv:2402.05098