Beyond Squared Error: Exploring Loss Design for Enhanced Training of Generative Flow Networks

Authors: Rui Hu, Yifan Zhang, Zhuoran Li, Longbo Huang
Primary area: generative models
Keywords: gflownet, generative models, f-divergence, loss function

Abstract:
Generative Flow Networks (GFlowNets) are a novel class of generative models designed to sample from unnormalized distributions and have found applications in various important tasks, attracting great research interest in their training algorithms. In general, GFlowNets are trained by fitting the forward flow to the backward flow on sampled training objects. Prior work focused on the choice of training objects, parameterizations, sampling and resampling strategies, and backward policies, aiming to enhance credit assignment, exploration, or exploitation of the training process. However, the choice of regression loss, which can highly influence the exploration and exploitation behavior of the under-training policy, has been overlooked. Due to the lack of theoretical understanding for choosing an appropriate regression loss, most existing algorithms train the flow network by minimizing the squared error of the forward and backward flows in log-space, i.e., using the quadratic regression loss. In this work, we rigorously prove that distinct regression losses correspond to specific divergence measures, enabling us to design and analyze regression losses according to the desired properties of the corresponding divergence measures. Specifically, we examine two key properties: zero-forcing and zero-avoiding, where the former promotes exploitation and higher rewards, and the latter encourages exploration and enhances diversity. Based on our theoretical framework, we propose three novel regression losses, namely, Shifted-Cosh, Linex(1/2), and Linex(1). We evaluate them across three benchmarks: hyper-grid, bit-sequence generation, and molecule generation. Our proposed losses are compatible with most existing training algorithms, and significantly improve the performances of the algorithms concerning convergence speed, sample diversity, and robustness.

Decision: accept (spotlight)
Reviewer scores: 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730539630429

thread_type:
official_review

<thread_object>
creator: reviewer_8npt

<content>
summary:
The authors propose to modify the loss function of Gflownet (which has been completely overlooked by prior work). They show that dinstinct losses lead to different divergences. They propose three new loss functions, evaluate them extensively on diverse benchmarks.

strengths:
Contributions:
- generalizing the objective function of gflownet
- derive impact of loss function on the gradient
- define zero-forcing (encourage epxloitation) and zero-avoiding (encourage exploration) as two key properties induced by certain loss functions
- They create 3 new losses (alonside the existing quadratic loss) to tackle all 4 possible combinations (with/without zero-avoiding, with/without zero-forcing)
- Linex(1) corresponds to the KL divergence
- experiments on 3 datasets
	- Non-zero-forcing losses (Linex(1) and Linex(0.5)) converge faster on hyper-grid
	- Linex(1) obtains all the modes almost always the fastest, but spearman corr between train and test is highest for shifted-cos on bit-sequence
	- Linex(1) tends increase diversity while quadratic and shifted-cos give higher quality (high average rewards) on molecule generation

Paper is well written.

weaknesses:
- it would be nice to have a few more choices of losses, taking inspiration let say from f-divergences (Reverse-KL, JSD, etc.)
- there may be more than just zero-forcing and zero-avoiding to the key properties of loss functions hence why studying more losses would be helpful
- it would be nice to let say consider hybrid methods with some kind of annealing. For example, why not use Linex(1) for its fast convergence to a large number of nodes, before then transitioning to shifted-cos for higher rewards around those now-discovered modes.

So to me, the paper is great, but its kind of stopping too quickly, it feels like its only just tapping the surface. These kinds of ideas would be easy to test out and add to the papers. 

If the authors add bit more meat to the paper, i.e. extra loss functions and hybrid-annealing (as discussed above), I would likely increase my score. Like I said, its just missing a little bit of filling to make it a great paper.

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730752584053

thread_type:
official_review

<thread_object>
creator: reviewer_ehxp

<content>
summary:
This paper presents a novel theoretical finding for GFlowNets regarding their objective function. Using f-divergence theories, they connect existing objectives of GFlowNets and show that they are special cases of the squared loss. They design a new loss structure that combines both properties together: (1) zero forcing (as considered in existing losses) and (2) zero avoiding, which compensates for exploration. Their new loss function seems to have empirical benefits.

strengths:
1. The paper is well-written and easy to follow.


2. The theories are insightful.

weaknesses:
1. The empirical results seem to be weak; they are only varied in synthetic tasks

questions:
Areas for improvement and suggestions:

1. Making a literature connection with f-GAN, which also uses f-divergence in GANs, might be insightful to readers.


