Abstract
Statistical matching combines partially overlapping datasets that share covariates X but observe the target Y and auxiliary variables Z separately. Classical approaches typically invoke the conditional independence assumption (CIA), which makes the problem identifiable but fundamentally implies that the imported auxiliary variable provides no additional predictive power for Y once X is known. To capture this latent Y--Z dependence, we propose a novel dependency-aware Schrödinger bridge for predictive statistical matching. Our approach couples the two separated databases by tilting the conservative CIA baseline with a transportation-based compatibility cost, recovering an informative joint distribution. The resulting statistical learning framework yields full probabilistic posterior rules for bidirectional imputation. Theoretically, we establish a sufficient condition under which the learned bridge strictly improves over the CIA baseline, alongside an exact joint recovery guarantee in the Gaussian setting under an appropriate cost. Across synthetic benchmarks and real-world datasets (CelebA and Adult), we demonstrate that our dependency-aware completion consistently improves downstream predictive utility, proving especially beneficial in settings like data recoding where the underlying population exhibits strong Y--Z dependence.
Explore similar work
Jul 30, 2026cs.LG
Flow matching assumes fully observed training data, which many real-world applications rarely provide. We propose Missing-Data Flow Matching, which treats the missing coordinates of training samples as latent variables and averages the flow matching loss over the values they could take. We first prove the correction is exact rather than approximate. Under missing completely at random with true completions, the incomplete-data objective equals the complete-data objective, so missingness changes nothing about what flow matching learns and the entire difficulty relocates to the completion model. Our finite-sample analysis then answers design questions that the algorithm leaves open, and the answers are not the ones intuition suggests. Missingness transfers estimator variance rather than adding it, one completion per example already matches complete-data variance exactly, and under a fixed evaluation budget one completion is optimal. A learned completion model contributes a single irreducible bias, which we bound by its expected conditional Wasserstein distance to the true completion law. Experiments numerically validate the theoretical predictions, show that deterministic rather than frozen imputation is what collapses the generated distribution, and place our method alongside strong classical and deep imputation baselines on real tabular data.
Fairoz Nower Khan, Nabuat Zaman Nahim, Peizhong Ju
May 6, 2026stat.ME
The standard constraint-based paradigm for causal discovery with incomplete data -- impute first, test second -- is frequently miscalibrated: any consistent conditional independence (CI) test rejects a true null with probability approaching 1 when imputation error induces spurious conditional dependence. We introduce PAIR-CI, a nonparametric CI test that restores calibration by integrating multiple imputation directly into the inferential procedure via a paired permutation design. PAIR-CI compares cross-validated models that include and exclude the candidate variable while receiving the same imputed conditioning set, forcing imputation error to cancel in their loss difference rather than contaminate the test statistic. A provably consistent variance estimator jointly accounts for uncertainty arising from cross-validation and multiple imputation -- to our knowledge, the first formal unification of these two inferential frameworks. In simulations, existing imputation-based CI tests exhibit false positive rates of 28--45% when data are missing not at random (MNAR), whereas PAIR-CI averages below the nominal 5% level across data-generating processes and missingness mechanisms. These gains are largest in nonlinear settings and grow with causal graph size: when integrated into the PC algorithm, PAIR-CI reduces structural Hamming distance by 8% on 10-variable nonlinear graphs, 15% on 30-variable equivalents, and up to 44% on the 56-variable HAILFINDER network, with stable performance in all settings.
Thomas S. Robinson, Ranjit Lall
Jul 8, 2026cs.LG
Missing data is prevalent in practical applications, making effective imputation an essential preprocessing step for downstream analysis. Real-world datasets often exhibit complex latent structures composed of multiple subgroups with distinct distributions. However, existing methods often overlook such population heterogeneity. Without explicit structural guidance, these methods tend to produce generic estimates that blur subgroup boundaries and lack instance-level fidelity. While incorporating subgroup information offers a remedy, it faces a circular dependency: reliable subgroup identification requires complete data, while data completion is the imputation objective itself. To resolve this, we propose CAGI (Cluster-Aware Generative Imputation), a framework that reformulates clustering and imputation as a mutually reinforcing co-optimization process. CAGI employs a ``Partition-Guide-Restore'' strategy where dynamic cluster assignments act as local priors to condition a Generative Adversarial Network. An iterative feedback loop is established to progressively refine both cluster structures and imputed values toward faithful subgroup distributions. To ensure distributional stability, CAGI further employs a multi-level optimization objective combining instance-level reconstruction with distribution-level regularization. Extensive experiments on 14 benchmark datasets with 15 representative baselines demonstrate the superiority of CAGI. The source code is available at: https://github.com/supercocachii/CAGI
Chuyao Zhang, E Li, Taochen Chen +5