Optimal transport (OT) is widely used to align distributions and transfer labels between unpaired source and target data, but source-target composition shifts can cause rare classes to be systematically under-predicted even with high transport confidence. We identify a source of this failure in the transport marginal: under the standard uniform weighting of source points, each class has transport capacity proportional to its source frequency, which can be insufficient to meet its target demand even when separated well in the embedding space. We show that displaced mass is preferentially absorbed by high-capacity classes rather than geometrically similar classes, producing confident errors that are difficult to detect from the transport plan. Downstream read-outs, entropic smoothing, and marginal relaxation provide limited recovery, whereas matching the transport marginal to the target composition substantially improves label transfer. On single-cell data, the oracle composition-matched marginal raises macro-F1 from 0.322 to 0.623, and the best deployable estimator closes 69% of the oracle gap. Correcting the marginal also separates a coverage-dependent component of error from residual error associated with class separability, showing that transport capacity and representation geometry are distinct sources of transfer error.
Figures & tables
Figure 1: Capacity and correction. (a) c4 is rare in the source yet well separated; c1,c2 overlap. (b) Under the default μ= source frequencies, c4 is starved and its mass is absorbed by the abundant c1 . (c) Replacing it with μ′=π^/ns , π^ estimating the target composition, restores c4 ’s budget. (d) Correct assignments return; the c1↔c2 off-diagonal is unchanged. (e) Recovered mass is the capacity term; the residual tracks separability.
Source marginal
B_Plasma
T_NK
Myeloid
Malignant
Epithelial
Endothelial
Stromal
source frequencies
.65/.63
.95/.19
.30/.53
.64/.55
.36/.53
.05/1.00
.09/.89
uniform per class
.73/.66
.80/.51
.43/.69
.66/.63
.72/ .26
.64/ .29
.60/.73
composition-matched
.78/.68
.75/.71
.61/.62
.61/.78
.45/.48
.46/.72
.69/.72
Table 1: Per-class recall / precision on single-cell data (Jint), natural target composition, 3 seeds. The default marginal produces the capacity signature: a capacity-limited class has low recall but high precision, while the largest class has high recall but low precision.
single-cell (Jint)
Office-Home
partial ρ (absorbed share, capacity ∣ distance)
+0.566
+0.329
partial ρ (absorbed share, distance ∣ capacity)
+0.144
−0.206
n (class, absorber) pairs / units
198/27
14,080/12
Table 2: Where displaced mass goes. Partial Spearman correlations of absorbed share with class capacity and embedding distance, each controlling for the other.
intervention
rk<0.15
rk≥0.15
marginal shape (matched to target)
+0.366
+0.016
marginal shape (uniform)
+0.274
+0.038
τ=0.01 (most relaxed)
+0.208
+0.053
τ=∞ (balanced Sinkhorn)
−0.021
−0.013
ε ( 3× default)
+0.000
−0.012
permuted marginal (control)
−0.015
−0.079
Table 3: Effect of the three OT parameters in Eq. 1 , stratified by capacity. Entries are changes in macro recall relative to the default source marginal on Office-Home.
single-cell (Jint)
Office-Home
estimate of π
macro
F1
micro
below
macro
F1
micro
%
source frequencies
0.433
0.322
0.393
0.030
0.392
0.359
0.437
n/a
permuted (control)
n/a
n/a
n/a
n/a
0.333
0.284
0.360
n/a
uniform
0.654
0.512
0.557
0.533
0.501
0.396
0.448
6%
k NN pre-pass
0.648
0.556
0.616
0.463
0.487
0.428
0.518
48%
JCPOT
0.585
0.505
0.595
0.280
0.467
0.434
0.537
59%
Table 4: Label-transfer performance under different source marginals and target-composition estimates. Results are shown for single-cell data (Jint) and Office-Home under the natural target composition.
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
transport
k NN vote
classes resolvable
65/65
64/65
classes with a finite r⋆
36/65
47/65
median R2
0.276
0.142
median r⋆
0.142
0.108
median ceiling a
0.657
0.606
fits declined at the a≤1 bound
0/36
0/47
Appendix
Table 5: The saturating law fitted on Office-Home, transport read-out against a k NN vote. Same protocol as the single-cell fits: per class, recall(t)≈art/(r⋆+rt) , ceiling constrained to a≤1 , and any fit landing within 10−3 of that bound declined. No fit is pinned here (the largest ceiling is 0.9951 ), so nothing is declined for that reason. The form does not transfer: the fits are poor, a threshold is undefined for roughly half the classes under transport, the transport-to-vote ratio is close to unity with a range spanning both sides of it, and no class sits at chance. Single-cell fits are not tabulated alongside, because the logged Janesick fits predate the a≤1 convention; Sec. 3.1 rests the single-cell claim on the fit-free resolvability count instead, which no fitting convention can affect.
Figure 2: The failure is in the coupling, and the marginal repairs it. Entropic unbalanced OT couplings between the source (rows) and the target (columns), both sorted by class. Each class is allocated the same number of bins on each axis, so a class holding 0.3% of the source is as legible as one holding half of it, and colour is the share of each target point’s transported mass falling in a given source bin, on one scale shared by both panels and all classes. Left: the source marginal is the source’s own class frequencies. Right: the marginal is rebuilt as μj=π^yj/nyjs from a k NN pre-pass estimate of the target composition, which uses no target labels. Everything else is identical. The horizontal band on the T_NK rows in the left panel is the absorption mechanism of Sec. 3.2 made visible: target points of a starved class are transported to a class that has capacity. Endothelial, the one class below capacity here, is starred and its diagonal block is ringed; uncorrected it sends 46.9% of its mass to T_NK and 3.9% to itself, and corrected it sends 51.3% to itself. Aggregated over the whole target, the T_NK source rows absorb 47.8% of all transported mass before correction and 14.4% after.
Figure 3: The correction acts on class correspondence, not on domain alignment. UMAP of the source together with the target mapped into source space by each row’s transport plan through barycentric projection. The reducer is fit once, on the source alone, and both rows are transformed into that fixed embedding, so the source points are identical in every panel and only the target moves; any difference between rows is therefore attributable to the plan and not to a refit. Top row: source-frequency marginal. Bottom row: composition-matched marginal from the k NN pre-pass. Left column: points coloured by domain, which measures alignment. Right column: points coloured by class, which checks that the biology is preserved. Domain mixing, the fraction of each point’s fifteen nearest neighbours drawn from the other domain on the plotted coordinates, where 0.5 is perfect interleaving, improves only from 0.346 to 0.387 : the two domains were already reasonably interleaved, and alignment is not what the marginal repairs. The class panels show what does change, with the Endothelial population diffuse and largely mislabelled before correction and consolidated after, and below-capacity recall rising from 0.036 to 0.543 .
class
r
recall before
recall after
diagonal mass
Endothelial (below capacity)
0.141
0.036
0.543
3.9%→51.3%
Stromal
0.320
0.094
0.619
8.7%→55.8%
Epithelial
1.903
0.238
0.514
23.7%→48.6%
Malignant
2.305
0.561
0.896
53.1%→84.9%
Myeloid
3.135
0.314
0.442
32.5%→43.3%
B_Plasma
5.943
0.711
0.853
69.9%→83.9%
Appendix
Table 6: Jint fold G1, per class, ordered by coverage ratio. Diagonal mass is the share of a class’s target mass landing on its own source block, read off the plotted matrices of Fig. 2 . Endothelial is the only class below capacity on this fold.
family
read-out
Jint
Janesick rep1
Janesick rep2
hard
argmax
0.040
0.071
0.211
tempered
softmax T=1
0.040
0.071
0.211
tempered
softmax T=0.05
0.040
0.071
0.211
soft
soft posterior
0.040
0.071
0.211
propagating
OTP, 1 step
0.062
0.085
0.254
propagating
OTP, 20 steps
0.063
0.086
0.259
Appendix
Table 7: Read-out battery, below-capacity recall under the natural marginal. Every arm except k NN consumes the same converged plan. Hard, tempered and soft read-outs are identical to four decimals; propagation buys a little and twenty steps buy nothing over one; the only plan read-outs that help are the two that divide by source class count, which is the marginal correction relocated after the solve. n=12 (Jint), 2,412 (rep1), 1,926 (rep2).
dataset
coupling mass
expected recall
argmax recall
extracted
Jint
0.0041
0.041
0.040
96%
Janesick rep1
0.0190
0.077
0.071
93%
Janesick rep2
0.0160
0.183
0.211
115%
Appendix
Table 8: The coupling-mass ceiling, natural marginal. Coupling mass is the share of total transported mass landing on a below-capacity class and expected recall is what a perfect read-out of that plan could achieve; both are read-out-independent properties of the coupling. The last column is what argmax actually extracts as a fraction of that ceiling.
Janesick
Jint
Janesick
τ
macro
macro
absorption
Endothelial / Stromal
identified r⋆
0.01 (most relaxed)
0.451
0.586
0.656
0.217 / 0.435
6/14
0.1
0.399
0.519
0.863
0.091 / 0.222
4/14
1 (default)
0.306
0.426
0.974
0.020 / 0.058
0/14
10
0.267
0.407
0.985
0.012 / 0.039
0/14
∞ (balanced)
0.262
0.404
0.987
0.012 / 0.036
0/14
Appendix
Table 9: The τ sweep on both datasets. Identified r⋆ is the number of the fourteen Janesick classes yielding a fit that is not declined; at the default τ and tighter, every fit is declined with its ceiling pinned and the note that the curve does not saturate within the observed window.
domain
marginal
τ=0.01
0.1
1
10
∞
span
z
Office-Home
default
0.450
0.359
0.217
0.193
0.190
0.260
11.8
deployable
0.477
0.451
0.431
0.426
0.426
0.051
1.8
oracle
0.528
0.542
0.541
0.538
0.537
0.014
0.5
single-cell
default
0.408
0.203
0.057
0.044
0.042
0.366
8.7
deployable
0.547
0.542
0.552
0.555
0.555
0.013
0.2
oracle
0.549
0.572
0.634
0.643
0.644
0.096
1.3
Appendix
Table 10: Marginal ×τ , below-capacity macro-F1. Deployable is JCPOT on Office-Home and the k NN pre-pass on single-cell data, each the leading estimator in its own panel; oracle uses the true composition and is an upper bound rather than a method. Span is best minus worst within the row, with its z . On a wrong marginal τ is a large lever; on a corrected one no setting is distinguishable from any other. Office-Home is 4 domain pairs ×3 seeds and single-cell is 3 folds ×3 seeds, both complete. A further arm at τ=0.001 leaves every conclusion unchanged (App. E ).
ε
degenerate
below-cap. recall
macro
sel.
conf. correct
conf. wrong
0.01×
yes
0.000
0.143
1.00
—
—
0.03×
yes
0.132
0.312
1.15
—
—
0.1×
no
0.027
0.441
2.94
0.974
0.874
0.3×
no
0.027
0.442
2.89
—
—
1×
no
0.024
0.438
3.28
0.927
0.744
3×
no
0.002
0.408
5.15
0.817
0.554
Appendix
Table 11: The ε sweep, as a multiple of the calibrated default. No setting has both meaningful below-capacity recall and a detector clearing the pre-registered bar. The two smallest values produce a numerically degenerate plan. Sel. is the selectivity of the best output-side detector; conf. are mean confidences on correct covered points and on wrong below-capacity points. Six scorable units: one fold has no below-capacity class, so the detector suite is undefined there and it contributes no rows. The bar is cleared on at most three of six units at any ε , never all six.
prior for π
removable ∼logr
removable ∼ sep
residual ∼ sep
residual ∼logr
oracle (published)
−0.813
+0.135
−0.338
+0.034
uniform
−0.573
+0.316
−0.502
−0.157
k NN pre-pass
−0.517
+0.263
−0.374
−0.435
JCPOT
−0.516
+0.188
−0.313
−0.287
BBSE
−0.541
+0.282
−0.422
−0.296
RLLS
−0.538
+0.289
−0.429
−0.305
Appendix
Table 12: The decomposition under seven priors, Office-Home, n=780 . Cells are partial Spearman ρ with the other predictor controlled. Removable stays coverage-governed and residual stays separability-governed on every arm, while the near-zero cells at the oracle become non-zero under every estimate.
Figure 4: Dose-response under a deployable class prior. Coverage deciles as in Fig. 8 , with the oracle marginal, the best deployable estimator and uniform balancing overlaid. The shape, the sign change and the cost at high coverage survive; the amplitude at the starved end and the break-even coverage do not.
single-cell (Jint)
Office-Home
α
macro-F1
below F1
absorption
macro-F1
below F1
absorption
0.0
0.322
0.062
0.971
0.349
0.211
0.846
0.1
0.383
0.232
0.852
0.386
0.298
0.754
0.2
0.429
0.338
0.763
0.413
0.368
0.669
0.3
0.466
0.418
0.685
0.434
0.413
0.599
0.4
0.498
0.476
0.620
0.453
0.452
0.534
Appendix
Table 13: Performance along the path from the source composition to the target. Single-cell is 3 folds ×3 seeds and Office-Home is 4 domain pairs ×3 seeds. Rows are the eight values of α measured on both domains. Absorption falls monotonically throughout, so the mechanism tracks the marginal rather than the metric.
Figure 5: Office-Home. Performance along the interpolation path, plotted against α (left) and against the divergence remaining to the target composition (right), with the divergence axis reversed so both panels read from the default marginal to the oracle. Bands are standard errors over 12 units. Circles mark where each estimator of Sec. 4.3 sits: its horizontal position is the divergence of its own estimate from the target, and its vertical position is its measured below-capacity macro-F1. The two transport-native estimators land on the curve, so for them the divergence alone predicts what the correction will deliver.
Figure 6: Single-cell. As Fig. 5 . The horizontal axis of the right panel spans roughly twice the range, because the single-cell composition shift is about twice as large in divergence. The k NN pre-pass is the best estimator on both axes at once: it achieves the lowest divergence to the target composition and the highest below-capacity macro-F1 of any arm.
domain
estimate of π
divergence
path equivalent
below-capacity F1
Office-Home
JCPOT
0.252
0.264
0.426
k NN pre-pass
0.257
0.258
0.407
RLLS
0.273
0.234
0.437
MLLS
0.351
0.236
0.460
uniform
0.423
0.071
0.465
single-cell
k NN pre-pass
0.585
0.408
0.583
Appendix
Table 14: Where each estimator of Sec. 4.3 falls relative to the path. Divergence is the mean KL from the estimate to the target composition, and path equivalent is the α at which the interpolation reaches the same divergence. Estimators whose divergence exceeds the endpoint separation lie off the swept range and are omitted.
solver
macro
Endothelial
absorption
detector (side)
median sel. / pass
k NN vote (reference)
0.638
0.295
0.541
plan confidence (out)
2.43 / 6/27
UOT + centering
0.426
0.020
0.974
split conformal (out)
2.38 / 1/27
moscot
0.424
0.014
0.981
reject-sink OT (out)
1.51 / 0/27
UOT, no centering
0.320
0.009
0.994
transport entropy (out)
0.42 / 0/27
coverage rule (source)
24.50∗ / 3/27
Appendix
Table 15: Solver invariance and output-side detectors, single-cell (Jint). Left: four independent entropic-OT implementations on identical embeddings, points, labels, folds and seeds, with a k NN vote on the same embedding for reference. Absorption is the fraction of a starved class’s target mass received by a single other class. Fixed-per-class target composition, so this is a within-run contrast across solvers rather than a cross-design comparison; the corrected natural-composition comparison is Table 17 . Right: detectors for below-capacity points, 9 folds ×3 seeds, against a bar fixed in advance of selectivity at least 3× at coverage at least 90% ; suite verdict 0/5 , with 10/135 unit evaluations passing. ∗ over the units where it fires.
Figure 7: The double dissociation. What the composition-matched marginal removes (top) and what survives it (bottom), against coverage (left) and class separability (right), on Office-Home. The two null panels are shown deliberately: what is removed is explained by coverage and not by separability, and what survives is explained by separability and not by coverage. Points are class-observations ( 4 domain pairs ×3 seeds), the line is a binned mean with standard error, and the band at residual =1.0 is classes the correction leaves at zero recall. The panels are interpretable only because the predictors were orthogonalized by construction; in the single-cell data they are confounded at ρ=−0.618 .
Figure 8: Dose-response of the correction. Change in macro recall from the composition-matched marginal against coverage, binned. Monotone on both domains, crossing zero and turning negative on the best-covered classes. The single-cell panel has 21 class-observations in 5 bins against Office-Home’s 780 in 10 and is a companion rather than an equally powered replication, and the x -ranges differ, so the two crossing points are not comparable thresholds.
foundation
PCA
cut
n
shape
τ
ε
confident errors
33/33
33/33
r<0.05
56
0.472
0.243
0.006
absorber = nearest
0/33
0/33
r<0.10
136
0.419
0.234
0.014
absorber = largest
30/33
33/33
r<0.15
235
0.366
0.208
0.013
confidence ratio
0.874
0.906
r<0.30
389
0.296
0.176
0.015
macro recall
0.445
0.495
detector suite
0/5
0/5
Appendix
Table 16: Left: backbone invariance, 9 folds ×3 seeds each on a 40,000 -point source, fixed-per-class target, so a within-run contrast. Counts differ slightly from Sec. 3.2 because that run uses natural composition and this one does not. Right: the below-capacity knob ordering of Table 3 at four cut points, as Δ macro recall against the default marginal.
read-out
macro
macro-F1
micro
below capacity
transport, default
0.433
0.322
0.393
0.030
k NN vote, uncorrected
0.627
0.553
0.625
0.351
linear probe, uncorrected
0.268
0.212
0.215
0.014
k NN, prior-corrected
0.654
0.571
0.640
0.430
k NN, composition-matched
0.635
0.577
0.652
0.414
linear probe, corrected
0.452
0.391
0.377
0.278
Appendix
Table 17: Matched corrections, single-cell (Jint), 3 folds ×3 seeds, natural composition. Composition-matched arms, transport and vote alike, use the true target composition and are therefore matched oracles; prior-corrected arms divide by the source class prior and use no target information, so they are matched to each other at a lower information level. The oracle-to-oracle comparison is transport 0.546 against vote 0.414 below capacity, where below is the below-capacity stratum, three (fold, class) cells averaged over seeds.
read-out
macro
macro-F1
micro
below capacity
transport, default
0.388
0.356
0.440
0.136
k NN vote, uncorrected
0.367
0.327
0.506
0.154
linear probe, uncorrected
0.199
0.136
0.313
0.004
k NN, prior-corrected
0.553
0.483
0.580
0.526
k NN, composition-matched
0.460
0.461
0.644
0.566
linear probe, class-balanced
0.590
0.533
0.641
0.542
Appendix
Table 18: Matched corrections on Office-Home, 4 domain pairs ×3 seeds, deep-starvation configuration. Below is the below-capacity stratum, n=235 class-observations.
replicate 1
replicate 2
estimate of π
macro
micro
macro
micro
source frequencies
0.141
0.224
0.334
0.477
uniform
0.206
0.149
0.591
0.346
k NN pre-pass
0.180
0.167
0.479
0.437
JCPOT
0.179
0.157
0.461
0.389
BBSE
0.161
0.316
0.346
0.436
Appendix
Table 19: Janesick under natural target composition, macro / micro recall by marginal, two replicates ×3 seeds.
Cross-city transfer improves prediction in label-scarce cities by leveraging labeled data from other cities, but it becomes challenging when cities adopt incompatible partitions and no ground-truth region correspondences exist. Existing approaches either rely on heuristic region matching, which is often sensitive to anchor choices, or perform distribution-level alignment that leaves correspondences implicit and can be unstable under strong heterogeneity. We propose SCOT, a cross-city representation learning framework that learns explicit soft correspondences between unequal region sets via Sinkhorn-based entropic optimal transport. SCOT further sharpens transferable structure with an OT-weighted contrastive objective and stabilizes optimization through a cycle-style reconstruction regularizer. For multi-source transfer, SCOT aligns each source and the target to a shared prototype hub using balanced entropic transport guided by a target-induced prototype prior. Across real-world cities and tasks, SCOT consistently improves transfer accuracy and robustness, while the learned transport couplings and hub assignments provide interpretable diagnostics of alignment quality.
Yuyao Wang, Min Yang, Meng Chen +3
Department of Mathematics and Statistics, Boston University, Boston, MA, USA · School of Software, Shandong University, Jinan, China · School of Geography, University of Leeds, Leeds, UK
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 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
Austin J. Brockmeier, Bilal Riaz
Department of Electrical and Computer Engineering, University of Delaware, Newark, DE, 19711 · Department of Computer and Information Sciences, University of Delaware, Newark, DE, 19711