Technical Summary: Linear Regression with Unknown Truncation Beyond Gaussian Features

1. Problem Definition

The paper addresses the problem of truncated linear regression where the truncation (or "survival") set is unknown. In this setting, a learner observes samples $(x, y)$ generated as follows:

1. A feature vector $x \in \mathbb{R}^d$ is drawn from an unknown distribution $\mathcal{D}$.
2. A label $y$ is computed as $y = x^\top w^\star + \xi$, where $w^\star$ is an unknown regressor and $\xi \sim \mathcal{N}(0, 1)$ is Gaussian noise.
3. The sample is observed if and only if $y \in S^\star$, where $S^\star \subseteq \mathbb{R}$ is an unknown "survival set." Otherwise, the sample is discarded.

The goal is to estimate $w^\star$ to $\varepsilon$-accuracy. The primary challenge is that standard estimators (like Ordinary Least Squares) fail due to selection bias, and prior efficient algorithms for this problem required either knowledge of $S^\star$ or strong assumptions on the feature distribution (specifically, that $\mathcal{D}$ is Gaussian).

The authors focus on the practically relevant case where $S^\star$ is unknown and is a union of at most $k$ intervals. They aim to achieve this under the weaker assumption that the feature distribution $\mathcal{D}$ is sub-Gaussian, rather than strictly Gaussian.

2. Methodology

The proposed algorithm (Algorithm 1) operates in two distinct phases:

Phase I: Learning the Truncation Set $S^\star$

Since $S^\star$ is unknown, the algorithm first learns an approximation $\hat{S} \approx S^\star$.

• Positive-Only Learning: The problem of learning $S^\star$ is framed as learning a concept (the set $S^\star$) from positive examples only (samples where $y \in S^\star$) without access to negative examples.
• Smoothness Assumption: To make this information-theoretically feasible, the authors utilize a "smoothed" learning framework. They construct an auxiliary distribution $\mathcal{D}$ that is "smooth" with respect to the true untruncated marginal distribution of $y$ (denoted $\mathcal{D}^\star_+$).
• Key Lemma (Smoothness): The authors prove that if the feature distribution $\mathcal{D}$ is sub-Gaussian, then for any parameter vector $w$ close to $w^\star$, the distribution of $y$ generated by $w$ (denoted $\mathcal{D}_y(w)$) is smooth with respect to the true distribution $\mathcal{D}_y(w^\star)$. This allows the construction of the necessary auxiliary distribution using a "warm start" $b_w$ (an initial estimate of $w^\star$).
• Algorithm: They employ a greedy procedure (Algorithm 3) to learn the union of $k$ intervals. This involves sorting positive samples, identifying intervals between them, and discarding the intervals containing the most samples from the auxiliary distribution. This runs in polynomial time for one-dimensional sets.

Phase II: Parameter Estimation via Projected Stochastic Gradient Descent (PSGD)

Once an approximation $\hat{S}$ is obtained, the algorithm estimates $w^\star$.

• Perturbed Objective: Since $\hat{S} \neq S^\star$, the standard negative log-likelihood (NLL) is not well-defined or convex in the usual sense. The authors minimize a perturbed NLL objective, $\mathcal{L}_{\hat{S}}(w)$, which is defined using the intersection of the estimated set and the true set for the data generation, but the estimated set for the model likelihood.
• Projection Set: To ensure convergence, the optimization is restricted to a projection set $\mathcal{P}$ (a ball of constant radius around the warm start $b_w$).
  • The authors prove that $w^\star$ lies within $\mathcal{P}$ with high probability.
  • They establish that $\mathcal{L}_{\hat{S}}$ is strongly convex within $\mathcal{P}$, provided the set estimation error is sufficiently small.
• Biased Gradient Sampling: A critical technical hurdle is sampling gradients for $\mathcal{L}_{\hat{S}}$. Exact sampling is intractable. The authors design an approximate gradient sampler that produces biased gradients. They prove that the bias is bounded and can be controlled by making the set estimation error $\varepsilon$ sufficiently small.
• Convergence: They analyze PSGD with these biased gradients, showing that the algorithm converges to a point close to the minimizer of $\mathcal{L}_{\hat{S}}$, which in turn is close to $w^\star$.

3. Key Contributions

1. First Polynomial-Time Algorithm for Unknown Truncation: The paper presents the first algorithm that learns $w^\star$ in $\text{poly}(d, k, 1/\varepsilon)$ time when the survival set $S^\star$ is unknown and a union of $k$ intervals.
2. Relaxation of Distributional Assumptions: Unlike prior work (e.g., Lee, Mehrotra, and Zampetakis [LMZ24]) which required the feature distribution to be Gaussian, this algorithm only requires the features to be sub-Gaussian. This significantly broadens the applicability to real-world data where Gaussianity is often an unrealistic assumption.
3. Efficient Positive-Only Learning: The authors provide an efficient algorithm for PAC learning unions of intervals from positive examples only, under a smoothness condition. This subroutine is of independent interest to the learning theory community.
4. Handling Biased Gradients: The work introduces a novel analysis for PSGD with biased gradient estimators in the context of truncated regression, demonstrating that convergence is achievable even when exact gradient sampling is impossible.

4. Results

• Theorem 3.1 (Main Result): Under Assumptions 1 (Survival Probability), 2 (Sub-Gaussianity), and 3 (Identifiability), there exists an algorithm that, given $n = \text{poly}(d, k, 1/\varepsilon)$ samples, outputs an estimate $\hat{w}$ such that $\|\hat{w} - w^\star\|_2 \leq \varepsilon$ with probability $1-\delta$.
• Complexity: Both the sample complexity and the running time scale polynomially in the dimension $d$, the number of intervals $k$, and the inverse accuracy $1/\varepsilon$.
• Comparison: In the special case where features are Gaussian, the authors improve the running time of the state-of-the-art [LMZ24] from $d^{\text{poly}(k/\varepsilon)}$ (exponential in $d$) to $\text{poly}(d, k, 1/\varepsilon)$.
• Simulations: Experiments on synthetic data (10-dimensional mixture of Gaussians, 5 intervals) show that the proposed algorithm converges to the true $w^\star$ at a rate comparable to an idealized algorithm that knows $S^\star$ a priori, and significantly outperforms OLS and PSGD with incorrect truncation sets.

5. Significance and Claims

The paper claims to resolve a long-standing open problem in truncated regression: achieving computational efficiency without the restrictive assumption of Gaussian features when the truncation set is unknown.

• Practical Relevance: The authors emphasize that real-world truncation sets (e.g., in astronomy or econometrics) are often complex unions of intervals and are rarely known explicitly. Furthermore, feature distributions in these domains are rarely strictly Gaussian.
• Theoretical Advancement: By moving beyond Gaussian features to sub-Gaussian ones, the work bridges the gap between theoretical tractability and practical data characteristics.
• Modesty: The authors acknowledge that their algorithm relies on the assumption that $S^\star$ is a union of a bounded number of intervals ($k$). They note that if $S^\star$ is known, the exponential dependence on $d$ can be avoided entirely, but the unknown case necessitates a dependence on $d$ (specifically in the exponent of the sample complexity for set learning, though the runtime remains polynomial). They also note that their approach is specific to truncation on the response variable $y$, distinct from prior work that handles high-dimensional truncation on both $x$ and $y$ but requires Gaussianity.

The work is presented as a foundational step toward handling systematic selection biases in high-dimensional regression without relying on parametric assumptions about the covariates.