Sharp Risk Bounds for Early-Stopping in Gaussian Linear Regression

1. Problem Statement

The paper addresses the problem of high-dimensional Gaussian linear regression where the number of features $d$ may exceed the number of samples $n$. The goal is to minimize the in-sample mean squared error (MSE) (prediction risk) under arbitrary convex constraints on the parameter vector $\alpha$.

• Model: $y = X\alpha^* + \xi$, where $X \in \mathbb{R}^{n \times d}$ is the design matrix, $\alpha^* \in \mathbb{R}^d$ is the ground truth satisfying $\alpha^* \in K_\tau = \tau K$ (where $K$ is a convex body containing the origin), and $\xi \sim \mathcal{N}(0, I_n)$.
• Method: The authors study Early-Stopped Mirror Descent (ESMD). Instead of running the optimization to convergence, the algorithm is stopped at a specific time $t^*$ to act as a regularizer.
• Gap: While the statistical performance of the Least Squares Estimator (LSE) under convex constraints is well-understood via local Gaussian width, similar sharp, localized risk bounds for iterative methods like Mirror Descent in high-dimensional, general-geometry settings were previously lacking. Existing bounds were often unlocalized, restricted to low dimensions, or limited to specific geometries (e.g., $\ell_2$).

2. Methodology

The authors bridge the gap between iterative optimization and statistical risk bounds by leveraging the framework of offset Rademacher complexities (Kanade et al., 2023) and adapting the local Gaussian width analysis used for LSE (Chatterjee, 2014; Bellec, 2016).

Key Technical Components:

1. Mirror Descent (MD):

   • The algorithm uses a potential function (mirror map) $\psi: \mathbb{R}^d \to \mathbb{R}$.
   • Continuous-time: $\frac{d}{dt}\alpha_t = -(\nabla^2 \psi(\alpha_t))^{-1} \nabla \hat{R}(\alpha_t)$.
   • Discrete-time: $\nabla \psi(\alpha_{t+1}) = \nabla \psi(\alpha_t) - \eta \nabla \hat{R}(\alpha_t)$.
   • The potential $\psi$ dictates the geometry of the optimization path.

2. Assumption A (Potential Conditions):
   The core contribution is establishing sufficient conditions on the potential $\psi$ relative to the constraint set $K$. These conditions ensure the optimization path stays within a region where statistical bounds hold. The conditions involve:

   • Smoothness/Convexity: $\psi$ must be differentiable (and twice differentiable for continuous time) with $\nabla \psi(0)=0$.
   • Square-root Convexity: $\sqrt{\psi}$ must be convex.
   • Strong Convexity: $\psi$ must be strongly convex (discrete time) or strictly convex (continuous time).
   • Approximation of Minkowski Functional: $\psi$ must approximate the squared Minkowski functional $\phi_K^2$ of the convex body $K$. Specifically, there must exist constants $c_l, c_u$ such that:
     $$\phi_K(\alpha) \leq c_l \sqrt{\psi(\alpha)} \quad \text{and} \quad \sqrt{\psi(\alpha)} \leq c_u \tau \quad \forall \alpha \in K_\tau$$
   • The product $c_a = c_l c_u$ is the "approximation constant."

3. Localization Argument:
   The proof demonstrates that under Assumption A, the iterates of early-stopped mirror descent remain within a Bregman ball that is contained within a scaled version of the constraint set ($3c_a K_\tau$). This allows the authors to bound the risk using the stationary radius of the local Gaussian width, similar to the LSE analysis.

3. Key Contributions

A. Main Theoretical Result (Theorem 1)

The paper proves that the in-sample risk of ESMD is bounded by the stationary radius of the local Gaussian width, scaled by the approximation constant $c_a$.
$$\min_{0 \leq t \leq T} R(\alpha_t) \lesssim \frac{r_0^2(X\alpha^*, XK_{3c_a\tau})}{n} + \frac{\log(1/\delta)}{n} + \epsilon$$
where $r_0$ is the stationary radius defined by the equation $w((K-\alpha) \cap rB_2) \leq r^2/2$.

