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.