2. Include a discussion connecting with off-policy exploration methods. Are your loss and off-policy search orthogonal? Which means, is your loss function combined with an off-policy method (e.g., local search) better than the TB loss combined with an off-policy method?


3. It's good to see the categorization of prior GFlowNet works. Can you include this recent work [1] that uses genetic search as an off-policy method for training GFlowNets and provide some discussion?



[1] Hyeonah Kim et al., "Genetic-guided GFlowNets for Sample Efficient Molecular Optimization," NeurIPS 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730948197080

thread_type:
official_review

<thread_object>
creator: reviewer_xe2i

<content>
summary:
This paper presents a novel framework for GFlowNet objective functions, unifying existing training algorithms and clarifying key components. By establishing a connection between objective functions and divergence measures, it offers valuable insights into designing effective training objectives. The authors investigate key regression properties—zero-forcing and zero-avoiding—and propose three new loss functions (Linex(1), Linex(1/2), and Shifted-Cosh) to balance exploration and exploitation. Extensive experiments on benchmarks, including hyper-grid, bit-sequence, and molecule generation, show that these losses outperform the common squared loss in convergence speed, diversity, quality, and robustness.

strengths:
This paper introduces a systematic framework for designing regression losses in GFlowNet training, linking each loss to specific divergence measures for targeted properties. Resulting in three new losses—Shifted-Cosh, Linex(1/2), and Linex(1)—that enhance exploration and exploitation balance.

weaknesses:
Broader exploration of other potential divergence-based losses would offer a more comprehensive understanding of the effects of different divergence properties on GFlowNet training. Although the novelty lies in extending GFlowNet loss functions, there are similar attempts in reinforcement learning and generative models. Although the paper derives theoretical properties of zero-forcing and zero-avoiding, it lacks direct theoretical comparison with existing GFlowNet training algorithms.

questions:
Could the authors clarify if they’ve noticed stability shifts in higher-dimensional or complex tasks and if adjustments might bolster robustness? Additionally, what drives the choice of a limited set of losses—are there theoretical or practical reasons for omitting other f-divergences, like Hellinger? Lastly, insights into each loss function’s hyperparameter sensitivity and effects on convergence guarantees would further clarify their resilience and adaptability.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732709584261

thread_type:
public_comment

<thread_object>
creator: ~eliezer_de_souza_da_silva1

<content>
title:
Some relevant references

comment:
We appreciated the authors' contributions to this domain and wanted to highlight some missing relevant references.

In our recent work, published at NeurIPS 2024, we investigated the properties of $\alpha$-divergences (including forward and reverse Kullback-Leibler (KL), R\'enyi-$\alpha$, and Tsallis-$\alpha$ divergences) in the context of training GFlowNets. In particular, we also presented the trade-offs between zero-forcing and zero-avoiding behaviors for this family of learning objectives.  

While this paper approaches the problem differently, we believe that acknowledging our work would provide additional context for understanding the theoretical and practical aspects of GFlowNet training.  

Additionally, we noticed that Heiko Zimmermann et al.'s work exploring the relationship between GFlowNets and Variational Inference (VI), appears to be missing from the related discussion. Including this reference would enrich the paper's perspective and be appropriate for a complete overview of recent works.

Thank you for considering our feedback. We hope this fosters a richer discussion and continued progress in the field! 

References:
1. Tiago Silva, Eliezer de Souza da Silva, and Diego Mesquita. *On Divergence Measures for Training GFlowNets.* NeurIPS, 2024. [Link](https://openreview.net/forum?id=N5H4z0Pzvn)  
2. Heiko Zimmermann et al. *A Variational Perspective on Generative Flow Networks.* Transactions on Machine Learning Research, 2023. [Link](https://openreview.net/forum?id=AZ4GobeSLq)
</content>
</thread_object>

</thread>


Metareview:
metareview:
In the paper, the authors addressed the issue of choosing regression loss in Generative Flow Networks (GFlowNets). They demonstrated a connection between regression losses and specific divergence measures. This connection enables the systematic design and evaluation of regression losses, tailored to the unique properties of their corresponding divergence measures. 

All the reviewers agree that the theoretical results are novel and insightful. The presentation of the paper is clear and easy to follow. After the rebuttal, most of the remaining concerns of the reviewers are addressed and all the reviewers are happy with the current stage of the paper. 

While there are some concerns about the limited experimental results, I believe that the current novelty and contribution of the paper are sufficient for ICLR. Therefore, I recommend accepting the paper at the current stage. The authors are encouraged to incorporate the feedbacks and comments of the reviewers into the camera-ready version of their paper.