Learning Spherical Polynomials by Learnable Channel Attention

We study the problem of learning a spherical polynomial of an unknown degree ${\ell }_0=\Theta (1)\ge 1$ by training a ReLU network with channel attention by gradient descent (GD) with finite training sample and finite network width guarantee. The target functions are general spherical polynomials with bounded $L^2$ norm and a fixed, unknown degree bound ${\ell }_0\ge 1$. The training data are $n$ i.i.d. uniformly distributed inputs on the unit sphere in ${\text{RR}}^d$ with i.i.d. centered noise of finite variance in their responses. The network width $m$ is the number of ReLU channels and the size of the diagonal channel-attention matrix. The training process of the network has two stages. In stage one, one ordinary GD step learns channel attention parameters and strictly enlarges the prediction space on the first half of the training data. In stage two, finitely many output GD steps on the second half of the training data achieve expected regression risk $\text{cO}((\log d{)}^{-2})\to 0$ as $d\to \infty$. The feature kernel of the ReLU attention network evolves during training and the learnable channel attention provably reduceds the regression risk, demonstrating the provable benefit of feature learning. With $n\asymp d^{{\ell }_0}(\log d{)}^4$ and network width $m\asymp d^{{\ell }_0}(\log d{)}^2$, the sample size and network width $m$ are only of logarithmic factors away from their corresponding lower bounds for learning degree-${\ell }_0$ spherical polynomials. Under sub-Gaussian noise, the same algorithm also satisfies a high-probability risk bound with logarithmic confidence dependence. To the best of our knowledge, this is the first guarantee of vanishing regression risk for a finite ReLU network trained by GD with learnable channel attention, uniformly over bounded-$L^2$ spherical polynomials of degree at most ${\ell }_0$, under these sample-size and width budgets.