Skip to content

Statistical Matching via Schrödinger Bridge beyond Conditional Independence

Conference: NeurIPS 2026
arXiv: 2606.22770
Area: Image Generation Keywords: Statistical Matching, Schrödinger Bridge, Conditional Independence Assumption, Entropic Optimal Transport, Bidirectional Imputation

TL;DR

Addressing the fundamental limitation of classical statistical matching that enforces the conditional independence assumption, this paper introduces a dependency-aware Schrödinger bridge framework that exponentially tilts the conservative baseline via a compatibility cost, recovering informative joint distributions with exact Gaussian recovery guarantees and bidirectional probabilistic posterior completion.

Background & Motivation

In data integration and cross-domain collaborative analysis, merging disparate databases that lack unique entity identification keys is ubiquitous due to privacy constraints and fragmented data collection. In typical statistical matching (SM) settings, a recipient database records shared covariates \(X\) and target variables \(Y\), while a separate donor database records \(X\) and auxiliary attributes \(Z\). For example, linking retail loyalty transaction records with credit card histories can reveal shopping patterns as effective proxies for creditworthiness, substantially improving risk estimation for populations lacking formal credit scores. The core objective of SM is to impute the unobserved auxiliary features into the recipient database to reconstruct a synthetic joint dataset that approximates the true joint population distribution.

To render this ill-posed identification problem tractable, classical statistical matching methodologies universally invoke the Conditional Independence Assumption (CIA), positing that \(Y\) and \(Z\) are conditionally independent given \(X\) (\(Y \perp Z \mid X\)). Although CIA turns an under-identified joint recovery problem into a feasible conditional transfer task, it imposes severe restrictions on downstream utility: under CIA, \(P_{Y|X,Z} = P_{Y|X}\), implying that the imputed auxiliary variable provides zero additional predictive power for \(Y\) once \(X\) is known. Consequently, conventional CIA-based integration is inherently incapable of extracting latent \(Y\)–\(Z\) interactions when substantial unobserved residual dependence exists, such as in data recoding or multimodal feature synthesis.

Existing unregularized optimal transport approaches (such as linear-programming-based discrete matching) attempt to align variables by minimizing transport costs, but they suffer from combinatorial instability, non-unique solutions, strict discrete support restrictions, and the inability to provide full probabilistic posterior distributions for uncertainty quantification. To address these bottlenecks, a principled generative framework is required that can both respect the empirical marginal constraints and systematically synthesize informative latent dependencies. The core idea is to frame predictive statistical matching as a dependency-aware Schrödinger bridge problem anchored to the conservative CIA product measure, tilting it via a structural compatibility cost and optimizing neural potential functions to deliver principled probabilistic posterior imputation.

Method

Overall Architecture

Statistical matching seeks a joint probability measure \(Q^\star\) within the admissible coupling set \(\Pi(P_{X,Y}, P_{X,Z})\) matching the observed marginals \(P_{X,Y}\) and \(P_{X,Z}\). The proposed framework casts this problem into the mathematical foundation of entropic optimal transport (EOT) and Schrödinger bridges (SB).

Locally, for any given covariate value \(x \in \mathcal{X}\), let \(\rho_x := P_{Y|X=x} \otimes P_{Z|X=x}\) denote the independent product coupling induced by the CIA. To recover latent dependencies beyond \(\rho_x\), a local EOT problem is formulated with a non-negative compatibility cost \(c(x,y,z)\): $$ \pi_x^\star \in \arg\min_{\pi_x \in \Pi(P_{Y|X=x}, P_{Z|X=x})} \left{ \mathbb{E}_{\pi_x}[c(x,Y,Z)] + \lambda \mathrm{KL}(\pi_x \parallel \rho_x) \right} $$ where \(\lambda > 0\) balances cost minimization and preservation of the conservative CIA reference. Defining the exponentially tilted reference measure \(dR_x(y,z) = \alpha(x)^{-1} \exp\{-c(x,y,z)/\lambda\} d\rho_x(y,z)\) transforms the optimization into minimizing \(\mathrm{KL}(\pi_x \parallel R_x)\).

