MLPs at the EOC: Concentration of the NTK

We study the concentration of the Neural Tangent Kernel (NTK) Kθ : ℝm0 × ℝm0 → ℝml × ml of l-layer Multilayer Perceptrons (MLPs) N : ℝm0 × Theta → ℝml equipped with activation functions φ(s) = a s + b vert s vert for some a,b ∈ ℝ with the parameter θ ∈ Theta 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θ(xi1,xi2) for i1,i2 ∈ [1:n] over a dataset x1,cdots,xn subset ℝm0 concentrate simultaneously via maximal inequalities, we prove that the NTK matrix K(θ) = [frac1n Kθ(xi1,xi2) : i1,i2 ∈ [1:n]] ∈ ℝnml × nml concentrates around its infinitely wide limit oversetscriptscriptstyle∞K ∈ ℝnml × nml 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 mk = k2 m for some m ∈ mathbbN+1 for sufficient concentration. For such MLPs, we obtain the concentration bound mathbbP( Vert K(θ) - oversetscriptscriptstyle∞K Vert ≤ O((Deltaφ-2 + ml^frac12 l) kappaφ2 m^-frac12)) ≥ 1-O(m-1) modulo logarithmic terms, where we denoted Deltaφ = fracb2a2+b2 and kappaφ = fracvert a vert + vert b vert√a2 + b2. This reveals in particular that the absolute value (Deltaφ=1, kappaφ=1) beats the ReLU (Deltaφ=frac12, kappaφ=√2) in terms of the concentration of the NTK.