MLPs at the EOC: Concentration of the NTK
Dávid Terjék, Diego González-Sánchez
TL;DR
This work derives finite-width concentration bounds for the Neural Tangent Kernel (NTK) of deep MLPs initialized at the Edge Of Chaos with $(a,b)$-ReLU activations, without relying on gradient independence. By introducing a flexible parameterization with width patterns $m_k=\gamma_k m$ and a common scale $q$ that interpolates between kernel ($q=0$) and rich ($q=1$) regimes, the authors decompose the NTK into layerwise terms and prove high-probability concentration around the infinite-width limit $\overset{\infty}{K}$, provided hidden layers grow quadratically as $m_k=k^2 m$. A key technical achievement is bounding the finite-width gradient-dependence error via propagation of activation cosines, establishing approximate gradient independence at finite width. The main result delivers an explicit concentration bound: $\mathbb{P}\left( \left\Vert K(\theta) - \overset{\infty}{K} \right\Vert \leq O\left( \overline{\tau}^2 \left( \Delta_\phi^{-2} + (\log(l) + m_l^{1/2}) l \right) \sqrt{\log(\ln) \log(m)} \kappa_\phi^2 m^{-1/2} \right) \right) \geq 1 - O(m^{-1})$, where $\overline{\tau}=\max_i \|x_i\|$, $\Delta_\phi=\frac{b^2}{a^2+b^2}$, and $\kappa_\phi=\frac{|a|+|b|}{\sqrt{a^2+b^2}}$. The results imply that absolute-value activations can yield stronger NTK concentration than ReLU within this framework and lay groundwork for analyzing kernel versus feature-learning behavior in finite-width networks.
Abstract
We study the concentration of the Neural Tangent Kernel (NTK) $K_θ: \mathbb{R}^{m_0} \times \mathbb{R}^{m_0} \to \mathbb{R}^{m_l \times m_l}$ of $l$-layer Multilayer Perceptrons (MLPs) $N : \mathbb{R}^{m_0} \times Θ\to \mathbb{R}^{m_l}$ equipped with activation functions $φ(s) = a s + b \vert s \vert$ for some $a,b \in \mathbb{R}$ with the parameter $θ\in Θ$ being initialized at the Edge Of Chaos (EOC). Without relying on the gradient independence assumption that has only been shown to hold asymptotically in the infinitely wide limit, we prove that an approximate version of gradient independence holds at finite width. Showing that the NTK entries $K_θ(x_{i_1},x_{i_2})$ for $i_1,i_2 \in [1:n]$ over a dataset $\{x_1,\cdots,x_n\} \subset \mathbb{R}^{m_0}$ concentrate simultaneously via maximal inequalities, we prove that the NTK matrix $K(θ) = [\frac{1}{n} K_θ(x_{i_1},x_{i_2}) : i_1,i_2 \in [1:n]] \in \mathbb{R}^{nm_l \times nm_l}$ concentrates around its infinitely wide limit $\overset{\scriptscriptstyle\infty}{K} \in \mathbb{R}^{nm_l \times nm_l}$ without the need for linear overparameterization. Our results imply that in order to accurately approximate the limit, hidden layer widths have to grow quadratically as $m_k = k^2 m$ for some $m \in \mathbb{N}+1$ for sufficient concentration. For such MLPs, we obtain the concentration bound $\mathbb{P}( \Vert K(θ) - \overset{\scriptscriptstyle\infty}{K} \Vert \leq O((Δ_φ^{-2} + m_l^{\frac{1}{2}} l) κ_φ^2 m^{-\frac{1}{2}})) \geq 1-O(m^{-1})$ modulo logarithmic terms, where we denoted $Δ_φ= \frac{b^2}{a^2+b^2}$ and $κ_φ= \frac{\vert a \vert + \vert b \vert}{\sqrt{a^2 + b^2}}$. This reveals in particular that the absolute value ($Δ_φ=1$, $κ_φ=1$) beats the ReLU ($Δ_φ=\frac{1}{2}$, $κ_φ=\sqrt{2}$) in terms of the concentration of the NTK.