To scale to high-dimensional continuous covariate spaces where pointwise local solutions are intractable, the problem is lifted to a unified global Schrödinger bridge: defining the global CIA product measure \(P_\rho(dx,dy,dz) = P_X(dx) P_{Y|X}(dy|x) P_{Z|X}(dz|x)\) and global reference measure \(dR(x,y,z) \propto \exp\{-c(x,y,z)/\lambda\} dP_\rho(x,y,z)\), the global bridge is defined by: $$ Q^\star \in \arg\min_{Q \in \Pi(P_{X,Y}, P_{X,Z})} \mathrm{KL}(Q \parallel R) $$ Because \(Q\) and \(R\) share identical \(P_X\) marginals, the global solution \(Q^\star\) decomposes into conditional kernels that almost surely coincide with the optimal local bridges.

%%{init: {'theme': 'base', 'themeVariables': {'fontSize': '12px', 'lineColor': '#2B6CB0'}}}%%
flowchart TD
    A["Separated Databases Input<br/>Recipient D^A(X,Y) and Donor D^B(X,Z)"] --> B["Structural Compatibility Cost Design<br/>Covariate Alignment & Conditional Residuals"]
    B --> C["Exponential Tilting of Reference Measure<br/>Injecting Transport Compatibility onto CIA Baseline"]
    C --> D["Neural Potential Parameterization<br/>Statistical Risk Minimization via Importance Sampling"]
    D --> E["Bidirectional Stochastic Posterior Sampling<br/>Completing Missing Attributes to Form D^T"]
    E --> F["Downstream Predictive Evaluation<br/>Enhanced Inference Leveraging Latent Y-Z Dependencies"]

Key Designs

1. Exponential Tilting of Global Reference Measure: Bridging Conservative Baseline and Active Coupling Classical methods either strictly enforce independence or solve brittle unregularized linear programs resulting in deterministic point assignments. The proposed formulation anchors the joint space to the conservative CIA product law \(P_\rho\), modulated by Boltzmann-type exponential factors \(\exp\{-c(x,y,z)/\lambda\}\). By variational duality, the optimal joint density factors as \(q^\star(x,y,z) = A(x,y) B(x,z) \exp\{-c(x,y,z)/\lambda\}\), where positive potentials \(A(x,y)\) and \(B(x,z)\) correspond to the Lagrange multipliers enforcing marginal feasibility. This entropic formulation ensures strict convexity, statistical stability, and continuous support across both continuous and categorical variables.

2. Structural Compatibility Cost Design: Covariate Projection Alignment and Conditional Residual Modeling Because joint pairs \((Y,Z)\) are unobserved across disparate databases, engineering informative surrogate costs is essential. For general benchmark settings, the empirical alignment cost measures discrepancies between backwards conditional expectations of the shared covariates, \(c_{\mathrm{align}}(y,z) = \|\mathbb{E}[X \mid Y=y] - \mathbb{E}[X \mid Z=z]\|_2^2\). For complex regression where covariates already explain substantial variance, a conditional residual cost is constructed: neural estimators fit conditional means \(\hat{m}_Y(x) = \mathbb{E}[u_Y(Y) \mid X=x]\) and \(\hat{m}_Z(x) = \mathbb{E}[u_Z(Z) \mid X=x]\) on standardized variables, computing normalized residual distances \(c(x,y,z) = \left[\{u_Y(y)-\hat{m}_Y(x)\} - \{u_Z(z)-\hat{m}_Z(x)\}\right]^2 / K_{\mathrm{res}}\) to isolate pure residual associations from confounding \(X\).

