Technical Summary: Exploring and Exploiting Stability in Latent Flow Matching

1. Problem Statement

Training generative models, particularly Diffusion Models and their Flow Matching (FM) counterparts, remains computationally expensive. These models typically require massive datasets, extensive compute resources, long training durations, and, in conditional setups, extensive annotations. While Flow Matching offers a deterministic formulation for sample generation via Ordinary Differential Equations (ODEs) that often reduces sampling steps compared to Stochastic Differential Equations (SDEs), the cost of training these models on full-scale datasets remains a bottleneck.

The central question addressed by this work is whether dataset size and model capacity can be significantly reduced without sacrificing generation quality. Specifically, the authors investigate the stability of Latent Flow Matching (LFM) models under major perturbations, including dataset subsampling (pruning), architectural changes, and altered training configurations. They posit that if LFM transport trajectories are inherently stable, one could exploit this property to derive more efficient training and inference algorithms.

2. Methodology

2.1. Theoretical Foundation: Stability in Flow Matching

The work builds upon the observation that Rectified Flow admits a closed-form expression for the velocity field. For a finite dataset, the optimal velocity field is a weighted sum of training samples, where the weights are determined by a softmax function over the squared distances between the current state and scaled training samples. The authors note that this softmax becomes "peaked" early in the transport phase, meaning a single training sample often dominates the velocity field for a given trajectory.

This theoretical insight suggests that the transport trajectory is largely determined by the presence of specific "dominating" samples. Consequently, removing non-dominant samples (pruning) should leave the majority of trajectories unchanged. The authors formalize this by defining an assignment metric: for a trajectory starting at noise $x_0$, they determine which training sample $x_i$ it ends at. They hypothesize that this assignment remains invariant even when a large fraction of the dataset is removed, provided the dominating sample for that trajectory is retained.

2.2. Data Pruning Strategies

To leverage this stability, the authors propose three data pruning criteria adapted from discriminative models but reformulated for the FM objective (which utilizes shared noise paths and timesteps):

1. Gradient-based Scoring (G): Samples are ranked by the magnitude of the gradient norm of the loss.
2. Loss-based Scoring (L): Samples are ranked by the magnitude of the loss computed along shared noise paths.
3. Clustering-based Scoring (C): Samples are clustered in a pretrained embedding space (e.g., CLIP). The authors propose two specific variants:
   • Proportional ($C_p$): Selects samples proportional to cluster size.
   • Balanced ($C_b$): Selects an equal number of samples from each cluster to mitigate distributional imbalances.
   • Kernel-based ($C_b^\kappa$): Selects samples greedily to match the mean of the cluster's Gaussian kernel in latent space.

The authors also evaluate the "inverse" of these criteria (selecting the lowest scores) to test the robustness of the stability claim.

2.3. Coarse-to-Fine (C2F) Inference

Exploiting the stability of trajectories across model capacities, the authors propose a two-stage inference strategy:

1. Coarse Model: A lightweight architecture (e.g., DiT-S/2) is trained on a pruned subset of the data. It handles the first phase of the trajectory ($t \in [0, t_0)$).
2. Fine Model: A high-capacity architecture (e.g., DiT-XL/2), pretrained on the full dataset, handles the second phase ($t \in [t_0, 1]$).

To ensure a smooth transition at the "seam" point $t_0$, the Coarse model is fine-tuned with a continuity loss. This loss minimizes the difference between the velocity fields of the Coarse and Fine models at $t_0$. The Fine model's trajectory is used to generate the target state $x_{t_0}$ via ODE inversion (backward integration), ensuring the Coarse model learns to align with the Fine model's path.

3. Key Contributions

1. Empirical Demonstration of LFM Stability: The authors show that independently trained LFM models converge to highly consistent transport trajectories under dataset subsampling, architectural changes, and different random initializations. This stability is linked to the theoretical properties of the FM objective.
2. Practical Pruning Algorithms: They introduce three pruning criteria (Gradient, Loss, Clustering) suitable for FM. They demonstrate that aggressive pruning (e.g., removing 50-95% of data) does not degrade performance and can even improve convergence speed and generation quality under limited compute budgets.
3. Coarse-to-Fine Inference Acceleration: They propose a two-model generation strategy that splits the trajectory. This approach achieves a significant inference speedup (approx. 2.15$\times$) while maintaining generation quality, provided the stability assumption holds across the model architectures.
4. Bias Mitigation via Balanced Pruning: The authors show that balanced clustering pruning ($C_b$) can mitigate biases in generated data (e.g., gender, skin tone) without requiring explicit label-aware balancing, by equalizing the representation of semantic clusters in the training set.

4. Experimental Results

The methods were evaluated on CelebA-HQ, FFHQ, and ImageNet-1K.

• Stability under Pruning: On CelebA-HQ, models trained on disjoint subsets (pruned at 50%) produced visually similar images when starting from the same noise seed $x_0$. The ArcFace cosine similarity between pruned and unpruned models ranged from 0.79 to 0.83, significantly higher than random pairs (0.37).
• Performance Metrics:
  • CelebA-HQ: Pruning 50% of the data using balanced clustering ($C_b$) slightly improved the Fréchet Inception Distance (FID) from 24.24 (unpruned) to 22.80.
  • ImageNet: Pruning 75% of the data ($pr=0.75$) resulted in faster convergence and maintained or slightly improved FID compared to the full dataset. Even at 95% pruning, the model converged faster initially, though performance degraded after a certain number of iterations.
• Robustness to Perturbations:
  • Architecture: Increasing model capacity (DiT-S/2 to DiT-XL/2) resulted in minimal trajectory changes (similarity $s=0.81$). Switching to a U-Net architecture reduced similarity ($s=0.55$) but preserved coarse attributes.
  • Latent Space: The stability is sensitive to changes in the latent space parametrization (e.g., changing the VAE seed or applying invertible transforms), which broke stability.
  • Objective: Switching from Flow Matching to score-based diffusion completely altered the trajectories, indicating stability is specific to the FM objective.
• Coarse-to-Fine Efficiency: On CelebA-HQ, the C2F approach with $t_0=0.7$ reduced inference time from 93.95 ms/image (Fine only) to 43.53 ms/image (approx. 2.15$\times$ speedup) while maintaining FID.
• Fairness: Balanced pruning ($C_b$) reduced KL divergence to a uniform distribution across attributes (gender, age, skin tone) more effectively than unpruned models, despite being label-agnostic.

5. Significance and Claims

The authors claim that the observed stability is a characteristic property of the Rectified Flow Matching paradigm under fixed latent representations, rather than an artifact of specific training algorithms.

• Efficiency: The primary significance is the ability to reduce training time and data requirements without sacrificing quality. The authors demonstrate that for large datasets like ImageNet, pruning can accelerate convergence, and for smaller datasets like CelebA-HQ, it can improve generation quality.
• Inference Speedup: The C2F approach offers a principled way to reduce inference costs by offloading the early, coarse stages of generation to lightweight models, relying on the stability of the trajectory to ensure the high-capacity model only needs to refine the final details.
• Bias Control: The work suggests that dataset pruning, specifically via balanced clustering, can serve as a mechanism to control and mitigate distributional imbalances in generated data, offering a path toward fairer generative models without complex label-aware reweighting.

The paper concludes that while stability is robust to data and model capacity changes, it is sensitive to changes in the latent space structure and the generative objective itself. This insight allows practitioners to design more resource-efficient pipelines while maintaining high-fidelity generation.