PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Keyi Li, Yuval Kluger, Boris Landa

arXiv:2608.16506v1stat.MLcs.LGstat.APstat.ME

TL;DR

Dataset alignment can be misleading when geometrically similar datasets have substantially different sampling densities. The paper introduces density-reweighted EOT to control this influence, proving convergence and recovering more geometrically faithful correspondences in simulations.

  • Problem

    EOT-based dataset alignment can produce geometrically misleading correspondences when geometrically similar manifolds have substantially different sampling densities.

  • Method

    Density-reweighted EOT discounts local sampling-density effects through a tunable factor θ, interpolating between standard EOT and geometry-driven alignment.

  • Results

    The method converges to population-level plans with explicit density dependence and outperforms standard EOT and UOT in recovering geometrically faithful correspondences.

  • Takeaways & Limitations

    The framework provides an interpretable way to reduce sampling-density bias in dataset alignment and obtain correspondences driven more by underlying geometry.

  • Takeaways & Limitations

    UOT performance varies substantially across hyperparameter configurations, limiting its practical utility when unsupervised alignment offers no natural criterion for choosing marginal relaxation.

Abstract

from arXiv · show

Dataset alignment is a central step in data analysis across science and engineering, where the goal is to match observations between datasets. Entropic Optimal Transport (EOT) offers a computationally tractable framework for this task by encoding cross-dataset affinities in a transport plan. However, when two datasets are sampled from geometrically similar low-dimensional structures with substantially different sampling densities, the EOT plan may match points by relative sampling density rather than geometric proximity, yielding geometrically misleading correspondences. To address this issue, we propose a density-reweighted EOT framework in which the influence of sampling density on the transport plan can be discounted to a desired degree, ranging from standard EOT to alignment driven purely by underlying geometry. Under suitable regularity conditions, we establish convergence of the reweighted EOT plan to a family of population-level plans whose dependence on sampling density is made explicit. Through simulations, we show that our approach recovers geometrically faithful correspondences, improving over related EOT-based frameworks when datasets exhibit substantial sampling density disparity.

1 Introduction

The paper addresses geometrically misleading EOT correspondences caused by sampling-density mismatch between datasets with similar underlying geometry. It introduces density-reweighted EOT to interpolate between density-sensitive and geometry-driven alignment, with convergence guarantees under regularity conditions.

  • Dataset alignment establishes correspondences across datasets and is fundamental in scientific and engineering applications.
  • Standard EOT can produce biased alignment when uniform marginals and mass conservation interact with sampling-density mismatch.
  • The framework provides a principled interpolation between purely geometry-driven alignment and the density-sensitive alignment of standard EOT.
  • The proposed density-reweighted EOT mitigates density dependence by reweighting the Gaussian kernel with local sampling density and imposing reweighted marginal constraints.
  • Under suitable regularity conditions, the reweighted plan converges with high probability to a population-level EOT plan whose density dependence is governed explicitly by θ ∈ [0, 1].At θ = 1, it converges to the population-level EOT plan under uniform measures on both manifolds.

2 Related Work

This section situates density-reweighted EOT among Unbalanced Optimal Transport, Diffusion Maps, and EOT Eigenmaps. It emphasizes controlling how sampling density influences alignment while preserving geometric information.

  • Unbalanced Optimal Transport: Unbalanced Optimal Transport relaxes hard marginal constraints with soft f-divergence penalties, accommodating variation in total transported mass.The KL-divergence formulation is widely adopted.
  • Unbalanced Optimal Transport: As λ1, λ2 →∞, UOT recovers EOT; as λ1, λ2 →0, it reduces to a Gaussian kernel encoding cross-data similarities through raw Euclidean distances.The parameters λ1 and λ2 control the degree of marginal relaxation.
  • Unbalanced Optimal Transport: For dataset alignment, UOT’s density-related behavior remains implicit because prior work does not characterize how λ1 and λ2 balance geometry against sampling densities.The proposed approach addresses this limitation by making the density dependence explicit.
  • Manifold Learning: Diffusion Maps motivate density discounting through α-normalization, which controls the sampling-density dependence of the recovered differential operator.At α = 0, density dependence is retained; at α = 1, it is completely removed.
  • Manifold Learning: EOT Eigenmaps use EOT-plan singular vectors to jointly embed datasets from a shared low-dimensional manifold, whereas this work targets differing sampling densities.The proposed reweighting can support joint manifold learning and embedding robust to dataset-specific sampling-density variation.

