Density-Reweighted Entropic Optimal Transport: Decoupling Geometry from Sampling Density

Introduction

• EOT provides a tractable way to establish cross-dataset correspondences, whereas exact OT scales cubically with sample count.EOT’s transport plan encodes pairwise affinities and can be computed efficiently using Sinkhorn scaling.Source: §1 · pp. 1–2

Our Contribution

• With uniform empirical marginals, EOT can match geometrically similar structures by relative sampling density rather than proximity; UOT allows total-mass variation but adds tuning and offers little intuition about density–geometry effects.Such mismatches can arise from technical acquisition biases in genomics.Source: §1 · p. 2
• Reweighting the Gaussian kernel and marginal constraints gives sparse-region points greater importance; discount factor θ ∈[0, 1] interpolates from standard EOT at θ = 0 to density-independent geometry-driven alignment at θ = 1.Intermediate θ values trade off density-sensitive and geometry-driven alignment.Source: §1 · p. 3
• Across multiple hyperparameter configurations, density-reweighted EOT outperforms standard EOT and UOT in recovering geometrically faithful correspondences under substantial sampling-density disparity.The experiments also discuss practical considerations and guidance for deployment.Source: §1 · p. 3
• Under suitable manifold-geometry and sampling-density regularity conditions, W(θ) converges with high probability, up to global scaling, to a population plan whose density dependence is governed by θ.At θ = 1, the limit is the population EOT plan under uniform measures; the results cover exact and kernel-estimated densities with explicit convergence rates.Source: §1 · p. 3

Unbalanced Optimal Transport

• UOT relaxes hard marginal constraints, but its influence on geometry versus sampling density is implicit; density-reweighted EOT makes that tradeoff explicit through θ.With KL penalties, UOT approaches EOT as λ1, λ2 → ∞ and the Gaussian kernel as λ1, λ2 → 0.Source: §2 · p. 3

Manifold Learning

• Diffusion Maps use α-normalization to control density influence, motivating EOT’s θ, which ranges from standard EOT at θ = 0 to geometry-only alignment at θ = 1.At θ = 1, the transport plan is invariant to the sampling measure; intermediate values balance density influence and geometry.Source: §2 · pp. 3–4
• Unlike EOT Eigenmaps [29], which assumes both datasets share a sampling density, density-reweighted EOT targets differing densities and can support robust joint manifold learning.EOT Eigenmaps uses singular vectors of the EOT plan to jointly embed datasets sampled from a shared low-dimensional manifold.Source: §2 · pp. 4–5

Density reweighting

• Inverse-density reweighting in the transport cost, entropy, and marginal constraints assigns more mass to sparse-region points and less to dense-region points; a scaling factor conserves total mass.Source: §3 · p. 5

Theoretical guarantees

• The convergence guarantees assume smooth, compact, boundaryless embedded manifolds, strictly positive C3 sampling densities, and sample sizes satisfying n ≥ m ≥ n^γ for a constant γ in (0, 1).Source: §3 · p. 5

Population-level plans

• As θ increases from 0 to 1, the reweighted densities interpolate from the sampling densities to uniform manifold densities, defining population EOT plans doubly stochastic under those measures.The associated positive continuous scaling functions exist and are unique up to a positive multiplicative constant.Source: §3 · p. 6

Theoretical guarantees

• m^-1/2 up to logarithmic factors is the stated error rate with exact densities; at θ=1, the plan depends only on manifold geometry, and Figure 3 shows both normalized ℓ1 and ℓ∞ errors decay approximately at that rate.The convergence result holds with high probability under the stated assumptions, up to a global scaling constant.Source: §3 · pp. 5–6, Figure 3 · p. 6

Estimated densities

• With unknown sampling densities, mean-normalized KDE estimates yield plans converging to the population plan up to global scaling; the bound depends on bandwidths, sample sizes, and intrinsic dimensions.The estimator is motivated by invariance of the correspondence structure to global multiplicative scaling of density estimates.Source: §3 · p. 7

Experiments

• The test uses two translated manifolds, each formed from a unit circle and a radius-2 semicircle, with opposing component sampling probabilities of 0.8/0.2 and 0.2/0.8.Additional density variation is imposed within each component.Source: §4 · p. 8
• R_k measures top-k coupling neighbors against true geometric neighbors, macro-averaging equally across components and symmetrically across both transfer directions.True neighbors are computed after translating Y back by the known shift and finding Euclidean nearest neighbors.Source: §4 · p. 8
• Across k=5–100, θ=1 achieved R_k consistently about 2× or more above the best UOT configuration and nearly an order of magnitude above standard EOT, averaged over 10 trials.Standard EOT performed worst across all neighborhood sizes.Source: §4 · p. 8, Figure 4 · p. 9
• UOT partially mitigates standard EOT’s failure, but its performance varies substantially across λ configurations, complicating unsupervised selection of marginal relaxation.The passage notes that dataset alignment typically offers no natural criterion for choosing the degree of marginal relaxation.Source: §4 · p. 9

