Ergodic Generative Flows

Dear HCTd,

We thank you for your detailed review. Before going into detail, allow us to emphasize that the present submission is an "Exercise in Style": how far can we go with "pure" GFN without a separate energy model? We understand it may be restrictive, but this is the game we decided to play.

1. You are right to point to Sendera et al., 2024, which deals with GFN generalization of diffusion models. To the best of my knowledge, Sendera et al. are restricted to reward and energy model objectives. As such, we should cite them together with Zhang et al., 2022.

2. Experimental Designs Or Analyses (1): We agree that DB or TB losses comparison is important, we made the initial choice to restrict ourselves to FM-losses for consistency. Also, the main purpose of the RL section is to check the basics of EGF in a standard GFN setting and confirm the 0-flow problem noted by Brunswic et al. DB or TB losses require changes in our implementation, but we will try to incorporate stability results by the end of the rebuttal discussion.

3. Experimental Designs Or Analyses (2): We fail to understand the comment as Figure 2 provides such a comparison.

4. Tractability of TB. Since TB objective does not enforce flow-matching property (there is no flow, only policies) it's not a "pure" GFN. Furthermore, to the best of our knowledge, TB requires a separate energy model since it has no flow. On the other hand, DB may be in the scope, but the expectancy of the DB loss is the same as the expectancy of the FM loss if the backward policy is chosen according to formulas (6-7), except that DB has a higher variance.

5. The other reviewers also made several suggestions and remarks regarding rewording. Please see our answers to the other reviews for the intended changes.

6. We wanted to emphasize the distinction between "target distribution with density" and "target distribution with samples." We felt that generative modeling in AI covers many practices that intersect both worlds (the work of Sendera et al., for instance), and we struggled to find the right terminology. Although imperfect choices, "Imitation Learning" carries the idea of "examples" while "Reinforcement Learning" is strongly associated with "reward" in the community. We cannot change this choice for the present submission, but we are open to suggestions for terminology that you may provide.

Dear mQ7G,

We thank you for your particularly detailed review. To begin with, we will add clear references to proofs in the appendix.

3. Theoretical Claims:

• Your interpretation of Definition 3.3 is correct. We may move the paragraph after Theorem 3.4 to after Definition 3.3 and expand it to better explain this definition by referring to the mixing properties of the Markov kernel.

We agree that understanding Theorem 3.8 would benefit from a specialization on graphs. Your suggestion is in line with our answer to reviewer 6zNQ regarding neural network approximations. We shall insert a paragraph answering both points in the RL setting.

5. Scientific Literature:

• We apologize for not citing Zhang et al. We will cite it properly.

6. Strengths And Weaknesses:

• Equation 7, you are right; this equation should be introduced more gently. It may be derived from $\pi^*_\leftarrow (x\rightarrow A) = \frac{d (F_{\rightarrow}^*\otimes \pi^*_\rightarrow)(A\rightarrow \cdot)}{d F_{\leftarrow}^*}(x)$ taken from section 3.2 of Brunswic et al. https://arxiv.org/pdf/2312.15246. This equation indeed formalizes a detailed balance condition in the measurable setting. Also, Proposition 1 in the appendix clarifies the link between bi-measures and Markov kernels. We will add a sentence to account for this interpretation and add proof in the appendix. More generally, we will enhance the introduction of Theorems of definitions with comments on their intuition.

7. Questions:

• The main point of Theorem 3.4 is that EGFs reduce the generative problem of finding a scalar field on the state space. Theorems 3.5 and 3.6 provide sufficient conditions for L2-exponential mixing, then an L2 error on $f_{\rightarrow}^*$ translates into an increase in $\delta$ in Corollary 3.9 and an increase of the TV, which is easy to relate to the L1 error of $f_{\rightarrow}^*$.
  We agree that the "expressive enough" formulation is vague and requires a comment.

• Regarding Algorithm 1 and importance densities. The densities were used as is for the results presented in the submission. In later experiments, we added a small constant (1e-2) to stabilize training; however, no ablation studies were conducted to assess its relevance or effectiveness. Flow collapse was the most common problem encountered but seemed more related to too strong a regularization than to collapsing densities. A lower bound on $\int f_{init}$ proved sufficient to force exploration.

• Your comment on the regularization echoes that of moVi. The regularization is inspired by Brunswic et al. regularization; their Theorem states that a decaying regularization may be used to enforce convergence toward an acyclic flow on graphs when using a stable FM loss. From an optimization viewpoint, the regularization pushes the derivative of the loss to be positive in the direction of any 0-flow, hence leading to a stable loss in the sense of definition 3 of Brunswic et al. In our setting, any 0-flow component is expected to vanish (but the theorem mentioned above does not apply to our setting due to their enumerable state space assumption). A side effect is reducing the sampling length. We will enhance the RL section with a discussion of this regularization.

