AlignFlow: Improving Flow-based Generative Models with Semi-Discrete Optimal Transport

ICLR 2026 Conference SubmissionAnonymous Authors
Flow-based generative model; flow matching; Semi-discrete optimal transport
Abstract:

Flow-based Generative Models (FGMs) effectively transform noise into a data distribution, and coupling the noise and data in the training of FGM by Optimal Transport (OT) improves the straightness of the flow paths. However, existing OT- based couplings are difficult to combine with modern models and/or to scale to large datasets due to the curse of dimensionality in the sample complexity of (batch) OT. This paper introduces AlignFlow, a new approach using Semi-Discrete Optimal Transport (SDOT) to enhance FGM training by establishing explicit alignment between noise and data pairs. SDOT computes a transport map by partitioning the noise space into Laguerre cells, each mapped to a corresponding data point. During the training of FGM, i.i.d.-sampled noise is matched with corresponding data by the SDOT map. AlignFlow bypasses the curse of dimensionality and scales effectively to large datasets and models. Our experiments demonstrate that AlignFlow improves a wide range of state-of-the-art FGM algorithms and can be integrated as a plug-and-play solution with negligible additional cost.

Disclaimer
This report is AI-GENERATED using Large Language Models and WisPaper (A scholar search engine). It analyzes academic papers' tasks and contributions against retrieved prior work. While this system identifies POTENTIAL overlaps and novel directions, ITS COVERAGE IS NOT EXHAUSTIVE AND JUDGMENTS ARE APPROXIMATE. These results are intended to assist human reviewers and SHOULD NOT be relied upon as a definitive verdict on novelty.
NOTE that some papers exist in multiple, slightly different versions (e.g., with different titles or URLs). The system may retrieve several versions of the same underlying work. The current automated pipeline does not reliably align or distinguish these cases, so human reviewers will need to disambiguate them manually.
If you have any questions, please contact: mingzhang23@m.fudan.edu.cn

Overview

Overall Novelty Assessment

AlignFlow introduces Semi-Discrete Optimal Transport (SDOT) to align noise-data pairs during flow-based generative model training, partitioning noise space into Laguerre cells mapped to data points. The paper resides in the 'Minibatch and Semi-Discrete Optimal Transport' leaf, which contains four papers including the original work. This leaf sits within the broader 'Optimal Transport Coupling Strategies for Flow Matching' branch, indicating a moderately populated research direction focused on scalable OT coupling schemes. The taxonomy reveals this is an active but not overcrowded area, with sibling leaves addressing conditional OT, rectified flows, and direct OT map learning.

The taxonomy structure shows AlignFlow's leaf neighbors include 'Conditional Optimal Transport for Flow-Based Models' (five papers addressing conditioning variables) and 'Rectified and Straight-Path Flow Construction' (three papers on iterative trajectory straightening). The parent branch excludes domain-specific applications and theoretical analyses, focusing purely on coupling algorithm design. Nearby branches cover core flow matching frameworks (simulation-free objectives, stochastic interpolants) and theoretical foundations (equivalence studies, gradient flow connections). AlignFlow's SDOT approach bridges computational efficiency concerns (addressed in the 'Computational and Scalability Enhancements' branch) with coupling quality, distinguishing it from purely batch-level or continuous OT formulations in sibling nodes.

Among thirty candidates examined across three contributions, the core 'AlignFlow: Semi-Discrete Optimal Transport' contribution shows one refutable candidate out of ten examined, suggesting some prior work overlap in the limited search scope. The 'Bypassing curse of dimensionality' contribution found zero refutable candidates among ten examined, indicating relative novelty in addressing sample complexity. The 'Plug-and-play integration' contribution similarly shows zero refutable candidates from ten examined. These statistics reflect a focused literature search rather than exhaustive coverage, with the single refutable match likely representing closely related minibatch or semi-discrete OT work within the same taxonomy leaf.

Based on the limited thirty-candidate search, AlignFlow appears to occupy a moderately explored niche within OT-coupled flow matching. The taxonomy reveals sufficient prior work in minibatch and semi-discrete methods to contextualize the contribution, yet the low refutation rate (one of thirty candidates) suggests meaningful differentiation. The analysis does not cover potential overlaps beyond top-K semantic matches or recent preprints outside the search scope, leaving open questions about exhaustive novelty assessment.

Taxonomy