3 Density-Reweighted EOT and Theoretical Guarantees

This section defines density-reweighted EOT to discount sampling-density effects by reweighting costs, regularization, and marginal masses, with θ interpolating between sampling-density and uniform reference measures. It establishes population-level convergence under regularity assumptions and extends the guarantees to unknown densities estimated by mean-normalized KDE.

  • Density-reweighted EOT: Reweighting by 1/(f̂_iĝ_j)^θ in costs and regularization, and inversely scaling marginal masses, assigns more mass to sparse regions and less to dense regions.This compensates for nonuniform sampling while preserving total mass across the reweighted marginals.
  • Theoretical guarantees: For fixed θ and exact sampling densities, the finite-sample plan converges uniformly up to a global scaling constant with high probability, at rate m^-1/2 up to logarithmic factors.The bound holds simultaneously over all point pairs and depends on the smaller dataset size m.
  • Theoretical guarantees: Under smooth, compact manifold and positive C3 density assumptions, the population plan is doubly stochastic with respect to the density-reweighted measures, with positive continuous scaling functions unique up to a positive multiplicative constant.The sample-size assumption requires n ≥ m ≥ n^γ for γ ∈ (0, 1).
  • Population-level formulation: As θ increases from 0 to 1, the reweighted reference densities interpolate between the sampling densities and uniform densities on the respective manifolds.The resulting population-level plan is defined with respect to these density-reweighted measures.
  • Unknown-density setting: With unknown densities, mean-normalized KDE yields the same type of convergence guarantee up to global scaling, with bounds depending on bandwidths, sample sizes, and intrinsic dimensions.Mean normalization avoids requiring the intrinsic dimension and preserves correspondence structure under global multiplicative density-estimation errors.

4 Experiments

The experiments test alignment of translated manifolds with opposing sampling densities and show that density-reweighted EOT recovers geometrically faithful correspondences. It outperforms standard EOT and UOT, while UOT is sensitive to hyperparameter selection.

  • Experimental setup: The experiment aligns two translated manifolds consisting of a unit circle and a radius-2 semi-circle under substantially different sampling densities.The manifolds differ by a translation of (3, 5)⊤, while component sampling probabilities are reversed between datasets.
  • Evaluation: Alignment quality is measured using Rk, which macro-averages across components and symmetrizes transfer directions to avoid favoring either sampling density or direction.True neighbors are computed after translating Y back by the known shift.
  • Results: Across k ∈[5, 100], the density-reweighted method with θ = 1 outperforms every other method in Rk over 10 independent trials.Its performance is consistently higher at every examined neighborhood size.
  • Results: Rk for the proposed method exceeds the best UOT configuration by approximately 2× or more and standard EOT by nearly an order of magnitude.Standard EOT performs worst across all k, indicating that heterogeneous sampling marginal constraints can produce misleading correspondences.
  • Limitations: UOT performance varies substantially across hyperparameter configurations, and its sensitivity to λ may hinder unsupervised dataset alignment.Unsupervised alignment offers no natural criterion for selecting the degree of marginal relaxation.

Appendix A Details of Experiments and Additional Results · A.1 Experimental Setup for Figures 1, 2, and 3 · A.2 Estimation Performance of mean-normalized KDE