• We conducted comparisons to diffusion models (DDM) and normalizing flows (FFJORD), but we felt that the comparison to Moser flow already carried our point. We will add qualitative comparisons in the appendix on 2D toys on tiny 32x2 models. On the NASA datasets, we already include NLL comparisons to baselines for volcanoes; we will add the results on earthquakes and floods, also the results for Riemannian diffusion, as well as a comment on the relative size of each baseline.

  • The experiments presented in the paper were solved with manual hyperparameter tuning; the only caveats were:
    • too high regularization scale leads to flow collapse
    • too high WeakFM scale tends to make learning the optimum harder especially for small models.

So far, we have only trained EGFs up to dimension 20 on distributions (toy and sparse reward Markov decision tasks). However, we prefer to limit the scope of the present work to proof of concept as hyperparameter tuning proved less straightforward in higher dimensions. - Regarding training time, we did not benchmark the training speed. We shall include a comparison table in the appendix. Our wall clock evaluation led us to conclude that our training for NASA datasets was faster than that of Moser Flow by several orders of magnitude. Diffusion models tend to train slightly faster than EGFs on 2D toys, while FFJORD is significantly slower. However, hyperparameters have been shown to impact convergence speed significantly.

Dear 6zNQ,

We thank you for your detailed review. We acknowledge the need for higher dimensional experiments, it is part of an ongoing project to scaling up EGFs as well as training conditioned EGFs.

Regarding your questions:

1. Unfortunately, they do not directly. A straightforward strategy would be to use Wasserstein bounds together with dequantization.

   • By Theorem 3.8 and its corollaries, the KL-WeakFM loss controls the TV distance between the EGF sampling distribution and the target dequantized empirical distribution.
   • On compact domains W < diameter * TV, the loss then controls the Wasserstein distance between EGF sampling distribution and dequantized empirical target distribution.
   • Finally, use triangular inequality in W distance for the path "dequantized empirical target distribution -> empirical target distribution -> true target distribution." Bounds such as https://arxiv.org/pdf/1707.00087 to estimate the Wasserstein distance of the empirical distance to the target distribution are leveraged at this step.

   Such a computation is straightforward and may be added in appendix. The unbounded case requires extra care and assumption for the second step.

2. Regarding the curse of dimensionality,

   • to begin with, we refer to the answer given to reviewer moVi about the control of the sampling time (short answer yes in a tame way).
   • more to the point, the curse of dimensionality is likely to arise in the estimation of the weakFM loss since we need to enforce a flow-matching condition. We expect the spectral gap to reduce the asymptotic theoretical complexity of the Monte Carlo estimator compared to naive translation moves (non-ergodic), but more theoretical work is needed. A more sophisticated replay buffer is also probably required. We note this question for future theoretical investigations.

3. On neural networks approximations

• Neural networks approximate:

  • the star outflow $f_\rightarrow^*$,
  • the policy softmax part $\alpha_\rightarrow^i$.
  • the translation part of the affine transforms on tori (but with no inputs since they do not depend on the state)

• The answer to your question then depends on whether we are in the RL or IL setting. In the IL setting, Corollary 3.9, together with equation 16, ensures that the KL-WeakFM loss controls the TV error. In the RL setting, we should add a paragraph explaining the following:

  • an L1 flow matching error controls the TV bound, itself bounded the L2 flow matching error controlled by the loss, which in turn depends (in a tractable way) on the trainable models $f_\rightarrow^*$, $\alpha_\rightarrow^i$ and the fixed models $f_{init},f_{term}$. More abstractly, the FM loss we use is an estimator of a twisted L2 distance of the density of the approximated flow bimeasure $F_\rightarrow \otimes \pi_{\rightarrow}$ to the affine subspace of target flow-matching bimeasures, where $F_\rightarrow = (f_\rightarrow^* + f_{term})\lambda + f_{init}\delta_{s_0}$ and $\pi_{\rightarrow}(x) = \frac{ f_\rightarrow^*(x)}{ f_{term}(x) + f_\rightarrow^*(x)} \pi_{\rightarrow}^*(x) + \frac{ f_{term}(x)}{ f_{term}(x) + f_\rightarrow^*(x)}\delta_{s_f}$. This L2 distance controls the L1 distance if we assume the state space to have a finite volume. The L1 distance is $\delta F_{init} (\mathcal S)$. As a result, the stable FM-loss controls the right-hand side of Theorem 3.8.

Thus, we will add our answer to question 1 in the appendix and a paragraph in the main text to answer question 3.

