Shallow Neural Networks Learn Low-Degree Spherical Polynomials with Learnable Channel Attention

TL;DR

This work investigates nonparametric regression on the unit sphere where the target is a degree-$\ell_0$ spherical polynomial. It introduces an over-parameterized two-layer neural network with learnable channel attention and a two-stage training scheme: a one-step gradient descent to identify the ground-truth channel count, followed by standard gradient descent to train the second layer with a fixed activation. Through a novel analysis combining kernel complexity, NTK-based decomposition, and local Rademacher complexity, the authors prove a minimax-optimal regression rate of $\Theta(d^{\ell_0}/n)$, achievable with finite width, and show that the required sample size scales as $n\asymp d^{\ell_0}/\varepsilon$. This work highlights how channel attention enables targeted feature learning that, together with NTK arguments, yields sharp rates for learning structured functions on high-dimensional spheres.

Abstract

We study the problem of learning a low-degree spherical polynomial of degree $\ell_0 = Θ(1) \ge 1$ defined on the unit sphere in $\RR^d$ by training an over-parameterized two-layer neural network (NN) with channel attention in this paper. Our main result is the significantly improved sample complexity for learning such low-degree polynomials. We show that, for any regression risk $\eps \in (0,1)$, a carefully designed two-layer NN with channel attention and finite width of $m \ge Θ({n^4 \log (2n/δ)}/{d^{2\ell_0}})$ trained by the vanilla gradient descent (GD) requires the lowest sample complexity of $n \asymp Θ(d^{\ell_0}/\eps)$ with probability $1-δ$ for every $δ\in (0,1)$, in contrast with the representative sample complexity $Θ\pth{d^{\ell_0} \max\set{\eps^{-2},\log d}}$, where $n$ is the training daata size. Moreover, such sample complexity is not improvable since the trained network renders a sharp rate of the nonparametric regression risk of the order $Θ(d^{\ell_0}/{n})$ with probability at least $1-δ$. On the other hand, the minimax optimal rate for the regression risk with a kernel of rank $Θ(d^{\ell_0})$ is $Θ(d^{\ell_0}/{n})$, so that the rate of the nonparametric regression risk of the network trained by GD is minimax optimal. The training of the two-layer NN with channel attention consists of two stages. In Stage 1, a provable learnable channel selection algorithm identifies the ground-truth channel number $\ell_0$ from the initial $L \ge \ell_0$ channels in the first-layer activation, with high probability. This learnable selection is achieved by an efficient one-step GD update on both layers, enabling feature learning for low-degree polynomial targets. In Stage 2, the second layer is trained by standard GD using the activation function with the selected channels.

Figures (1)

• Figure : Learnable Channel Selection

Theorems & Definitions (56)

• Theorem 3.1
• Theorem 4.1
• Theorem 5.1
• Theorem 5.2
• Theorem 5.3
• Theorem 5.4
• Definition A.1
• Theorem A.2: bartlett2005
• Theorem A.3
• Definition A.4: Sub-root function,bartlett2005