Appendix A specifies the toy experiments used to study density effects and convergence, and states conditions for analyzing mean-normalized KDE estimation. The experiments contrast datasets with reversed sampling-density profiles while varying reweighting and sample size.

  • A.1 Experimental Setup for Figures 1, 2, and 3: The toy example uses two-dimensional datasets supported on distinct curves sharing the first-coordinate domain [0, 1].The source lies on xi2 = xi1, while the target lies on yj2 = 2 + yj1 + 0.5 yj1^2.
  • A.1 Experimental Setup for Figures 1, 2, and 3: Sampling densities differ substantially because X places probability 0.9 on [0, 0.5] and 0.1 on [0.5, 1], with probabilities reversed for Y.Each point’s second coordinate is determined by its dataset’s curve equation.
  • A.1 Experimental Setup for Figures 1, 2, and 3: Figure 1 uses n = 3,000, m = 2,000, and θ = 1 in Algorithm 1.These settings define the displayed experiment for the fully specified reweighting parameter.
  • A.1 Experimental Setup for Figures 1, 2, and 3: Figure 2 varies θ on a uniform grid over [0, 1] and computes W(θ) for each value using Algorithm 1.The sample sizes remain n = 3,000 and m = 2,000.
  • A.1 Experimental Setup for Figures 1, 2, and 3: Approximate population-level EOT plans are constructed by cubic-spline interpolation of the estimated functions in the first coordinate and then calculated entrywise.The interpolated quantities approximate uε and vε defined in (12).
  • A.1 Experimental Setup for Figures 1, 2, and 3: The convergence experiment varies m over 10 log-uniform values from 100 to 10,000, sets n = 2m, and averages 20 independent repetitions.Two error metrics are considered, though their definitions are not included in the supplied passage.
  • A.2 Estimation Performance of mean-normalized KDE: Lemma A.1 analyzes mean-normalized KDE on a smooth, compact, boundaryless d-dimensional manifold with a positive C3 probability density and iid samples.The manifold is isometrically embedded in RD, and the estimator uses bandwidth h.
  • A.2 Estimation Performance of mean-normalized KDE: Under the lemma’s sample-size and bandwidth conditions, positive constants h0(Mx, µ), κ(Mx, µ), C0(Mx, µ), and m0(h0, κ) exist.The supplied statement specifies m ≥ m0(h0, κ), while the remaining bandwidth inequality is truncated.

A.3 Experimental Setup for Figure 4

Figure 4 evaluates transport-plan alignment on circle and semi-circle components with deliberately varied sampling densities, using fixed sample sizes and repeated comparisons across EOT-based methods. KDE bandwidths are selected independently for the two datasets using the specified bootstrap–Lepski settings.

  • Sampling Density: Sampling densities are parameterized by ϕ using distinct wrapped Gaussian mixtures on the unit-circle and semi-circle components of Mx and My.The four component-specific distributions use different mixture weights, means, variances, and wrapping domains.
  • Sampling Density: The component densities are arbitrary and are not tuned to favor any particular method.Points are generated as c + r(cos ϕ, sin ϕ)⊤ on circles or semi-circles of radius r centered at c.
  • Figure 4 Evaluation: Figure 4 fixes m = n = 2,000 and compares standard EOT, the proposed method with θ = 1, and UOT with λ = γε for γ ∈ {0.1, 1, 10, 100}.Each transport plan is evaluated using Rk from (18), with results averaged over 10 independent replicates.

Appendix B Proofs and Related Background · B.1 Proof of Proposition 3.1

