Theoretical guarantees in KL for Diffusion Flow Matching

Flow Matching (FM) (also referred to as stochastic interpolants or rectified flows) stands out as a class of generative models that aims to bridge in finite time the target distribution $\nu^\star$ with an auxiliary distribution $\mu$, leveraging a fixed coupling $\pi$ and a bridge which can either be deterministic or stochastic. These two ingredients define a path measure which can then be approximated by learning the drift of its Markovian projection. The main contribution of this paper is to provide relatively mild assumptions on $\nu^\star$, $\mu$ and $\pi$ to obtain non-asymptotics guarantees for Diffusion Flow Matching (DFM) models using as bridge the conditional distribution associated with the Brownian motion. More precisely, we establish bounds on the Kullback-Leibler divergence between the target distribution and the one generated by such DFM models under moment conditions on the score of $\nu^\star$, $\mu$ and $\pi$, and a standard $L^2$-drift-approximation error assumption.

Paper

References (48)

Scroll for more · 36 remaining

Similar papers

Peer review

Reviewer igTf7/10 · confidence 3/52024-06-20

Summary

This paper proposes theoretical guarantees for diffusion flow matching, which is of the most popular generative models. Specifically, it provides an upper bound of the KL divergence between the target distribution and the one learned by the DFM. This work promotes the theoretical understanding of the discretized flow matching model.

Strengths

1. The paper is clearly written, and the proofs seem to be sound. 2. The authors consider the early-stopping case and remove the Lipschitz assumption of the velocity field.

Weaknesses

1. Assumption H1 is quite strong that requires both the source and target distributions to have finite 8-th moments. 2. The procedure is standard compared to some theoretical work on SGM. For example, the authors assume an L2-drift-approximation error and use Girsanov theorem for bounding the KL divergence.

Questions

1. If you in addition assume the velocity field is Lipschitz, will you be able to relax assumption H1 to have finite second moments?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Reviewer 83K85/10 · confidence 3/52024-07-12

Summary

This paper gives theoretical guarantees for Diffusion Flow Matching similar in spirit to existing results for Denoising Diffusion Probabilistic Models and Probability Flow ODE. This work requires weaker assumptions than prior work, replacing the Lipschitzness of the score function with a relative Fisher information condition.

Strengths

This paper weakens this required assumptions of prior work, and is the first to develop guarantees for diffusion flow matching, instead of, for example, an analysis of DDPM or probability flow ODE. Unlike these other works, the authors further exploit properties of the heat kernel that appears in the estimation procedure. Their results also easily extend to the early-stopping setting.

Weaknesses

One could make the argument that the present ideas are not as novel as its predecessors (e.g., Conforti et al. (2023) for the finite-fisher information criteria, or the line of works similar to Chen et al. (2023) in general). Moreover, the dimension-dependence in these results is quite poor (though the authors do mention this as an avenue for future work).

Questions

- The authors claim to tackle "all sources of error" in the abstract, though I would argue the statistical convergence rates remain open. - What is the primary obstacle in improving the dimension dependence? Do you conjecture that it should be the same for DDPM? Some typos I managed to catch and write down: - L28: "are important... " or "were an important" - L34: I don't know what the ":" is there - L63: "In the case where..." maybe? - L73: Reference to LWYql23 --- might be worth fixing the authors name in the bibliography - L77: "... an approximate sample ..." - L144: "... a solution of the ..." - L164: "... is therefore the solution ..." [the subsequent page has a few more instances of missing articles that are worth fixing] - L170: "through" (not trough) - L173: Fix the [Kry] citation (no year) - L190: "...gives an ideal..."

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

N/A

Reviewer TMKS8/10 · confidence 2/52024-07-14

Summary

This work provides theoretical guarantees for diffusion flow matching (DFM) models, which are a recent class of generative models similar to score-based generative models (SGM). Extensive background is given in sections 1 and 2, and section 3 contains the results, namely, bounds on the KL divergence from the target distribution to the distribution outputted by DFM. A detailed comparison with existing results is presented in section 3.2.

Strengths

The presentation of the whole paper is remarkably clear, notably the background explanations of section 2 are very clear and very welcome. Regarding the guarantees of section 3, they are (claimed to be) the first ones for DFMs, which are well-motivated as alternatives to SGMs and deterministic FMs; this makes these results quite valuable.

Weaknesses

