Distributionally Robust Linear Regression With Block Lewis Weights
For machine learning practitioners needing robust regression under group shifts, this provides a faster algorithm with theoretical guarantees.
The paper presents an algorithm for group distributionally robust least squares that achieves a (1+ε)-multiplicative optimal solution using Õ(min{rank(A),m}^{1/3}ε^{-2/3}) linear-system-solves, improving over interior point methods for moderate accuracy and matching state-of-the-art for ℓ∞ regression.
We present an algorithm for the group distributionally robust (GDR) least squares problem. Given $m$ groups, a parameter vector in $\mathbb{R}^d$, and stacked design matrices and responses $\mathbf{A}$ and $\mathbf{b}$, our algorithm obtains a $(1+\varepsilon)$-multiplicative optimal solution using $\widetilde{O}(\min\{\mathsf{rank}(\mathbf{A}),m\}^{1/3}\varepsilon^{-2/3})$ linear-system-solves of matrices of the form $\mathbf{A}^{\top}\mathbf{B}\mathbf{A}$ for block-diagonal $\mathbf{B}$. Our technical methods follow from a recent geometric construction, block Lewis weights, that relates the empirical GDR problem to a carefully chosen least squares problem and an application of accelerated proximal methods. Our algorithm improves over known interior point methods for moderate accuracy regimes and matches the state-of-the-art guarantees for the special case of $\ell_{\infty}$ regression. We also give algorithms that smoothly interpolate between minimizing the average least squares loss and the distributionally robust loss.