MLPs at the EOC: Spectrum of the NTK

TL;DR

This work analyzes the NTK of infinitely wide MLPs at the Edge Of Chaos for activations of the form φ(s)=a s + b|s|. It shows that NTK entries are effectively governed by inverse cosine distances between layer activations, and it develops a propagation framework using the cosine map and the convex inverse-cosine distance map ω to derive non-asymptotic bounds and tight spectral bounds on the limiting NTK. A central finding is that the nonlinear parameter Δ_0phi = b^2/(a^2+b^2) controls how quickly the NTK spectrum stabilizes with depth, with Δ_0phi=1 (absolute value) providing the best conditioning compared to Δ_0phi=1/2 (ReLU). The authors provide a closed-form expression for the limiting NTK in terms of depth-$l$ cosine-map compositions and offer explicit eigenvalue bounds via inverse cosine-distance matrices, highlighting the favorable conditioning afforded by more nonlinear activations. These insights pave the way for extensions to other architectures and activation families and for exploring NTK dynamics during training beyond the lazy, kernel regime.

Abstract

We study the properties of the Neural Tangent Kernel (NTK) $\overset{\scriptscriptstyle\infty}{K} : \mathbb{R}^{m_0} \times \mathbb{R}^{m_0} \to \mathbb{R}^{m_l \times m_l}$ corresponding to infinitely wide $l$-layer Multilayer Perceptrons (MLPs) taking inputs from $\mathbb{R}^{m_0}$ to outputs in $\mathbb{R}^{m_l}$ equipped with activation functions $φ(s) = a s + b \vert s \vert$ for some $a,b \in \mathbb{R}$ and initialized at the Edge Of Chaos (EOC). We find that the entries $\overset{\scriptscriptstyle\infty}{K}(x_1,x_2)$ can be approximated by the inverses of the cosine distances of the activations corresponding to $x_1$ and $x_2$ increasingly better as the depth $l$ increases. By quantifying these inverse cosine distances and the spectrum of the matrix containing them, we obtain tight spectral bounds for the NTK matrix $\overset{\scriptscriptstyle\infty}{K} = [\frac{1}{n} \overset{\scriptscriptstyle\infty}{K}(x_{i_1},x_{i_2}) : i_1, i_2 \in [1:n]]$ over a dataset $\{x_1,\cdots,x_n\} \subset \mathbb{R}^{m_0}$, transferred from the inverse cosine distance matrix via our approximation result. Our results show that $Δ_φ= \frac{b^2}{a^2+b^2}$ determines the rate at which the condition number of the NTK matrix converges to its limit as depth increases, implying in particular that the absolute value ($Δ_φ=1$) is better than the ReLU ($Δ_φ=\frac{1}{2}$) in this regard.

Figures (1)

• Figure 1: Condition number of $\overset{\infty}{K}$ for $(a,b)$-ReLUs as a function of depth for different values of $\Delta_\phi$ with the EOC parameterization \ref{['def:eoc_parameterization']}. Each curve is the average of $100$ runs, where in each run we sampled $n=32$ points $x_1,\cdots,x_{32}$ uniformly from the unit sphere in $\mathbb{R}^{16}$ to form the dataset and calculated the corresponding $\overset{\infty}{K}$ in closed form via Proposition \ref{['prop:cosine_map']} and Proposition \ref{['prop:limiting_ntk_at_eoc']}.

Theorems & Definitions (20)

• Definition 1: $(a,b)$-ReLU
• Definition 2: Limiting activation and derivative inner products
• Definition 3: Limiting NTK
• Definition 4: EOC initialization for $(a,b)$-ReLUs
• Definition 5: Dual function
• Remark 6: $\vert \cdot \vert$ and $\mathop{\mathrm{sgn}}\nolimits$
• Proposition 7: Limiting inner products as dual functions
• Proposition 8: Limiting inner products at the EOC
• Proposition 9: Cosine map
• Proposition 10: $\overset{\infty}{K}$ at the EOC