Learning High-dimensional Gaussians from Censored Data
Arnab Bhattacharyya, Constantinos Daskalakis, Themis Gouleakis, Yuhao Wang
TL;DR
This work tackles learning a high-dimensional Gaussian $\mathcal{N}(\bm{\mu}^*, \mathbf{\Sigma}^*)$ from data where coordinates may be censored in a MNAR fashion. It introduces two censoring regimes—self-censoring and linear-thresholding—and develops computationally efficient algorithms with non-asymptotic guarantees. For self-censoring, the approach reduces to 1D and 2D truncation subproblems to recover $\bm{\mu}^*$ and $\mathbf{\Sigma}^*$, achieving $d_{TV}$-accurate recovery with sample complexity $\tilde{O}\left(\frac{d^2 (\lambda_{\max}/\lambda_{\min})^2}{\alpha \varepsilon^2}\right)$. For linear-thresholding with known $\mathbf{\Sigma}$, the MissingDescent method uses anchored missingness and projected SGD with Langevin-based gradient estimates to recover $\bm{\mu}^*$ in polynomial time with complexity depending on $\alpha, \beta, \gamma$ and spectral properties of $\mathbf{\Sigma}$. The results advance finite-sample MNAR density estimation in high dimensions and delineate clear directions when the covariance is unknown, highlighting a rigorous pathway for robust mean and covariance recovery under censoring.
Abstract
We provide efficient algorithms for the problem of distribution learning from high-dimensional Gaussian data where in each sample, some of the variable values are missing. We suppose that the variables are missing not at random (MNAR). The missingness model, denoted by $S(y)$, is the function that maps any point $y$ in $R^d$ to the subsets of its coordinates that are seen. In this work, we assume that it is known. We study the following two settings: (i) Self-censoring: An observation $x$ is generated by first sampling the true value $y$ from a $d$-dimensional Gaussian $N(μ*, Σ*)$ with unknown $μ*$ and $Σ*$. For each coordinate $i$, there exists a set $S_i$ subseteq $R^d$ such that $x_i = y_i$ if and only if $y_i$ in $S_i$. Otherwise, $x_i$ is missing and takes a generic value (e.g., "?"). We design an algorithm that learns $N(μ*, Σ*)$ up to total variation (TV) distance epsilon, using $poly(d, 1/ε)$ samples, assuming only that each pair of coordinates is observed with sufficiently high probability. (ii) Linear thresholding: An observation $x$ is generated by first sampling $y$ from a $d$-dimensional Gaussian $N(μ*, Σ)$ with unknown $μ*$ and known $Σ$, and then applying the missingness model $S$ where $S(y) = {i in [d] : v_i^T y <= b_i}$ for some $v_1, ..., v_d$ in $R^d$ and $b_1, ..., b_d$ in $R$. We design an efficient mean estimation algorithm, assuming that none of the possible missingness patterns is very rare conditioned on the values of the observed coordinates and that any small subset of coordinates is observed with sufficiently high probability.