Finite-data fixed-time NTK stability under integrably bounded training directions

Regard the network as a finite composition of $C^1$ maps on its Euclidean parameter space; [F1] gives the shared-parameter gradients. Let $J_f(\theta (t))$ be the Jacobian of $(f_{\theta }(x_i){)}_{i=1}^N$. It is continuous in $\theta$ because $\sigma$ is $C^2$, and it is bounded on the compact image of the continuous path $\theta ([0,T])$. The flow equation and $d\in L^1([0,T];{\mathbb{R}}^N)$ therefore imply $\dot{\theta }\in L^1([0,T])$. By [F8], applied to each of the finitely many parameter coordinates, $\theta$ is absolutely continuous and its increments equal the integrals of its derivative. Define the terminal array by $D_{1i}^L(t)=d_i(t)$, and obtain each hidden backpropagated array $D^{\mathrm{\ell }+1}(t)\in {\mathbb{R}}^{n_{\mathrm{\ell }+1}\times N}$ by weighting the derivative of each data-point output with $d_i(t)$. If $Z^{\mathrm{\ell }}(t)\in {\mathbb{R}}^{n_{\mathrm{\ell }}\times N}$ is the activation array, the weight and bias equations are ${\dot{W}}^{(\mathrm{\ell })}=-D^{\mathrm{\ell }+1}(Z^{\mathrm{\ell }}{)}^T\mathrm{/}(N\sqrt{n_{\mathrm{\ell }}})$ and ${\dot{b}}^{(\mathrm{\ell })}=-\beta D^{\mathrm{\ell }+1}\mathbf{1}\mathrm{/}N$, while hidden backpropagation is $D^{\mathrm{\ell }}=\mathrm{diag}(\sigma \mathrm{\prime }(U^{\mathrm{\ell }}))(W^{(\mathrm{\ell })}{)}^T{D}^{\mathrm{\ell }+1}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$. Here division by $N$ acts on the data sum.