A.1 Experimental Setup for Figures 1, 2, and 3

• The toy datasets lie on a line and a curved manifold over the shared first-coordinate domain [0, 1], with opposite 90%/10% sampling probabilities across the two halves.Each dataset’s second coordinate is determined by its respective curve equation.Source: §A.1 · p. 9
• Standard EOT uses uniform empirical marginal constraints, while the proposed method computes W(θ) with ground-truth sampling densities; both use ε = 5 × 10−2.Source: §A.1 · p. 9
• Figure 1 uses n = 3,000 and m = 2,000, with θ = 1 in Algorithm 1.Source: §A.1 · p. 9
• Figure 2 uses n = 3,000 and m = 2,000 and computes W(θ) over a uniform grid of θ values on [0, 1].Source: §A.1 · p. 9
• For Figure 3, standard EOT on an equally arc-length-spaced 10,000-point grid supplies scaling vectors, which are spline-interpolated to approximate continuous scaling functions.The approximations are used to calculate the population-level plan entrywise.Source: §A.1 · pp. 9–10
• The convergence study varies m across 10 log-uniform values from 100 to 10,000, sets n = 2m, and averages 20 independent repetitions; manifold volumes are computed by arc-length integration.The passage says two error metrics are considered but does not provide their names in the supplied text.Source: §A.1 · p. 10

A.2 Estimation Performance of mean-normalized KDE

• The KDE lemma assumes iid samples from a positive C^3 density on a smooth, compact, boundaryless manifold and introduces positive manifold- and density-dependent constants and a sufficiently large sample-size threshold.The supplied passage truncates the bandwidth condition, so its precise guarantee is not recoverable here.Source: §A.2 · p. 10

A.3 Experimental Setup for Figure 4

• Mx’s unit-circle and semi-circle components use distinct wrapped Gaussian-mixture sampling densities.For the unit circle, ϕ = ˜ϕ mod 2π with ˜ϕ ∼ 0.3 · N(0.5π, 0.8²) + 0.7 · N(1.5π, 1.0²); for the semi-circle, ϕ = ˜ϕ mod π with ˜ϕ ∼ 0.6 · N(0.25π, 0.3²) + 0.4 · N(0.65π, 0.8²). The component densities are chosen arbitrarily, not to favor a method.Source: §A.3 · p. 10
• My’s unit-circle and semi-circle components likewise use distinct wrapped Gaussian-mixture sampling densities.For the unit circle, ϕ = ˜ϕ mod 2π with ˜ϕ ∼ 0.5 · N(0.2π, 0.5²) + 0.5 · N(1.3π, 0.9²); for the semi-circle, ϕ = ˜ϕ mod π with ˜ϕ ∼ 0.25 · N(0.1π, 0.4²) + 0.75 · N(0.7π, 0.6²). The component densities are chosen arbitrarily, not to favor a method.Source: §A.3 · p. 10
• KDE bandwidths for X and Y are selected independently with Algorithm 2 using k = 20, m = 50, α = 0.05, and C = 1.The passage also specifies 2h2 max = 1.Source: §A.3 · p. 10
• Each method’s ε is selected by maximizing the number of MKNN pairs (K = 50) via golden-section search over [0, 5], stopping below interval width 0.10 or after 20 iterations.The search range was chosen based on the data scale.Source: §A.3 · p. 11
• With m = n = 2,000, Figure 4 compares standard EOT, the proposed method with θ = 1, and UOT at four λ = γε settings using R_k averaged over 10 replicates.The UOT settings use γ ∈ {0.1, 1, 10, 100}; remaining hyperparameters are selected as described in the setup.Source: §A.3 · p. 11

B.1 Proof of Proposition 3.1

• The feasible set’s nonemptiness and compactness, together with the objective’s strict convexity and continuity, guarantee a unique minimizer.The proof then identifies W(θ) as that minimizer by feasibility and first-order conditions.Source: §B.1 · p. 13
• A rank-one plan built from h_ij, r_i, and c_j satisfies both marginal constraints, establishing that the feasible set is nonempty.The proof states that a direct calculation using S(θ) yields the needed normalization.Source: §B.1 · p. 13
• The feasible set is closed as an intersection of hyperplanes with the nonnegative orthant, and bounded by entrywise bounds, so it is compact.Each entry satisfies 0 ≤ W_ij ≤ h_ij min{n r_i, m c_j}.Source: §B.1 · p. 13
• Substituting W(θ), defined in (6), into the Sinkhorn scaling conditions yields exactly the marginal constraints in (9), proving feasibility.The proof uses the definition of M(θ) in (4).Source: §B.1 · p. 13
• Introducing row and column multipliers and setting the Lagrangian derivatives with respect to each plan entry to zero gives the first-order conditions.These conditions are sufficient to identify the unique minimizer, as stated in the proof overview.Source: §B.1 · p. 13
• The proof notes that W(θ) factors in the same way as the plan characterized by the first-order conditions.This factorization links W(θ) to the stationarity argument.Source: §B.1 · p. 13

