Optimal Transport and Distributional Alignment

• Optimal Transport and Distributional Alignment is a framework that finds cost-optimal couplings between distributions to align their complete structures.
• It underpins applications in machine learning such as domain adaptation, generative modeling, and time-series analysis by robustly matching statistical features.
• Recent advances incorporate entropic regularization, adaptive strategies, and multi-marginal techniques to enhance scalability, interpretability, and practical performance.

Optimal transport (OT) is a mathematical and algorithmic framework for aligning probability distributions by finding a transport plan or map that minimizes the expected value of a given cost function. The optimal transport perspective goes beyond traditional pointwise and moment-matching approaches, enabling the comparison, transformation, and synthesis of full distributions. Distributional alignment, within the OT paradigm, refers to constructing maps or couplings between source and target distributions so that certain probabilistic, geometric, or functional criteria are optimized. This framework underpins a broad spectrum of modern methodologies in machine learning, generative modeling, domain adaptation, time-series analysis, statistical inference, and beyond.

1. Mathematical Foundations and Formulations

At its core, OT seeks a coupling $\gamma$ between two measures $\mu$ and $\nu$ on spaces $\mathcal X$, $\mathcal Y$, minimizing a cost integral: $$W_c(\mu,\nu) = \min_{\gamma \in \Pi(\mu,\nu)} \int_{\mathcal X \times \mathcal Y} c(x,y)\, d\gamma(x,y)$$ where $\Pi(\mu,\nu)$ denotes the set of couplings (joint distributions with marginals $\mu,\nu$) and $c(x,y)$ is a specified cost, e.g., squared Euclidean distance. In the Monge formulation, for atomless $\mu$, one seeks a measurable map $\mu$0 with $\mu$1 (the pushforward), minimizing $\mu$2. For empirical measures, the problem reduces to a (linear or entropic-regularized) assignment or network flow.

Generalizations pertinent to distributional alignment include:

• Multi-marginal OT, where several distributions are fit via a common or joint coupling (Karimi et al., 2021);
• Semi-discrete OT, with one continuous and one empirical measure, leading to explicit mappings via Laguerre cells (power diagrams) (Kong et al., 16 Oct 2025);
• Partial and adaptive OT, where only a subset of mass is transported due to support mismatch or outliers, and transport is adaptively tuned to data geometry (Yang et al., 7 Mar 2025);
• Unbalanced OT frameworks, penalizing mass creation/destruction in distribution support mismatch (Lee et al., 16 Mar 2026).

Key duality results, convex relaxations, and regularized variants (e.g., entropic, group-structured, information-theoretic) underpin scalable optimization schemes and extended modeling capabilities (Chuang et al., 2022, Courty et al., 2015).

2. Distributional Alignment in Machine Learning Applications

OT-based alignment frameworks have emerged as universally applicable tools across multiple domains. Notable instantiations include:

• Domain adaptation: Courty et al. pioneered OT-based domain adaptation, seeking mappings that align source and target feature distributions, with regularizers enforcing class coherence and computationally-expedient entropic smoothing (Courty et al., 2015). Extensions handle deep representations, multi-domain scenarios, and partial alignments robust to outliers (Lin et al., 2020, Yang et al., 7 Mar 2025).
• Flow-based generative modeling: AlignFlow utilizes semi-discrete OT (SDOT), partitioning the latent (noise) space into Laguerre cells mapped deterministically to data samples, providing explicit, low-variance couplings for training flow-matching generative models and yielding straight transport trajectories with provable convergence and empirical improvements in FID, speed, and NFE (Kong et al., 16 Oct 2025).
• Hierarchical and low-rank alignment: For multimodal data or data with latent structure, hierarchical OT (Lee et al., 2019) and latent OT (Lin et al., 2020) approaches incorporate cluster structure or anchor-based low-rank factorization, enhancing interpretability, noise robustness, and efficiency, especially for complex, high-dimensional, or partially-overlapping distributions.
• Large-scale and efficient approximations: Slicing-based approaches (min-sliced transport plans) minimize OT cost over projections, enabling closed-form, 1D matching combined with amortized learning and rapid inference; theoretical results prove transferability of slicing parameters under distributional shift (Liu et al., 24 Nov 2025).

3. Advanced Distributional Alignment Objectives and Regularizations

Recent advances augment classical geometric alignment by integrating additional constraints, objectives, and statistical structure:

• Information-theoretic regularization: InfoOT incorporates mutual information maximization, encouraging transport plans that align points with coherent, statistically-dependent features, robustifying alignments against outliers and fostering transferability to unseen samples (Chuang et al., 2022).
• Partial/Adaptive OT: Adaptive OT relaxes marginal constraints, endowing OT with intrinsic mass selection: data-driven adaptive-mass transport is optimal for partial or noisy dataset alignment without pre-specifying transported mass (Yang et al., 7 Mar 2025).
• Preference and safety alignment in LLMs: Distributional OT objectives express token-level alignment and stochastic-dominance constraints for preference alignment (PLOT (Zhu et al., 2 Apr 2026); AOT (Melnyk et al., 2024)), as well as safety-driven push–pull distributional alignment via dual-reference OT optimization (SOT (Wang et al., 12 Jan 2026)). These methods connect global distributional geometry with alignment desiderata not achievable by instance-level or aggregate objectives.
• Unbalanced OT and variational codebook optimization: Distributional alignment in open-ended evaluation or explainability leverages unbalanced OT metrics (UOT) on compacted representations of high-dimensional data, such as value codebooks for cultural alignment assessment (Lee et al., 16 Mar 2026) or group-level counterfactuals (You et al., 2024).