For an array $R\in {\mathbb{R}}^{m\times N}$ write $\mathrm{\parallel }R{\mathrm{\parallel }}_{F,N}=\mathrm{\parallel }R{\mathrm{\parallel }}_F\mathrm{/}\sqrt{N}$. Let $c=\sup \mathrm{\mid }{\sigma }^{\mathrm{\prime }}\mathrm{\mid }$ and $M_2=\sup \mathrm{\mid }{\sigma }^{\mathrm{\prime }\mathrm{\prime }}\mathrm{\mid }$, both finite by the hypotheses (taking a difference-quotient limit in the global Lipschitz inequality gives $c\le \mathrm{Lip}(\sigma )$), and set $w_{\mathrm{\ell }}(t)=\mathrm{\parallel }{W}^{(\mathrm{\ell })}(t){\mathrm{\parallel }}_{op}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$ and $a_{\mathrm{\ell }}(t)=\mathrm{\parallel }{Z}^{\mathrm{\ell }}(t){\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$ for $1\le \mathrm{\ell }<L$. Let $K^{\mathrm{\ell }}$ be the tangent-kernel block for the preactivation array $U^{\mathrm{\ell }}$, acting on direction arrays by $({\mathcal{K}}^{\mathrm{\ell }}D{)}_{ai}=N^{-1}{\sum }_{b,j}{K}_{(i,a),(j,b)}^{\mathrm{\ell }}{D}_{bj}$, and put $B_{\mathrm{\ell }}=\mathrm{\parallel }{\mathcal{K}}^{\mathrm{\ell }}{\mathrm{\parallel }}_{op}$ using $\mathrm{\parallel }\cdot {\mathrm{\parallel }}_{F,N}$. The initial block and layerwise chain rule give $B_1\le \mathrm{\parallel }X{\mathrm{\parallel }}_F^2\mathrm{/}(n_{0}N)+{\beta }^2$ and $B_{\mathrm{\ell }+1}\le c^2{w}_{\mathrm{\ell }}^2{B}_{\mathrm{\ell }}+a_{\mathrm{\ell }}^2+{\beta }^2$ for $1\le \mathrm{\ell }<L$: the current-layer weight-and-bias Gram has operator norm at most $a_{\mathrm{\ell }}^2+{\beta }^2$, while the preceding Jacobian is composed with $W^{(\mathrm{\ell })}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$ and diagonal activation derivatives of norm at most $c$. The backpropagation identity in step 1.1 also gives $\mathrm{\parallel }{D}^{\mathrm{\ell }}(t){\mathrm{\parallel }}_{F,N}\le c^{L-\mathrm{\ell }}({\prod }_{k=\mathrm{\ell }}^{L-1}{w}_k(t))\mathrm{\parallel }d(t){\mathrm{\parallel }}_N$. Recursively, each $B_{\mathrm{\ell }}$ is bounded by a polynomial in the finitely many $a_k,w_k,c,\beta$, and each backpropagation norm is bounded by such a polynomial times $\mathrm{\parallel }d(t){\mathrm{\parallel }}_N$.

Cauchy–Schwarz in the data index and step 1.1 give, for $1\le \mathrm{\ell }<L$, $\mathrm{\parallel }{\dot{W}}^{(\mathrm{\ell })}{\mathrm{\parallel }}_{op}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}\le a_{\mathrm{\ell }}\mathrm{\parallel }{D}^{\mathrm{\ell }+1}{\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$. For hidden layers the chain rule gives ${\dot{U}}^{\mathrm{\ell }}=-{\mathcal{K}}^{\mathrm{\ell }}{D}^{\mathrm{\ell }}$, so $\mathrm{\parallel }{\dot{U}}^{\mathrm{\ell }}{\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}\le B_{\mathrm{\ell }}\mathrm{\parallel }{D}^{\mathrm{\ell }}{\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$; for $\mathrm{\ell }=1$ this is the first-layer parameter Gram. The hidden-bias equation gives $\beta \mathrm{\parallel }{\dot{b}}^{(\mathrm{\ell }-1)}{\mathrm{\parallel }}_2\mathrm{/}\sqrt{n_{\mathrm{\ell }}}\le {\beta }^2\mathrm{\parallel }{D}^{\mathrm{\ell }}{\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$, and the first-layer matrix gives $\mathrm{\parallel }{\dot{W}}^{(0)}{\mathrm{\parallel }}_F\mathrm{/}\sqrt{n_1}\le C\mathrm{\parallel }{D}^1{\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_1}$ with $C=\mathrm{\parallel }X{\mathrm{\parallel }}_F\mathrm{/}\sqrt{n_{0}N}$. The output bias is excluded because its derivative contributes only the constant feature $\beta$ and changes no later feature or parameter derivative. By [F7], $\mathrm{\mid }{\sigma }^{\mathrm{\prime }}(u)-\sigma \mathrm{\prime }(v)\mathrm{\mid }\le M_2\mathrm{\mid }u-v\mathrm{\mid }$; global Lipschitz continuity gives $\mathrm{\parallel }{\dot{Z}}^{\mathrm{\ell }}{\mathrm{\parallel }}_{F,N}\le c\mathrm{\parallel }{\dot{U}}^{\mathrm{\ell }}{\mathrm{\parallel }}_{F,N}$. Whenever $a_{\mathrm{\ell }}(t),w_{\mathrm{\ell }}(t)\le R$ for every hidden layer, the recurrence in step 2.1 gives $B_{\mathrm{\ell }}\le P_{\mathrm{\ell }}(R)$, where $P_1(R)=\mathrm{\parallel }X{\mathrm{\parallel }}_F^2\mathrm{/}(n_{0}N)+{\beta }^2$ and $P_{\mathrm{\ell }+1}(R)=c^2{R}^2{P}_{\mathrm{\ell }}(R)+R^2+{\beta }^2$; backpropagation gives $\mathrm{\parallel }{D}^{\mathrm{\ell }}{\mathrm{\parallel }}_{F,N}\le (cR{)}^{L-\mathrm{\ell }}\mathrm{\parallel }d{\mathrm{\parallel }}_N$. Substituting these bounds in the displayed rates shows that each normalized activation, preactivation, hidden-bias, and weight change rate is at most a polynomial in $R$ times $\mathrm{\parallel }d(t){\mathrm{\parallel }}_N\mathrm{/}\sqrt{n_{\min }}$, with $n_{\min }={\min }_{1\le k<L}{n}_k$.

For $1\le \mathrm{\ell }<L$ define ${\tilde{a}}_{\mathrm{\ell }}=\mathrm{\parallel }{Z}^{\mathrm{\ell }}(t)-Z^{\mathrm{\ell }}(0){\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$, ${\tilde{u}}_{\mathrm{\ell }}=\mathrm{\parallel }{U}^{\mathrm{\ell }}(t)-U^{\mathrm{\ell }}(0){\mathrm{\parallel }}_{F,N}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$, and ${\tilde{w}}_{\mathrm{\ell }}=\mathrm{\parallel }{W}^{(\mathrm{\ell })}(t)-W^{(\mathrm{\ell })}(0){\mathrm{\parallel }}_{op}\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$; for $0\le \mathrm{\ell }<L-1$ put ${\tilde{b}}_{\mathrm{\ell }}=\beta \mathrm{\parallel }{b}^{(\mathrm{\ell })}(t)-b^{(\mathrm{\ell })}(0){\mathrm{\parallel }}_2\mathrm{/}\sqrt{n_{\mathrm{\ell }+1}}$, and put ${\tilde{w}}_0=\mathrm{\parallel }{W}^{(0)}(t)-W^{(0)}(0){\mathrm{\parallel }}_F\mathrm{/}\sqrt{n_1}$. Define $A(t)={\sum }_{\mathrm{\ell }=1}^{L-1}(a_{\mathrm{\ell }}(0)+w_{\mathrm{\ell }}(0)+{\tilde{a}}_{\mathrm{\ell }}+{\tilde{u}}_{\mathrm{\ell }}+{\tilde{w}}_{\mathrm{\ell }}+{\tilde{b}}_{\mathrm{\ell }-1})+{\tilde{w}}_0$. Since $a_{\mathrm{\ell }}(t)\le a_{\mathrm{\ell }}(0)+{\tilde{a}}_{\mathrm{\ell }}(t)$ and $w_{\mathrm{\ell }}(t)\le w_{\mathrm{\ell }}(0)+{\tilde{w}}_{\mathrm{\ell }}(t)$, all $a_{\mathrm{\ell }},w_{\mathrm{\ell }}$ are bounded by $A(t)$. Each term of $A$ is absolutely continuous by [F8] and the norm/Lipschitz inequalities; the preceding estimates and the polynomial bounds in step 3.1 give, with $n_{\min }={\min }_{1\le \mathrm{\ell }<L}{n}_{\mathrm{\ell }}$, $\mathrm{\mid }\dot{A}(t)\mathrm{\mid }\le Q(A(t))\mathrm{\parallel }d(t){\mathrm{\parallel }}_N\mathrm{/}\sqrt{n_{\min }}$ for a polynomial $Q$ with nonnegative coefficients depending only on $L,N,c,M_2,\beta$ and the fixed inputs. Initially $A(0)$ is stochastically bounded in the sequential limit: each $a_{\mathrm{\ell }}(0{)}^2$ is an empirical second moment of a Lipschitz function of a finite Gaussian vector, and at each width limit [F4] applies after conditioning on earlier layers; for each fixed next-layer width, $W^{(\mathrm{\ell })}{W}^{(\mathrm{\ell })T}\mathrm{/}{n}_{\mathrm{\ell }}\to I$ entrywise by [F4], so $w_{\mathrm{\ell }}(0)$ is bounded in that iterated limit.

Fix probability tolerance $\epsilon >0$. Stochastic boundedness of $A(0)$ and of ${\int }_0^T\mathrm{\parallel }d(t){\mathrm{\parallel }}_{N}dt$ lets us restrict, except on a set of probability at most $2\epsilon$, to $A(0)\le R$ and this integral at most $M$ for finite $R,M$. On this event, the differential inequality in step 4.1 implies $A(t)\le R+1$ for all $t\le T$ once $n_{\min }$ is large: on $[0,R+1]$, $Q$ is bounded by some $Q_{\ast }$, so $\mathrm{\mid }A(t)-A(0)\mathrm{\mid }\le Q_{\ast }M\mathrm{/}\sqrt{n_{\min }}<1$; a first-exit argument justifies the bound on the whole interval. Since the normalized changes are nonnegative summands of $A(t)-A(0)$, each tends uniformly to zero.

For each layer, the exact tangent-kernel block recursion is $K_{(i,c),(j,d)}^{\mathrm{\ell }+1}={\delta }_{cd}(n_{\mathrm{\ell }}^{-1}{\sum }_a{Z}_{ai}^{\mathrm{\ell }}{Z}_{aj}^{\mathrm{\ell }}+{\beta }^2)+n_{\mathrm{\ell }}^{-1}{\sum }_{a,b}{W}_{ca}^{(\mathrm{\ell })}{W}_{db}^{(\mathrm{\ell })}{\sigma }^{\mathrm{\prime }}(U_{ai}^{\mathrm{\ell }}){\sigma }^{\mathrm{\prime }}(U_{bj}^{\mathrm{\ell }})K_{(i,a),(j,b)}^{\mathrm{\ell }}$. It follows by applying the chain rule to each coordinate and splitting parameter derivatives into the current layer and earlier shared parameters, as in [F1]. The first term varies uniformly by the activation and hidden-bias estimates in step 5.1. For the second term, for each output index $c$ and data index $i$ use the vector $v_{c,i}(t)\in {\mathbb{R}}^{n_{\mathrm{\ell }}}$ with entries $W_{ca}^{(\mathrm{\ell })}(t){\sigma }^{\mathrm{\prime }}(U_{ai}^{\mathrm{\ell }}(t))$. Their normalized Euclidean norms stay bounded. Their changes tend to zero in normalized Euclidean norm: the weight-change part is bounded by $c\mathrm{\parallel }\mathrm{\Delta }{W}_c{\mathrm{\parallel }}_2\mathrm{/}\sqrt{n_{\mathrm{\ell }}}$, while Cauchy–Schwarz for the initial-weight part gives $n_{\mathrm{\ell }}^{-1}{\sum }_a\mathrm{\mid }{W}_{ca}^{(\mathrm{\ell })}(0){\mathrm{\mid }}^2\mathrm{\mid }\mathrm{\Delta }{\sigma }_{ai}^{\mathrm{\prime }}{\mathrm{\mid }}^2\le (n_{\mathrm{\ell }}^{-1}{\sum }_a\mathrm{\mid }{W}_{ca}^{(\mathrm{\ell })}(0){\mathrm{\mid }}^4{)}^{1\mathrm{/}2}(n_{\mathrm{\ell }}^{-1}{\sum }_a\mathrm{\mid }\mathrm{\Delta }{\sigma }_{ai}^{\mathrm{\prime }}{\mathrm{\mid }}^4{)}^{1\mathrm{/}2}$. The first factor is stochastically bounded by [F3]–[F4]; the second tends uniformly to zero because [F7] gives $\mathrm{\mid }\mathrm{\Delta }{\sigma }^{\mathrm{\prime }}\mathrm{\mid }\le M_2\mathrm{\mid }\mathrm{\Delta }U\mathrm{\mid }$ and $\mathrm{\mid }\mathrm{\Delta }{\sigma }^{\mathrm{\prime }}\mathrm{\mid }\le 2c$, hence $\mathrm{\mid }\mathrm{\Delta }{\sigma }^{\mathrm{\prime }}{\mathrm{\mid }}^4\le (2c{)}^2{M}_2^2\mathrm{\mid }\mathrm{\Delta }U{\mathrm{\mid }}^2$; normalized preactivation changes tend to zero by step 5.1. For each fixed next-layer width, the finitely many output-index blocks have uniformly bounded subnetwork Gram operators by step 2.1, and induction from the constant first-layer Gram makes their changes tend uniformly to zero; applying $\mathrm{\mid }⟨v,K{v}^{\mathrm{\prime }}⟩-⟨w,K_0{w}^{\mathrm{\prime }}⟩\mathrm{\mid }\le \mathrm{\parallel }v-w\mathrm{\parallel }\mathrm{\parallel }K\mathrm{\parallel }\mathrm{\parallel }{v}^{\mathrm{\prime }}\mathrm{\parallel }+\mathrm{\parallel }w\mathrm{\parallel }\mathrm{\parallel }K-K_0\mathrm{\parallel }\mathrm{\parallel }{v}^{\mathrm{\prime }}\mathrm{\parallel }+\mathrm{\parallel }w\mathrm{\parallel }\mathrm{\parallel }{K}_0\mathrm{\parallel }\mathrm{\parallel }{v}^{\mathrm{\prime }}-w^{\mathrm{\prime }}\mathrm{\parallel }$ at each of the finitely many layers proves uniform convergence of the finite-data kernel blocks. The output layer has fixed scalar dimension, so this is exactly ${\sup }_{t\le T}\mathrm{\parallel }{K}_t-K_0{\mathrm{\parallel }}_{op}\to 0$ in probability.