The proof establishes that the feasible set is nonempty and compact, then verifies that W(θ) is feasible and satisfies sufficient first-order optimality conditions. Strict convexity and continuity therefore imply that W(θ) is the unique minimizer.

  • B.1 Proof of Proposition 3.1: The proof uses two steps: establish a nonempty compact feasible set, then identify W(θ) through feasibility and first-order conditions.These conditions, combined with strict convexity and continuity, establish existence and uniqueness.
  • B.1 Proof of Proposition 3.1: The proof introduces notation to simplify the expressions used in constructing and verifying the transport plan.The notation defines h_ij in terms of the reweighted quantities (ˆf_i)^θ and (ˆg_j)^θ.
  • B.1 Proof of Proposition 3.1: A direct calculation from S(θ) shows that a rank-one plan satisfies both marginal constraints, proving feasibility.The plan is defined by W_ij = h_ij r_i c_j / Σ_k r_k.
  • B.1 Proof of Proposition 3.1: The feasible set is closed and bounded because it is formed by finitely many hyperplanes and the nonnegative orthant, with bounded entries.Each entry satisfies 0 ≤ W_ij ≤ h_ij min{n r_i, m c_j}, so the feasible set is compact.
  • B.1 Proof of Proposition 3.1: Using (6) and the definition of M(θ) in (4), the constructed matrix W(θ) is shown to be feasible.The proof labels this verification as Step 2, Feasibility.
  • B.1 Proof of Proposition 3.1: Substitution into the Sinkhorn scaling conditions in (5) yields exactly the marginal constraints in (9).This connects the scaling representation to feasibility of W(θ).
  • B.1 Proof of Proposition 3.1: Introducing row and column multipliers λ and µ produces the Lagrangian used to derive first-order optimality conditions.Differentiation with respect to the plan entries and setting the derivative to zero gives the required stationarity equations.
  • B.1 Proof of Proposition 3.1: The factorization of W(θ) matches the factorization obtained from the stationarity conditions, identifying it as the optimizer.This factorization agreement completes the optimality argument.

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

This section defines diagonal matrix scaling and records existence, uniqueness, and stability results for scaling factors. It also introduces high-probability order notation used throughout the remaining proofs.

  • Matrix scaling: Matrix scaling diagonally rescales a positive matrix A to prescribed row sums r and column sums c, with scaling factors α and β.The target sums must satisfy ∥r∥1 = ∥c∥1.
  • Existence and uniqueness: For a positive matrix with matching total row and column sums, valid scaling factors exist and are unique up to the transformation (cα, c−1β) for c > 0.This is stated as Lemma B.1.
  • Stability of scaling factors: Lemma B.2 states that approximate scaling factors are close to true scaling factors for the corresponding row and column sums under bounded-entry and approximation conditions.The stability bounds use a = mini,j Ai,j, b = maxi,j Ai,j, δ ∈ (0, 1), C1, C2, and s = ∥r∥1 = ∥c∥1.
  • Order in probability: Definition B.3 introduces e O(ε) notation for random variables satisfying bounds with probability at least 1 − n−t when n and m meet the stated growth conditions.When n0 and C are independent of ε, the notation is written e Om,n(f(m, n)).
  • Order in probability: The proofs use the union bound and Taylor expansion properties associated with Definition B.3 without further restatement.These properties are applied throughout the remaining proofs.

B.3 Proof of Theorem 3.2

The proof controls the empirical scaling factors by first bounding sampling fluctuations, then applying a scaling-stability lemma under positivity and compactness assumptions. These bounds yield the multiplicative estimate that implies the additive bound in (13), completing Theorem 3.2.

  • Scaling-factor comparison: The proof compares the exact scaling factors α(θ), β(θ) with approximate factors that nearly enforce the prescribed row sums r and column sums c.The scaling factors diagonally scale M(θ) to satisfy the transport constraints, while the approximate factors are used in the perturbation argument.
  • Concentration bounds: Hoeffding’s inequality and union bounds control the relevant empirical quantities simultaneously across all rows and columns under Assumption 2.The argument conditions on individual sample points, derives analogous row and column bounds, and then applies union bounds.
  • Uniform boundedness: Positivity and continuity of ε, µ, and ν on compact manifolds provide uniform upper and lower bounds, which bound the coefficients required by Lemma B.2.These bounds establish positive constants controlling the quantities in the scaling-stability argument.
  • Completion of the proof: The scaling-stability lemma produces a pair of factors satisfying the prescribed constraints, and the resulting multiplicative estimate implies the additive bound in (13).Choosing n0(ε, τ) properly ensures the required conditions, while continuity and boundedness of ε convert the multiplicative estimate into the final additive bound.

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

