Improved Guarantees for Heterogeneous Treatment-Effect Estimation via Matrix Completion
2026-05-28 • Artificial Intelligence
Artificial IntelligenceData Structures and AlgorithmsMachine Learning
AI summaryⓘ
The authors focus on estimating how different treatments affect each individual unit over time using panel data, rather than just the average effect. They model the problem as completing a matrix of treatment effects where each row corresponds to a unit's effects over time. Traditional methods can't accurately estimate effects for each unit, so the authors propose a new computational method that works without knowing treatment assignment probabilities and relies on the matrix having a low-rank structure. Their approach also provides new theoretical guarantees for how well each individual row can be approximated, improving on previous error bounds.
causal inferenceheterogeneous treatment effectspanel datamatrix completionlow-rank approximationtreatment effect estimationpropensity scoresrow-wise errorperturbation theory
Authors
Anay Mehrotra, Phuc Tran, Van H. Vu, Manolis Zampetakis
Abstract
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.