B.2 Auxiliary Definition and Lemma for the Proof of Theorem 3.2 and 3.4

• A positive matrix can be diagonally scaled to prescribed row and column sums when their total masses agree; the diagonal vectors are its scaling factors.The setup follows Section SM5 of [29].Source: §B.2 · p. 13
• For a positive matrix and equal-total-mass marginals, valid scaling factors exist and are unique up to reciprocal scalar multiplication.Source: §B.2 · p. 13
• If positive vectors approximately scale a positive matrix to given marginals, Lemma B.2 guarantees true scaling factors close to them under its stated conditions [29].The lemma’s quantitative bounds depend on the matrix entry minimum and maximum and the row and column sums.Source: §B.2 · pp. 13–14
• Definition B.3 uses a probability-order notation in which the stated bound holds with probability at least 1 − n^-t for sufficiently large t, n, and m under Assumption 2.The section also invokes union bounds for polynomially many such variables and Taylor expansion for smooth transformations.Source: §B.2 · p. 14

B.3 Proof of Theorem 3.2

• For fixed θ, approximate scaling factors are used to scale M(θ) to the prescribed row sums r and column sums c.Source: §B.3 · pp. 14–15
• Hoeffding bounds and union bounds under Assumption 2 make the row and column estimates hold simultaneously across indices.Source: §B.3 · pp. 15–16
• Positivity and continuity of ε, µ, and ν on compact manifolds bound them away from zero and infinity, allowing the proof to bound coefficients in Lemma B.2.Source: §B.3 · p. 16

B.4 Auxiliary Lemmas for the Proof of Lemma A.1

• Under smoothness and compactness assumptions, the Gaussian-kernel expansion includes a geometric potential term and the Laplace–Beltrami operator.The result applies to a smooth, compact, boundaryless manifold and a C3 function; continuity makes the first-order term uniformly bounded by a constant multiple of h.Source: §B.4 · p. 17
• Lemma B.5 provides a high-probability kernel-density estimation bound for samples from a positive C3 density on a compact, boundaryless manifold.Its stated probability is at least 1−n^-2 for sufficiently large n and bandwidths in the specified admissible range.Source: §B.4 · pp. 17–18

B.5 Proof of Lemma A.1

• The proof of Lemma A.1 uses bandwidth choices that place the estimator within Lemma B.5’s admissible range.The bandwidth conditions are selected for sufficiently large m, so Lemma B.5 applies to the m-sample estimator.Source: §B.5 · p. 18
• The leave-one-out estimate at each sample point is related exactly to the full kernel-density estimate by removing that point’s unit self-kernel contribution.This identity is used to bound the difference between the two estimates.Source: §B.5 · p. 18
• A uniform event from Lemma B.5 and a concentration event are combined by a union bound, yielding probability at least 1−2m^-2.The proof uses Hoeffding’s inequality for the concentration step.Source: §B.5 · pp. 18–19
• The proof then bounds the normalization quantity away from zero and combines the preceding estimates to conclude Lemma A.1.The argument uses the bandwidth conditions and bounds on the deviation of the normalization quantity from the squared L2 norm of the density.Source: §B.5 · p. 19

• On the density-estimation event, the estimates are bounded away from zero and infinity, and fixed compact intervals contain both estimates and their targets independently of sample sizes and bandwidths.The positive continuous densities on compact manifolds provide the bounds used to construct these intervals.Source: §B.6 · p. 20
• For θ ∈ [0, 1], the mean value theorem controls transformed density-estimate errors uniformly over all sample points, with constants depending only on the manifolds, densities, and θ.The derivative of ϕ(z) = z^-θ is bounded on the fixed intervals containing the estimates and targets.Source: §B.6 · pp. 20–21
• For δ := τ C(ε)(ρ1 + ρ2) ≤ 1/2, achieved by choosing n0(ε, τ) and h′1, h′2 properly, positive vectors (α(θ)′, β(θ)′) scale M(θ) to row sums r and column sums c.Source: §B.6 · p. 22
• With probability at least 1 - n^-τ - 4n^-2γ, the multiplicative estimate implies the additive bound (17), completing the proof.The probability bound follows from intersecting the density-estimation event with the events underlying (89)–(99).Source: §B.6 · p. 23