• Significance: This bound holds for any convex body $K$ and any design matrix $X$ in the high-dimensional regime ($d \gg n$), provided the potential $\psi$ satisfies Assumption A.

B. Comparison with LSE (Corollary 1)

The authors show that if the stationary radius and the critical radius of the LSE are comparable (a condition known as Assumption B, which holds for many regular classes), then the worst-case risk of ESMD is bounded by a constant factor of the LSE's risk.
$$\sup_{\alpha^*} \mathbb{E}[R(\alpha_{t^*})] \lesssim C \cdot c_a \cdot \sup_{\alpha^*} \mathbb{E}[R(\alpha_{LSE})]$$
This establishes that ESMD can achieve the same statistical performance as the LSE without explicitly solving a constrained optimization problem.

C. Minimax Optimality (Corollary 2)

Sufficient conditions are derived for ESMD to be minimax optimal. If the local entropy of the constraint set satisfies specific growth conditions (Assumption C) and $c_a \approx 1$, ESMD achieves the minimax rate.

D. Construction of Potentials (Lemma 1 & Section 4)

A major practical contribution is showing how to construct valid potentials for non-smooth or non-strongly convex constraints (like $\ell_1$) using Moreau envelopes or other smoothing techniques. This allows the theory to apply to standard sparsity-inducing norms.

4. Results and Applications

The paper applies the general theory to specific settings, deriving sharp statistical rates:

1. $\ell_p$-norms ($1 < p < 2$):

   • Using $\psi(\alpha) = \|\alpha\|_p^2$, ESMD achieves rates matching the known minimax rates for column-normalized and Gaussian designs.
   • The authors prove a minimax lower bound for a worst-case fixed design matrix, showing that the derived upper bound is tight (up to constants) for column-normalized designs.

2. $\ell_1$-norm (Sparsity/LASSO):

   • The constrained LSE corresponds to the LASSO. Previous bounds for early-stopped mirror descent with hyperbolic entropy had a $\log d$ gap compared to the optimal LASSO rate.
   • By using adjusted potentials (e.g., squared Moreau envelope of $\ell_1$, adjusted hypentropy, or sigmoidal potentials), the authors close this gap.
   • Result: ESMD with these potentials achieves the rate $\tau \sqrt{\frac{\log d}{n}}$, matching the minimax optimal rate of the LASSO.

3. M-Convex Hulls:

   • The theory is extended to convex hulls of $M$ points. A smoothed potential based on the log-sum-exp function is constructed, yielding sharp rates dependent on the geometry of the hull.

4. Computational-Statistical Trade-off:

   • The paper analyzes the relationship between the strong convexity parameter $\rho$ of the potential and the stopping time $T$. Since $T \propto 1/\rho$, potentials with vanishing strong convexity (like squared $\ell_p$ for large $d$) may require more iterations to reach the statistical optimum, highlighting a trade-off between statistical tightness and computational cost.

5. Significance

• Unification: The work unifies the analysis of iterative regularization (early stopping) with the sharp statistical theory of constrained estimation (LSE). It proves that "implicit regularization" via early-stopped mirror descent is not just a heuristic but can be rigorously shown to match the performance of explicit regularization (LSE) under general geometries.
• Generality: Unlike previous works restricted to $\ell_2$ or Hilbert spaces, this framework applies to arbitrary convex bodies and high-dimensional settings.
• Practical Impact: It provides a systematic recipe for choosing mirror descent potentials. Instead of guessing, practitioners can construct potentials (via smoothing the Minkowski functional) that guarantee minimax optimality.
• Resolution of Open Problems: It resolves the logarithmic gap in the $\ell_1$-constrained setting, demonstrating that early-stopped mirror descent can be as statistically efficient as the LASSO, provided the potential is chosen correctly.

In summary, the paper establishes that Early-Stopped Mirror Descent is a statistically optimal method for high-dimensional regression over convex sets, provided the algorithm's potential function is carefully designed to approximate the geometry of the constraint set.