Heavy-Tailed Flow Matching via Random Clocks

Technical Summary: Heavy-Tailed Flow Matching via Random Clocks

Problem Statement
Heavy-tailed data, characterized by rare events with disproportionate impact (e.g., financial crashes, weather extremes, imbalanced image classes), pose a significant challenge for modern generative modeling. Standard diffusion and flow-matching models typically rely on Gaussian source distributions or noise processes. While Gaussian sources offer tractable training targets, they provide a poor inductive match for heavy-tailed data. When the data distribution exhibits significantly heavier tails than the source, the model must learn the tail behavior entirely through the learned transport map, often leading to inaccurate reproduction of rare events.

Existing attempts to address this by replacing Gaussian noise with heavy-tailed alternatives (e.g., $\alpha$-stable Lévy processes, Student-t kernels) face limitations. These approaches often require non-local fractional-score machinery, are restricted to specific tail families (e.g., finite-variance regimes for Student-t), or rely on robust aggregation losses that break the $L^2$ projection structure essential for efficient training. There is a lack of a unified framework that preserves the low-number-of-function-evaluations (NFE) sampling efficiency of flow matching while accommodating diverse heavy-tailed source laws.

Methodology: Heavy-Tailed Flow Matching via Random Clocks (HTFM)
The authors propose HTFM, a framework that models heavy-tailed sources as mixtures of clock-conditioned Gaussian sources. The core innovation is the introduction of a random clock, a non-decreasing, path-valued latent variable $T = (T_t)_{t \in [0,1]}$ that controls the conditional covariance of the source distribution.

1. Clock-Conditioned Dynamics:

   • Conditioning on a realized clock path $T_{[0,1]}$, the source distribution $X_0$ and the flow between the source and data are Gaussian.
   • Specifically, $X_0 | T_{[0,1]} \sim \mathcal{N}(0, \Sigma(T_{[0,1]}))$.
   • The flow follows an affine path: $X_t^T = \alpha_t^T X_1 + \beta_t^T X_0^T$.
   • Because the conditional problem is Gaussian, the endpoint-conditioned velocity target $v_t^T(x | x_1)$ retains the same tractable algebraic form as standard Gaussian flow matching. This allows the use of standard $L^2$ regression objectives pathwise.

2. Marginalization and Heavy Tails:

   • The heavy-tailed nature of the marginal source distribution arises only after integrating (marginalizing) over the random clock $T$.
   • By choosing specific laws for the clock, the framework recovers various heavy-tailed families:
     • Deterministic Clock ($T_t = t$): Recovers standard Gaussian flow matching.
     • $\alpha$-Stable Subordinator Clock: Yields $\alpha$-stable source marginals.
     • Inverse-Gamma Clock: Yields Student-t marginals for any degrees of freedom $\nu > 0$, including the infinite-variance regime ($\nu \le 2$) often excluded by other methods.

3. Signature-Based Conditioning:

   • Since the clock path $T_{[0,1]}$ is infinite-dimensional, passing it directly to the neural network is impractical.
   • The authors encode the realized clock path using truncated logsignature features ($\ell_m(T)$), derived from rough path theory. These features provide a finite-dimensional, universal representation of the path.
   • The neural velocity field $u_\theta$ is conditioned on these features alongside time $t$ and state $X_t$. This allows the network to adapt to the specific conditional Gaussian space defined by the realized clock with negligible computational overhead.

Key Contributions

1. Unified Random-Clock Framework: HTFM provides a single flow-matching objective that can generate Gaussian, $\alpha$-stable, and Student-t sources by simply varying the clock law. It avoids family-specific score formulas or unavailable density functions by working in a conditionally Gaussian space.
2. Signature-Based Conditioning: The paper introduces a practical interface for clock-conditioned vector fields using truncated logsignatures. This enables the velocity field to adapt to the realized conditional space without the cost of operator learning.
3. Theoretical Guarantees: The framework establishes a pathwise conditional-marginal equivalence (Proposition 3.1), ensuring that minimizing the clock-conditioned objective minimizes the marginal heavy-tailed objective, provided integrability conditions are met.

Experimental Results
The authors evaluate HTFM on three benchmarks: 2D imbalanced $\alpha$-stable mixtures, CIFAR10-LT (long-tailed image dataset), and HRRR weather fields (Vertically Integrated Liquid).

• Performance: HTFM consistently outperforms Gaussian flow matching (HTFM-G) and competitive heavy-tailed baselines (LIM, DLPM, t-Flow) in terms of mode coverage, sample quality (FID), and tail-statistic recovery (Kurtosis Ratio, Skewness Ratio, KS statistic).
• Low-NFE Efficiency: A primary advantage is the retention of flow matching's low-NFE sampling efficiency. HTFM substantially outperforms $\alpha$-stable diffusion baselines (like DLPM) at low NFE (e.g., 20–100 steps), where diffusion models often struggle with slow convergence.
• Tail Calibration: The experiments demonstrate that the "heaviness" of the generated tails can be calibrated by varying only the clock law or its tail parameter, without changing the network architecture.
• Ablation Studies:
  • Clock Family: Different clock families (e.g., $\alpha$-stable vs. Student-t) perform differently even at matched polynomial tail-decay rates, indicating the clock family is a meaningful design knob beyond just tail heaviness.
  • Signature Conditioning: Explicitly conditioning the velocity field on the clock signature (using $m \ge 1$) yields better results than using a heavy-tailed source alone without clock conditioning ($m=0$), confirming the value of adapting to the conditional space.

Significance and Claims
The paper claims that HTFM bridges the gap between the tractability of Gaussian flow matching and the necessity of heavy-tailed modeling for real-world extreme events. By decoupling the tail behavior (controlled by the random clock) from the conditional regression problem (which remains Gaussian), the method avoids the computational and theoretical pitfalls of previous heavy-tailed diffusion approaches.

The authors emphasize that HTFM offers a practical tail-control interface: users can calibrate the "heaviness" of generated samples across different distribution families simply by adjusting the clock law. Furthermore, the use of logsignature features allows the model to leverage the full path information of the clock without incurring the high costs associated with infinite-dimensional operator learning, making it a scalable solution for high-dimensional heavy-tailed generative modeling.