SD-FM: Semidiscrete Flow Matching

• Semidiscrete Flow Matching (SD-FM) is a generative modeling framework that uses semidiscrete optimal transport to pair continuous noise with finite data points.
• It optimizes an entropic dual potential vector via stochastic methods, reducing flow curvature and improving computational and memory efficiency.
• SD-FM is applied in unconditional, conditional, and super-resolution tasks, demonstrating enhanced integration with modern generative modeling frameworks.

Semidiscrete Flow Matching (SD-FM) denotes a family of flow-based generative models that leverage semidiscrete optimal transport couplings between a continuous source measure (typically a standard Gaussian) and a finite target dataset or data-supported empirical distribution. SD-FM was developed to address the computational and geometric limitations of independent and batch-optimal transport couplings in flow matching, enabling the efficient alignment of noise and data points in high-dimensional synthesis, supervised and conditional generation, super-resolution, and other contexts. The method parameterizes the coupling through an entropic dual potential vector, supports efficient maximum inner product search for training pair assignment, and can be incorporated into contemporary flow matching objectives, consistency models, and guidance schemes. This paradigm yields significant improvements in flow curvature, sample quality, and scalability relative to previous flow matching approaches.

1. Motivation and Theoretical Foundations

SD-FM formalizes flow matching within the semidiscrete optimal transport (SD-OT) regime, in which the data set or target measure $\nu$ is finite and the source measure $\mu$ (noise) is continuous. Given the goal of training a velocity field $v_\theta(x, t)$ that transforms a simple distribution to data via an ODE or SDE, early approaches independently sample $(x_0, x_1)$ and minimize the misalignment between $v_\theta(x, t)$ and $(x_1 - x_0)$ along the associated path. This results in high flow curvature and inefficient sampling, especially in high dimensions.

OT-based flow matching (OT-FM) employs batch couplings using the Sinkhorn algorithm or similar solvers on batch sizes $n$, with computational complexity $O(n^2/\varepsilon^2)$ per pairing and quadratic scaling with batch size and regularization parameter. These approaches are feasible only for moderate $n$, are increasingly limited by sample bias, and frequently incur quadratic memory overhead.

SD-FM resolves these bottlenecks by reformulating the OT coupling: The finite empirical support of $\nu$ is encoded by a dual potential vector $g \in \mathbb{R}^N$ (with $N$ the number of data points). The soft c-transform, central to entropic SD-OT, is defined as:

$$f(x) = -\varepsilon \log \left[ \sum_{j=1}^N b_j \exp\left( \frac{g_j - c(x, y_j)}{\varepsilon} \right) \right],$$

with cost function $c(x, y)$ (commonly $-x^T y$) and weights $b_j$ (often uniform). In the zero-regularization case ($\varepsilon = 0$), this reduces to the hard maximum:

$f(x) = - \max_j [ g_j - c(x, y_j) ].$

Pairing during training uses maximum inner product search (MIPS) or its regularized variant.

2. Methodology and Implementation Framework

The SD-FM pipeline consists of two key stages:

1. Dual Potential Estimation: The dual potential $g$ is found by optimizing the semidual OT objective via stochastic gradient descent (SGD):

$$F_\varepsilon(g) = \langle b, g \rangle - \varepsilon \mathbb{E}_{x \sim \mu} \left[ \log \left( \sum_{j=1}^N b_j \exp \left( \frac{g_j - c(x, y_j)}{\varepsilon} \right) \right) \right]$$

With gradients given by $$\partial F_\varepsilon / \partial g_j = b_j - m_j(g)$$, where $m_j(g)$ is the expected soft assignment weight to data point $y_j$ over the source measure.

1. Pair Assignment via Soft c-Transform: At each training iteration, newly sampled $x_0 \sim \mu$ is paired by either:
   • Sampling from the categorical distribution defined by $s_{\varepsilon,g}(x_0)_j$ for $\varepsilon > 0$, or
   • Assigning $x_0$ to $\arg\max_j [ g_j + \langle x_0, y_j \rangle ]$ for $\varepsilon = 0$. This operation is highly efficient ($O(N)$ or faster with approximate MIPS), decoupling training from batch size and regularization cost.

These pairings are now used in the flow matching objective:

$$\mathbb{E}_{t \sim U[0,1], (x_0, x_1) \text{ from SD-OT coupling}} \left[ \| v_\theta(x_t, t) - (x_1 - x_0) \|^2 \right]$$

