Fast algorithms for learning a Gaussian under halfspace truncation with optimal sample complexity
For machine learning theorists, this solves an open problem by providing the first optimal sample and time complexity algorithm for Gaussian estimation under halfspace truncation, improving upon prior polynomial-time methods.
The paper presents an efficient algorithm for learning a high-dimensional Gaussian truncated to an unknown halfspace, achieving optimal sample complexity of Õ(d²/ε²) and runtime dominated by computing the empirical covariance matrix, matching the complexity of learning without truncation.
We study the fundamental problem of learning a high-dimensional Gaussian truncated to an unknown halfspace. Lee, Mehrotra and Zampetakis (FOCS'24) recently obtained the first polynomial time algorithm for this problem, but their resulting sample and time complexity bounds are not optimal. Under non-trivial truncation, for any target accuracy $\varepsilon > 0$ and dimension $d$ we give an efficient algorithm that uses $n = \tilde{O}(d^2/\varepsilon^2)$ samples and learns the underlying Gaussian to error $\varepsilon$ in total variation distance. Our algorithm is also fast: its runtime is dominated by the cost of computing the empirical covariance matrix. Both our sample and time complexity are optimal in terms of $d$ and $\varepsilon$ even without truncation: in this regard, we can learn a Gaussian under halfspace truncation for free. The key ingredient behind our result is a novel reinterpretation of the low-degree moments of the truncated Gaussian in terms of a relative truncation parameter. This relative truncation parameter uniquely determines the parameters of the untruncated Gaussian and enables direct parameter recovery. This reinterpretation allows us to circumvent the time intensive projected stochastic gradient descent procedure that is widely used in learning under truncation.