Dynamic Neural Optimal Transport for Subset Alignment
Authors: Austin J. Brockmeier, Bilal Riaz
Organizations: Department of Electrical and Computer Engineering, University of Delaware, Newark, DE, 19711 · Department of Computer and Information Sciences, University of Delaware, Newark, DE, 19711
We study optimal transport between distributions when only a subset of the target population is relevant to the source. The goal is to map the source distribution to a reweighted target distribution whose density ratio with respect to the full target lies in [0,c]. This constraint enables the transported distribution to select an appropriate target subpopulation while limiting how strongly any region of the target can be emphasized. We propose a dynamic neural optimal transport framework for this setting using two neural networks. The first parameterizes a time-dependent potential whose spatial gradient defines the velocity field and whose terminal value enforces the density-ratio constraint. The second parameterizes an interpolation between the source distribution and the learned terminal distribution. Unlike prior neural dynamic optimal transport methods, which typically assume a fixed terminal distribution, our formulation learns the terminal distribution within the support of the target. We demonstrate the method on positive-unlabeled learning, unpaired domain translation, and trajectory inference, where meaningful alignments exist between source and target distributions, but direct distribution matching is inappropriate because the target contains additional or irrelevant components
Figures & tables
Figure 1: Subset alignment between two uniform distributions in R2 at t=1 , obtained using Algorithm 1 with c∈{1,2,4,8} . The black contour line indicates the zero-level set of φ1(x) .
Figure 2: Source, target, and transported source histograms with sampling densities for 1D subset alignment between a uniform distribution and a semicircular density. φ1 is shown in red.
Method
c=1
c=2
c=4
c=8
static ( Gazdieva et al., 2023 )
46.93
75.33
82.44
87.32
dynamic (one step)
64.85
75.16
93.57
95.84
dynamic (ODE)
58.80
70.47
92.68
95.00
Table 1: Classification accuracies of translated images for MNIST → EMNIST.
Figure 3: Image translation outputs for MNIST → EMNIST.
Dataset
PUSB
NTC-MI
static
dynamic
sorted
sign
sorted
sign
Abalone
0.544 ± 0.060
0.575 ± 0.025
0.561 ± 0.029
0.503 ± 0.008
0.578 ± 0.025
0.540 ± 0.054
Banknote
0.829 ± 0.050
0.922 ± 0.019
0.883 ± 0.037
0.882 ± 0.039
0.952 ± 0.022
0.939 ± 0.018
Breast-w
0.766 ± 0.145
0.870 ± 0.028
0.930 ± 0.028
0.941 ± 0.027
0.927 ± 0.024
0.931 ± 0.026
Diabetes
0.546 ± 0.042
0.700 ± 0.039
0.635 ± 0.044
0.629 ± 0.044
0.629 ± 0.043
0.630 ± 0.046
Haberman
0.513 ± 0.023
0.532 ± 0.066
0.539 ± 0.066
0.540 ± 0.067
0.544 ± 0.074
0.540 ± 0.075
Table 2: Comparison of average balanced accuracies of 20 models trained using dynamic subset alignment methods with PUSB and NTC-MI reported by Teisseyre et al. (2025) , and our formulation of incomplete extremal optimal transport from Gazdieva et al. (2023) (labeled static). The best methods are colored red and second best blue.
Class
Man
Woman
Young
15K
23K
Old
7K
3.5K
Table 3: Division of FFHQ train images.
Figure 4: FFHQ old → young using dynamic subset alignment with Euler integration (100 steps).
c=1
c=2
c=4
c=8
UOT
UOT
U
UOTM
UOTM
UOTM
Task
Accuracy
static
dynamic
static
dynamic
static
dynamic
static
dynamic
-SD
-GAN
-LOT
KL-KL
χ2 -KL
lin.– χ2
Young → Old
Target
84.09
83.45
70.47
79.23
47.63
74.51
27.07
73.93
87.33
84.25
81.78
71.06
72.68
80.78
Class
70.43
71.55
91.41
84.31
95.47
90.03
96.06
90.30
45.71
73.85
84.49
90.17
88.68
78.77
Old → Young
Target
96.06
93.36
92.55
86.65
85.04
85.03
72.93
83.33
97.39
95.88
87.79
88.67
89.70
94.19
Class
71.77
71.92
89.46
87.69
94.43
89.84
95.45
90.88
49.30
74.74
89.48
89.64
88.19
80.68
Man → Woman
Target
96.11
94.53
92.18
91.74
85.66
90.96
77.55
90.29
98.16
97.38
90.23
90.63
91.19
94.51
Table 4: Target and source accuracy (%) for different domain translations on the FFHQ dataset. Dynamic subset selection is evaluated using Euler integration with 100 steps.
Figure 5: Accuracy curves for c∈{1,2,4,8} , in comparison to results from LOT ( Gazdieva et al., 2024 ) , UOT-GAN ( Yang & Uhler, 2019 ) , UOT-SD and UOTM ( Choi et al., 2023 ; Choi et al., 2024a ) .
Model
MMD
W2
Model
MMD
W2
SB-FBSDE ( Chen et al., 2022 )
1.86e−2
6.23
Subset Alignment ( c=1 , 476/490)
5.21e−3
4.83
uDSB (no growth) ( Pariset et al., 2023 )
1.86e−2
6.27
Subset Alignment ( c=2.0 , 247/490)
6.10e−3
4.68
uDSB (w/growth) ( Pariset et al., 2023 )
1.75e−2
6.11
Subset Alignment ( c=4.0 , 148/490)
9.24e−3
5.40
WLF-OT no growth ( Neklyudov et al., 2024 )
5.04e−3
5.20
Subset Alignment ( c=8.0 , 103/490)
1.48e−2
6.60
WLF-UOT λ=0.1 ( Neklyudov et al., 2024 )
9.16e−3
5.01
Table 5: Results for the Melanoma dataset ( Pariset et al., 2023 , Table 1) with three marginals and train/test splits. MMD and W2 are measured between νt=48 and ρt=48 obtained by running dynamics from μt=8 . For subset alignment, dynamics are run from the selected subset of μt=8 (cells with φt=8≤0 ; number of selected test cells in parentheses), and results are reported at the iteration with the lowest MMD on a validation set; note that the validation cells are a subset of the training cells.
Figure 6: CelebA → Anime.
Metric
One-Step
ODE
Time/sample (s)
1.44×10−4
2.84×10−3
Table 6: Average evaluation time per sample for FFHQ domain translation tasks from Table 4. Evaluation is performed with batch size 128.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Temporal snapshots at t∈{0.0,0.25,0.50,0.75,1.0} for c∈1,2,4,8 . (a) Evolution between two 2D uniform distributions shown in Figure 1 ; the black contour line indicates the zero-level set. (b) Evolution from 1D uniform to semicircular density is shown in Figure 2 with support subset alignment. φt(x) is denoted by red curve.
Figure 8: Temporal snapshots at t∈{0.0,0.25,0.50,0.75,1.0} for c∈1,2,4,8 . (a) Evolution between two 2D bimodal distributions; the black contour line indicates the zero-level set of φt(x) . (b) Evolution from 1D uniform to another 1D uniform density, with dynamic subset subset alignment. φt(x) is denoted by red curve.
Figure 9: Static subset alignment results for c∈{1,2,4,8} . (a) Translation between two 2D uniform distributions shown in Figure 1 . (b) 2D bimodal distributions shown in Figure 8(a) . Red curves indicate the zero-level sets of η(x) .
Figure 10: Static subset alignment results for c∈{1,2,4,8} . (a) Transport from 1D uniform distribution to semicircular density shown in Figure 7(b) . (b) Transport from 1D uniform distribution to a uniform distribution shown in Figure 8(b) . Red curves indicate the η(x) .
Dataset
data-set size
dimension
π
batch size
Abalone
4177
8
0.16
20
Banknote
1372
4
0.44
10
Breast-w
699
9
0.34
10
Diabetes
768
8
0.35
6
Haberman
306
3
0.26
6
Heart
270
13
0.44
6
Appendix
Table 7: UCI datasets for PU learning, along with the total number of data points (n), dimension (dim), positive prior ( π ) and batch sizes employed in training the corresponding dynamic models.
Figure 11: (a) Normalized and (b) unnormalized confusion matrices for the EMNIST classifier discussed in Section 4.2 .
Figure 12: Confusion matrices for MNIST → EMNIST domain translation using static subset selection. Accuracy is computed as the ratio between the trace and the sum of all entries in the confusion matrices.
Figure 13: Confusion matrices for MNIST → EMNIST domain translation using dynamic subset selection.
Figure 15: FFHQ young → old translation using (a) static and (b) dynamic subset selection. Dynamic subset selection is evaluated using Euler integration with 100 steps.
Figure 16: FFHQ woman → man translation using (a) static and (b) dynamic subset selection. Dynamic subset selection is evaluated using Euler integration with 100 steps.
Figure 17: FFHQ man → woman translation using (a) static and (b) dynamic subset selection. Dynamic subset selection is evaluated using Euler integration with 100 steps.
Optimal Transport (OT) provides a principled framework for learning transformations between probability distributions from unpaired samples. In many applications, however, a single transformation must map several source distributions to a common target distribution. For example, image restoration might require handling different types of degradation without knowing the degradation of each input at inference time. Simple approaches of pooling the source distributions only encourage alignment with the target at the aggregate level and may leave individual sources misaligned. In our paper, we consider the simultaneous OT problem which formalizes the task of learning a shared transport map that minimizes the average transport cost while aligning each source distribution with a prescribed target. We propose a neural method for solving the simultaneous OT problem by learning a shared transport map that minimizes the average transport cost while aligning each source distribution with a prescribed target. We derive a max-min formulation for learning this map. We illustrate its application to image restoration, where a single model handles multiple degradation types using a common collection of clean target images.
Milena Gazdieva, Kirill Sokolov, Jiawei Chen +2
Applied AI Institute AXXX Moscow, Russia · Lomonosov Moscow State University Moscow, Russia · MIRIAI Moscow, Russia
We propose an implicit neural formulation of optimal transport that eliminates adversarial min--max optimization and multi-network architectures commonly used in existing approaches. Our key idea is to parameterize a single potential in the Kantorovich dual and reformulate the associated c-transform as a proximal fixed-point problem. This yields a stable single-network framework in which dual feasibility is enforced exactly through proximal optimality conditions rather than adversarial training. Despite the inner fixed-point computation, gradients can be computed without differentiating through the fixed-point iterations, enabling efficient training without requiring implicit differentiation. We further establish convergence of stochastic gradient descent. The resulting framework is efficient, scalable, and broadly applicable: it simultaneously recovers forward and backward transport maps and naturally extends to class-conditional settings. Experiments on high-dimensional Gaussian benchmarks, physical datasets, and image translation tasks demonstrate strong transport accuracy together with improved training stability and favorable computational and memory efficiency.
Yesom Park, Eric Gelphman, Stanley Osher +1
Department of Mathematics, University of California, Los Angeles · Department of Applied Mathematics and Statistics, Colorado School of Mines
In recent years, the machine learning community has increasingly embraced the optimal transport (OT) framework for modeling distributional relationships. In this work, we introduce a sample-based neural solver for computing the Wasserstein geodesic between a source and target distribution, along with the associated velocity field. Building on the dynamical formulation of the optimal transport (OT) problem, we recast the constrained optimization as a minimax problem, using deep neural networks to approximate the relevant functions. This approach not only provides the Wasserstein geodesic but also recovers the OT map, enabling direct sampling from the target distribution. By estimating the OT map, we obtain velocity estimates along particle trajectories, which in turn allow us to learn the full velocity field. The framework is flexible and readily extends to general cost functions, including the commonly used quadratic cost. We demonstrate the effectiveness of our method through experiments on both synthetic and real datasets.
Hailiang Liu, Yan-Han Chen
Department of Mathematics, Iowa State University, Ames, IA 50011, USA. · Department of Statistics, Iowa State University, Ames, IA 50011, USA.