4. Statistical Guarantees, Optimization Algorithms, and Computational Considerations

OT and distributional alignment frameworks benefit from a mature theory and diverse algorithmic landscape:

• Duality and efficient solvers: Kantorovich duality facilitates convex optimization, LP relaxations, and scalable entropic-regularized solvers with Sinkhorn iterations dominating applied workflows (Courty et al., 2015, Chuang et al., 2022).
• Power diagram and Laguerre cell methods: SDOT computes explicit transport maps as partitions of the continuous source space—enabling exact, deterministic mappings with guaranteed mass assignment and convergence (Kong et al., 16 Oct 2025).
• Multi-marginal and barycentric optimization: Multi-marginal OT is tractable for Gaussian or low-rank structures, with explicit SDPs or Burer–Monteiro factorizations (Dandapanthula et al., 3 Dec 2025, Karimi et al., 2021).
• Minibatch, sliced, and amortized algorithms: Sliced OT and min-STP reduce memory and computational costs, and amortized learning generalizes alignment strategies across closely-related distribution pairs with statistical error guarantees (Liu et al., 24 Nov 2025).
• Statistical consistency and convergence rates: Theoretical bounds for OT-based estimators, sample complexities for regularized or sliced variants, and asymptotic consistency (often $\mu$3 in 1D) underpin the deployment of OT in high-dimensional, finite-sample regimes (Melnyk et al., 2024, Bateni et al., 12 Nov 2025).

5. Specialized Settings: Time-Series, Denoising, and Causal Alignment

OT-based methodologies naturally extend to temporal and causal settings:

• Distributional time-series analysis: Autoregressive OT models regress optimal maps along Wasserstein geodesics; stationary solutions and least-squares estimation are analyzed via contraction principles, with applications in spatial-temporal data and climate analysis (Zhu et al., 2021).
• Causal inference with heterogeneous sites: Fused Gromov–Wasserstein frameworks synthesize counterfactual treatment effect distributions by optimally aligning feature–outcome distributions across sites and applying learned transport maps to intervention groups; convergence to the oracle distribution is guaranteed under regularity conditions (Bateni et al., 12 Nov 2025).
• Optimal transport denoisers: A hierarchy of denoisers with increasing higher-order scores (score-based transport maps) bridges basic MMSE denoisers and the Monge map between noisy and signal distributions, providing plug-in estimators with precise OT guarantees (Liang, 10 Dec 2025).

6. Interpretability, Robustness, and Limitations

OT-based alignment also advances interpretability and model insight:

• Anchor-based and hierarchical interpretations: Factorization via anchors or hierarchical structures clarifies how mass flows between data substructures, improving model explainability and trust (Lin et al., 2020, Lee et al., 2019). Visualization of transport plans exposes data geometry and the correspondence between clusters or latent factors.
• Robustness to outliers/noise: Structured OT methods, adaptive regularization, and information-augmenting objectives all address classical OT’s sensitivity to support mismatch, imbalance, and outlier points.
• Limitations and open problems: Challenges remain in scaling OT to massive datasets (necessitating stochastic, low-rank, or sliced approximations), adapting to non-Euclidean cost structures (graphs, manifolds), and automatically selecting regularization or model-complexity parameters. Theoretical understanding of sample complexity in high-dimensional, adaptive, or structured OT remains an active area of research.

7. Summary Table: Major Classes of OT-Based Distributional Alignment

Method / Objective	Key Property	Typical Strengths
Entropic/Sinkhorn OT	Regularized/smooth plan	Efficiency, scalable computation
SDOT (Laguerre mapping)	Explicit deterministic	Low variance, provable convergence
Hierarchical / Anchor	Structured alignment	Robustness, interpretability
Sliced / 1D projections	Low memory / compute	Scalability, transferability
InfoOT (MI reg.)	Coherence preservation	Robustness to outliers, semantics
Adaptive OT	Data-driven mass select	Noise/outlier resilience
Multi-marginal OT	Alignment of families	Barycentric estimates, regression

In summary, optimal transport and distributional alignment form a unified, geometrically and statistically principled paradigm for complex data transformation and integration. OT-based approaches underpin state-of-the-art advances in generative modeling, domain/task adaptation, explainability, and robust machine learning (Kong et al., 16 Oct 2025, Courty et al., 2015, Lin et al., 2020, Chuang et al., 2022, Liu et al., 24 Nov 2025). They provide well-founded mechanisms for matching distributions under varied structural, statistical, and computational constraints, with growing impact across foundational and applied research.