The auxiliary lemmas establish uniform kernel-expansion and empirical-concentration bounds on compact, boundaryless manifolds. The proof of Lemma A.1 applies these results under an admissible bandwidth range, controls leave-one-out density estimates and normalization, and concludes with high probability.

  • B.4 Auxiliary Lemmas for the Proof of Lemma A.1: For smooth f ∈ C3(M), Lemma B.4 gives a uniform kernel expansion on compact, boundaryless M for sufficiently small h.The first-order error is uniformly O(h), using continuity of the potential, f, and Δ_Mf on compact M.
  • B.4 Auxiliary Lemmas for the Proof of Lemma A.1: For positive µ ∈ C3(M), Lemma B.5 supplies a high-probability uniform empirical kernel bound when n ≥ 2h0(M, µ)^−1/2 and 4n^−2 ≤ h ≤ h0(M, µ).The bound is obtained by combining a deterministic bias estimate with a Gaussian-kernel empirical-process bound.
  • B.4 Auxiliary Lemmas for the Proof of Lemma A.1: The empirical-process constant Cgc depends only on d, diam(M), and uniform upper and lower bounds for µ.These quantities are finite because M is compact and µ is positive and continuous.
  • B.5 Proof of Lemma A.1: In Lemma A.1, the bandwidth condition (κ log m/m)^(1/d) ≤ h ≤ h0 ensures the admissible range of Lemma B.5 for all sufficiently large m.The proof chooses m0 so this condition also implies h ≥ 4m^−2 and h ≤ h⋆.
  • B.5 Proof of Lemma A.1: The leave-one-out estimator satisfies ˆp−i(xi) = m/(m−1) ˆp(xi) − 1/((m−1)(πh)^(d/2)).This identity follows from Kh(xi, xi) = 1 and is used to compare leave-one-out and full-sample estimates.
  • B.5 Proof of Lemma A.1: With probability at least 1 − 2m^−2, the proof simultaneously controls the kernel-estimation event and the empirical normalization event via a union bound.Uniformity over Mx makes the kernel estimate hold simultaneously at every sample point xi.
  • B.5 Proof of Lemma A.1: The normalization ¯Z is bounded away from zero because m h^d ≥ κ log m and its deviation from (∥µ∥L2)^2 is at most half that quantity.The argument combines the previously established concentration bounds with the triangle inequality.

B.6 Proof of Theorem 3.4

The proof establishes uniform control of density estimates and approximate scaling factors, then applies matrix-scaling stability to obtain the theorem’s multiplicative and additive bounds. The resulting estimate holds with probability at least 1 − n^−τ − 4n^−2γ.

  • Density-estimation control: On the density-estimation event E, the estimated densities are bounded away from zero and infinity uniformly over both datasets.Positivity and continuity of µ and ν on compact manifolds provide fixed bounds, while the intervals I_x and I_y are independent of m, n, h_1, and h_2.
  • Density-estimation control: Uniform derivative and mean-value bounds control the errors in the density-dependent scaling quantities for every i ∈[m] and j ∈[n].The absorbed constants depend only on M_x, M_y, µ, ν, and θ.
  • Scaling-factor control: Hoeffding’s inequality and union bounds yield uniform control of the row- and column-scaling discrepancies under Assumption 2.The argument bounds the relevant averages and kernel-weighted terms while preserving constants depending only on the manifolds, densities, and θ.
  • Matrix-scaling stability: Matrix-scaling stability converts the uniform discrepancies into positive scaling vectors that satisfy the prescribed row sums r and column sums c.The stability lemma is applied with δ := τ C(ε)(ρ_1 + ρ_2), after choosing n_0(ε, τ), h′_1, and h′_2 so that δ ≤1/2.
  • Conclusion: The final multiplicative estimate implies the additive bound in (17), and the bound holds with probability at least 1 −n^−τ −4n^−2γ.Continuity and boundedness of ε on M_x ×M_y enable the conversion, while the probability follows from intersecting the controlling events and applying the union bound.
Loading 2608.16506v1…