Dear moVi,

We thank you for your detailed review and valuable feedback. Please find our responses to your comments below:

1. Methods and Evaluation Criteria Stable Loss: We appreciate your observation regarding the difference between our stable loss and the one proposed by Brunswic et al. (2024) in equation (19). We would like to clarify that Brunswic et al. introduce a family of stable losses in Section 4.2, denoted by $\Delta_{f,g,\nu}(\alpha, \beta)$. Our stable loss, as given in Eq. (3), is an instance of this family. Since our focus in this work is on the foundational aspects of Ergodic Generative Flows (EGF), we opted for the simplest stable loss, as it seemed more relevant for our goals. Furthermore, we experimented with Brunswic et al.'s loss on non-acyclic GFlowNets and observed little improvement, with additional difficulties in hyperparameter tuning.

Regularization: Brunswic et al. introduced a cycle-killing regularization schedule in Theorem 1 to ensure convergence of non-acyclic flows toward an acyclic flow. Although they do not explicitly suggest using this regularization to stabilize losses in non-acyclic settings, it seemed a natural step for us. We note that a similar approach was also taken by Morosov et al. in their work "Revisiting Non-Acyclic GFlowNets in Discrete Environments."

While we believe the first point does not require further discussion in the main paper, we will include a clarification about our architectural and hyperparameter choices in the appendix. We will also ensure the regularization approach is more clearly explained in the main text. Also there was a typo on the Bengio loss (the square was missing).

2. Experimental Designs or Analyses Purpose of RL Experiment: The RL experiment serves a dual purpose: A) testing the applicability of EGF in the original context of GFlowNets, and B) addressing limitation 4. While limitation 4 is secondary to the primary discussion on non-acyclic GFlowNets, we think it is an interesting point that warrants clarification.

Additional Datasets: We agree that the inclusion of flood and earthquake results would improve the completeness of our experiments. We are currently running updated experiments to incorporate these benchmarks. However, we were unable to locate the original fire dataset due to changes in the NASA repository.

3. Supplementary Material We have corrected the issues in the supplementary material as per your suggestions.

4. Other Comments or Suggestions We will make the necessary corrections and reformulations to the notations and definitions as suggested.

5. Sampling Time and Performance

• Spectral Gap and Sampling Time: The estimation of sampling time can be derived from the spectral gap and the $L_2$ norm of the target distribution. We intend to discuss these relations in a broader follow-up work focused on spectral gap control for EGF. Although this theoretical framework is still in development, the following informal results provide some insight:

  • A spectral gap $\eta>0$ guarantees exponential decay in $L^2$ deviation from the mean: $||(f-\int f)\cdot(\pi^*_\rightarrow)^t|| = O((1-\eta)^t).$

  • To transition from an initial distribution $f_{init}$ to a terminal distribution $f_{term}$ with error $\epsilon$, we require: $||(f_{init}-f_{term})(\pi^*_\rightarrow)^t||\leq \epsilon$.

  • Thus, the required sampling time $\tau$ is expected to scale as: $\tau = O\left(\frac{\log(\epsilon)-\log(||f_{init}-f_{term}|| )}{\log(1-\eta)}\right)$.

  • One could then leverage general lower bounds on spectral gaps such as those of https://arxiv.org/abs/2306.12358 to guarantee a spectral gap scaling as $\eta \sim 1/dim$.

  • Furthermore, the difficulty of the objective is controlled by $||f_{init} - f_{term}||$, the co-dimension $q$ of the manifold supporting the distribution would then leads (with a bit of gaussian dequantization) to $\log||f_{init} - f_{term}|| \sim \log Vol($Ball of radius $r$ in dim $q)$ so that one would end up with $\tau = O\left(-\log(\epsilon)dim+q\log(q)dim\right)$.

  • However, this bound is probably very loose: in the small volume limit, the "effective" spectral gap is bounded away from 0 (we may discuss this in more details if you wish) so the term $q\log(q)dim$ becomes $q\log(q)$.

  • We stress the need for a proper analysis, hence a later submission.

• Trainable Parameters and Sampling Time: the simplest answer is via Theorem 1 as the sampling time is directly controlled by $\int f^*_{\rightarrow}$ and the WeakFM loss via equation 17. The master Theorem 3.4 provides a control: The proof controls the $L^2$ norm of the target $\||f^*_{\rightarrow}-\int f^*_{\rightarrow} \|| = O(\||f_{init}-f_{term}\||\sum_t(1-\eta)^t) = O(\||f_{init}-f_{term}\||\frac{1}{1-\eta})$. One could link some measure of expressivity of neural network to such a bound, we did not do it at this stage as we would like to reserve it for a future work (as mentioned above).