Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Mirror Flow Matching with Heavy-Tailed Priors for Generative Modeling on Convex Domains

Yunrui Guan, Krishnakumar Balasubramanian, Shiqian Ma

Abstract

We study generative modeling on convex domains using flow matching and mirror maps, and identify two fundamental challenges. First, standard log-barrier mirror maps induce heavy-tailed dual distributions, leading to ill-posed dynamics. Second, coupling with Gaussian priors performs poorly when matching heavy-tailed targets. To address these issues, we propose Mirror Flow Matching based on a \emph{regularized mirror map} that controls dual tail behavior and guarantees finite moments, together with coupling to a Student-$t$ prior that aligns with heavy-tailed targets and stabilizes training. We provide theoretical guarantees, including spatial Lipschitzness and temporal regularity of the velocity field, Wasserstein convergence rates for flow matching with Student-$t$ priors and primal-space guarantees for constrained generation, under $\varepsilon$-accurate learned velocity fields. Empirically, our method outperforms baselines in synthetic convex-domain simulations and achieves competitive sample quality on real-world constrained generative tasks.

Mirror Flow Matching with Heavy-Tailed Priors for Generative Modeling on Convex Domains

Abstract

We study generative modeling on convex domains using flow matching and mirror maps, and identify two fundamental challenges. First, standard log-barrier mirror maps induce heavy-tailed dual distributions, leading to ill-posed dynamics. Second, coupling with Gaussian priors performs poorly when matching heavy-tailed targets. To address these issues, we propose Mirror Flow Matching based on a \emph{regularized mirror map} that controls dual tail behavior and guarantees finite moments, together with coupling to a Student- prior that aligns with heavy-tailed targets and stabilizes training. We provide theoretical guarantees, including spatial Lipschitzness and temporal regularity of the velocity field, Wasserstein convergence rates for flow matching with Student- priors and primal-space guarantees for constrained generation, under -accurate learned velocity fields. Empirically, our method outperforms baselines in synthetic convex-domain simulations and achieves competitive sample quality on real-world constrained generative tasks.

Paper Structure

This paper contains 22 sections, 12 theorems, 79 equations, 6 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Challenges and Solutions
  3. Ingredients for Designing Mirror Flow Matching
  4. Ingredient 1: The Mirror map
  5. Ingredient 2: The Prior Distribution
  6. Mirror Flow Matching
  7. Theoretical Results
  8. Guarantees for Euclidean Flow Matching with t-Distribution Priors
  9. Primal Space Guarantee for Mirror Flow Matching
  10. Experiments
  11. Numerical simulation
  12. Real-data application: Watermarked image generation
  13. Conclusion
  14. Visual Illustration of Methodological Challenges
  15. Examples verifying Proposition \ref{['P_Wasserstein_Primal_Dual_Tail_Bound']}
  16. ...and 7 more sections

Key Result

Lemma 1

Let $Y$ be a random variable in $\mathbb{R}^{d}$ with law $P$. Then, (i) if $P(\|Y\| \ge R) \ge {C}/{R^{p}}$ for some constant $C>0$, then $\mathbb{E}[\|Y\|^{p}]$ does not exist, and (ii) if $P(\|Y\| \ge R) \le {C}/{R^{\beta}}$ with $\beta > p$, then $\mathbb{E}[\|Y\|^{p}]$ is finite.

Figures (6)

  • Figure 1: Dual space distribution comparison between the log barrier and our mirror map ($\kappa = 0.5$). The primal distribution is a truncated Gaussian mixture within a polytope (see Appendix \ref{['sec:visualmetchallenge']}). The log barrier yields a heavy-tailed distribution, while our mirror map produces a much lighter tail.
  • Figure 2: Visualization of interpolations in primal and dual spaces -- Straight line interpolation in the dual space (Figure (b)) corresponds to curved "geodesic" interpolation in primal space Figure (a)).
  • Figure 3: Samples of generated watermarked images from the AFHQv2 $64 \times 64$ dataset. Constraint satisfaction were checked with built-in functions of liu2023mirror.
  • Figure 4: Figure \ref{['N_Numerical2_a']} shows the ground-truth reference distribution. Figures \ref{['N_Numerical2_b']} and \ref{['N_Numerical2_c']} illustrate that the log-barrier method performs poorly (both with G or t-flow), while Figure \ref{['N_Numerical2_d']} demonstrates that G-flow (with our mirror map) fails to capture the mode centered near $(-10,0)$. In contrast, Figure \ref{['N_Numerical2_e']} shows that t-flow with our mirror map covers the target distribution better. All results are obtained with discretization step size $h=0.1$. See also Figure \ref{['N_Numerical2_bdry']} for a zoomed-in illustration near the boundary.
  • Figure 5: We generate a total of $10{,}000$ samples, but for visualization we only display those lying in the boundary region $x \in [-14, -12], y \in [0, 2]$. Figure \ref{['N_Numerical2_bdry_a']} shows the ground-truth reference distribution. Figures \ref{['N_Numerical2_bdry_b']} and \ref{['N_Numerical_bdry_c']} demonstrate that, near the boundary, t-flow provides a closer approximation to the ground truth than G-flow.
  • ...and 1 more figures

Theorems & Definitions (16)

  • Lemma 1
  • Proposition 2
  • Example 3: Uniform distribution on the cube
  • Example 4
  • Proposition 5
  • Proposition 6
  • Theorem 7: Discretization Error of t-Flow
  • Lemma 8
  • Theorem 9
  • Example 10
  • ...and 6 more