Theoretical guarantees in KL for Diffusion Flow Matching

Marta Gentiloni Silveri,Giovanni Conforti,Alain Durmus

2024-09-12

Abstract: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.

Machine Learning,Probability

What problem does this paper attempt to address?

### Problems Addressed by the Paper

This paper primarily explores the theoretical guarantees of the **Diffusion Flow Matching (DFM) model** in generative modeling. Specifically, the paper aims to analyze the performance of the DFM model when using a Brownian bridge as the intermediary and provides an upper bound estimate of the Kullback-Leibler (KL) divergence between the target distribution and the DFM model's generated distribution.

#### Main Contributions:

1. **KL Divergence Upper Bound without Early Stopping**:
- Without using early stopping, the paper establishes a clear and simple upper bound on the KL divergence between the target distribution and the DFM model's final time point distribution. This result assumes only that the target distribution ν⋆, the base distribution µ, and the coupling π satisfy certain moment and integral conditions, and that there exists an L2 drift approximation error.

2. **KL Divergence Upper Bound with Early Stopping**:
- With early stopping, the paper also provides an explicit upper bound on the KL divergence between a smoothed version of the target distribution and the early-stopped version of the DFM model. This result assumes that π is an independent coupling of µ and ν⋆, and that µ satisfies absolute continuity and log-derivative integral conditions.

### Abstract and Background

The paper investigates the Flow Matching (FM) method, a generative model that connects the target distribution ν⋆ with an auxiliary distribution µ over a finite time. This method uses a fixed coupling π and a bridge (which can be deterministic or stochastic) to define a path measure and approximates this path measure by learning the drift of its Markov projection. The main contribution of this paper is to provide non-asymptotic guarantees for the Diffusion Flow Matching (DFM) model using the conditional distribution of Brownian motion as the bridge under mild conditions.

### Conclusion

This paper fills a gap in the existing literature regarding the theoretical analysis of the DFM model, especially when considering all sources of error (drift approximation error and time discretization error). Compared to previous studies, the approach in this paper does not rely on smoothness assumptions of the flow velocity field or its estimator and is applicable to a broader range of data distributions. These theoretical results pave the way for further understanding and improvement of generative modeling techniques.