Core-task Taxonomy Papers
50
3
Claimed Contributions
30
Contribution Candidate Papers Compared
1
Refutable Paper

Research Landscape Overview

Core task: improving flow-based generative models with optimal transport coupling. The field has evolved around several major branches that reflect different emphases in leveraging optimal transport (OT) to enhance flow matching and continuous normalizing flows. At the highest level, one finds branches dedicated to OT coupling strategies—ranging from minibatch and semi-discrete schemes (e.g., Minibatch Optimal Transport[3], Semidiscrete Couplings[12]) to sliced and all-to-all variants—alongside frameworks that establish the theoretical underpinnings connecting OT, stochastic interpolants, and rectified flows (Flow Matching[2], Rectified Flow[28], Stochastic Interpolants[13]). Other branches address domain-specific applications (molecular generation, biological trajectory inference, visuomotor policies), advanced training objectives and regularization techniques, specialized inference and sampling methods, and computational enhancements that scale OT-based training to large datasets. Emerging cross-domain extensions further illustrate the breadth of this taxonomy, spanning speech enhancement, point cloud modeling, and beyond. Within the landscape of OT coupling strategies, a particularly active line of work explores how to efficiently compute transport plans over minibatches or semi-discrete settings, balancing computational cost against the quality of learned trajectories. AlignFlow[0] sits squarely in this cluster, focusing on minibatch and semi-discrete optimal transport to align source and target distributions more effectively during training. Nearby efforts such as Minibatch Optimal Transport[3] and Semidiscrete Couplings[12] share a similar emphasis on scalable coupling schemes, while works like Faster Flow Inference[1] and All-to-All Flow[6] investigate alternative pairing strategies or inference speedups. A key trade-off across these methods concerns the granularity of the coupling—whether to match samples in small batches, use sliced projections (Sliced Wasserstein Flows[5]), or adopt global transport maps—and how these choices affect both training stability and sample quality. AlignFlow[0] contributes to this dialogue by proposing refined minibatch OT techniques that aim to preserve geometric structure without prohibitive computational overhead, positioning it as a natural extension of recent semi-discrete and minibatch approaches.

Claimed Contributions

AlignFlow: Semi-Discrete Optimal Transport for Flow-based Generative Models

The authors propose AlignFlow, a method that uses Semi-Discrete Optimal Transport to compute a deterministic mapping between noise samples and dataset points. This approach partitions the noise space into Laguerre cells, each mapped to a corresponding data point, enabling more efficient training of flow-based generative models.

10 retrieved papers
Can Refute
Bypassing the curse of dimensionality in optimal transport

Unlike prior batch OT-based approaches that suffer from exponential sample complexity in high dimensions, AlignFlow computes the SDOT plan between the known empirical dataset (discrete distribution) and the known noise distribution (continuous distribution), avoiding the need to estimate the unknown data distribution and thus circumventing the curse of dimensionality.

10 retrieved papers
Plug-and-play integration with state-of-the-art FGM algorithms

AlignFlow is designed as a modular component that can be readily combined with existing flow-based generative model training techniques, including flow matching, shortcut models, MeanFlow, and consistency training, with minimal computational overhead (less than 1% extra cost).

10 retrieved papers

Core Task Comparisons

Comparisons with papers in the same taxonomy category

Contribution Analysis

Detailed comparisons for each claimed contribution

Contribution

AlignFlow: Semi-Discrete Optimal Transport for Flow-based Generative Models

The authors propose AlignFlow, a method that uses Semi-Discrete Optimal Transport to compute a deterministic mapping between noise samples and dataset points. This approach partitions the noise space into Laguerre cells, each mapped to a corresponding data point, enabling more efficient training of flow-based generative models.

Contribution

Bypassing the curse of dimensionality in optimal transport

Unlike prior batch OT-based approaches that suffer from exponential sample complexity in high dimensions, AlignFlow computes the SDOT plan between the known empirical dataset (discrete distribution) and the known noise distribution (continuous distribution), avoiding the need to estimate the unknown data distribution and thus circumventing the curse of dimensionality.

Contribution

Plug-and-play integration with state-of-the-art FGM algorithms

AlignFlow is designed as a modular component that can be readily combined with existing flow-based generative model training techniques, including flow matching, shortcut models, MeanFlow, and consistency training, with minimal computational overhead (less than 1% extra cost).