3. Neural Potential Parameterization and Risk Minimization: Unbiased Importance Sampling over Dual Integrals The system of dual potential equations requires satisfying \(A(x,y) \int B(x,z) e^{-c/\lambda} dz = p_{X,Y}(x,y)\) and its symmetric counterpart. Setting \(A_p(x,y) = A(x,y)/p_{X,Y}(x,y)\) and density ratio \(d(x,y,z) = p_{Y|X}(y \mid x)/p_{Z|X}(z \mid x)\), an auxiliary proposal distribution \(q\) enables importance sampling to convert continuous integral constraints into a tractable statistical risk minimization problem: $$ \mathcal{L}(\theta,\phi) = \mathbb{E}{P}}\left[\left{A_p(X,Y;\theta) \mathbb{E{Z\sim q}\left[\frac{B(X,Z;\phi)}{q(Z)} e^{-\frac{c(X,Y,Z)}{\lambda}}\right] - 1\right}^2\right] + \mathbb{E}\right] - 1\right}^2\right] $$ Multi-layer perceptrons parameterize }}\left[\left{B(X,Z;\phi) \mathbb{E}_{Y\sim q}\left[\frac{A_p(X,Y;\theta)}{q(Y)} d(X,Y,Z) e^{-\frac{c(X,Y,Z)}{\lambda}\(A_p\) and \(B\) (with unit \(L^2(P_{X,Z})\) normalization on \(B\) for identifiability), optimized via alternating gradient descent over empirical mini-batches and Monte Carlo inner draws.

4. Bidirectional Generative Posterior Sampling: Full Probabilistic Imputation for Uncertainty Quantification Departing from deterministic nearest-neighbor or hot-deck donor matching, the learned Schrödinger bridge yields explicit closed-form posterior conditional densities: recipient imputation draws from \(q^\star_{Z|X,Y}(z \mid x,y) \propto B(x,z) \exp\{-c(x,y,z)/\lambda\}\), while donor imputation draws from \(q^\star_{Y|X,Z}(y \mid x,z) \propto A(x,y) \exp\{-c(x,y,z)/\lambda\}\). Missing variables are stochastically sampled from candidate donor pools weighted by these posterior probabilities, yielding completed databases \(\hat{D}^A\) and \(\hat{D}^B\). This stochastic sampling naturally aligns with Rubin's multiple imputation paradigm, enabling downstream models to capture imputation uncertainty faithfully.

Theoretical Analysis & Trade-offs

The theoretical guarantees of the framework establish both informativeness over the baseline and exact joint recoverability:

  1. Sufficient Condition for Informative Gain (Theorem 3.1): Quantifying recipient-side predictive improvement by expected KL reduction \(\Delta_Z(Q) = \mathbb{E}_{P_{X,Y}}[\mathrm{KL}(P_{Z|X,Y} \parallel P_{Z|X}) - \mathrm{KL}(P_{Z|X,Y} \parallel Q_{Z|X,Y})]\), any non-negative cost with finite second moment under \(P_\rho\) satisfies: $$ \Delta_Z(Q^\star) \ge \frac{\delta_c}{\lambda} - \frac{C}{\lambda^2} + C_\lambda $$ where \(\delta_c = (\mathbb{E}_{P_\rho} - \mathbb{E}_P)[c(X,Y,Z)]\). When the cost assigns lower expected penalties to true dependent triples than to independent pairs (\(\delta_c > 0\)), the leading order term dominates for sufficiently large \(\lambda\), strictly guaranteeing positive information gain \(\Delta_Z(Q^\star) > 0\).

  2. Exact Joint Recovery under Gaussianity (Theorem 3.2): For jointly Gaussian distributions \((X,Y,Z) \sim \mathcal{N}(m, \Sigma)\), using the oracle cost \(c_{\mathrm{oracle}}(y,z) = -\mathrm{sgn}(\sigma_{YZ|X}) y z\), or the empirical alignment cost \(c_{\mathrm{align}}(y,z)\) under geometric consistency \(\sigma_{YZ|X}\Sigma_{YX}\Sigma_{XZ} > 0\), there exists an optimal regularization parameter: $$ \lambda^\star = \frac{\beta |\Sigma_{YZ|X}|}{\sigma_{YZ|X}} > 0 $$ such that the global bridge exactly recovers the true joint law, \(Q^\star = P\). This confirms that the proposed framework serves not merely as an empirical heuristic, but as an exact reconstructor under structured populations.

Key Experimental Results

Main Results

Evaluations were conducted across Gaussian simulations, synthetic non-linear regressions (SYN1 linear, SYN2 sine-quadratic, SYN3 logarithmic-exponential), and real-world benchmarks (CelebA and Adult). Baselines include NAIVE (CIA baseline predicting from \(X\) alone), ORACLE (empirical upper bound trained on fully observed triples), GO (unregularized optimal transport), and five mixed-method variants (Best MM).

1. Gaussian Benchmark

In a trivariate Gaussian with \(\sigma_{XY}=\sigma_{XZ}=0.5\), evaluating across varying latent correlation \(\sigma_{YZ} \in \{0.9, 0.7, 0.5, 0.3\}\) for downstream \(R^2\) (Down. \(R^2\)) and donor imputation error (\(\mathrm{RMSE}^B\)):

\(\sigma_{YZ}\) Metric NAIVE (CIA) ORACLE SB (Strong) SB (Weak) GO Best MM
0.9 Down. \(R^2\) \(0.255 \pm 0.008\) \(0.730 \pm 0.013\) \(0.725^\dagger \pm 0.014\) \(0.723 \pm 0.013\) \(0.697 \pm 0.019\) \(0.264 \pm 0.009\) (MM5)
\(\mathrm{RMSE}^B\) — — \(0.572 \pm 0.009\) \(0.527 \pm 0.005\) \(0.558 \pm 0.023\) \(0.866 \pm 0.008\) (MM2)
0.7 Down. \(R^2\) \(0.256 \pm 0.010\) \(0.485 \pm 0.009\) \(0.452^\dagger \pm 0.008\) \(0.441 \pm 0.009\) \(0.371 \pm 0.064\) \(0.261 \pm 0.008\) (MM3)
\(\mathrm{RMSE}^B\) — — \(0.896 \pm 0.012\) \(0.872 \pm 0.007\) \(0.800 \pm 0.031\) \(0.866 \pm 0.009\) (MM2)
0.5 Down. \(R^2\) \(0.259 \pm 0.004\) \(0.331 \pm 0.008\) \(0.265^\dagger \pm 0.009\) \(0.238 \pm 0.016\) \(0.080 \pm 0.097\) \(0.260 \pm 0.007\) (MM2)
\(\mathrm{RMSE}^B\) — — \(1.038 \pm 0.003\) \(1.018 \pm 0.008\) \(0.975 \pm 0.046\) \(0.868 \pm 0.010\) (MM2)
0.3 Down. \(R^2\) \(0.260 \pm 0.006\) \(0.262 \pm 0.007\) \(0.207 \pm 0.010\) \(0.209 \pm 0.010\) \(-0.191 \pm 0.095\) \(0.260 \pm 0.006\) (MM2)
\(\mathrm{RMSE}^B\) — — \(1.099 \pm 0.001\) \(1.096 \pm 0.010\) \(1.110 \pm 0.044\) \(0.867 \pm 0.009\) (MM2)

Under strong dependencies (\(\sigma_{YZ}=0.9\) and \(0.7\)), SB (Weak) recovers 98.5% and 80.8% of the gap to ORACLE respectively, while GO achieves only 93.1% and 50.2%. Under weak correlations (\(\sigma_{YZ} \le 0.5\)), SB gracefully converges near NAIVE, while unregularized GO collapses into negative \(R^2\) (\(-0.191\)).

2. Real-World Datasets (CelebA and Adult)

CelebA maps 8-dimensional VAE embeddings as \(X\), Attractive attribute as \(Y\), and a 32-state discrete facial feature as \(Z\). Adult predicts high income ($>$50K) with 32-cluster occupation/education groups as \(Z\).

Dataset Metric NAIVE (CIA) ORACLE SB (Strong) SB (Weak) GO Best MM
CelebA Down. AUC \(0.752 \pm 0.001\) \(0.853 \pm 0.002\) \(0.830^\dagger \pm 0.002\) \(0.830 \pm 0.001\) \(0.713 \pm 0.020\) \(0.814 \pm 0.007\) (MM3)
\(\mathrm{AUC}^B\) — — \(0.832^\dagger \pm 0.002\) \(0.832 \pm 0.002\) \(0.714 \pm 0.006\) \(0.788 \pm 0.006\) (MM5)
Adult Down. AUC \(0.834 \pm 0.002\) \(0.888 \pm 0.002\) \(0.845^\dagger \pm 0.002\) \(0.844 \pm 0.004\) \(0.516 \pm 0.039\) \(0.841 \pm 0.003\) (MM3)
\(\mathrm{AUC}^B\) — — \(0.843^\dagger \pm 0.001\) \(0.843 \pm 0.001\) \(0.507 \pm 0.015\) \(0.813 \pm 0.004\) (MM1)

On CelebA, SB (Weak) closes 77.2% and 74.5% of the downstream and donor gaps to ORACLE. On Adult, GO deteriorates toward near-random guessing (AUC 0.516), whereas SB achieves 0.844, showcasing consistent robustness.

Ablation Study

The residual non-linear modeling capability was ablated across SYN1, SYN2, and SYN3 under varying noise variance \(\sigma^2_{\mathrm{noise}}\) and residual correlation \(\rho_{\mathrm{noise}}\).

Setting (\(\sigma^2_{\mathrm{noise}}, \rho_{\mathrm{noise}}\)) Task NAIVE ORACLE SB (Strong) SB (Weak) GO Best MM
\((0.5, 0.9)\) SYN1 (Linear) \(0.907 \pm 0.001\) \(0.948 \pm 0.004\) \(0.948^\dagger \pm 0.004\) \(0.944 \pm 0.004\) \(0.776 \pm 0.027\) \(0.908 \pm 0.001\) (MM2)
\((0.5, 0.9)\) SYN2 (Nonlinear I) \(0.821 \pm 0.004\) \(0.937 \pm 0.005\) \(0.930 \pm 0.005\) \(0.930 \pm 0.005\) \(0.361 \pm 0.042\) \(0.657 \pm 0.013\) (MM1)
\((0.5, 0.9)\) SYN3 (Nonlinear II) \(0.713 \pm 0.002\) \(0.881 \pm 0.007\) \(0.879^\dagger \pm 0.007\) \(0.875 \pm 0.007\) \(0.598 \pm 0.042\) \(0.716 \pm 0.010\) (MM3)
\((0.5, 0.6)\) SYN2 (Nonlinear I) \(0.821 \pm 0.003\) \(0.872 \pm 0.002\) \(0.861^\dagger \pm 0.003\) \(0.854 \pm 0.001\) \(0.232 \pm 0.043\) \(0.641 \pm 0.008\) (MM1)
\((0.1, 0.9)\) SYN2 (Nonlinear I) \(0.958 \pm 0.001\) \(0.980 \pm 0.002\) \(0.978^\dagger \pm 0.002\) \(0.969 \pm 0.001\) \(0.445 \pm 0.040\) \(0.759 \pm 0.008\) (MM1)

The ablations confirm that as non-linearity intensifies (SYN2/SYN3) and residual correlation grows (\(\rho_{\mathrm{noise}}=0.9\)), the relative gain of SB becomes prominent (e.g. elevating \(R^2\) from 0.821 to 0.930 on SYN2, whereas GO drops to 0.361).

Key Findings

  1. High Sample Efficiency of Weak Auxiliary Guidance: Hyperparameter selection with only \(n_C=50\) triples in SB (Weak) yields performance virtually identical to SB (Strong), highlighting immediate practical applicability.
  2. Dependence Strength Governs Matching Utility: When substantial latent correlation exists, SB reliably closes the performance gap to the fully observed benchmark; when variables are truly independent, entropic regularization safeguards against spurious matching.

Highlights & Insights

  • Conceptual Paradigm Shift: Pioneers the application of dependency-aware Schrödinger bridges and entropic optimal transport to statistical matching, lifting decades-old conditional independence barriers.
  • Rigorous Theoretical Guarantees: Supplies both an information-theoretic lower bound for posterior predictive gains and exact full joint recoverability theorems in multivariate Gaussian regimes.
  • Algorithmic Elegance: By reformulating dual integral equations via importance sampling and density ratios into neural mean-squared risk minimization, the method attains scalability and numerical robustness across continuous and categorical data.

Limitations & Future Work

  • Risk of Amplifying Spurious Biases: In synthesizing latent associations from unlinked data, existing historical biases in marginals may inadvertently be reinforced, highlighting the need for algorithmic fairness audits.
  • Gaussian Boundaries for Exact Recovery: Exact theoretical recovery is currently restricted to Gaussian regimes; extending recovery guarantees to non-Gaussian manifolds remains an open research question.
  • Dependency on Prior Validation Samples: Regularization tuning currently benefits from small auxiliary validation sets; developing completely assumption-free, self-supervised tuning mechanisms represents a promising future avenue.