[Paper Review] Learning Polynomials of Few Relevant Dimensions

Sitan Chen, Raghu Meka|arXiv (Cornell University)|Apr 28, 2020

Advanced X-ray Imaging Techniques75 references4 citations

TL;DR

This paper proposes a novel algorithm for learning low-rank polynomials—degree-$d$ polynomials depending on only $r \ll n$ relevant dimensions—under Gaussian covariates. It combines filtered PCA for a warm start and geodesic stochastic Riemannian gradient descent to achieve sample complexity $O_{r,d}(n\log^2(1/\epsilon)(\log n)^d)$ and runtime $O_{r,d}(Nn^2)$, significantly improving over the standard $n^d$ complexity when $r$ is small.

ABSTRACT

Polynomial regression is a basic primitive in learning and statistics. In its most basic form the goal is to fit a degree $d$ polynomial to a response variable $y$ in terms of an $n$-dimensional input vector $x$. This is extremely well-studied with many applications and has sample and runtime complexity $Θ(n^d)$. Can one achieve better runtime if the intrinsic dimension of the data is much smaller than the ambient dimension $n$? Concretely, we are given samples $(x,y)$ where $y$ is a degree at most $d$ polynomial in an unknown $r$-dimensional projection (the relevant dimensions) of $x$. This can be seen both as a generalization of phase retrieval and as a special case of learning multi-index models where the link function is an unknown low-degree polynomial. Note that without distributional assumptions, this is at least as hard as junta learning. In this work we consider the important case where the covariates are Gaussian. We give an algorithm that learns the polynomial within accuracy $ε$ with sample complexity that is roughly $N = O_{r,d}(n \log^2(1/ε) (\log n)^d)$ and runtime $O_{r,d}(N n^2)$. Prior to our work, no such results were known even for the case of $r=1$. We introduce a new filtered PCA approach to get a warm start for the true subspace and use geodesic SGD to boost to arbitrary accuracy; our techniques may be of independent interest, especially for problems dealing with subspace recovery or analyzing SGD on manifolds.

Motivation & Objective

• Address the challenge of polynomial regression with high ambient dimension $n$ but low intrinsic dimension $r$, where the polynomial depends only on $r$ relevant projections of the input.
• Overcome the $n^d$ sample and runtime complexity of standard polynomial regression by exploiting low-rank structure in the data.
• Develop an efficient algorithm that achieves near-linear dependence on $n$ and polynomial dependence on $r$ and $d$, under Gaussian covariate assumptions.
• Achieve accurate recovery of both the hidden subspace $U^*$ and the link polynomial $p$ within $\epsilon$-accuracy.
• Provide theoretical guarantees on sample complexity, runtime, and convergence for both subspace and coefficient estimation in the presence of noise and finite samples.

Proposed method

• Introduce a filtered PCA approach to obtain a warm start by identifying a subspace close to the true $r$-dimensional hidden subspace $U^*$, using trimmed singular value decomposition on the data matrix.
• Apply geodesic stochastic Riemannian gradient descent (GeoSGD) on the Grassmann manifold to refine the subspace estimate, leveraging the Riemannian geometry of subspaces to ensure convergence.
• Use a two-stage optimization: first recover the subspace $V^{(t)}$ via SubspaceDescent, then re-estimate the polynomial coefficients $\mathbf{c}$ using RealignPolynomial on the current subspace.
• Introduce a novel Procrustes-based distance measure $d_P(V, V^*)$ to track subspace error and ensure convergence to the true subspace.
• Employ martingale concentration inequalities and Hermite polynomial expansions to control error in gradient estimates and ensure high-probability convergence guarantees.
• Design a termination condition that dynamically adapts to subspace accuracy, ensuring that if the subspace is already close enough, the algorithm terminates early without further refinement.

Experimental results

Research questions

• RQ1 Can polynomial regression be performed with sample and runtime complexity sub-quadratic in the ambient dimension $n$ when the underlying polynomial depends on only $r \ll n$ relevant dimensions?
• RQ2 Is it possible to achieve efficient learning of low-rank polynomials under Gaussian covariates without distributional assumptions on the link function beyond being a low-degree polynomial?
• RQ3 How can one simultaneously recover both the hidden subspace $U^*$ and the link polynomial $p$ with high accuracy using a single algorithmic framework?
• RQ4 What is the optimal trade-off between sample complexity, runtime, and accuracy in the low-rank polynomial regression setting?
• RQ5 Can geodesic SGD on the Grassmann manifold be effectively used to recover low-dimensional subspaces in high-dimensional data with provable convergence guarantees?

Key findings

• The proposed algorithm achieves sample complexity $O_{r,d}(n\log^2(1/\epsilon)(\log n)^d)$, which is polynomial in $n$ and independent of $n^d$, significantly improving over the standard $n^d$ complexity.
• The runtime is $O_{r,d}(Nn^2)$, where $N$ is the sample complexity, making it efficient for small $r$ and moderate $d$.
• The algorithm guarantees that the estimated subspace $V^{(T)}$ satisfies $d_P(V^{(T)}, V^*) \leq \epsilon$ with high probability, using $T = \alpha^{-1} \log(1/\epsilon)$ iterations.
• The coefficient vector $\mathbf{c}$ of the link polynomial is estimated within $\|\mathbf{c}^{(T)} - \mathbf{c}^*\|_2 \leq \epsilon$ after $T$ iterations, ensuring accurate polynomial recovery.
• Theoretical guarantees are established via local smoothness and curvature analysis of the Riemannian optimization landscape, with high-probability concentration bounds on gradient estimates.
• The method is robust to initialization: even if the initial subspace is not close to $V^*$, the algorithm converges to the true subspace at a linear rate $1 - \alpha$, where $\alpha$ depends on the condition number and problem parameters.