(I can only comment on the presentation and context as I am not familiar with the guarantees for SGMs and FMs, sorry.) - Two things remained unclear to me after reading section 2; see "Questions". - The discussion of the previous work on SGMs seems a bit difficult to compare, since this paper is about DFMs, which (if I understand) are a different thing although similar.

Questions

- Are SGMs a sub-case of DFMs, for a well-chosen coupling and bridge, and perhaps up to a change of time variable? - Is the Markovian projection of an Ito process always unique, and/or is the Markovian projection of (3) unique? - How do DFMs compare to SGMs and deterministic FMs in practice, say in generalization, speed, difficulty of implementation, robustness w.r.t. hyperparameters, ...? - Are there settings where it makes sense to want to choose something else than $\pi$ being the independent coupling and the bridge being Brownian bridge?

Rating

8

Confidence

2

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors have adequately addressed the limitations.

Reviewer NqRW7/10 · confidence 3/52024-07-15

Summary

The paper derives theoretical guarantees for a flow matching procedure which constructs a diffusion-like process which allows for samples obtained from a source distribution $\mu$ to a target distribution $\nu^*$. The goal is to build a suitable simple sampling procedure where a source sample is updated through a series of relatively simple Markovian procedure to an approximate sample from $\nu^*$. In the typical setting of diffusion models, this takes the form of a score function augmented with additional Gaussian noise. However, this suffers from two main drawbacks -- firstly, the choice of base distribution is fixed (the base function is always Gaussian) and secondly, the diffusion is truncated after a finite number of timesteps which incurs error in the approximation of the intermediate distribution by a Gaussian. Flow matching addresses both of these by allowing for transformations of arbitrary distributions to each other. To define a flow-matching process as in the paper, one obtains samples $(X, Y)$ generated from a coupling (joint distribution) $\pi (X, Y)$ where the marginals of $X$ and $Y$ are $\mu$ and $\nu^*$ (the source and target distribution respectively). Given these samples, an interpolating distribution, $\{X_t\}_{t \in [0, 1]}$, is constructed with $X_0 \coloneqq X$ and $X_1 \coloneqq Y$ and the intermediate values are distributed based on the choice of an appropriate bridge distribution. This paper utilizes the Brownian bridge which is the conditional distribution of Brownian motion conditioned on fixing its endpoints. The main issue with this approach is that such an interpolating distribution is not \emph{Markovian} which makes sampling from it in a diffusion-like process challenging. This may be addressed through the notion of a Markovian projection which constructs another \emph{Markovian} process with the same marginals as the interpolating process for which a score-function-like analogue which enables efficient Markovian sampling. The paper then shows that under some relatively mild assumptions on the approximation of the score function and the moments of the coupled distributions, this procedure guarantees closeness to the target distribution in Total Variation distance. All previous approaches either require additional assumptions on the estimated score function, and/or do not account for the discretization error. One of the main challenges of the proof is in analyzing the discretization error of the scheme where the continuous process is approximated with one with discrete time steps. The paper observes that the terms controlling this error may be rather large at the end points of the interval. Therefore, the paper considers a truncated process on $[\delta, 1 - \delta]$ and show that on the truncated set, the gradients of the densities with control the discretization error are bounded in ($L^8$) norm. Overall, the results of the paper are novel and interesting. The observation that the truncated distribution is well-behaved with respect to the discretization error is valuable for future theoretical and practical work on flow models. This also illustrates the drawbacks of prior work which require additional Lipschitz assumptions on the scores/estimated scores and presents a natural approach towards circumventing them. This is a nice contribution and I recommend acceptance.

Strengths

See main review

Weaknesses

See main review

Questions

See main review

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

See main review

Reviewer igTf2024-08-07

Response to the rebuttal

I appreciate the authors' comprehensive response. My questions and concerns are sufficiently addressed. Thus, I would like to increase the rating.

Authorsrebuttal2024-08-08

We thank the reviewer for considering our response.

Reviewer TMKS2024-08-08

Thank you for your response. To improve the presentation, I would recommend adding a few words early in the paper on the fact that SGMs are a sub-case of DFMs. I maintain my rating (and low confidence score).

Reviewer NqRW2024-08-08

I acknowledge the authors' rebuttal and will retain my current evaluation.

Area Chair U9Fa2024-08-14

Discussions

Dear reviewers, Thanks very much for your great efforts in the review process. Please read the authors' response and confirm it. You are encouraged to discuss with all of us if you have any concerns. Thanks, AC

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC