Causal Matrix Completion under Multiple Treatments via Mixed Synthetic Nearest Neighbors

Minrui Luo; Zhiheng Zhang

cs.LGMar 12, 2026

Causal Matrix Completion under Multiple Treatments via Mixed Synthetic Nearest Neighbors

Authors: Minrui Luo, Zhiheng Zhang

Organizations: Institute for Interdisciplinary Information Sciences, Tsinghua University, Beijing, China · School of Statistics and Data Science, Shanghai University of Finance and Economics, Shanghai, China

Abstract

Synthetic Nearest Neighbors (SNN) provides a principled solution to causal matrix completion under missing-not-at-random (MNAR) by exploiting local low-rank structure through fully observed anchor submatrices. However, its effectiveness critically relies on sufficient data availability within each treatment level, a condition that often fails in settings with multiple or complex treatments. In this work, we propose Mixed Synthetic Nearest Neighbors (MSNN), a new entry-wise causal identification estimator that integrates information across treatment levels. We show that MSNN retains the finite-sample error bounds and asymptotic normality guarantees of SNN, while enlarging the effective sample size available for estimation. Empirical results on synthetic and real-world datasets illustrate the efficacy of the proposed approach, especially under data-scarce treatment levels.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison between SNN and MSNN. The leftmost subfigure illustrates the SNN algorithm with $K_{\rm SNN}=2$ : it requires $\boldsymbol{S}^{(k)}$ , $q^{(k)}$ and $x^{(k)}$ are all fully observed at treatment level as the same of the estimated treatment, which is rare under data-scarce levels (e.g. the “red” level in the second subfigure). The rest four subfigures explain the procedure of MSNN for a specific subgroup $k$ : given entry $(i,j)$ and estimated treatment level (here is “red”), one need to find a fully observed $x^{(k)}$ under same “red” level, but the $\boldsymbol{S}^{(k)}$ and $q^{(k)}$ can be integrated from other treatments (here: “blue” and “green” level). The only requirement is that for each column of $\boldsymbol{S}^{(k)}$ (namely, $s_{i}^{(k)}$ ), its treatments should be as the same as the treatment of corresponding $q_{i}^{(k)}$ , see the third and fourth subfigures.

Algorithm	Low ( $p(d)=0.01$ )		Medium ( $p(d)=0.025$ )		High ( $p(d)=0.05$ )
Algorithm	FR % ( $\uparrow$ )	MRE ( $\downarrow$ )	FR % ( $\uparrow$ )	MRE ( $\downarrow$ )	FR % ( $\uparrow$ )	MRE ( $\downarrow$ )
SNN	$0.03\pm 0.00$	$0.806\pm 0.240$	$1.20\pm 0.36$	$0.577\pm 0.110$	$11.34\pm 0.82$	$0.515\pm 0.046$
MSNN (ours)	$\bf 4.69\pm 1.11$	$\bf 3.91\pm 1.09\;\times 10^{-2}$	$\bf 63.73\pm 5.67$	$\bf 1.18\pm 0.33\;\times 10^{-3}$	$\bf 99.29\pm 0.81$	$\bf 7.05\pm 0.21\;\times 10^{-4}$

Table 1: Performance comparison of MSNN and SNN on MCAR data.

Panel of Figure 2 — Figure 2: Selected prediction results of MSNN on Proposition 99 study in Abadie et al. (2010) . The three states Kansas, Arizona, New Jersey belong to treatment group of control, program, taxes respectively. The dashed lines are estimation results, while the solid lines indicates real-world observation. Before the year of 1989 (illustrated by the vertical dotted gray lines) all states are in control group so the solid lines are black-colored, after which their color varies. The dotted line indicates the time of Proposition 99 assignment. The dashed lines and solid lines of same color at the same time periods are close to each other, indicating successful validation and thus the correctness of applying our model on this real-world dataset.

Causal Effect Estimation

Causal effect estimation measures how outcomes would change under interventions, using experiments or observational data and explicit assumptions.

See all 263 papers

Improved Guarantees for Heterogeneous Treatment-Effect Estimation via Matrix Completion
May 28, 2026Anay Mehrotra, Phuc Tran, Van H. Vu +1Causal Effect Estimation Conditional Average Treatment Effect Estimation
Vin University · The University of Hong Kong · Yale University
Learning Causal Normalizing Flows from Incomplete Data via Observed-Data Likelihood
Sep 29, 2026Trung-Dung Hoang, Alceu Bissoto, Tim Flühmann +4Normalizing Flows Learning with Missing Data
Department of Digital Medicine, University of Bern, Switzerland · Department of Diabetes, Endocrinology, Nutritional Medicine and Metabolism (UDEM), Inselspital, Bern University Hospital, University of Bern, Switzerland · Graduate School for Cellular and Biomedical Sciences (GCB), University of Bern, Switzerland +3
MatrixFormer: A Foundation Model for Matrix Completion
Oct 5, 2026Dwaipayan Saha, Jacob Feitelberg, Kyuseong Choi +2Incomplete Data Imputation Transformer
Cornell Tech

May 28, 2026J/K move · Enter open · S save

Improved Guarantees for Heterogeneous Treatment-Effect Estimation via Matrix Completion

Anay Mehrotra, Phuc Tran, Van H. Vu, Manolis Zampetakis

Vin University · The University of Hong Kong · Yale University

Causal Effect Estimation Conditional Average Treatment Effect Estimation Low-Rank Approximation Individualized Treatment Effect Estimation Matrix Completion

Open paper AskPDF arXiv

A central goal of modern causal inference is estimating heterogeneous treatment effects to answer questions like "how does an intervention affect each unit," rather than only on average. We study this problem with panel-data where we observe

n

units across

m

times under unknown, non-uniform treatment assignments. The data in this setting is naturally represented as a matrix of all unit--time treatment effects. Estimating heterogeneous treatment effects can then be expressed as obtaining a good estimation of each row's average in this matrix. This allows us to formulate the problem as matrix completion, which can be solved under natural low-rankness assumptions. However, existing matrix-completion guarantees are not powerful enough to get meaningful bounds for the per-row guarantee required for estimating the heterogeneous treatment effect; roughly speaking, they are only useful for estimating average treatment effect bounds, as also illustrated in a recent line of work. We give a simple, computationally efficient estimator that, without knowledge of the propensities and under standard low-rankness and regularity assumptions, achieves a row-wise

\ell_2

error of

\tilde{O}(\sqrt{\frac{1}{n} + \frac{n}{m^2}})

. Technically, our analysis establishes the first sharp row-wise

\ell_2

-perturbation bound for low-rank approximation, complementing existing spectral-, Frobenius-, and entrywise perturbation theory.

Abstract

Causal Effect Estimation

Explore similar work

Improved Guarantees for Heterogeneous Treatment-Effect Estimation via Matrix Completion

Learning Causal Normalizing Flows from Incomplete Data via Observed-Data Likelihood

MatrixFormer: A Foundation Model for Matrix Completion