where $x_t$ is interpolated between $x_0$ and $x_1$.

3. Comparative Advantages over I-FM and OT-FM

SD-FM presents major improvements relative to both I-FM (independent coupling) and OT-FM:

• Computational Scalability: The cost to estimate $g$ is linear in $N$ and decoupled from training batch size and regularization parameter, unlike OT-FM with $O(n^2/\varepsilon^2)$ per batch.
• Pairing Optimality: SD-OT pairing with precomputed $g$ closely approximates the continuous-to-discrete OT, outperforming independent and batch pairings when measured by chi-squared divergence and flow curvature.
• Integration Efficiency: Empirical evaluations show SD-FM models require fewer numerical function evaluations (NFEs) during ODE inference, due to the improved straightness of the learned transport trajectories.
• Memory Efficiency: Storing $g \in \mathbb{R}^N$ costs $O(N)$ memory, far less than the quadratic overhead in batch OT-FM.

The following table summarizes key complexity features:

Method	Pairing Cost per Step	Memory Cost	Pairing Bias
I-FM	$O(1)$	$O(1)$	High
OT-FM	$O(n^2/\varepsilon^2)$	$O(n^2)$	Moderate
SD-FM	$O(N)$ (MIPS)	$O(N)$	Low

4. Empirical Results and Applications

SD-FM demonstrates improved sample fidelity and efficiency in several generative modeling tasks:

• Unconditional Generation: On ImageNet (32x32 and 64x64), PetFace, and similar datasets, SD-FM achieves notably lower FID scores compared to I-FM, OT-FM, and mean-flow models; training with PCA-reduced data further expedites pair assignment.
• Conditional Generation: SD-FM extends to class-conditional and structured settings by augmenting the cost function to include label or auxiliary conditioning (e.g., $c((x, z), (x', z')) = c_{x,x'} + \beta c_{z,z'}$).
• Super-Resolution: SD-FM supports continuous-conditional pairing (e.g. for input-conditioned high-resolution synthesis), improving PSNR and SSIM.
• Consistency and Guidance Models: The semidiscrete coupling enables precise Tweedie correction in consistency and guided sampling, further reducing curvature and improving sample precision.

Across all settings, qualitative and quantitative evaluations demonstrate superior sample quality, integration speed, and computational cost-effectiveness.

5. Integration with Contemporary Frameworks

SD-FM is consistent with and synergistic to several modern directions in generative learning:

• Stable Autonomous Flow Matching: Autonomous (time-independent) vector fields and pseudo-time augmentations from control theory (Sprague et al., 8 Feb 2024) can be incorporated into SD-FM to ensure stability in physically consistent or discrete-time settings.
• Model-Aligned Coupling (MAC): Although MAC (Lin et al., 29 May 2025) focuses on model-learnable couplings and efficient error-driven mini-batch assignments, SD-FM achieves global optimality and scalability through data-wide dual potential precomputation.
• Unified Bridge Algorithms: SD-FM naturally extends the unified SDE bridge framework (Kim, 27 Mar 2025), supporting cases where one marginal is discrete.

SD-FM's pair assignment protocol and objective formulation also facilitate easy adaptation to guidance, consistency, diffusion alignment, and multi-subject control strategies in image synthesis.

6. Limitations and Future Directions

SD-FM scaling is ultimately limited by the need to search or enumerate the finite target set ($O(N)$ time per datum), which presents challenges for extremely large datasets. The use of PCA or locality-sensitive hashing-based MIPS can mitigate this, but further algorithmic refinement is needed. In addition, convergence properties of the SGD-based dual potential estimation for very large $N$ or highly structured target distributions warrant deeper theoretical and empirical analysis.

Extensions to:

• Hybrid settings, with both source and target measures being semidiscrete or partially continuous.
• Multi-scale pairing and hierarchical models, leveraging SD-FM for blockwise or attention-based composition.

7. Conclusion

Semidiscrete Flow Matching (SD-FM) enables flow matching models to utilize data-wide optimal transport pairings in a computational and memory-efficient manner by exploiting the finite support of the target dataset. The method employs entropic dual potential optimization and fast assignment (MIPS), yielding reduced transport curvature, improved sample quality, and scalability to contemporary generative tasks. SD-FM outperforms independent and batch OT-based couplings across unconditional, conditional, super-resolution, consistency, and guidance scenarios on standard benchmarks. Continuing work aims to optimize SD-FM for extremely large targets and to further integrate these principles into broader generative modeling and control-theoretic frameworks.