Introduction
It has been widely demonstrated that stochasticity in the training of neural networks (NNs)—mostly introduced by mini-batch sampling in stochastic gradient descent (SGD)—is often beneficial for optimization and generalization [e.g., 61, 46, 45, 104, 101]. While theoretical analyses establish a clear separation in favor of stochasticity in convex optimization [5 , 4 ] , and highlight the benefit of injected noise in escaping saddle points [56 , 57 ] , non-asymptotic convergence guarantees for overparameterized neural networks do not yet fully reflect this empirical picture.
Much prior work focuses on settings where the parameters remain close to their initialization throughout training. In several of these works, the convergence rate is governed by the minimum eigenvalue of the Neural Tangent Kernel (NTK), which is typically controlled by the near-initialization condition [40 , 105 , 85 , 14 ] . This condition is poorly aligned with noisy dynamics, which tend to move the parameters away from their initialization. Consistent with this, much of the existing theory focuses on deterministic algorithms [91 , 17 , 83 , 79 ] . Related analyses that do attempt to account for noise often treat it primarily as a nuisance [24 , 72 , 89 , 2 , 124 , 116 , 23 ] . As a result, finite-width convergence guarantees for noisy neural network training beyond the near-initialization regime remain limited.
In this work, we treat noise not merely as a feature of realistic training, but as a central tool in the analysis. We study continuous-time Langevin dynamics (CLD), in which isotropic Gaussian noise is added to gradient flow, and derive a change-of-measure inequality inspired by the second law of thermodynamics. This inequality allows us to control key quantities under the evolving noisy dynamics, in particular, the minimum eigenvalue of the NTK, without requiring the parameters to remain near initialization. This leads to a non-asymptotic global convergence guarantee for feedforward neural networks.
Our analysis differs conceptually from many existing works on Langevin dynamics [e.g., 96, 117, 64]: Rather than proving convergence by analyzing the log-Sobolev inequality (LSI) constant of the Gibbs distribution, we obtain a direct optimization bound along the dynamics. Avoiding this route is crucial, since the LSI constant typically deteriorates exponentially as the noise temperature decreases. In turn, this yields linear convergence with a decay exponent that is independent of the noise temperature.
Main contributions.
We prove a general change-of-measure inequality (Lemma 2.2 ) inspired by the second law of thermodynamics [35 , 78 , 47 ] . This bounds expectations under the time-t t t
We strengthen the lower bound on the minimum eigenvalue of the NTK at initialization introduced by Banerjee et al. [9 ] , by sharpening their high-probability guarantees (Lemma 5.3 ).
We apply the “second-law change-of-measure” to prove that the NTK stays well-conditioned during training, yielding a global convergence bound on the expected loss under regularized CLD (Theorem 5.1 ) with non-annealing, isotropic noise, and provide explicit dependence on width, temperature, and parameterization. We explicitly show our result applies to standard and NTK parameterizations. For shallow nets, we also prove convergence to zero loss if the noise is eventually removed.
We empirically demonstrate that the resulting Langevin dynamics evolve the Jacobian away from its initial value under standard parameterization, suggesting that the proof technique does not implicitly enforce lazy training.
Notation. We denote [ n ] = { 1 , … , n } [n] = \{1, \ldots, n\} [ n ] = { 1 , … , n } e.g. , v \mathbf{v} v X \mathbf{X} X f : R d → R f : \mathbb{R}^{d} \to\mathbb{R} f : R d → R ∇ f ∈ R d \nabla f \in\mathbb{R}^{d} ∇ f ∈ R d ∇ 2 f ∈ R d × d \nabla^{2}f \in\mathbb{R}^{d \times d} ∇ 2 f ∈ R d × d Δ f = Tr ( ∇ 2 f ) \Delta f = \operatorname{Tr}(\nabla^{2}f) Δ f = Tr ( ∇ 2 f ) f : R d → R n f : \mathbb{R}^{d} \to\mathbb{R}^{n} f : R d → R n ∇ f ∈ R n × d \nabla f \in\mathbb{R}^{n \times d} ∇ f ∈ R n × d
We use ≲ \lesssim ≲ i.e. , a ≲ b a \lesssim b a ≲ b c > 0 c > 0 c > 0 a ≤ c b a \le cb a ≤ c b c c c 1 a ≳ b a \gtrsim b a ≳ b c > 0 c > 0 c > 0 a = c b a = cb a = c b O , Ω , Θ O, \Omega, \Theta O , Ω , Θ O ~ , Ω ~ , Θ ~ \widetilde{O}, \widetilde{\Omega}, \widetilde{\Theta} O , Ω , Θ c c c numerical subscript or decoration (e.g. , c 1 c_{1} c 1 c ~ \tilde{c} c ~ C C C
We denote for a random variable (RV) X X X L p L^{p} L p ∥ X ∥ L p ≜ ( E [ ∣ X ∣ p ] ) 1 / p \lVert X\rVert_{L^{p}} \triangleq\left(\mathbb{E}\left[|X|^{p}\right]\right)^{1/p} ∥ X ∥ L p ≜ ( E [ ∣ X ∣ p ] ) 1/ p X X X ∥ X ∥ ψ 2 \lVert X\rVert_{\psi_{2}} ∥ X ∥ ψ 2 ∥ X ∥ ψ 1 \lVert X\rVert_{\psi_{1}} ∥ X ∥ ψ 1 p p p q q q p ≪ q p \ll q p ≪ q p p p q q q i.e. , for any event E \mathcal{E} E q ( E ) = 0 q(\mathcal{E}) = 0 q ( E ) = 0 p ( E ) = 0 p(\mathcal{E}) = 0 p ( E ) = 0 p ≪ q p \ll q p ≪ q q ≪ p q \ll p q ≪ p equivalent , denoted as p ∼ q p \sim q p ∼ q
Key technical result: parameter control via a change of measure inequality
We begin by introducing a change of measure inequality, inspired by the second law of thermodynamics, as a technical tool that allows us to control properties of Markov processes. In order to state our result, we require the following definition.
Definition 2.1 (Gibbs distribution). Let p p p q q q X \mathcal{X} X p ≪ q p \ll q p ≪ q p p p q q q Ψ : X → R \Psi:\mathcal{X}\to\mathbb{R} Ψ : X → R
Z ≔ ∫ e − Ψ d q < ∞ Z \coloneqq\int e^{-\Psi}\mathrm{d}q < \infty Z : = ∫ e − Ψ d q < ∞ and the Radon–Nikodym derivative of p p p q q q
d p d q = Z − 1 e − Ψ . \frac{\mathrm{d}p}{\mathrm{d}q}=Z^{-1}e^{-\Psi}. d q d p = Z − 1 e − Ψ . We consider a time-invariant Markov process on a state space X \mathcal{X} X p t p_t p t t t t p ∞ p_\infty p ∞ ν \nu ν D α D_\alpha D α p t p_t p t ν \nu ν D α D_\alpha D α α → 1 \alpha\to1 α → 1 D ∞ D_\infty D ∞ p 0 = ν p_0=\nu p 0 = ν [11 Theorem 8] , to obtain our main technical result.
Lemma 2.2 (Second law of thermodynamics change of measure inequality). Consider a time-invariant Markov process over the state space X \mathcal{X} X p t p_t p t t ≥ 0 t\geq0 t ≥ 0 p t ≪ p 0 p_t\ll p_0 p t ≪ p 0 p ∞ p_\infty p ∞ p 0 p_0 p 0 Ψ \Psi Ψ p ∞ ∼ p 0 p_\infty\sim p_0 p ∞ ∼ p 0 φ : X → R \varphi:\mathcal{X}\to\mathbb{R} φ : X → R α > 1 \alpha>1 α > 1 t ≥ 0 t\geq0 t ≥ 0
E p t φ ( x ) ≤ ( E p 0 [ e 2 ( α − 1 ) Ψ ] ) 1 2 α ( E p 0 [ φ ( x ) α α − 1 ] ) α − 1 α . \mathbb{E}_{p_t}\varphi(x)\leq\left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\Psi}\right]\right)^{\frac{1}{2\alpha}}\left(\mathbb{E}_{p_0}\left[\varphi(x)^{\frac{\alpha}{\alpha-1}}\right]\right)^{\frac{\alpha-1}{\alpha}}. E p t φ ( x ) ≤ ( E p 0 [ e 2 ( α − 1 ) Ψ ] ) 2 α 1 ( E p 0 [ φ ( x ) α − 1 α ] ) α α − 1 . We provide the proof in App. B. Lemma 2.2 enables us to control specific properties of a Markov process throughout its mixing phase using only properties at initialization. In the following section, we demonstrate its applicability to regularized Langevin dynamics.
Regularized Langevin dynamics
We study the evolution of parameters θ ∈ R N p \theta\in\mathbb{R}^{N_p} θ ∈ R N p θ 0 ∼ N ( 0 N p , Λ − 1 ) \theta_0\sim\mathcal{N}(0_{N_p},\Lambda^{-1}) θ 0 ∼ N ( 0 N p , Λ − 1 ) Λ ∈ R N p × N p \Lambda\in\mathbb{R}^{N_p\times N_p} Λ ∈ R N p × N p p ∞ p_\infty p ∞ p 0 p_0 p 0 L 2 L_2 L 2 θ \theta θ
d θ t = − ∇ θ L ( θ t ) d t − β − 1 Λ θ t d t + 2 β − 1 d w t , (1) \mathrm{d}\theta_t=-\nabla_\theta\mathcal{L}(\theta_t)\,\mathrm{d}t-\beta^{-1}\Lambda\theta_t\mathrm{d}t+\sqrt{2\beta^{-1}}\mathrm{d}w_t, \tag*{(1)} d θ t = − ∇ θ L ( θ t ) d t − β − 1 Λ θ t d t + 2 β − 1 d w t , (1) where L ( θ ) ≥ 0 \mathcal{L}(\theta)\geq0 L ( θ ) ≥ 0 θ \theta θ β \beta β d w t \mathrm{d}w_t d w t R N p \mathbb{R}^{N_p} R N p L ≥ 0 \mathcal{L}\geq0 L ≥ 0 L 2 L_2 L 2
Z = ∫ R N p e − β L ( θ ) − 1 2 θ ⊤ Λ θ d θ < ∞ . Z=\int_{\mathbb{R}^{N_p}}e^{-\beta\mathcal{L}(\theta)-\frac{1}{2}\theta^{\top}\Lambda\theta}\mathrm{d}\theta<\infty. Z = ∫ R N p e − β L ( θ ) − 2 1 θ ⊤ Λ θ d θ < ∞. Therefore, denoting by φ Λ \varphi_\Lambda φ Λ N ( 0 N p , Λ − 1 ) \mathcal{N}(0_{N_p},\Lambda^{-1}) N ( 0 N p , Λ − 1 ) ∝ e − β L ( θ ) φ Λ ( θ ) \propto e^{-\beta\mathcal{L}(\theta)}\varphi_\Lambda(\theta) ∝ e − β L ( θ ) φ Λ ( θ ) (1) (Proposition F.1).2 p ∞ ∝ e − β L ( θ ) φ Λ ( θ ) p_\infty\propto e^{-\beta\mathcal{L}(\theta)}\varphi_\Lambda(\theta) p ∞ ∝ e − β L ( θ ) φ Λ ( θ ) φ Λ \varphi_\Lambda φ Λ Ψ = β L ( θ ) \Psi=\beta\mathcal{L}(\theta) Ψ = β L ( θ ) (1) , with θ 0 ∼ N ( 0 N p , Λ − 1 ) \theta_0\sim\mathcal{N}(0_{N_p},\Lambda^{-1}) θ 0 ∼ N ( 0 N p , Λ − 1 )
Corollary 3.1. With θ \theta θ p 0 = N ( 0 N p , Λ − 1 ) p_0=\mathcal{N}(0_{N_p},\Lambda^{-1}) p 0 = N ( 0 N p , Λ − 1 ) (1) , we have for any non-negative measurable function φ : R N p → R \varphi:\mathbb{R}^{N_p}\to\mathbb{R} φ : R N p → R t ≥ 0 t\geq0 t ≥ 0 α > 1 \alpha>1 α > 1
E p t φ ( θ t ) ≤ ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α ∥ φ ( θ 0 ) ∥ L α α − 1 . \mathbb{E}_{p_t}\varphi(\theta_t)\leq\left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}}\|\varphi(\theta_0)\|_{L^{\frac{\alpha}{\alpha-1}}}. E p t φ ( θ t ) ≤ ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1 ∥ φ ( θ 0 ) ∥ L α − 1 α . Remark 3.2. Under our model, as defined in Section 4, we are mostly focused on regimes where
( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α = O ( 1 ) , \left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}}=O(1), ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1 = O ( 1 ) , which simplifies the analysis greatly.
Itô’s formula. To analyze the dynamics of the loss under regularized CLD (Eq. (1) ), we can employ Itô’s formula [52 ] , [75 Chapter 1, Theorem 6.4] .3
d L ( θ t ) = ∇ L ( θ t ) ⊤ ( − ∇ L ( θ t ) − β − 1 Λ θ t ) d t + 2 β − 1 ∇ L ( θ t ) ⊤ d w t + 1 2 Tr ( 2 β − 1 ∇ 2 L ( θ t ) ) d t . \mathrm{d}\mathcal{L}(\theta_t)=\nabla\mathcal{L}(\theta_t)^{\top}\left(-\nabla\mathcal{L}(\theta_t)-\beta^{-1}\Lambda\theta_t\right)\mathrm{d}t+\sqrt{2\beta^{-1}}\nabla\mathcal{L}(\theta_t)^{\top}\mathrm{d}w_t+\frac{1}{2}\operatorname{Tr}\left(2\beta^{-1}\nabla^2\mathcal{L}(\theta_t)\right)\mathrm{d}t. d L ( θ t ) = ∇ L ( θ t ) ⊤ ( − ∇ L ( θ t ) − β − 1 Λ θ t ) d t + 2 β − 1 ∇ L ( θ t ) ⊤ d w t + 2 1 Tr ( 2 β − 1 ∇ 2 L ( θ t ) ) d t . By taking the expected value over both the initialization and the Brownian noise, we arrive at
d d t E L ( θ t ) = E [ − ∥ ∇ L ( θ t ) ∥ 2 + β − 1 Δ L ( θ t ) − β − 1 ∇ L ( θ t ) ⊤ Λ θ t ] , (2) \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)=\mathbb{E}\left[-\left\|\nabla\mathcal{L}(\theta_t)\right\|^2+\beta^{-1}\Delta\mathcal{L}(\theta_t)-\beta^{-1}\nabla\mathcal{L}(\theta_t)^{\top}\Lambda\theta_t\right], \tag*{(2)} d t d E L ( θ t ) = E [ − ∥ ∇ L ( θ t ) ∥ 2 + β − 1 Δ L ( θ t ) − β − 1 ∇ L ( θ t ) ⊤ Λ θ t ] , (2) where Δ L ( θ t ) \Delta\mathcal{L}(\theta_t) Δ L ( θ t ) (2) , thereby obtaining bounds on the expected loss for feedforward neural networks trained via regularized CLD.
Model: deep neural network
We consider a scalar-output feedforward NN trained with the mean squared error (MSE) loss [e.g., 40, 6, 79].
Formally, consider a training set D = { ( x i , y i ) } i ∈ [ n ] \mathcal{D}=\{(x_i,y_i)\}_{i\in[n]} D = {( x i , y i ) } i ∈ [ n ] x i ∈ R d x_i\in\mathbb{R}^d x i ∈ R d y i ∈ R y_i\in\mathbb{R} y i ∈ R X ∈ R n × d X\in\mathbb{R}^{n\times d} X ∈ R n × d i i i x i ⊤ x_i^{\top} x i ⊤ y ∈ R n y\in\mathbb{R}^n y ∈ R n i i i y i y_i y i f f f L + 1 L+1 L + 1 L L L m l m_l m l l ∈ [ L ] l\in[L] l ∈ [ L ] α ( 0 ) ( x ) = x \alpha^{(0)}(x)=x α ( 0 ) ( x ) = x ∀ x ∈ R d \forall x\in\mathbb{R}^d ∀ x ∈ R d
α ( l ) ( x ) = ϕ ( s l W ( l ) α ( l − 1 ) ( x ) ) , l ∈ [ L ] , \alpha^{(l)}(x)=\phi\left(s_l\mathbf{W}^{(l)}\alpha^{(l-1)}(x)\right),\quad l\in[L], α ( l ) ( x ) = ϕ ( s l W ( l ) α ( l − 1 ) ( x ) ) , l ∈ [ L ] , f ( θ ; x ) = α ( L + 1 ) ( x ) = s L + 1 v ⊤ α ( L ) ( x ) , (3) f(\theta;x)=\alpha^{(L+1)}(x)=s_{L+1}\mathbf{v}^{\top}\alpha^{(L)}(x), \tag*{(3)} f ( θ ; x ) = α ( L + 1 ) ( x ) = s L + 1 v ⊤ α ( L ) ( x ) , (3) where W ( l ) ∈ R m l × m l − 1 \mathbf{W}^{(l)}\in\mathbb{R}^{m_l\times m_{l-1}} W ( l ) ∈ R m l × m l − 1 l ∈ [ L ] l\in[L] l ∈ [ L ] v ∈ R m L \mathbf{v}\in\mathbb{R}^{m_L} v ∈ R m L ϕ ( ⋅ ) \phi(\cdot) ϕ ( ⋅ ) s l s_l s l l ∈ [ L + 1 ] l\in[L+1] l ∈ [ L + 1 ] s l = 1 s_l=1 s l = 1 l ∈ [ L + 1 ] l\in[L+1] l ∈ [ L + 1 ]
θ ≜ ( vec ( W ( 1 ) ) ⊤ , … , vec ( W ( L ) ) ⊤ , v ⊤ ) ⊤ ∈ R N p , \theta\triangleq\left(\operatorname{vec}\left(\mathbf{W}^{(1)}\right)^{\top},\ldots,\operatorname{vec}\left(\mathbf{W}^{(L)}\right)^{\top},\mathbf{v}^{\top}\right)^{\top}\in\mathbb{R}^{N_p}, θ ≜ ( vec ( W ( 1 ) ) ⊤ , … , vec ( W ( L ) ) ⊤ , v ⊤ ) ⊤ ∈ R N p , with m 0 ≜ d m_0\triangleq d m 0 ≜ d N p ≜ ∑ l = 1 L m l m l − 1 + m L N_p\triangleq\sum_{l=1}^{L}m_lm_{l-1}+m_L N p ≜ ∑ l = 1 L m l m l − 1 + m L i ∈ [ n ] i\in[n] i ∈ [ n ] f i ( θ ) ≜ f ( θ ; x i ) ∈ R f_i(\theta)\triangleq f(\theta;x_i)\in\mathbb{R} f i ( θ ) ≜ f ( θ ; x i ) ∈ R f ( θ ) ≜ f ( θ ; X ) = ( f 1 ( θ ) , … , f n ( θ ) ) ⊤ ∈ R n f(\theta)\triangleq f(\theta;X)=(f_1(\theta),\ldots,f_n(\theta))^{\top}\in\mathbb{R}^n f ( θ ) ≜ f ( θ ; X ) = ( f 1 ( θ ) , … , f n ( θ ) ) ⊤ ∈ R n ∇ f ( θ ) ≜ ∂ f ( θ ; X ) ∂ θ ∈ R n × N p \nabla f(\theta)\triangleq\frac{\partial f(\theta;X)}{\partial\theta}\in\mathbb{R}^{n\times N_p} ∇ f ( θ ) ≜ ∂ θ ∂ f ( θ ; X ) ∈ R n × N p K N T K ( θ ) ≜ ∇ f ( θ ) ∇ f ( θ ) ⊤ K_{\mathrm{NTK}}(\theta)\triangleq\nabla f(\theta)\nabla f(\theta)^{\top} K NTK ( θ ) ≜ ∇ f ( θ ) ∇ f ( θ ) ⊤ r ( θ ) ≜ f ( θ ) − y r(\theta)\triangleq f(\theta)-y r ( θ ) ≜ f ( θ ) − y
L ( θ ) = 1 2 n ∥ f ( θ ) − y ∥ 2 = 1 2 n ∥ r ( θ ) ∥ 2 . (4) \mathcal{L}(\theta)=\frac{1}{2n}\left\|f(\theta)-y\right\|^2=\frac{1}{2n}\left\|r(\theta)\right\|^2. \tag*{(4)} L ( θ ) = 2 n 1 ∥ f ( θ ) − y ∥ 2 = 2 n 1 ∥ r ( θ ) ∥ 2 . (4) Weight initialization. For some σ 0 > 0 \sigma_0>0 σ 0 > 0 c ϕ , σ 0 ≜ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z ) ] = Θ ( 1 ) c_{\phi,\sigma_0}\triangleq\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(z)]=\Theta(1) c ϕ , σ 0 ≜ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z )] = Θ ( 1 ) v 0 2 ≜ σ 0 2 c ϕ , σ 0 = Θ ( 1 ) v_0^2\triangleq\frac{\sigma_0^2}{c_{\phi,\sigma_0}}=\Theta(1) v 0 2 ≜ c ϕ , σ 0 σ 0 2 = Θ ( 1 )
vec ( W ( l ) ) ∼ N ( 0 m l m l − 1 , v 0 2 a l I m l m l − 1 ) , l ∈ [ L ] ; v ∼ N ( 0 m L , 1 a L + 1 I m L ) , (5) \operatorname{vec}\left(\mathbf{W}^{(l)}\right)\sim\mathcal{N}\left(0_{m_lm_{l-1}},\frac{v_0^2}{a_l}I_{m_lm_{l-1}}\right),\quad l\in[L];\quad\mathbf{v}\sim\mathcal{N}\left(0_{m_L},\frac{1}{a_{L+1}}I_{m_L}\right), \tag*{(5)} vec ( W ( l ) ) ∼ N ( 0 m l m l − 1 , a l v 0 2 I m l m l − 1 ) , l ∈ [ L ] ; v ∼ N ( 0 m L , a L + 1 1 I m L ) , (5) where a l a_l a l l ∈ [ L + 1 ] l\in[L+1] l ∈ [ L + 1 ] ∀ l ∈ [ L ] \forall l\in[L] ∀ l ∈ [ L ] a l s l 2 = m l − 1 \frac{a_l}{s_l^2}=m_{l-1} s l 2 a l = m l − 1 [48 ] .4 a L + 1 a_{L+1} a L + 1 κ ≜ s L + 1 m L a L + 1 \kappa\triangleq s_{L+1}\sqrt{\frac{m_L}{a_{L+1}}} κ ≜ s L + 1 a L + 1 m L κ \kappa κ [32 , 115 ] . Equivalently, Eq. (5) may be presented as
θ 0 ∼ N ( 0 N p , Λ − 1 ) , Λ = diag ( a 1 v 0 2 I m 1 m 0 , … , a L v 0 2 I m L m L − 1 , a L + 1 I m L ) . (6) \theta_0\sim\mathcal{N}\left(0_{N_p},\Lambda^{-1}\right),\quad\Lambda=\operatorname{diag}\left(\frac{a_1}{v_0^2}I_{m_1m_0},\ldots,\frac{a_L}{v_0^2}I_{m_Lm_{L-1}},a_{L+1}I_{m_L}\right). \tag*{(6)} θ 0 ∼ N ( 0 N p , Λ − 1 ) , Λ = diag ( v 0 2 a 1 I m 1 m 0 , … , v 0 2 a L I m L m L − 1 , a L + 1 I m L ) . (6) Activation function. We impose the following assumptions on the activation function:
Assumption 4.1 (Bounded activation). The activation ϕ \phi ϕ ∣ ϕ ∣ ≤ M ϕ \lvert\phi\rvert\le M_{\phi} ∣ ϕ ∣ ≤ M ϕ
Assumption 4.2 (Activation derivative). The activation is twice continuously differentiable and its derivatives are bounded (∣ ϕ ′ ∣ ≤ α ϕ \lvert\phi'\rvert\le\alpha_{\phi} ∣ ϕ ′ ∣ ≤ α ϕ ∣ ϕ ′ ′ ∣ ≤ β ϕ \lvert\phi''\rvert\le\beta_{\phi} ∣ ϕ ′′ ∣ ≤ β ϕ 5
Assumption 4.3 (Activation at origin). The activation ϕ \phi ϕ ϕ ( 0 ) = 0 \phi(0)=0 ϕ ( 0 ) = 0
We note that under our analysis, Assumption 4.1 could be relaxed to apply only to the activation function used in the last hidden layer. Taken together, these assumptions are satisfied in practice by a broad class of activations, including tanh, erf, and soft-root-sign [68 ] , as well as by GELU [49 ] and Swish [97 ] in all hidden layers except the last. Activations satisfying similar assumptions have also recently been shown to be beneficial, for example in the context of self-normalization [21 ] .
Data scaling. We make mild data-scaling assumptions that are common in such analyses [e.g., 91, 105, 14] and can be achieved in practice by normalization.
Assumption 4.4 (Normalized data). Every input sample x i ∈ R d x_i \in\mathbb{R}^{d} x i ∈ R d i ∈ [ n ] i \in[n] i ∈ [ n ] ∥ x i ∥ 2 2 = c ϕ , σ 0 d \lVert x_i\rVert_2^2 = c_{\phi,\sigma_0}d ∥ x i ∥ 2 2 = c ϕ , σ 0 d 6
Assumption 4.5 (Bounded labels). For all i ∈ [ n ] i \in[n] i ∈ [ n ] ∣ y i ∣ ≤ R y \lvert y_i\rvert\le R_y ∣ y i ∣ ≤ R y R y = Θ ( 1 ) R_y=\Theta(1) R y = Θ ( 1 )
Data separability. Similarly to Banerjee et al. [9 ] , our analysis relies on the following assumption:
Assumption 4.6 (Data separability). We assume
λ 1 ≜ λ min ( E g ∼ N ( 0 d , 1 α 1 2 I d ) [ ϕ ( s 1 X g ) ϕ ( s 1 X g ) T ] ) > 0. \lambda_1 \triangleq\lambda_{\min}\left(\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}(0_d,\frac{1}{\alpha_1^2}I_d)}\left[\phi(s_1X\boldsymbol{g})\phi(s_1X\boldsymbol{g})^{\mathsf{T}}\right]\right)>0. λ 1 ≜ λ m i n ( E g ∼ N ( 0 d , α 1 2 1 I d ) [ ϕ ( s 1 X g ) ϕ ( s 1 X g ) T ] ) > 0. Assumption 4.6 is satisfied under common data assumptions, e.g., it is satisfied with high probability (over the data) for i.i.d. data on the sphere [85 Theorem 3.3] , and for δ \delta δ [91 ] with ReLU activation. See also the discussion in Banerjee et al. [9 Remark 4.2] . We note that λ 1 \lambda_1 λ 1 λ 1 ≤ α ϕ 2 ν ϕ 2 c ϕ , σ 0 = O ( 1 ) \lambda_1 \le\alpha_{\phi}^2\nu_{\phi}^2c_{\phi,\sigma_0}=O(1) λ 1 ≤ α ϕ 2 ν ϕ 2 c ϕ , σ 0 = O ( 1 ) (17) ). Crucially, this condition is imposed only on the inputs; for the labels, we assume only boundedness.
For a summary of the assumptions, their role in the analysis, and possible relaxations, see Table 4 .
Table 4. Assumptions grouped by their role in the analysis, with where each enters and whether it can be relaxed.
Global convergence of deep networks under regularized CLD
We now present our main theorem.
Theorem 5.1 (Global convergence of deep networks via regularized CLD). Consider the deep network defined in § 4 , under the assumptions therein, with equal widths m l = m m_l=m m l = m l ∈ [ L ] l\in[L] l ∈ [ L ] L = Θ ( 1 ) L=\Theta(1) L = Θ ( 1 ) (1) , with Λ \Lambda Λ (6) . Denoting κ ˉ ≜ max ( 1 , κ ) \bar{\kappa}\triangleq\max(1,\kappa) κ ˉ ≜ max ( 1 , κ )
m = Ω ( κ ˉ 2 β n λ 1 ln n λ 1 ln ( κ ˉ 4 β λ 1 ) ) , β = Ω ( n λ 1 ) . m=\Omega\left(\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\ln\left(\frac{\bar{\kappa}^{4}\beta}{\lambda_1}\right)\right),\quad\beta=\Omega\left(\frac{n}{\lambda_1}\right). m = Ω ( κ ˉ 2 β λ 1 n ln λ 1 n ln ( λ 1 κ ˉ 4 β ) ) , β = Ω ( λ 1 n ) . Then, for all t ≥ 0 t\ge0 t ≥ 0
E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c 1 s L + 1 2 + λ 1 m n t + L ∞ , (7) \mathbb{E}\mathcal{L}(\theta_t)\le\left(\mathbb{E}\mathcal{L}(\theta_0)-\mathcal{L}_{\infty}\right)e^{-c_1s_{L+1}^{2}+\frac{\lambda_1^{m}}{n}t}+\mathcal{L}_{\infty}, \tag*{(7)} E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c 1 s L + 1 2 + n λ 1 m t + L ∞ , (7) where c 1 c_1 c 1
L ∞ ≲ ( 1 + κ 4 β ) exp ( − c 2 m κ 2 β n λ 1 ln n λ 1 ) + β − 1 n λ 1 ( 1 + κ 2 s 1 2 d s L + 1 2 m + κ 2 s L + 1 2 ∑ l = 2 L s l 2 ) + β − 1 n 2 λ 1 2 κ 2 κ ˉ 2 ( s 1 4 d 2 s L + 1 4 m 2 + 1 s L + 1 4 ( ∑ l = 2 L s l 2 ) 2 + 1 s L + 1 4 ∑ l = 2 L s l 4 ) + β − 2 n 2 λ 1 2 1 κ 4 . \mathcal{L}_{\infty}\lesssim(1+\kappa^4\beta)\exp\left(-c_2\frac{m}{\kappa^2\beta}\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)+\beta^{-1}\frac{n}{\lambda_1}\left(1+\kappa^2\frac{s_1^2d}{s_{L+1}^2m}+\frac{\kappa^2}{s_{L+1}^2}\sum_{l=2}^{L}s_l^2\right)
+\beta^{-1}\frac{n^2}{\lambda_1^2}\kappa^2\bar{\kappa}^2\left(\frac{s_1^4d^2}{s_{L+1}^4m^2}+\frac{1}{s_{L+1}^4}\left(\sum_{l=2}^{L}s_l^2\right)^2+\frac{1}{s_{L+1}^4}\sum_{l=2}^{L}s_l^4\right)+\beta^{-2}\frac{n^2}{\lambda_1^2}\frac{1}{\kappa^4}. L ∞ ≲ ( 1 + κ 4 β ) exp ( − c 2 κ 2 β m λ 1 n ln λ 1 n ) + β − 1 λ 1 n ( 1 + κ 2 s L + 1 2 m s 1 2 d + s L + 1 2 κ 2 l = 2 ∑ L s l 2 ) + β − 1 λ 1 2 n 2 κ 2 κ ˉ 2 s L + 1 4 m 2 s 1 4 d 2 + s L + 1 4 1 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 1 l = 2 ∑ L s l 4 + β − 2 λ 1 2 n 2 κ 4 1 . For shallow networks, i.e., L = 1 L = 1 L = 1
\begin{aligned}
\mathcal{L}_{\infty}^{(L=1)} \lesssim\left(1+\kappa^{4}\beta\right)\exp\left(-\frac{c_{2}m}{\kappa^{2}\beta\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)+\beta^{-1}\frac{n}{\lambda_{1}}\left(1+\kappa^{2}\frac{s_{1}^{2}d}{s_{2}^{2}m}\right)} \\
&+\beta^{-1}\frac{n^{2}}{\lambda_{1}^{2}}\kappa^{2}\kappa^{2}\frac{s_{1}^{4}d^{2}}{s_{2}^{4}m^{2}}+\beta^{-2}\frac{n^{2}}{\lambda_{1}^{2}}\left(\frac{1}{\kappa^{4}}+\frac{s_{1}^{4}d^{2}}{s_{2}^{4}m^{2}}\right). \tag*{(8)}
\end{aligned}
We provide the proofs for the shallow and deep cases in App. G and App. H, respectively. The floor L ∞ \mathcal{L}_{\infty} L ∞ β \beta β (1) is Gibbs rather than a point mass, and in Table 1 we show it can be made smaller than any fixed target by the choice of m m m β \beta β
Table 1. Convergence bounds from Theorem 5.1 under standard and NTK parameterizations.
Proposition 5.2 (Convergence to zero loss for shallow networks). For shallow networks, i.e., L = 1 L = 1 L = 1 τ = Ω ~ ( n s 2 2 λ 1 m ) \tau= \widetilde{\Omega}\left(\frac{n}{s_{2}^{2}\lambda_{1}m}\right) τ = Ω ( s 2 2 λ 1 m n )
d d t θ t = − ∇ L ( θ t ) . (9) \frac{\mathrm{d}}{\mathrm{d}t}\theta_t=-\nabla\mathcal{L}(\theta_t). \tag*{(9)} d t d θ t = − ∇ L ( θ t ) . (9) Let K > 2 K > 2 K > 2
1-\eta\exp\left(-\frac{c_{4}m}{c_{\beta}\kappa^{2}\beta\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)-\frac{2}{K},}
provided that parameters are chosen such that
L ∞ ( L = 1 ) ≲ s 2 2 m 2 K n 2 d 2 s 1 4 ( 1 d s 1 2 + κ 2 m s 2 2 ) ( κ 2 s 1 2 d s 2 2 + n m ) . (10) \mathcal{L}_{\infty}^{(L=1)}\lesssim\frac{s_{2}^{2}m^{2}}{Kn^{2}d^{2}s_{1}^{4}\left(\frac{1}{ds_{1}^{2}}+\frac{\kappa^{2}}{ms_{2}^{2}}\right)\left(\frac{\kappa^{2}s_{1}^{2}d}{s_{2}^{2}}+\sqrt{nm}\right)}. \tag*{(10)} L ∞ ( L = 1 ) ≲ K n 2 d 2 s 1 4 ( d s 1 2 1 + m s 2 2 κ 2 ) ( s 2 2 κ 2 s 1 2 d + nm ) s 2 2 m 2 . (10) E.g., under the standard parameterization with d = o ( n ) d = o(\sqrt{n}) d = o ( n ) K = Θ ( n / d ) K = \Theta(\sqrt{n}/d) K = Θ ( n / d ) β = Ω ~ ( n ) \beta= \widetilde{\Omega}(n) β = Ω ( n ) m = Ω ~ ( n 2 ) m = \widetilde{\Omega}(n^{2}) m = Ω ( n 2 ) 1 − o ( 1 ) 1-o(1) 1 − o ( 1 )
In Table 1 we provide concrete examples for simplified forms of L ∞ \mathcal{L}_{\infty} L ∞ L ∞ ( L = 1 ) \mathcal{L}_{\infty}^{(L=1)} L ∞ ( L = 1 ) d = O ( n 3 / 2 ) d = O(n^{3/2}) d = O ( n 3/2 )
Proof outline
By substituting the MSE loss definition, (4) , into (2) , we arrive at the following ODE for the expected loss:
d d t E L ( θ t ) = 1 n E [ − n − 1 r ( θ t ) ⊤ K N T K ( θ t ) r ( θ t ) + β − 1 Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)=\frac{1}{n}\mathbb{E}\left[-n^{-1}\mathbf{r}(\theta_t)^{\top}\mathbf{K}_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)+\beta^{-1}\operatorname{Tr}\left(\nabla f(\theta_t)^{\top}\nabla f(\theta_t)\right)
\right. d t d E L ( θ t ) = n 1 E [ − n − 1 r ( θ t ) ⊤ K NTK ( θ t ) r ( θ t ) + β − 1 Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) + β − 1 ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) − β − 1 r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] . (11) \left.
+\beta^{-1}\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)-\beta^{-1}\mathbf{r}(\theta_t)^{\top}\nabla f(\theta_t)\Lambda\theta_t
\right]. \tag*{(11)} + β − 1 i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) − β − 1 r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] . (11) The proof then proceeds by upper bounding each of the four terms in (11) for all times t ≥ 0 t \ge0 t ≥ 0
The first term is the descent term, for which we require a lower bound on the minimum eigenvalue of the NTK at initialization, observing that for any λ ⋆ > 0 \lambda_{\star} > 0 λ ⋆ > 0
− n − 1 E [ r ( θ t ) ⊤ K N T K ( θ t ) r ( θ t ) ] ≤ − n − 1 E [ ∥ r ( θ t ) ∥ 2 2 λ min ( K N T K ( θ t ) ) ] = − 2 E [ L ( θ t ) λ min ( K N T K ( θ t ) ) ] ≤ − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) max ( 0 , λ ⋆ − λ min ( K N T K ( θ t ) ) ) ] . (12) \begin{aligned}
-n^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}\mathbf{K}_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)\right]
\le-n^{-1}\mathbb{E}\left[\|\mathbf{r}(\theta_t)\|_{2}^{2}\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_t)\right)\right] \\
&= -2\mathbb{E}\left[\mathcal{L}(\theta_t)\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_t)\right)\right] \\
&\le-2\lambda_{\star}\mathbb{E}\left[\mathcal{L}(\theta_t)\right]+2\mathbb{E}\left[\mathcal{L}(\theta_t)\max\left(0,\lambda_{\star}-\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_t)\right)\right)\right]. \tag*{(12)}
\end{aligned} − n − 1 E [ r ( θ t ) ⊤ K NTK ( θ t ) r ( θ t ) ] ≤ − n − 1 E [ ∥ r ( θ t ) ∥ 2 2 λ m i n ( K NTK ( θ t ) ) ] = − 2 E [ L ( θ t ) λ m i n ( K NTK ( θ t ) ) ] ≤ − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ t ) ) ) ] . (12) In App. A.4, we review existing bounds for the NTK minimum eigenvalue, and find that existing bounds are not suitable for our setting, with the closest match being Banerjee et al. [9 ] . However, they analyze only the NTK parameterization, and their bound has failure probability polynomially small in m m m § 4 , and has exponentially low failure probability. To obtain such a bound, we tighten the analysis of Banerjee et al. [9 ] , yielding the following lemma, proven in App. D.
Lemma 5.3. For m = Ω ( n λ 1 ln n λ 1 ln n ) m=\Omega\left(\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\ln n\right) m = Ω ( λ 1 n ln λ 1 n ln n ) 1 − exp ( − c ~ 1 m λ 1 n λ 1 ) 1-\exp\left(-\widetilde{c}_1\frac{m}{\lambda_1\frac{n}{\lambda_1}}\right) 1 − exp ( − c 1 λ 1 λ 1 n m ) { W ( l ) } l ∈ [ L ] \{W^{(l)}\}_{l\in[L]} { W ( l ) } l ∈ [ L ]
λ min ( K N T K ( θ ) ) ≥ c 0 s L + 1 2 m λ 1 . \lambda_{\min}\left(K_{\mathrm{NTK}}(\theta)\right)\ge c_0s_{L+1}^{2}m\lambda_1 . λ m i n ( K NTK ( θ ) ) ≥ c 0 s L + 1 2 m λ 1 . Then, we choose λ ⋆ = c 0 s L + 1 2 m λ 1 \lambda_\star=c_0s_{L+1}^{2}m\lambda_1 λ ⋆ = c 0 s L + 1 2 m λ 1 (12) using our change-of-measure inequality Corollary 3.1.
Similarly, we upper bound the remaining terms in Eq. (11) using Corollary 3.1 and known properties at initialization, relying heavily on properties of subgaussian and subexponential components. Finally, we arrive at a (possibly complicated) constant a > 0 a>0 a > 0
d d t E L ( θ t ) ≤ 1 n ( − c λ ⋆ E L ( θ t ) + a ) . \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)\le\frac{1}{n}\left(-c\lambda_\star\mathbb{E}\mathcal{L}(\theta_t)+a\right). d t d E L ( θ t ) ≤ n 1 ( − c λ ⋆ E L ( θ t ) + a ) . Then, to conclude the proof of Theorem 5.1, Grönwall’s inequality yields, for all t ≥ 0 t\ge0 t ≥ 0
E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c λ ⋆ t / n + L ∞ . \mathbb{E}\mathcal{L}(\theta_t)\le\left(\mathbb{E}\mathcal{L}(\theta_0)-L_\infty\right)e^{-c\lambda_\star t/n}+L_\infty. E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c λ ⋆ t / n + L ∞ . Convergence to zero loss by gradient flow (Proposition 5.2 ) is then established by adapting the proximity-to-initialization claims of Song et al. [105 ] . Using Theorem 5.1 and the change-of-measure inequality, we show that their conditions hold with positive probability at a sufficiently large deterministic stopping time for CLD, without requiring the parameters to stay near θ 0 \theta_0 θ 0 7
Analysis of the convergence result
Each term of L ∞ L_\infty L ∞ (11) . Reading Eq. (8) from left to right: The exponent comes from the NTK term; the (unavoidable) β − 1 n λ 1 \beta^{-1}\frac{n}{\lambda_1} β − 1 λ 1 n β − 1 Tr ( ∇ f T ∇ f ) \beta^{-1}\operatorname{Tr}\left(\nabla f^\mathsf{T}\nabla f\right) β − 1 Tr ( ∇ f T ∇ f ) β − 1 n 2 λ 1 2 \beta^{-1}\frac{n^2}{\lambda_1^2} β − 1 λ 1 2 n 2 β − 1 ∑ i r i Δ f i \beta^{-1}\sum_i r_i\Delta f_i β − 1 ∑ i r i Δ f i β − 2 n 2 λ 1 2 \beta^{-2}\frac{n^2}{\lambda_1^2} β − 2 λ 1 2 n 2
In order to interpret our bounds in Theorem 5.1 , we discuss the requirements for achieving L ∞ ≤ L g L_{\infty} \leq L_g L ∞ ≤ L g L g < 1 L_g < 1 L g < 1 λ 1 = Θ ( 1 ) \lambda_1 = \Theta(1) λ 1 = Θ ( 1 ) L ∞ L_{\infty} L ∞ L g < 1 L_g < 1 L g < 1 m = Ω ~ ( κ ˉ 2 β n ) m = \widetilde{\Omega}(\bar{\kappa}^{2}\beta n) m = Ω ( κ ˉ 2 β n ) κ ˉ ≥ 1 \bar{\kappa} \geq1 κ ˉ ≥ 1 m = Ω ~ ( β n ) m = \widetilde{\Omega}(\beta n) m = Ω ( β n ) β − 1 n λ 1 \beta^{-1}\frac{n}{\lambda_1} β − 1 λ 1 n β = Ω ( n / L g ) \beta= \Omega(n/L_g) β = Ω ( n / L g ) 8
Taken together, these demands point to a tradeoff in the role of noise. On one hand, injected noise is essential in our analysis because, through the second-law change-of-measure argument, it allows control of quantities under the evolving training distribution without requiring the parameters to remain in a small neighborhood of initialization. On the other hand, the same noise leads to an unavoidable loss floor, since it is isotropic and non-diminishing. This is also evident empirically (see § 6 ). The bound thus makes explicit how width, temperature, and parameterization interact: larger width improves kernel conditioning, while larger inverse temperature reduces the optimization floor but weakens the stochastic mechanism of the proof. Moreover, any scaling κ > 1 \kappa> 1 κ > 1 [32 ] .
Width requirements in context. Under standard parameterization, our requirement m = Ω ~ ( β n ) m = \widetilde{\Omega}(\beta n) m = Ω ( β n ) m = Ω ~ ( n 2 ) m = \widetilde{\Omega}(n^2) m = Ω ( n 2 ) m = Ω ~ ( n 3 ) m = \widetilde{\Omega}(n^3) m = Ω ( n 3 ) m = Ω ( n 5 / 3 ) m = \Omega(n^{5/3}) m = Ω ( n 5/3 ) [105 ] 9 m = Ω ( n 2 ) m = \Omega(n^2) m = Ω ( n 2 ) [85 ] ; see Table 2 for a broader comparison. Our width requirement may be tightened for deep networks, but the shallow gap is intrinsic to our approach.
Table 2. Convergence analyses of gradient descent and stochastic gradient descent.
From a stopped schedule to noise annealing. Proposition 5.2 can be viewed as the simplest nontrivial noise schedule : fixed inverse temperature β \beta β t = τ t = \tau t = τ β = ∞ \beta= \infty β = ∞ β ( t ) \beta(t) β ( t ) [62 ] and empirical deep-learning variants [81 , 38 ] . Under our approach, a rigorous treatment of general annealing schedules would first require extending Lemma 2.2 to time-dependent potentials, with their family of instantaneous Gibbs stationary measures; we leave this to future work.
Experiments
Our motivation for performing experiments is twofold. First, we wish to investigate a matter that our theory does not fully answer: whether the dynamics fall within the lazy regime under standard parameterization. While our analysis does not rely on proximity to initialization, which is often associated with lazy training, it does not rule out lazy behavior either. Second, we wish to empirically validate aspects of our theory, such as the obtained loss floor and the preservation of a lower bound of λ min ( K N T K ) \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}) λ m i n ( K NTK ) β \beta β β = ∞ \beta= \infty β = ∞ § 18 for full details.
Figure 1 a shows Jacobian drift during training, defined as the cosine distance between the network Jacobian (evaluated for all datapoints and concatenated) at time t t t § 18 .§ 2 , we provide additional evidence that the noisy neural-network dynamics are not fully captured by the frozen-Jacobian linearization around initialization.
Figure 1. Training statistics as a function of the number of updates for β ∈ { 10 n , 50 n , 100 n , ∞ } \beta\in\{10n, 50n, 100n, \infty\} β ∈ { 10 n , 50 n , 100 n , ∞ } ( β , seed ) (\beta, \text{seed}) ( β , seed ) β = ∞ \beta= \infty β = ∞
Figure 1 b shows training-loss plateaus consistent with our theory, scaling proportionally to β − 1 \beta^{-1} β − 1 Figure 1 c shows the minimum eigenvalue of the feature Gram matrix ϕ ( W X ) ϕ ( W X ) T \phi(WX)\phi(WX)^{\mathsf{T}} ϕ ( W X ) ϕ ( W X ) T λ 1 \lambda_1 λ 1 β \beta β β \beta β
Discussion
Convergence of Langevin dynamics. A common route to optimization guarantees for Langevin dynamics first proves convergence to the Gibbs law and then converts this distributional convergence into optimization bounds. Non-asymptotic analyses typically rely on functional inequalities for the Gibbs law, such as the log-Sobolev or Poincaré inequalities, whose constants control the rate [110 , 28 , 80 ] . A prominent work of this kind is Raginsky et al. [96 ] , where indeed the convergence rate is governed by the log-Sobolev constant associated with the Gibbs measure. As the authors note, this constant is bounded by the spectral gap and, unavoidably, worsens exponentially as the inverse temperature β \beta β § 9 .§ 1 and Table 3 for an extended discussion.
Table 3. Convergence analyses of noisy training dynamics.
Comparison to the generalized NTK analysis of Chen et al. [23 ] . Chen et al. [23 ] study two-layer networks through an infinite-width mean-field description and show that, under a sufficiently large output scaling, noisy gradient descent with weight decay can still exhibit kernel-like convergence. Their analysis replaces the usual Euclidean closeness-to-initialization argument by a distributional one: the parameter law remains close to the initialization law, in Wasserstein/KL-type distances, which keeps the minimum eigenvalue of the NTK bounded away from zero and yields linear convergence up to a noise floor. Thus, as in our work, convergence is ultimately obtained from a kernel lower bound beyond the conventional lazy-training argument. When our finite-width bounds are specialized to the infinite-width limit (m → ∞ m \to\infty m → ∞ [23 ] . Their analysis is tailored to this limiting setting: the deterministic NTK lower bound is imposed and maintained at the level of the limiting parameter law. In contrast, our result applies directly at finite width, where the empirical NTK is random and its initialization-level nondegeneracy must be transferred to the trained parameters with high probability.10 [23 ] , the regularization/noise parameter is part of the dynamics, but not a constructive ingredient in preserving the kernel lower bound; rather, increasing it worsens both their convergence guarantee and their KL-divergence control. In our analysis, increasing the Langevin noise strengthens the divergence control, as in Harel et al. [47 ] , which is precisely what propagates the initialization-level NTK lower bound along the finite-width trajectory. Together, these perspectives suggest a more complete picture of how favorable initialization properties may be maintained during training in practice.
Is a bound with β = o ( n ) \beta= o(n) β = o ( n ) A natural question is whether the analysis can be tightened. Our result requires β = Ω ( n ) \beta= \Omega(n) β = Ω ( n ) β = o ( n ) \beta= o(n) β = o ( n ) [122 ] . However, consider the generalization result of Harel et al. [47 ] : under similar assumptions, with probability at least 1 − δ 1-\delta 1 − δ S ∼ D n S \sim D^n S ∼ D n
E p t [ L D ( θ t ) − L S ( θ t ) ∣ S ] ≤ β E p 0 [ L S ( θ 0 ) ∣ S ] + ln ( n / δ ) 2 n . \mathbb{E}_{p_t}\left[\mathcal{L}_D(\theta_t)-\mathcal{L}_S(\theta_t)\mid S\right]\leq\sqrt{\frac{\beta\mathbb{E}_{p_0}\left[\mathcal{L}_S(\theta_0)\mid S\right]+\ln(n/\delta)}{2n}}. E p t [ L D ( θ t ) − L S ( θ t ) ∣ S ] ≤ 2 n β E p 0 [ L S ( θ 0 ) ∣ S ] + ln ( n / δ ) . Thus, for any β = o ( n ) \beta= o(n) β = o ( n ) β = Θ ( n 0.99 ) \beta= \Theta(n^{0.99}) β = Θ ( n 0.99 ) n → ∞ n \to\infty n → ∞ E p t [ L S ( θ t ) ] \mathbb{E}_{p_t}[\mathcal{L}_S(\theta_t)] E p t [ L S ( θ t )] [55 ] , or by including a corrective term, as in Arora et al. [6 ] .
Neural network parameterizations. NN parameterizations, i.e. , the scaling of weights, outputs, and learning rate, largely determine the limiting training dynamics. Under NTK and standard parameterizations, infinite-width training with GF is well approximated by kernel gradient descent with a fixed kernel, suppressing feature learning. Mean-field (MF) parameterization instead captures genuine feature evolution through the evolution of the empirical distribution of neurons [76 , 31 , 103 , 100 ] . However, finite-width MF guarantees remain limited and typically rely on restrictive assumptions [54 , 73 , 37 ] ; moreover, MF dynamics are largely limited to two or three layers, as they may freeze at greater depth [82 ] . Thus, despite its ability to model feature learning, it remains unclear whether MF analysis explains the practical success of deep networks, whose default implementation often follows the standard parameterization, as in PyTorch [92 ] . We were also unable to extend our proof mechanism to the MF parameterization; see App. § 9 .§ 2 for further discussion.
Why study CLD? We study CLD as a tractable model for the role of stochasticity in neural-network training. Although CLD differs from common training methods such as SGD, it is the continuous-time limit of other discrete-time algorithms, such as stochastic gradient Langevin dynamics (SGLD) that is motivated both empirically and theoretically [62 , 81 , 38 , 56 , 57 , 114 , 123 ] . Technically, CLD admits a Gibbs stationary distribution, enabling our change-of-measure argument. Related Gibbs-like stationary behavior and SDE approximations have been extensively studied for SGD [74 , 118 , 8 , 102 , 69 , 70 , 51 , 26 , 65 , 71 ] . Thus, CLD is not meant to reproduce SGD exactly, but to isolate the effect of noise in a mathematically tractable and practically motivated setting, with extensions to anisotropic state-dependent noise left for future work. We elaborate further in App. § 9 .§ 3 .
Future work. Our work could be extended in several natural directions to better capture practical training, for instance by incorporating noise annealing or state-dependent noise, such as noise induced by the SGD covariance, as discussed above, and by deriving analogous bounds for discretized Langevin dynamics. Another direction is to investigate how noise in the training process may promote feature learning, using our method. More broadly, our change-of-measure lemma may apply to other settings requiring rare-event control or finite-time analysis when convergence times are unavailable.
Acknowledgments and Disclosure of Funding
The research of DS was Funded by the European Union (ERC, A-B-C-Deep, 101039436). Views and opinions expressed are however those of the author only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency (ERCEA). Neither the European Union nor the granting authority can be held responsible for them. The research of DS was also supported by the Council for Higher Education in Israel under the Moonshot Project. NS was partially supported by the NSF TRIPOD Institute on Data Economics Algorithms and Learning (IDEAL) and an NSF-IIS award.
Additional related work
Existing work on the convergence of Langevin dynamics
Langevin diffusion has been used for global optimization for decades [29 ] , alongside its classical role in sampling from the Gibbs distribution [12 , 99 ] .
Works that study the convergence rate of Langevin dynamics from the sampling perspective typically focus on establishing distributional convergence to the Gibbs posterior, i.e. , convergence of a divergence measure such as the KL divergence or Rényi divergence to the Gibbs law. Representative examples include Vempala and Wibisono [110 ] and subsequent works [28 , 80 ] ; these works assume that the Gibbs distribution satisfies a suitable functional inequality, such as a log-Sobolev or Poincaré inequality. These inequalities rely on unspecified constants, which can be arbitrarily bad.
Many works on convergence rates of Langevin dynamics for optimization proceed through the same route, namely by converting distributional convergence to the Gibbs distribution into optimization guarantees, driven by the fact that the Gibbs distribution concentrates around minimizers. A prominent example is Raginsky et al. [96 ] , where the authors decompose the generalization error under SGLD and bound its constituent terms: the error controlled by convergence of SGLD toward the Gibbs distribution, the generalization gap of the Gibbs distribution itself, and the empirical loss under the Gibbs distribution. The resulting convergence rate depends on the log-Sobolev constant of the Gibbs distribution, which, as the authors discuss, is controlled by the spectral gap and necessarily deteriorates exponentially with the inverse temperature β \beta β 11 N p log ( β + 1 ) β \frac{N_p \log(\beta+1)}{\beta} β N p l o g ( β + 1 ) N p N_p N p N p > β N_p > \beta N p > β [117 ] , Kumar et al. [64 ] , employ similar analyses and suffer from the same limitations, applying only in the low-noise regime with a convergence rate that is exponentially slow. A conceptually different recent contribution is Chen et al. [19 ] , who replace the standard sampling-to-optimization argument with a direct Lyapunov/hitting-time analysis of Langevin dynamics; nevertheless, their main quantitative guarantees still rely on a Poincaré inequality with an unspecified constant for the Gibbs measure.
Neural-network-specific works. More recent works on the convergence of Langevin dynamics that focus specifically on neural network training still mostly rely on the Gibbs functional-inequality route, or are confined to special regimes such as lazy training. Lugosi and Nualart [72 ] study a Langevin-type continuous-time approximation of SGD with anisotropic noise covariance matching that of the stochastic gradient. However, their analysis requires the parameters to remain close to initialization, preventing its extension to settings in which parameters drift away from initialization. Their work extends the noiseless analysis of Chatterjee [18 ] , assumes a non-standard initialization scheme, and importantly yields its strongest bounds in the zero-loss regime, thereby treating the noise essentially as a nuisance. Oberweis and Cayci [89 ] analyze similar noise dynamics to those of Lugosi and Nualart [72 ] , but their analysis is explicitly confined to the lazy-training regime induced by output scaling, similarly to Chizat et al. [32 ] .
Another major line of research is mean-field Langevin dynamics (MFLD), where convergence is studied in an infinite-width lifted measure-space regime. Quantitative exponential-rate results in this line commonly rely on uniform log-Sobolev-type assumptions [e.g., 30, 86, 108, 20, 88]; see Section A.2.2 for further details.
The above works help distinguish our contribution: we prove empirical-loss convergence directly for finite-width neural networks. Our proof uses noise as a crucial mechanism, together with a network-specific NTK lower bound that is preserved along Langevin dynamics, thereby avoiding the derivation of optimization guarantees from global functional inequalities of the Gibbs measure, and their associated unspecified constants. Table 3 places the analyses discussed here alongside ours, recording for each the width required, whether the parameters may leave initialization, the objective that is contracted, the rate, and the loss floor.
L 2 L_2 L 2 Weight decay, i.e., L 2 L_2 L 2 L → ∞ \mathcal{L} \to\infty L → ∞ ∥ θ ∥ → ∞ \lVert\theta\rVert\to\infty ∥ θ ∥ → ∞ e − β L e^{-\beta\mathcal{L}} e − β L L \mathcal{L} L 12 [96 ] , L 2 L_2 L 2 dissipativity of the loss, which is leveraged to derive a log-Sobolev inequality for the Gibbs distribution. Dissipativity is likewise assumed in later analyses of Langevin dynamics [e.g., 117, 19], and L 2 L_2 L 2 [64 ] , as well as in the MFLD literature [86 , 30 , 108 , 88 ] .
In our setting, we introduce L 2 L_2 L 2 ν = p 0 \nu= p_0 ν = p 0 Lemma 2.2 . Other regularization strengths are also possible, at the price of an additional divergence term in the change-of-measure inequality; see Remark B.8.
Neural network parameterizations
Introduction
The problem of neural network optimization has been analyzed under several parameterizations, that is, under varying choices of neuron scaling and initialization. Importantly, these choices have a profound impact on the optimization and generalization behavior of the network, as observed in prior works such as Chizat et al. [32 ] , Woodworth et al. [115 ] , Yang and Hu [119 ] .
The parameterization most commonly used in practice—also the default in PyTorch—is the so-called standard parameterization , in which no explicit scaling is applied within the network, and the parameters are initialized by He initialization [48 , 92 ] . However, exploring alternative parameterizations, and particularly different scalings with respect to network width, has led to fundamentally distinct perspectives on overparameterized neural network training dynamics. In particular, two main theoretical frameworks have emerged, each corresponding to a different scaling regime: the neural tangent kernel (NTK) approach [53 , 41 , 32 ] , and the mean-field (MF) approach [76 , 31 , 103 , 100 ] .
NTK. In the infinite-width limit, the NTK analysis shows that the network training dynamics are equivalent to gradient descent in a linearized model with a fixed kernel [53 , 67 , 121 ] . In this regime, the parameters remain close to their initialization throughout training, a phenomenon known as lazy learning [32 ] . These properties extend to standard parameterization in the limit of infinite width [119 ] . Under these assumptions, one can prove linear convergence of deep networks to zero empirical loss with width scaling polynomially in the number of samples [40 , 2 , 124 , 105 ] . However, since the parameters move only slightly from their initial values, the network fails to learn new features. This significantly limits the expressive power observed in practice and may lead to weaker generalization performance [113 , 115 ] .
Mean-field. In contrast, under mean-field analysis, the network is viewed as a collection of infinite particles in the limit of infinite width. Then, the empirical distribution of neurons (particles) under gradient flow evolves in the space of measures according to gradient flow in the 2-Wasserstein metric. This PDE perspective captures feature learning dynamics, as demonstrated, for example, by Suzuki et al. [106 ] in the context of the k k k [73 ] , Glasgow et al. [43 ] for learning single-index models. However, as we elaborate in the following section, non-asymptotic global convergence analyses of mean-field parametrized networks under realistic assumptions are still an open problem.
Existing convergence results for mean-field parametrized networks
Preliminary mean-field analyses of noiseless training dynamics established asymptotic convergence to zero population loss in the infinite-width, infinite-time regime [76 , 31 , 103 , 100 ] . In contrast, existing non-asymptotic convergence results for shallow networks under the mean-field parameterization typically rely on restrictive assumptions on the data distribution, such as strong separability conditions [e.g., 22, 16, 37], or on highly simplified teacher models (e.g., concave [54 ] or quartic [73 ] objectives).
More broadly, convergence results under milder assumptions, which would accommodate more realistic learning problems, remain largely open. A major obstacle is controlling the finite-width deviation from the infinite-particle PDE limit, termed the propagation of chaos (PoC) error, where “chaos” refers to the asymptotic independence of neurons [107 ] . Available bounds on this error, typically obtained via Grönwall-type arguments, grow exponentially with training time [76 , 77 ] , and thus do not suffice for problems which require longer time horizons for feature learning, such as learning multi-index models [7 , 1 , 13 ] . Recently, Glasgow et al. [43 ] derived a bound that grows polynomially with training time and thus addresses this category of problems, but requires additional stability and local strong convexity assumptions, which do not necessarily hold in neural networks. More recently, Glasgow and Bruna [42 ] established uniform-in-time weak PoC bound, relying on a convergence assumption for the mean-field loss.
Meanwhile, analyses of mean-field Langevin dynamics (MFLD), in which Gaussian noise is added to the gradient updates (as in our work), establish linear convergence of both the infinite-width PDE limit [86 , 30 ] and the finite-width network [87 , 108 ] , together with uniform-in-time propagation-of-chaos bounds [20 , 88 ] . Notably, these results rely on the associated Gibbs distribution satisfying a uniform log-Sobolev inequality (LSI). However, existing bounds on the LSI constant deteriorate exponentially with the inverse temperature (β \beta β [96 , 30 , 27 ] , leading to vacuous guarantees in the low-noise regime for the convergence rates derived in the aforementioned works. More recently, Wang and Chizat [112 ] strengthened this line of research by showing that local convergence (i.e., convergence when initialized near the limiting distribution) can instead be established under the weaker Poincaré inequality, yielding faster rates. Importantly, both the LSI and Poincaré inequalities are abstract functional inequalities imposed on the Gibbs measure, and they are not verified to hold with usable constants in neural network models. It is also worth mentioning Hu et al. [50 ] , who establish asymptotic convergence of MFLD without relying on log-Sobolev inequalities, and briefly discuss non-asymptotic convergence under stronger assumptions that are not shown to hold for neural networks in general. Moreover, the corresponding convergence rate deteriorates as the noise level decreases, with no explicit characterization of this dependence. The MFLD rows of Table 3 summarize how these guarantees compare with ours.
Discussion: connection to practical neural network parameterizations
A natural question is to what extent these approaches explain training dynamics under the parameterizations used in practice, since neither the NTK parameterization (as introduced by Jacot et al. [53 ] ) nor the mean-field parameterization (averaging over neurons) is commonly implemented.
As noted above, the standard parameterization remains the dominant choice in practice. As Yang and Hu [119 ] point out, in the infinite-width limit it leads to the kernel (lazy) regime when trained with sufficiently small step sizes, and therefore does not exhibit feature learning. Small step sizes are also assumed in several convergence analyses of standardly parameterized neural networks, which show—and rely on—the parameters remaining close to their initialization throughout training [e.g., 91, 105].
However, in practice, neural networks trained under the standard parameterization do learn features and achieve strong generalization performance [63 , 48 , 122 ] . This discrepancy is plausibly due to the fact that practical networks have finite width and are typically trained with step sizes larger than those allowed by the conservative stability conditions in the infinite-width analysis of Yang and Hu [119 ] . Indeed, larger step sizes have been shown to improve generalization capabilities [36 , 95 ] .
To reconcile stable infinite-width training with feature learning for deep networks, Yang and Hu [119 ] propose the maximal-update parameterization (μ \mu μ μ \mu μ
Failure to extend our result to mean-field parameterization
Mean-field parameterization, as well as μ \mu μ κ = 1 m \kappa= \frac{1}{\sqrt{m}} κ = m 1 L ∞ L_{\infty} L ∞
β − 2 n 2 λ 1 2 1 κ 4 . \beta^{-2}\frac{n^2}{\lambda_1^2}\frac{1}{\kappa^4}. β − 2 λ 1 2 n 2 κ 4 1 . When κ = 1 m \kappa= \frac{1}{\sqrt{m}} κ = m 1 β = Ω ( n m ) \beta= \Omega(nm) β = Ω ( nm ) m = Ω ~ ( β n ) m = \widetilde{\Omega}(\beta n) m = Ω ( β n )
It is not surprising that this obstruction arises from regularization: we use regularization to preserve favorable initialization properties, and this may mark a boundary of our method. The second-law change-of-measure argument controls training by propagating initialization-time properties of the full parameter process. This is sufficient when convergence follows from preserving a favorable initialization, but mean-field training is expected to rely on the creation of new, label-dependent features through the particle dynamics. Since our assumptions allow arbitrary labels and our proof does not track label-dependent feature evolution, the framework may be poorly suited to proving a teacher-free mean-field convergence theorem.
Why study continuous-time Langevin dynamics for training neural networks?
Our work employs CLD as a tractable framework for studying the role of stochasticity in neural network training. However, CLD differs from practical training algorithms, especially mini-batch SGD, in two central respects: (i) it is formulated in continuous rather than discrete time, and (ii) it uses additive isotropic noise on top of the full-batch gradient, rather than the anisotropic, state-dependent noise induced by mini-batch sampling. Nevertheless, there are several theoretical and practical reasons to study CLD in the context of neural network training.
First, although CLD is not the continuous-time limit of SGD, it is the continuous-time limit of algorithms such as Langevin Monte Carlo (LMC)13 [60 ] . In particular, SGLD is well motivated both theoretically and empirically. Injected noise has been shown empirically to help optimization escape poor local minima in highly non-convex landscapes, as in simulated annealing [62 ] , and to improve both optimization and generalization in neural network training [81 , 38 ] . On the theoretical side, injected noise is also known to facilitate escape from saddle points [56 , 57 ] . Beyond its theoretical appeal, SGLD is used in practice in settings where explicit stochasticity is beneficial, such as Bayesian posterior sampling [114 ] and low-precision optimization [123 ] .
Second, the main technical reason we focus on CLD is that it admits the Gibbs stationary distribution, which enables us to apply our change-of-measure inequality. Importantly, Mandt et al. [74 ] , Yaida [118 ] , Azizian et al. [8 ] show that, under certain conditions, SGD admits a stationary distribution that resembles the Gibbs distribution, while works such as Vempala and Wibisono [110 ] , Chewi et al. [28 ] , Mousavi-Hosseini et al. [80 ] bound the Rényi divergence of LMC from the Gibbs distribution under functional inequalities such as the Poincaré inequality and the log-Sobolev inequality. These results suggest that tools similar to our change-of-measure inequality may apply beyond the exact CLD setting.
Finally, analyzing SGD via stochastic differential equations (SDEs) in continuous time is a standard approach. Shi et al. [102 ] study the effects of learning rate under an isotropic Gaussian noise model, as in our setting, while Li et al. [69 , 70 ] , Hu et al. [51 ] , Cheng et al. [26 ] , Lanconelli and Lauria [65 ] rigorously justify approximating SGD in the small-step regime by an SDE with anisotropic, state-dependent Gaussian noise. Li et al. [71 ] further support this perspective, both theoretically and empirically, by showing that such SDEs can accurately capture the behavior of practical SGD even at finite learning rates, provided that the equilibrium gradient noise is sufficiently large relative to the norm of the gradient. Extending our analysis to more realistic anisotropic noise models is a natural direction for future work.
Overall, the point of studying CLD is not that it reproduces SGD in every detail, but that it isolates the role of stochasticity in a mathematically tractable and still practically relevant setting. Moreover, the role of noise in our proof highlights its theoretical—and potentially practical—utility for global convergence.
Existing lower bounds on the NTK minimum eigenvalue at initialization
The eigenvalues of the empirical NTK matrix at initialization are well-studied [91 , 84 , 105 , 79 , 14 , 9 , 59 ] . Prior work varies substantially in the scope of their guarantees: they consider different assumptions on the initialization, output scaling, activations and data. Additionally, some are restricted to shallow networks or to polynomially-decaying failure probability. We first provide a detailed survey of prior work, and then discuss their relevance to our setting.
Oymak and Soltanolkotabi [91 ] consider shallow networks with Θ ( 1 ) \Theta(1) Θ ( 1 ) Θ ( m ) \Theta(\sqrt{m}) Θ ( m ) λ min ( K N T K ) = Ω ( m ) \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}) = \Omega(m) λ m i n ( K NTK ) = Ω ( m ) m = Ω ( n log 2 n ) m = \Omega\left(n \log^{2} n\right) m = Ω ( n log 2 n ) 1 / n 1/n 1/ n [79 ] consider a linearized model with output scaling of 1 / m 1/\sqrt{m} 1/ m Ω ( 1 ) \Omega(1) Ω ( 1 ) m d log ( m d ) > n md \log(md) > n m d log ( m d ) > n n = poly ( d ) n = \operatorname{poly}(d) n = poly ( d ) [84 ] consider deep networks with ReLU activations under any scaling and initialization variance. For standard parameterization, they show that λ min ( K N T K ) = Θ ( m + d ) \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}) = \Theta(m+d) λ m i n ( K NTK ) = Θ ( m + d ) m = Ω ( n log n ) m = \Omega(n \log n) m = Ω ( n log n ) [105 ] consider shallow networks with standard parameterization and bounded activations. They show a bound of Ω ( 1 ) \Omega(1) Ω ( 1 ) m = Ω ( d − 1 n 5 / 3 ) m = \Omega\left(d^{-1}n^{5/3}\right) m = Ω ( d − 1 n 5/3 ) [14 ] consider deep networks where the last layer has initialization variance of Θ ( 1 ) \Theta(1) Θ ( 1 ) [85 ] . Denoting the widths of the two last hidden layers as m L − 1 m_{L-1} m L − 1 m L − 2 m_{L-2} m L − 2 λ min ( K N T K ) = Ω ( m L − 1 m L − 2 ) \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}) = \Omega(m_{L-1}m_{L-2}) λ m i n ( K NTK ) = Ω ( m L − 1 m L − 2 ) m L − 1 m L − 2 = ω ( n log 8 n ) m_{L-1}m_{L-2} = \omega\left(n \log^{8} n\right) m L − 1 m L − 2 = ω ( n log 8 n ) m L − 1 = Ω ( n γ ) m_{L-1} = \Omega\left(n^{\gamma}\right) m L − 1 = Ω ( n γ ) γ > 0 \gamma> 0 γ > 0 exp ( − log 2 m ) \exp\left(-\log^{2}m\right) exp ( − log 2 m ) [9 ] consider deep networks with NTK parameterization, and show that λ min ( K N T K ) = Ω ( 1 ) \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}) = \Omega(1) λ m i n ( K NTK ) = Ω ( 1 ) m = Ω ( n log 2 n ) m = \Omega\left(n \log^{2} n\right) m = Ω ( n log 2 n ) 1 / m 1/m 1/ m [59 ] consider both shallow and deep networks with NTK parameterization and ReLU activations. They assume the data exhibits δ \delta δ λ min ( K N T K ) = Ω ( d − 3 δ 2 ) \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}) = \Omega\left(d^{-3}\delta^{2}\right) λ m i n ( K NTK ) = Ω ( d − 3 δ 2 ) m = Ω ( ∥ X ∥ 2 d 3 δ − 2 log n ) m = \Omega\left(\lVert X\rVert^{2}d^{3}\delta^{-2}\log n\right) m = Ω ( ∥ X ∥ 2 d 3 δ − 2 log n )
None of these results directly apply to our setting. The works of Oymak and Soltanolkotabi [91 ] , Montanari and Zhong [79 ] , Bombari et al. [14 ] , Banerjee et al. [9 ] , Karhadkar et al. [59 ] do not match our initialization and output scaling; Nguyen et al. [84 ] , Karhadkar et al. [59 ] focus on ReLU activations; Oymak and Soltanolkotabi [91 ] , Bombari et al. [14 ] , Banerjee et al. [9 ] do not provide sufficiently strong failure probabilities for our analysis; Song et al. [105 ] use different data scaling and assumptions on the output layer. For this reason we extend the work of Banerjee et al. [9 ] to arbitrary parameterizations and strengthen its probability guarantee.
Tabular comparison with prior convergence analyses
Tables 2 and 3 place our result among existing convergence analyses for overparameterized networks. The first collects analyses of gradient descent and its stochastic variant; the second collects analyses that model training as a diffusion, which is where our result sits. Several of the works listed are discussed in prose in Sections A.1 and A.2.2; here we place them, together with further analyses, side by side. Widths and floors are stated up to constants and logarithmic factors, under each work’s own conventions and omitting dependence on depth and data separability. The output scaling κ \kappa κ κ = 1 \kappa= 1 κ = 1
Every analysis in Table 2 contracts the empirical risk to zero at a rate set by width, output scaling and data separability, with the parameters confined to a neighborhood of initialization; the rows differ only in the architecture they cover and the width they demand. The analyses in Table 3 differ from one another in more ways, beginning with what their theorems contract: the empirical risk itself, the population risk, the law of the iterates, or an entropy-regularized free energy. Every floor is nonetheless reported on the empirical risk, so that the rows remain comparable; where the theorem contracts something else, the qualification below gives the conversion. They differ also in what sets the rate. For Lugosi and Nualart [72 ] , Oberweis and Cayci [89 ] and Chen et al. [23 ] it is the familiar combination of width, scaling and separability, and the noise plays no part; for the remaining rows it is a Poincaré or log-Sobolev constant, whose only general estimates degrade like e − Θ ( β ) e^{-\Theta(\beta)} e − Θ ( β ) [30 ] assumes to do so outright. Every floor in the table grows as the noise grows, and only Chen et al. [23 ] shares the property our argument turns on, a noise level that can be reduced to shrink the floor while the rate is left untouched. They have it only at infinite width, and only once the output scaling is large enough to hold the model in its linear regime. Among the finite-width analyses in either table, ours is therefore the only one that lets the parameters leave a neighborhood of initialization, treats the injected noise as the mechanism of the proof rather than as an obstacle, and still contracts the empirical risk at a rate independent of the noise.
Table 2: Convergence analyses of gradient descent and stochastic gradient descent. In every row the empirical risk contracts to zero at a linear rate, the parameters must remain within a neighborhood of initialization, and the randomness is either absent or an obstacle the analysis must absorb; the rows differ in the architecture they cover and in the width they require. Regime: the parameterization the analysis is stated in or, where the guarantee is driven by an explicit output scaling rather than by the scale of the initialization, the scaling regime it requires. Here n n n d d d m m m κ \kappa κ Table 3 .
Table 3: Convergence analyses of noisy training dynamics. Rows are grouped by the role the randomness plays in the proof. Regime: as in Table 2 . Leaves init.: whether the parameters may travel far from their initialization (✓) or must remain close (×). Objective: the quantity the theorem itself contracts, namely the empirical risk, the population risk, the law of the iterates towards a Gibbs measure, or the entropy-regularized free energy. Floor: the non-vanishing residual left in the empirical risk; where the objective is something else, the qualifications give the conversion. Rate: linear when the objective contracts at an exponential rate set by the data and the architecture; exponentially slow in β \beta β e − Θ ( β ) e^{-\Theta(\beta)} e − Θ ( β ) N p N_p N p β \beta β Table 2 . Rows marked † carry qualifications set out below.
Qualifications. The rows marked † carry conditions that a table cell cannot hold.
Montanari and Zhong [79 ] require the output scaling κ ≳ n / m d \kappa\gtrsim n/\sqrt{md} κ ≳ n / m d d ≤ n ≤ poly ( d ) d \le n \le\operatorname{poly}(d) d ≤ n ≤ poly ( d )
Lugosi and Nualart [72 ] study the population risk, under the diffusion driven by the covariance of the stochastic gradients rather than by injected isotropic noise. Their general theorem assumes a local gradient-dominance condition on ∥ ∇ f ∥ 2 / f \lVert\nabla f\rVert^{2}/f ∥ ∇ f ∥ 2 / f n n n
Oberweis and Cayci [89 ] keep the kernel non-degenerate throughout training at finite width, but only inside the lazy regime induced by the output scaling. Theirs is the one floor in the table not governed by the noise: their bound holds for every q > 1 q>1 q > 1 κ − 1 \kappa^{-1} κ − 1 κ → ∞ \kappa\to\infty κ → ∞ n n n
Chen et al. [23 ] work at infinite width, and assume that the Gram matrix of the initial law satisfies λ min ( H ( p 0 ) ) = Λ > 0 \lambda_{\min}(H(p_{0}))=\Lambda>0 λ m i n ( H ( p 0 )) = Λ > 0 κ = 1 \kappa=1 κ = 1 κ = m \kappa=\sqrt{m} κ = m κ \kappa κ β − 1 \beta^{-1} β − 1
L ( p t ) ≤ 2 e − 2 κ 2 Λ t / n + O ( d 2 n 2 β 2 κ 2 Λ 2 ) . L(p_t) \le2e^{-2\kappa^{2}\Lambda t/n}+O\left(\frac{d^{2}n^{2}}{\beta^{2}\kappa^{2}\Lambda^{2}}\right). L ( p t ) ≤ 2 e − 2 κ 2 Λ t / n + O ( β 2 κ 2 Λ 2 d 2 n 2 ) . At any fixed admissible κ \kappa κ κ \kappa κ β − 1 \beta^{-1} β − 1
Raginsky et al. [96 ] reach the empirical risk in two steps: their Proposition 3.3 bounds the Wasserstein distance between the law of the iterates and the Gibbs measure π z ∝ e − β F z \pi_z \propto e^{-\beta F_z} π z ∝ e − β F z e − κ n / ( β c L S ) e^{-\kappa n/(\beta c_{\mathrm{LS}})} e − κn / ( β c LS )
∫ F z d π z − min w F z ( w ) ≤ N p 2 β log ( e M m ( b β N p + 1 ) ) , \int F_z\,\mathrm{d}\pi_z-\min_w F_z(w)\le\frac{N_p}{2\beta}\log\left(\frac{eM}{m}\left(\frac{b\beta}{N_p}+1\right)\right), ∫ F z d π z − w min F z ( w ) ≤ 2 β N p log ( m e M ( N p b β + 1 ) ) , which is the floor listed. Driving it to o ( 1 ) o(1) o ( 1 ) β \beta β c L S c_{\mathrm{LS}} c LS β \beta β
Kumar et al. [64 ] prove convergence of the law of the iterates to the Gibbs measure μ s ∝ e − 2 L / s \mu_s \propto e^{-2\mathcal{L}/s} μ s ∝ e − 2 L / s empirical loss, in total variation and in q q q [96 ] quoted there for the step that makes the Gibbs measure an approximate empirical risk minimizer; the floor listed is that step’s, and it is sketched rather than stated as a theorem. Their rate is governed by the Poincaré constant of μ s \mu_s μ s λ ∥ W ∗ ∥ F 2 \lambda\lVert W^{*}\rVert_{F}^{2} λ ∥ W ∗ ∥ F 2
Nitanda et al. [86 ] and Chizat [30 ] contract the free energy F β = L + β − 1 Ent F_{\beta}=L+\beta^{-1}\operatorname{Ent} F β = L + β − 1 Ent F β ( μ t ) − min F β ≤ e − 2 ρ β t / β ( F β ( μ 0 ) − min F β ) F_{\beta}(\mu_t)-\min F_{\beta}\le e^{-2\rho_{\beta}t/\beta}(F_{\beta}(\mu_0)-\min F_{\beta}) F β ( μ t ) − min F β ≤ e − 2 ρ β t / β ( F β ( μ 0 ) − min F β ) ρ β \rho_{\beta} ρ β ρ β ≥ C 0 e − α ∗ β \rho_{\beta}\ge C_0e^{-\alpha^*\beta} ρ β ≥ C 0 e − α ∗ β β − 1 KL ( μ ∘ ∥ μ 0 ) \beta^{-1}\operatorname{KL}(\mu^\circ\Vert\mu_0) β − 1 KL ( μ ∘ ∥ μ 0 ) μ ∘ \mu^\circ μ ∘ L L L β t − 1 = α / log t \beta_t^{-1}=\alpha/\log t β t − 1 = α / log t log log t / log t \log\log t/\log t log log t / log t [88 ] add a uniform-in-time particle term of order 1 / m 1/m 1/ m
Our own floor, L ∞ L_{\infty} L ∞ m = Ω ~ ( β n / λ 1 ) m=\tilde{\Omega}(\beta n/\lambda_1) m = Ω ~ ( β n / λ 1 ) λ 1 \lambda_1 λ 1 β \beta β L ∞ L_{\infty} L ∞ β = Ω ~ ( n ) \beta=\tilde{\Omega}(n) β = Ω ~ ( n ) β = Ω ~ ( n 2 ) \beta=\tilde{\Omega}(n^2) β = Ω ~ ( n 2 ) Ω ~ ( n 2 ) \tilde{\Omega}(n^2) Ω ~ ( n 2 ) Ω ~ ( n 3 ) \tilde{\Omega}(n^3) Ω ~ ( n 3 ) n n n β = Ω ~ ( n 2 ) \beta=\tilde{\Omega}(n^2) β = Ω ~ ( n 2 ) Ω ~ ( n 3 ) \tilde{\Omega}(n^3) Ω ~ ( n 3 )
Change-of-measure inequality derivation
Throughout this section, we work in a density setting. Namely, all probability measures on the relevant spaces are assumed to admit densities with respect to fixed underlying reference measures, and all Markov transition kernels are assumed to admit corresponding transition densities. We use the same notation for a probability measure and its density. On product spaces, joint densities and marginal densities are taken with respect to the corresponding product and marginal reference measures. All ratios of densities are understood on the sets where the denominator is positive; their values elsewhere are irrelevant for the divergence quantities considered below.
Divergence upper bound
Second law of thermodynamics
Definition B.1 (Rényi divergence [98 , 109 ] ). Let p p p q q q X \mathcal{X} X α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ] \alpha\in(0,1) \cup(1,\infty] α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ] Rényi divergence of order α \alpha α p p p q q q p ≪ q p \ll q p ≪ q
D α ( p ∥ q ) ≜ 1 α − 1 ln ∫ ( d p d q ) α − 1 d p , α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ) ; D ∞ ( p ∥ q ) ≜ ln ess sup x ∼ p d p d q ( x ) . D_{\alpha}(p\|q) \triangleq\frac{1}{\alpha-1}\ln\int\left(\frac{\mathrm{d}p}{\mathrm{d}q}\right)^{\alpha-1}\mathrm{d}p,\quad\alpha\in(0,1)\cup(1,\infty);\qquad D_{\infty}(p\|q)\triangleq\ln\operatorname*{ess\,sup}_{x\sim p}\frac{\mathrm{d}p}{\mathrm{d}q}(x). D α ( p ∥ q ) ≜ α − 1 1 ln ∫ ( d q d p ) α − 1 d p , α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ) ; D ∞ ( p ∥ q ) ≜ ln x ∼ p ess sup d q d p ( x ) . If p ≪̸ q p \not\ll q p ≪ q D α ( p ∥ q ) ≜ ∞ D_{\alpha}(p\|q) \triangleq\infty D α ( p ∥ q ) ≜ ∞ α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ] \alpha\in(0,1) \cup(1,\infty] α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ]
We begin by showing a data processing inequality for the Rényi divergence, similar to Theorem 9 in Van Erven and Harremos [109 ] , and provide a proof for completeness and to fit our notation.
Proposition B.2 (Data processing inequality). Let p ( x , y ) p(x,y) p ( x , y ) q ( x , y ) q(x,y) q ( x , y ) X × Y \mathcal{X}\times\mathcal{Y} X × Y p X ( x ) p_X(x) p X ( x ) q X ( x ) q_X(x) q X ( x )
p X ( x ) = ∫ Y p ( x , y ) d y . p_X(x)=\int_{\mathcal{Y}}p(x,y)\,\mathrm{d}y. p X ( x ) = ∫ Y p ( x , y ) d y . We further denote by p ( y ∣ x ) p(y\mid x) p ( y ∣ x ) q ( y ∣ x ) q(y\mid x) q ( y ∣ x ) p ( x , y ) = p ( y ∣ x ) p X ( x ) p(x,y)=p(y\mid x)p_X(x) p ( x , y ) = p ( y ∣ x ) p X ( x ) q q q α > 1 \alpha>1 α > 1
D α ( p X ∥ q X ) ≤ D α ( p ∥ q ) . D_{\alpha}(p_X\|q_X)\le D_{\alpha}(p\|q). D α ( p X ∥ q X ) ≤ D α ( p ∥ q ) . Proof. Directly from Definition B.1, we have
D α ( p ∥ q ) = 1 α − 1 ln ( ∫ X × Y ( p ( x , y ) q ( x , y ) ) α − 1 d p ( x , y ) ) = 1 α − 1 ln ( ∫ X × Y ( p ( x , y ) q ( x , y ) ) α − 1 p ( x , y ) d x d y ) = 1 α − 1 ln ( ∫ X × Y ( p ( y ∣ x ) p X ( x ) q ( y ∣ x ) q X ( x ) ) α − 1 p ( y ∣ x ) p X ( x ) d x d y ) = 1 α − 1 ln ( ∫ X × Y ( p X ( x ) q X ( x ) ) α − 1 p X ( x ) ( p ( y ∣ x ) q ( y ∣ x ) ) α − 1 p ( y ∣ x ) d x d y ) = [Fubini-Tonelli] = 1 α − 1 ln ( ∫ X ( p X ( x ) q X ( x ) ) α − 1 p X ( x ) ( ∫ Y ( p ( y ∣ x ) q ( y ∣ x ) ) α − 1 p ( y ∣ x ) d y ) d x ) = 1 α − 1 ln ( ∫ X ( p X ( x ) q X ( x ) ) α − 1 p X ( x ) ⋅ exp ( ( α − 1 ) 1 α − 1 ln ( ∫ Y ( p ( y ∣ x ) q ( y ∣ x ) ) α − 1 d p ( y ∣ x ) ) ) d x ) . = 1 α − 1 ln ( ∫ X ( p X ( x ) q X ( x ) ) α − 1 p X ( x ) exp ( ( α − 1 ) D α ( p ( ⋅ ∣ x ) ∥ q ( ⋅ ∣ x ) ) ) d x ) \begin{aligned}
D_{\alpha}(p\|q)
&=\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}\times\mathcal{Y}}\left(\frac{p(x,y)}{q(x,y)}\right)^{\alpha-1}\mathrm{d}p(x,y)\right)\\
&=\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}\times\mathcal{Y}}\left(\frac{p(x,y)}{q(x,y)}\right)^{\alpha-1}p(x,y)\,\mathrm{d}x\mathrm{d}y\right)\\
&=\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}\times\mathcal{Y}}\left(\frac{p(y\mid x)p_X(x)}{q(y\mid x)q_X(x)}\right)^{\alpha-1}p(y\mid x)p_X(x)\,\mathrm{d}x\mathrm{d}y\right)\\
&=\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}\times\mathcal{Y}}\left(\frac{p_X(x)}{q_X(x)}\right)^{\alpha-1}p_X(x)\left(\frac{p(y\mid x)}{q(y\mid x)}\right)^{\alpha-1}p(y\mid x)\,\mathrm{d}x\mathrm{d}y\right)\\
&\mathrel{\phantom{=}} \text{[Fubini-Tonelli]}=\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}}\left(\frac{p_X(x)}{q_X(x)}\right)^{\alpha-1}p_X(x)\left(\int_{\mathcal{Y}}\left(\frac{p(y\mid x)}{q(y\mid x)}\right)^{\alpha-1}p(y\mid x)\,\mathrm{d}y\right)\mathrm{d}x\right)\\
&=\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}}\left(\frac{p_X(x)}{q_X(x)}\right)^{\alpha-1}p_X(x)\cdot\exp\left((\alpha-1)\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{Y}}\left(\frac{p(y\mid x)}{q(y\mid x)}\right)^{\alpha-1}\mathrm{d}p(y\mid x)\right)\right)\mathrm{d}x\right). \\
&= \frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}}\left(\frac{p_X(x)}{q_X(x)}\right)^{\alpha-1}p_X(x)\exp\left((\alpha-1)D_\alpha\left(p(\cdot\mid x)\mathbin\Vert q(\cdot\mid x)\right)\right)\mathrm{d}x\right)
\end{aligned} D α ( p ∥ q ) = α − 1 1 ln ( ∫ X × Y ( q ( x , y ) p ( x , y ) ) α − 1 d p ( x , y ) ) = α − 1 1 ln ( ∫ X × Y ( q ( x , y ) p ( x , y ) ) α − 1 p ( x , y ) d x d y ) = α − 1 1 ln ( ∫ X × Y ( q ( y ∣ x ) q X ( x ) p ( y ∣ x ) p X ( x ) ) α − 1 p ( y ∣ x ) p X ( x ) d x d y ) = α − 1 1 ln ( ∫ X × Y ( q X ( x ) p X ( x ) ) α − 1 p X ( x ) ( q ( y ∣ x ) p ( y ∣ x ) ) α − 1 p ( y ∣ x ) d x d y ) = [Fubini-Tonelli] = α − 1 1 ln ( ∫ X ( q X ( x ) p X ( x ) ) α − 1 p X ( x ) ( ∫ Y ( q ( y ∣ x ) p ( y ∣ x ) ) α − 1 p ( y ∣ x ) d y ) d x ) = α − 1 1 ln ( ∫ X ( q X ( x ) p X ( x ) ) α − 1 p X ( x ) ⋅ exp ( ( α − 1 ) α − 1 1 ln ( ∫ Y ( q ( y ∣ x ) p ( y ∣ x ) ) α − 1 d p ( y ∣ x ) ) ) d x ) . = α − 1 1 ln ( ∫ X ( q X ( x ) p X ( x ) ) α − 1 p X ( x ) exp ( ( α − 1 ) D α ( p ( ⋅ ∣ x ) ∥ q ( ⋅ ∣ x ) ) ) d x ) [ D α ≥ 0 , ∗ ] ≥ 1 α − 1 ln ( ∫ X ( p X ( x ) q X ( x ) ) α − 1 p X ( x ) exp ( 0 ) d x ) = 1 α − 1 ln ( ∫ X ( p X ( x ) q X ( x ) ) α − 1 d p X ( x ) ) = D α ( p X ∥ q X ) , \begin{aligned}
\left[D_\alpha\ge0,\ast\right]\ge\frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}}\left(\frac{p_X(x)}{q_X(x)}\right)^{\alpha-1}p_X(x)\exp(0)\mathrm{d}x\right) \\
&= \frac{1}{\alpha-1}\ln\left(\int_{\mathcal{X}}\left(\frac{p_X(x)}{q_X(x)}\right)^{\alpha-1}\mathrm{d}p_X(x)\right)=D_\alpha(p_X\mathbin\Vert q_X),
\end{aligned} [ D α ≥ 0 , ∗ ] ≥ α − 1 1 ln ( ∫ X ( q X ( x ) p X ( x ) ) α − 1 p X ( x ) exp ( 0 ) d x ) = α − 1 1 ln ( ∫ X ( q X ( x ) p X ( x ) ) α − 1 d p X ( x ) ) = D α ( p X ∥ q X ) , where the non-negativity of the Rényi divergence (*) was shown in Theorem 8 of Van Erven and Harremos [109 ] , and the sign is maintained since α > 1 \alpha>1 α > 1
From here, the proof of a form of the second law of thermodynamics for the Rényi divergence in time-invariant Markov processes follows along the lines of the proof of Lemma A.2 in Harel et al. [47 ] for the second law of thermodynamics for the KL divergence of discrete-time Markov processes, and then their Corollary A.4, which extends the result to continuous-time Markov processes.
Proposition B.3 (Second law of thermodynamics for discrete-time Markov processes). Given a time-invariant Markov process on state space X \mathcal{X} X p n p_n p n n ∈ N 0 n\in\mathbb{N}_0 n ∈ N 0 p ∞ p_\infty p ∞ n ∈ N 0 n\in\mathbb{N}_0 n ∈ N 0 α > 1 \alpha>1 α > 1
D α ( p n + 1 ∥ p ∞ ) ≤ D α ( p n ∥ p ∞ ) . D_\alpha(p_{n+1}\mathbin\Vert p_\infty)\le D_\alpha(p_n\mathbin\Vert p_\infty). D α ( p n + 1 ∥ p ∞ ) ≤ D α ( p n ∥ p ∞ ) . Proof. Fix any n ∈ N 0 n\in\mathbb{N}_0 n ∈ N 0 p n ≪̸ p ∞ p_n\not\ll p_\infty p n ≪ p ∞
D α ( p n ∥ p ∞ ) = ∞ , D_\alpha(p_n\mathbin\Vert p_\infty)=\infty, D α ( p n ∥ p ∞ ) = ∞ , and the claim is trivial. Hence assume p n ≪ p ∞ p_n\ll p_\infty p n ≪ p ∞
Let K ( ⋅ ∣ x ) K(\cdot\mid x) K ( ⋅ ∣ x ) X × X \mathcal{X}\times\mathcal{X} X × X
p ( x , y ) = p n ( x ) K ( y ∣ x ) , q ( x , y ) = p ∞ ( x ) K ( y ∣ x ) . p(x,y)=p_n(x)K(y\mid x),\qquad q(x,y)=p_\infty(x)K(y\mid x). p ( x , y ) = p n ( x ) K ( y ∣ x ) , q ( x , y ) = p ∞ ( x ) K ( y ∣ x ) . The Y Y Y
p Y ( y ) = ∫ p ( x , y ) d x = ∫ p n ( x ) K ( y ∣ x ) d x = p n + 1 ( y ) , p_Y(y)=\int p(x,y)\,\mathrm{d}x=\int p_n(x)K(y\mid x)\,\mathrm{d}x=p_{n+1}(y), p Y ( y ) = ∫ p ( x , y ) d x = ∫ p n ( x ) K ( y ∣ x ) d x = p n + 1 ( y ) , q Y ( y ) = ∫ p ∞ ( x ) K ( y ∣ x ) d x = p ∞ ( y ) . q_Y(y)=\int p_\infty(x)K(y\mid x)\,\mathrm{d}x=p_\infty(y). q Y ( y ) = ∫ p ∞ ( x ) K ( y ∣ x ) d x = p ∞ ( y ) . The identity p Y = p n + 1 p_Y=p_{n+1} p Y = p n + 1 q Y = p ∞ q_Y=p_\infty q Y = p ∞ p ∞ p_\infty p ∞
D α ( p n + 1 ∥ p ∞ ) ≤ D α ( p ∥ q ) . D_\alpha(p_{n+1}\mathbin\Vert p_\infty)\le D_\alpha(p\mathbin\Vert q). D α ( p n + 1 ∥ p ∞ ) ≤ D α ( p ∥ q ) . From Definition B.1 we have,
D α ( p ∥ q ) = 1 α − 1 ln ∫ ( p ( x , y ) q ( x , y ) ) α − 1 d p ( x , y ) = 1 α − 1 ln ∬ ( p n ( x ) K ( y ∣ x ) p ∞ ( x ) K ( y ∣ x ) ) α − 1 p ( x , y ) d x d y = 1 α − 1 ln ∬ ( p n ( x ) p ∞ ( x ) ) α − 1 p n ( x ) K ( y ∣ x ) d x d y \begin{aligned}
D_\alpha(p\mathbin\Vert q)=\frac{1}{\alpha-1}\ln\int\left(\frac{p(x,y)}{q(x,y)}\right)^{\alpha-1}\mathrm{d}p(x,y) \\
&=\frac{1}{\alpha-1}\ln\iint\left(\frac{p_n(x)K(y\mid x)}{p_\infty(x)K(y\mid x)}\right)^{\alpha-1}p(x,y)\,\mathrm{d}x\mathrm{d}y \\
&=\frac{1}{\alpha-1}\ln\iint\left(\frac{p_n(x)}{p_\infty(x)}\right)^{\alpha-1}p_n(x)K(y\mid x)\,\mathrm{d}x\mathrm{d}y
\end{aligned} D α ( p ∥ q ) = α − 1 1 ln ∫ ( q ( x , y ) p ( x , y ) ) α − 1 d p ( x , y ) = α − 1 1 ln ∬ ( p ∞ ( x ) K ( y ∣ x ) p n ( x ) K ( y ∣ x ) ) α − 1 p ( x , y ) d x d y = α − 1 1 ln ∬ ( p ∞ ( x ) p n ( x ) ) α − 1 p n ( x ) K ( y ∣ x ) d x d y = [ Fubini-Tonelli ] 1 α − 1 ln ( ∫ ( p n ( x ) p ∞ ( x ) ) α − 1 p n ( x ) ( ∫ K ( y ∣ x ) d y ) d x ) = 1 α − 1 ln ∫ ( p n ( x ) p ∞ ( x ) ) α − 1 d p n ( x ) = D α ( p n ∥ p ∞ ) . \begin{aligned}
\mathop{=}^{[\text{Fubini-Tonelli}]}\frac{1}{\alpha-1}\ln\left(\int\left(\frac{p_n(x)}{p_\infty(x)}\right)^{\alpha-1}p_n(x)\left(\int K(y\mid x)\,\mathrm{d}y\right)\mathrm{d}x\right) \\
&=\frac{1}{\alpha-1}\ln\int\left(\frac{p_n(x)}{p_\infty(x)}\right)^{\alpha-1}\mathrm{d}p_n(x) \\
&=D_\alpha(p_n\mathbin\Vert p_\infty).
\end{aligned} = [ Fubini-Tonelli ] α − 1 1 ln ( ∫ ( p ∞ ( x ) p n ( x ) ) α − 1 p n ( x ) ( ∫ K ( y ∣ x ) d y ) d x ) = α − 1 1 ln ∫ ( p ∞ ( x ) p n ( x ) ) α − 1 d p n ( x ) = D α ( p n ∥ p ∞ ) . Note that the value of the ratio on the set where q ( x , y ) = 0 q(x,y)=0 q ( x , y ) = 0 q q q
Corollary B.4 (Second law of thermodynamics for continuous-time Markov processes). Given a time-invariant Markov process on state space X \mathcal{X} X p t p_t p t t ≥ 0 t \ge0 t ≥ 0 p ∞ p_\infty p ∞ t ≥ 0 t \ge0 t ≥ 0 α > 1 \alpha> 1 α > 1
D α ( p t ∥ p ∞ ) ≤ D α ( p 0 ∥ p ∞ ) . D_\alpha(p_t\|p_\infty) \le D_\alpha(p_0\|p_\infty). D α ( p t ∥ p ∞ ) ≤ D α ( p 0 ∥ p ∞ ) . Proof. The case t = 0 t = 0 t = 0 t > 0 t > 0 t > 0 p ~ n = p n t \widetilde{p}_n = p_{nt} p n = p n t ∀ n ∈ N 0 \forall n \in\mathbb{N}_0 ∀ n ∈ N 0 p ~ ∞ = p ∞ \widetilde{p}_\infty= p_\infty p ∞ = p ∞ n = 0 n = 0 n = 0 Proposition B.3 .
Divergence bound derivation
We now define a “relative” version of the Rényi divergence, similarly to Harel et al. [47 ] .
Definition B.5 (Relative Rényi divergence). Let p p p q q q μ \mu μ X \mathcal{X} X μ ≪ q \mu\ll q μ ≪ q μ ≪ p \mu\ll p μ ≪ p α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ) \alpha\in(0,1) \cup(1,\infty) α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ) relative Rényi divergence of order α \alpha α p p p q q q μ \mu μ
D α μ ( p ∥ q ) ≔ 1 α − 1 ln ∫ ( d p d q ) α − 1 d μ D^\mu_\alpha(p\|q) \coloneqq\frac{1}{\alpha-1}\ln\int\left(\frac{\mathrm{d}p}{\mathrm{d}q}\right)^{\alpha-1}\mathrm{d}\mu D α μ ( p ∥ q ) : = α − 1 1 ln ∫ ( d q d p ) α − 1 d μ when p ≪ q p \ll q p ≪ q D α μ ( p ∥ q ) = ∞ D^\mu_\alpha(p\|q) = \infty D α μ ( p ∥ q ) = ∞ 14 p p p ({xref}Definition B.1`) is then D α ( p ∥ q ) = D α p ( p ∥ q ) D_\alpha(p\|q) = D^p_\alpha(p\|q) D α ( p ∥ q ) = D α p ( p ∥ q )
We proceed by showing a Rényi-divergence analogue of Claim 2.3 of Harel et al. [47 ] .
Claim B.6. Let p p p q q q μ 1 \mu_1 μ 1 μ 2 \mu_2 μ 2 p p p q q q Ψ \Psi Ψ Definition 2.1 ), p ∼ q p \sim q p ∼ q μ 1 , μ 2 ≪ p , q \mu_1,\mu_2 \ll p,q μ 1 , μ 2 ≪ p , q α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ ) \alpha\in(0,1) \cup(1,\infty) α ∈ ( 0 , 1 ) ∪ ( 1 , ∞ )
E μ 1 [ e − ( α − 1 ) Ψ ] < ∞ , E μ 2 [ e ( α − 1 ) Ψ ] < ∞ . \mathbb{E}_{\mu_1}\left[e^{-(\alpha-1)\Psi}\right] < \infty,\qquad\mathbb{E}_{\mu_2}\left[e^{(\alpha-1)\Psi}\right] < \infty. E μ 1 [ e − ( α − 1 ) Ψ ] < ∞ , E μ 2 [ e ( α − 1 ) Ψ ] < ∞. Then,
D α μ 1 ( p ∥ q ) + D α μ 2 ( q ∥ p ) = 1 α − 1 [ ln E μ 1 [ e − ( α − 1 ) Ψ ] + ln E μ 2 [ e ( α − 1 ) Ψ ] ] . D^{\mu_1}_\alpha(p\|q) + D^{\mu_2}_\alpha(q\|p) = \frac{1}{\alpha-1}\left[\ln\mathbb{E}_{\mu_1}\left[e^{-(\alpha-1)\Psi}\right] + \ln\mathbb{E}_{\mu_2}\left[e^{(\alpha-1)\Psi}\right]\right]. D α μ 1 ( p ∥ q ) + D α μ 2 ( q ∥ p ) = α − 1 1 [ ln E μ 1 [ e − ( α − 1 ) Ψ ] + ln E μ 2 [ e ( α − 1 ) Ψ ] ] . Proof. By Definition 2.1 , d p d q = Z − 1 e − Ψ \frac{\mathrm{d}p}{\mathrm{d}q} = Z^{-1}e^{-\Psi} d q d p = Z − 1 e − Ψ Z < ∞ Z < \infty Z < ∞
D α μ 1 ( p ∥ q ) + D α μ 2 ( q ∥ p ) = 1 α − 1 ln ( ∫ ( d p d q ) α − 1 d μ 1 ) + 1 α − 1 ln ( ∫ ( d q d p ) α − 1 d μ 2 ) = 1 α − 1 [ ln ( ∫ Z − ( α − 1 ) e − ( α − 1 ) Ψ d μ 1 ) + ln ( ∫ Z ( α − 1 ) e ( α − 1 ) Ψ d μ 2 ) ] = 1 α − 1 [ ln ( Z − ( α − 1 ) ) + ln ( ∫ e − ( α − 1 ) Ψ d μ 1 ) + ln ( Z ( α − 1 ) ) + ln ( ∫ e ( α − 1 ) Ψ d μ 2 ) ] = 1 α − 1 [ ln ( ∫ e − ( α − 1 ) Ψ d μ 1 ) + ln ( ∫ e ( α − 1 ) Ψ d μ 2 ) ] = 1 α − 1 [ ln E μ 1 [ e − ( α − 1 ) Ψ ] + ln E μ 2 [ e ( α − 1 ) Ψ ] ] . \begin{aligned}
D^{\mu_1}_\alpha(p\|q) + D^{\mu_2}_\alpha(q\|p)
&= \frac{1}{\alpha-1}\ln\left(\int\left(\frac{\mathrm{d}p}{\mathrm{d}q}\right)^{\alpha-1}\mathrm{d}\mu_1\right) + \frac{1}{\alpha-1}\ln\left(\int\left(\frac{\mathrm{d}q}{\mathrm{d}p}\right)^{\alpha-1}\mathrm{d}\mu_2\right) \\
&= \frac{1}{\alpha-1}\left[\ln\left(\int Z^{-(\alpha-1)}e^{-(\alpha-1)\Psi}\mathrm{d}\mu_1\right) + \ln\left(\int Z^{(\alpha-1)}e^{(\alpha-1)\Psi}\mathrm{d}\mu_2\right)\right] \\
&= \frac{1}{\alpha-1}\left[\ln\left(Z^{-(\alpha-1)}\right) + \ln\left(\int e^{-(\alpha-1)\Psi}\mathrm{d}\mu_1\right) + \ln\left(Z^{(\alpha-1)}\right) + \ln\left(\int e^{(\alpha-1)\Psi}\mathrm{d}\mu_2\right)\right] \\
&= \frac{1}{\alpha-1}\left[\ln\left(\int e^{-(\alpha-1)\Psi}\mathrm{d}\mu_1\right) + \ln\left(\int e^{(\alpha-1)\Psi}\mathrm{d}\mu_2\right)\right] \\
&= \frac{1}{\alpha-1}\left[\ln\mathbb{E}_{\mu_1}\left[e^{-(\alpha-1)\Psi}\right] + \ln\mathbb{E}_{\mu_2}\left[e^{(\alpha-1)\Psi}\right]\right].
\end{aligned} D α μ 1 ( p ∥ q ) + D α μ 2 ( q ∥ p ) = α − 1 1 ln ( ∫ ( d q d p ) α − 1 d μ 1 ) + α − 1 1 ln ( ∫ ( d p d q ) α − 1 d μ 2 ) = α − 1 1 [ ln ( ∫ Z − ( α − 1 ) e − ( α − 1 ) Ψ d μ 1 ) + ln ( ∫ Z ( α − 1 ) e ( α − 1 ) Ψ d μ 2 ) ] = α − 1 1 [ ln ( Z − ( α − 1 ) ) + ln ( ∫ e − ( α − 1 ) Ψ d μ 1 ) + ln ( Z ( α − 1 ) ) + ln ( ∫ e ( α − 1 ) Ψ d μ 2 ) ] = α − 1 1 [ ln ( ∫ e − ( α − 1 ) Ψ d μ 1 ) + ln ( ∫ e ( α − 1 ) Ψ d μ 2 ) ] = α − 1 1 [ ln E μ 1 [ e − ( α − 1 ) Ψ ] + ln E μ 2 [ e ( α − 1 ) Ψ ] ] . We now obtain an upper bound on the Rényi divergence of p t p_t p t ν \nu ν [47 ] .
Lemma B.7 (Divergence Bound). Consider a time-invariant Markov process with marginal distributions p t p_t p t t ≥ 0 t \ge0 t ≥ 0 p ∞ p_\infty p ∞ ν \nu ν Ψ \Psi Ψ p ∞ ∼ ν p_\infty\sim\nu p ∞ ∼ ν ∀ t ≥ 0 \forall t \ge0 ∀ t ≥ 0 p t ≪ ν , p ∞ p_t \ll\nu,p_\infty p t ≪ ν , p ∞ t ≥ 0 t \ge0 t ≥ 0 α > 1 \alpha> 1 α > 1
D α ( p t ∥ ν ) ≤ D ∞ ( p 0 ∥ ν ) + 1 2 ( α − 1 ) ln E p 0 [ e 2 ( α − 1 ) Ψ ] + 1 2 ( α − 1 ) ln E p t [ e − 2 ( α − 1 ) Ψ ] . D_\alpha(p_t\Vert\nu) \le D_\infty(p_0\Vert\nu) + \frac{1}{2(\alpha-1)}\ln\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\Psi}\right] + \frac{1}{2(\alpha-1)}\ln\mathbb{E}_{p_t}\left[e^{-2(\alpha-1)\Psi}\right]. D α ( p t ∥ ν ) ≤ D ∞ ( p 0 ∥ ν ) + 2 ( α − 1 ) 1 ln E p 0 [ e 2 ( α − 1 ) Ψ ] + 2 ( α − 1 ) 1 ln E p t [ e − 2 ( α − 1 ) Ψ ] . Proof. Applying the chain rule,
D α ( p t ∥ ν ) = 1 α − 1 ln ( ∫ ( d p t d ν ) α − 1 d p t ) = 1 α − 1 ln ( ∫ ( d p t d p ∞ d p ∞ d ν ) α − 1 d p t ) \begin{aligned}
D_\alpha(p_t\Vert\nu)=\frac{1}{\alpha-1}\ln\left(\int\left(\frac{\mathrm{d}p_t}{\mathrm{d}\nu}\right)^{\alpha-1}\mathrm{d}p_t\right) \\
&=\frac{1}{\alpha-1}\ln\left(\int\left(\frac{\mathrm{d}p_t}{\mathrm{d}p_\infty}\frac{\mathrm{d}p_\infty}{\mathrm{d}\nu}\right)^{\alpha-1}\mathrm{d}p_t\right)
\end{aligned} D α ( p t ∥ ν ) = α − 1 1 ln ( ∫ ( d ν d p t ) α − 1 d p t ) = α − 1 1 ln ( ∫ ( d p ∞ d p t d ν d p ∞ ) α − 1 d p t ) [Hölder: r = s = 2 r=s=2 r = s = 2 ≤ \le ≤
1 α − 1 ln ( ( ∫ ( d p t d p ∞ ) 2 ( α − 1 ) d p t ) 1 / 2 ( ∫ ( d p ∞ d ν ) 2 ( α − 1 ) d p t ) 1 / 2 ) = 1 2 ( α − 1 ) ln ( ∫ ( d p t d p ∞ ) 2 ( α − 1 ) d p t ) + 1 2 ( α − 1 ) ln ( ∫ ( d p ∞ d ν ) 2 ( α − 1 ) d p t ) = D 2 α − 1 ( p t ∥ p ∞ ) + D 2 α − 1 p t ( p ∞ ∥ ν ) \begin{aligned}
\frac{1}{\alpha-1}\ln\left(\left(\int\left(\frac{\mathrm{d}p_t}{\mathrm{d}p_\infty}\right)^{2(\alpha-1)}\mathrm{d}p_t\right)^{1/2}\left(\int\left(\frac{\mathrm{d}p_\infty}{\mathrm{d}\nu}\right)^{2(\alpha-1)}\mathrm{d}p_t\right)^{1/2}\right) \\
&=\frac{1}{2(\alpha-1)}\ln\left(\int\left(\frac{\mathrm{d}p_t}{\mathrm{d}p_\infty}\right)^{2(\alpha-1)}\mathrm{d}p_t\right)+\frac{1}{2(\alpha-1)}\ln\left(\int\left(\frac{\mathrm{d}p_\infty}{\mathrm{d}\nu}\right)^{2(\alpha-1)}\mathrm{d}p_t\right) \\
&=D_{2\alpha-1}(p_t\Vert p_\infty)+D^{p_t}_{2\alpha-1}(p_\infty\Vert\nu)
\end{aligned} α − 1 1 ln ( ∫ ( d p ∞ d p t ) 2 ( α − 1 ) d p t ) 1/2 ( ∫ ( d ν d p ∞ ) 2 ( α − 1 ) d p t ) 1/2 = 2 ( α − 1 ) 1 ln ( ∫ ( d p ∞ d p t ) 2 ( α − 1 ) d p t ) + 2 ( α − 1 ) 1 ln ( ∫ ( d ν d p ∞ ) 2 ( α − 1 ) d p t ) = D 2 α − 1 ( p t ∥ p ∞ ) + D 2 α − 1 p t ( p ∞ ∥ ν ) [Corollary B.4 ] ≤ \le ≤
D 2 α − 1 ( p 0 ∥ p ∞ ) + D 2 α − 1 p t ( p ∞ ∥ ν ) D_{2\alpha-1}(p_0\Vert p_\infty)+D^{p_t}_{2\alpha-1}(p_\infty\Vert\nu) D 2 α − 1 ( p 0 ∥ p ∞ ) + D 2 α − 1 p t ( p ∞ ∥ ν ) [Chain rule] = = =
1 2 α − 2 ln ( ∫ ( d p 0 d ν d ν d p ∞ ) 2 α − 2 d p 0 ) + D 2 α − 1 p t ( p ∞ ∥ ν ) \frac{1}{2\alpha-2}\ln\left(\int\left(\frac{\mathrm{d}p_0}{\mathrm{d}\nu}\frac{\mathrm{d}\nu}{\mathrm{d}p_\infty}\right)^{2\alpha-2}\mathrm{d}p_0\right)+D^{p_t}_{2\alpha-1}(p_\infty\Vert\nu) 2 α − 2 1 ln ( ∫ ( d ν d p 0 d p ∞ d ν ) 2 α − 2 d p 0 ) + D 2 α − 1 p t ( p ∞ ∥ ν ) [Hölder: r = ∞ r=\infty r = ∞ s = 1 s=1 s = 1 ≤ \le ≤
1 2 α − 2 ln ( ess sup p 0 ( d p 0 d ν ) 2 α − 2 ∫ ( d ν d p ∞ ) 2 α − 2 d p 0 ) + D 2 α − 1 p t ( p ∞ ∥ ν ) = D ∞ ( p 0 ∥ ν ) + D 2 α − 1 p 0 ( ν ∥ p ∞ ) + D 2 α − 1 p t ( p ∞ ∥ ν ) \begin{aligned}
\frac{1}{2\alpha-2}\ln\left(\mathop{\operatorname{ess\,sup}}_{p_0}\left(\frac{\mathrm{d}p_0}{\mathrm{d}\nu}\right)^{2\alpha-2}\int\left(\frac{\mathrm{d}\nu}{\mathrm{d}p_\infty}\right)^{2\alpha-2}\mathrm{d}p_0\right)+D^{p_t}_{2\alpha-1}(p_\infty\Vert\nu) \\
&=D_\infty(p_0\Vert\nu)+D^{p_0}_{2\alpha-1}(\nu\Vert p_\infty)+D^{p_t}_{2\alpha-1}(p_\infty\Vert\nu)
\end{aligned} 2 α − 2 1 ln ( ess sup p 0 ( d ν d p 0 ) 2 α − 2 ∫ ( d p ∞ d ν ) 2 α − 2 d p 0 ) + D 2 α − 1 p t ( p ∞ ∥ ν ) = D ∞ ( p 0 ∥ ν ) + D 2 α − 1 p 0 ( ν ∥ p ∞ ) + D 2 α − 1 p t ( p ∞ ∥ ν ) [Claim B.6 ] = = =
D ∞ ( p 0 ∥ ν ) + 1 2 α − 2 ( ln E p t [ e − ( 2 α − 2 ) Ψ ] + ln E p 0 [ e ( 2 α − 2 ) Ψ ] ) , D_\infty(p_0\Vert\nu)+\frac{1}{2\alpha-2}\left(\ln\mathbb{E}_{p_t}\left[e^{-(2\alpha-2)\Psi}\right]+\ln\mathbb{E}_{p_0}\left[e^{(2\alpha-2)\Psi}\right]\right), D ∞ ( p 0 ∥ ν ) + 2 α − 2 1 ( ln E p t [ e − ( 2 α − 2 ) Ψ ] + ln E p 0 [ e ( 2 α − 2 ) Ψ ] ) , where the integrability requirements of Claim B.6 are omitted since, when they do not hold, the RHS is infinite and the inequality is trivial. □ \square □
Remark B.8 (Avoiding Rényi-infinity divergence). Using r = 3 2 r=\frac{3}{2} r = 2 3 s = 3 s=3 s = 3 r = s = 2 r=s=2 r = s = 2 D ∞ ( p 0 ∥ ν ) D_\infty(p_0\Vert\nu) D ∞ ( p 0 ∥ ν )
D α ( p t ∥ ν ) ≤ D 3 α − 2 ( p 0 ∥ ν ) + 1 3 ( α − 1 ) ( ln E p 0 [ e 3 ( α − 1 ) Ψ ] + ln E p t [ e − 3 ( α − 1 ) Ψ ] ) , D_\alpha(p_t\Vert\nu)\le D_{3\alpha-2}(p_0\Vert\nu)+\frac{1}{3(\alpha-1)}\left(\ln\mathbb{E}_{p_0}\left[e^{3(\alpha-1)\Psi}\right]+\ln\mathbb{E}_{p_t}\left[e^{-3(\alpha-1)\Psi}\right]\right), D α ( p t ∥ ν ) ≤ D 3 α − 2 ( p 0 ∥ ν ) + 3 ( α − 1 ) 1 ( ln E p 0 [ e 3 ( α − 1 ) Ψ ] + ln E p t [ e − 3 ( α − 1 ) Ψ ] ) , which may be more useful when one does not assume ν = p 0 \nu=p_0 ν = p 0
Rényi change-of-measure
We now cite Theorem 8 from Bégin et al. [11 ] , stating a change of measure inequality based on the Rényi divergence, and provide the proof for completeness.
Lemma B.9 (Rényi change-of-measure [11 ] ). Let p p p q q q X \mathcal{X} X p ≪ q p \ll q p ≪ q ϕ : X → R \phi:\mathcal{X}\to\mathbb{R} ϕ : X → R α > 1 \alpha>1 α > 1
α α − 1 ln E p [ ϕ ( x ) ] ≤ D α ( p ∥ q ) + ln E q [ ϕ ( x ) α α − 1 ] . \frac{\alpha}{\alpha-1}\ln\mathbb{E}_{p}\left[\phi(x)\right]\leq D_{\alpha}(p\|q)+\ln\mathbb{E}_{q}\left[\phi(x)^{\frac{\alpha}{\alpha-1}}\right]. α − 1 α ln E p [ ϕ ( x ) ] ≤ D α ( p ∥ q ) + ln E q [ ϕ ( x ) α − 1 α ] . Proof. We rewrite the expectation over p p p q q q p ≪ q p\ll q p ≪ q r = α r=\alpha r = α s = α α − 1 s=\frac{\alpha}{\alpha-1} s = α − 1 α 1 r + 1 s = 1 \frac{1}{r}+\frac{1}{s}=1 r 1 + s 1 = 1
α α − 1 ln E p [ ϕ ( x ) ] = α α − 1 ln E q [ d p d q ϕ ( x ) ] \frac{\alpha}{\alpha-1}\ln\mathbb{E}_{p}\left[\phi(x)\right]=\frac{\alpha}{\alpha-1}\ln\mathbb{E}_{q}\left[\frac{dp}{dq}\phi(x)\right] α − 1 α ln E p [ ϕ ( x ) ] = α − 1 α ln E q [ d q d p ϕ ( x ) ] ≤ α α − 1 ln ( ( E q [ ( d p d q ) α ] ) 1 α ( E q [ ϕ ( x ) α α − 1 ] ) α − 1 α ) \leq\frac{\alpha}{\alpha-1}\ln\left(\left(\mathbb{E}_{q}\left[\left(\frac{dp}{dq}\right)^{\alpha}\right]\right)^{\frac{1}{\alpha}}\left(\mathbb{E}_{q}\left[\phi(x)^{\frac{\alpha}{\alpha-1}}\right]\right)^{\frac{\alpha-1}{\alpha}}\right) ≤ α − 1 α ln ( ( E q [ ( d q d p ) α ] ) α 1 ( E q [ ϕ ( x ) α − 1 α ] ) α α − 1 ) = 1 α − 1 ln ( E q [ ( d p d q ) α ] ) + ln ( E q [ ϕ ( x ) α α − 1 ] ) =\frac{1}{\alpha-1}\ln\left(\mathbb{E}_{q}\left[\left(\frac{dp}{dq}\right)^{\alpha}\right]\right)+\ln\left(\mathbb{E}_{q}\left[\phi(x)^{\frac{\alpha}{\alpha-1}}\right]\right) = α − 1 1 ln ( E q [ ( d q d p ) α ] ) + ln ( E q [ ϕ ( x ) α − 1 α ] ) = D α ( p ∥ q ) + ln ( E q [ ϕ ( x ) α α − 1 ] ) , =D_{\alpha}(p\|q)+\ln\left(\mathbb{E}_{q}\left[\phi(x)^{\frac{\alpha}{\alpha-1}}\right]\right), = D α ( p ∥ q ) + ln ( E q [ ϕ ( x ) α − 1 α ] ) , where in the last equality we used the definition of the Rényi divergence and rewrote the integral with respect to q q q p ≪ q p\ll q p ≪ q d p d q \frac{dp}{dq} d q d p
D α ( p ∥ q ) ≜ 1 α − 1 ln ( ∫ ( d p d q ) α − 1 d p ) D_{\alpha}(p\|q)\triangleq\frac{1}{\alpha-1}\ln\left(\int\left(\frac{dp}{dq}\right)^{\alpha-1}dp\right) D α ( p ∥ q ) ≜ α − 1 1 ln ( ∫ ( d q d p ) α − 1 d p ) = 1 α − 1 ln ( ∫ ( d p d q ) α − 1 d p d q d q ) =\frac{1}{\alpha-1}\ln\left(\int\left(\frac{dp}{dq}\right)^{\alpha-1}\frac{dp}{dq}dq\right) = α − 1 1 ln ( ∫ ( d q d p ) α − 1 d q d p d q ) = 1 α − 1 ln ( E q [ ( d p d q ) α ] ) . =\frac{1}{\alpha-1}\ln\left(\mathbb{E}_{q}\left[\left(\frac{dp}{dq}\right)^{\alpha}\right]\right). = α − 1 1 ln ( E q [ ( d q d p ) α ] ) . Conclusion
Recall Lemma 2.2 (Second law of thermodynamics change of measure inequality). Consider a time invariant Markov process over the state space X \mathcal{X} X p t p_t p t t ≥ 0 t\geq0 t ≥ 0 p t ≪ p 0 p_t\ll p_0 p t ≪ p 0 p ∞ p_\infty p ∞ p 0 p_0 p 0 Ψ \Psi Ψ p ∞ ∼ p 0 p_\infty\sim p_0 p ∞ ∼ p 0 ϕ : X → R \phi:\mathcal{X}\to\mathbb{R} ϕ : X → R α > 1 \alpha>1 α > 1 t ≥ 0 t\geq0 t ≥ 0
E p t ϕ ( x ) ≤ ( E p 0 [ e 2 ( α − 1 ) Ψ ] ) 1 2 α ( E p 0 [ ϕ ( x ) α α − 1 ] ) α − 1 α . \mathbb{E}_{p_t}\phi(x)\leq\left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\Psi}\right]\right)^{\frac{1}{2\alpha}}\left(\mathbb{E}_{p_0}\left[\phi(x)^{\frac{\alpha}{\alpha-1}}\right]\right)^{\frac{\alpha-1}{\alpha}}. E p t ϕ ( x ) ≤ ( E p 0 [ e 2 ( α − 1 ) Ψ ] ) 2 α 1 ( E p 0 [ ϕ ( x ) α − 1 α ] ) α α − 1 . Proof. We simply plug Lemma B.7 with ν = p 0 \nu=p_0 ν = p 0 Lemma B.9 with p = p t p=p_t p = p t q = p 0 q=p_0 q = p 0 Lemma B.7 .
Remark B.10 (Absolute continuity requirements). When imposing ν = p 0 = N ( 0 N p , Λ − 1 ) \nu=p_0=\mathcal{N}\left(0_{N_p},\Lambda^{-1}\right) ν = p 0 = N ( 0 N p , Λ − 1 ) ν \nu ν R N p \mathbb{R}^{N_p} R N p (4) ) we have Ψ = β L ( θ ) < ∞ \Psi=\beta L(\theta)<\infty Ψ = β L ( θ ) < ∞ R N p \mathbb{R}^{N_p} R N p
p ∞ ∝ e − Ψ ν p_\infty\propto e^{-\Psi}\nu p ∞ ∝ e − Ψ ν also has strictly positive Lebesgue density on R N p \mathbb{R}^{N_p} R N p p ∞ ∼ ν p_\infty\sim\nu p ∞ ∼ ν
By Proposition F.1 , the process is global, p ∞ p_\infty p ∞ p t ≪ p ∞ ∼ ν p_t\ll p_\infty\sim\nu p t ≪ p ∞ ∼ ν t ≥ 0 t\geq0 t ≥ 0 Lemma 2.2 hold directly.
Assumption table
Lower bound on the empirical NTK at initialization
Setting: deep fully connected network
We follow most of the notation from § 4 , with a feed-forward L + 1 L+1 L + 1 L L L
α ( 0 ) ( x ) = x , \boldsymbol{\alpha}^{(0)}(\mathbf{x})=\mathbf{x}, α ( 0 ) ( x ) = x , α ( l ) ( x ) = ϕ ( s l W ( l ) α ( l − 1 ) ( x ) ) , l ∈ [ L ] , \boldsymbol{\alpha}^{(l)}(\mathbf{x})=\phi\left(s_l\mathbf{W}^{(l)}\boldsymbol{\alpha}^{(l-1)}(\mathbf{x})\right),\quad l\in[L], α ( l ) ( x ) = ϕ ( s l W ( l ) α ( l − 1 ) ( x ) ) , l ∈ [ L ] , f ( θ ; x ) = α ( L + 1 ) ( x ) = s L + 1 v T α ( L ) ( x ) , f(\boldsymbol{\theta};\mathbf{x})=\boldsymbol{\alpha}^{(L+1)}(\mathbf{x})=s_{L+1}\mathbf{v}^{\mathsf T}\boldsymbol{\alpha}^{(L)}(\mathbf{x}), f ( θ ; x ) = α ( L + 1 ) ( x ) = s L + 1 v T α ( L ) ( x ) , where W ( l ) ∈ R m l × m l − 1 \mathbf{W}^{(l)}\in\mathbb{R}^{m_l\times m_{l-1}} W ( l ) ∈ R m l × m l − 1 l ∈ [ L ] l\in[L] l ∈ [ L ] v ∈ R m L \mathbf{v}\in\mathbb{R}^{m_L} v ∈ R m L ϕ ( ⋅ ) \phi(\cdot) ϕ ( ⋅ ) s l s_l s l l ∈ [ L + 1 ] l\in[L+1] l ∈ [ L + 1 ]
θ ≜ ( vec ( W ( 1 ) ) T , … , vec ( W ( L ) ) T , v T ) T ∈ R N p , \boldsymbol{\theta}\triangleq\left(\operatorname{vec}\left(\mathbf{W}^{(1)}\right)^{\mathsf T},\ldots,\operatorname{vec}\left(\mathbf{W}^{(L)}\right)^{\mathsf T},\mathbf{v}^{\mathsf T}\right)^{\mathsf T}\in\mathbb{R}^{N_p}, θ ≜ ( vec ( W ( 1 ) ) T , … , vec ( W ( L ) ) T , v T ) T ∈ R N p , with m 0 ≜ d m_0\triangleq d m 0 ≜ d N p ≜ ∑ l = 1 L m l m l − 1 + m L N_p\triangleq\sum_{l=1}^{L}m_lm_{l-1}+m_L N p ≜ ∑ l = 1 L m l m l − 1 + m L L = O ( 1 ) L=O(1) L = O ( 1 )
Recall that c ϕ , σ 0 ≜ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z ) ] c_{\phi,\sigma_0}\triangleq\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(z)] c ϕ , σ 0 ≜ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z )] ν 0 2 ≜ σ 0 2 c ϕ , σ 0 \nu_0^2\triangleq\frac{\sigma_0^2}{c_{\phi,\sigma_0}} ν 0 2 ≜ c ϕ , σ 0 σ 0 2 ∥ x i ∥ 2 2 = c ϕ , σ 0 d \|\mathbf{x}_i\|_2^2=c_{\phi,\sigma_0}d ∥ x i ∥ 2 2 = c ϕ , σ 0 d i ∈ [ n ] i\in[n] i ∈ [ n ] W ( l ) \mathbf{W}^{(l)} W ( l ) l ∈ [ L ] l\in[L] l ∈ [ L ]
vec ( W ( l ) ) ∼ N ( 0 m l m l − 1 , ν 0 2 a l I m l m l − 1 ) , \operatorname{vec}\left(\mathbf{W}^{(l)}\right)\sim\mathcal{N}\left(\mathbf{0}_{m_lm_{l-1}},\frac{\nu_0^2}{a_l}\mathbf{I}_{m_lm_{l-1}}\right), vec ( W ( l ) ) ∼ N ( 0 m l m l − 1 , a l ν 0 2 I m l m l − 1 ) , such that ∀ l ∈ [ L ] \forall l\in[L] ∀ l ∈ [ L ] a l s l 2 = m l − 1 \frac{a_l}{s_l^2}=m_{l-1} s l 2 a l = m l − 1 v \mathbf{v} v ϕ ( 0 ) = 0 \phi(0)=0 ϕ ( 0 ) = 0 ∣ ϕ ′ ∣ ≤ α ϕ |\phi'|\le\alpha_\phi ∣ ϕ ′ ∣ ≤ α ϕ λ 1 > 0 \lambda_1>0 λ 1 > 0 15
Let A ( l ) ∈ R n × m l \mathbf{A}^{(l)}\in\mathbb{R}^{n\times m_l} A ( l ) ∈ R n × m l A i , : ( l ) ≜ α ( l ) ( x i ) T \mathbf{A}^{(l)}_{i,:}\triangleq\boldsymbol{\alpha}^{(l)}(\mathbf{x}_i)^{\mathsf T} A i , : ( l ) ≜ α ( l ) ( x i ) T l ∈ [ L ] l\in[L] l ∈ [ L ] i ∈ [ n ] i\in[n] i ∈ [ n ] w j ( l ) ∈ R m l − 1 \mathbf{w}^{(l)}_j\in\mathbb{R}^{m_{l-1}} w j ( l ) ∈ R m l − 1 j j j W ( l ) \mathbf{W}^{(l)} W ( l ) l ∈ [ L ] l\in[L] l ∈ [ L ] m ≜ min l ∈ [ L ] m l m\triangleq\min_{l\in[L]}m_l m ≜ min l ∈ [ L ] m l
The analysis in this section closely follows that of Banerjee et al. [9 ] , with adjustments used to tighten the probability of their bound, apply it to our generalized parameterization, and fix a minor error we found in their proof.
Analysis
We begin with a supporting proposition:
Proposition D.1. For all l ∈ [ L ] l\in[L] l ∈ [ L ] j ∈ [ m l ] j\in[m_l] j ∈ [ m l ] i ∈ [ n ] i\in[n] i ∈ [ n ] ∣ α j ( l ) ( x i ) ∣ \left|\alpha_j^{(l)}(\mathbf{x}_i)\right| α j ( l ) ( x i ) A ( l − 1 ) \mathbf{A}^{(l-1)} A ( l − 1 )
∥ α j ( l ) ( x i ) ∥ ψ 2 ≤ c 2 α ϕ ν 0 s l a l ∥ α ( l − 1 ) ( x i ) ∥ 2 \left\|\alpha_j^{(l)}(\mathbf{x}_i)\right\|_{\psi_2}\le\frac{c\sqrt{2\alpha_\phi\nu_0s_l}}{\sqrt{a_l}}\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2 α j ( l ) ( x i ) ψ 2 ≤ a l c 2 α ϕ ν 0 s l α ( l − 1 ) ( x i ) 2 for some absolute constant c > 0 c>0 c > 0
Proof. We follow a similar proof to their proof of Lemma A.1.
Observe that for all j ∈ [ m l ] j\in[m_l] j ∈ [ m l ] α j ( l ) ( x i ) = ϕ ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ) \alpha_j^{(l)}(\mathbf{x}_i)=\phi\left(s_l(\mathbf{w}^{(l)}_j)^{\mathsf T}\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right) α j ( l ) ( x i ) = ϕ ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ) w j ( l ) ∼ N ( 0 m l − 1 , ν 0 2 a l I m l − 1 ) \mathbf{w}^{(l)}_j\sim\mathcal{N}\left(\mathbf{0}_{m_{l-1}},\frac{\nu_0^2}{a_l}\mathbf{I}_{m_{l-1}}\right) w j ( l ) ∼ N ( 0 m l − 1 , a l ν 0 2 I m l − 1 )
Note that if ∥ α ( l − 1 ) ( x i ) ∥ 2 = 0 \left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2=0 α ( l − 1 ) ( x i ) 2 = 0
α j ( l ) ( x i ) = ϕ ( 0 ) = 0 , \alpha_j^{(l)}(\mathbf{x}_i)=\phi(0)=0, α j ( l ) ( x i ) = ϕ ( 0 ) = 0 , and the claim is trivial. Hence assume this norm is positive. Therefore
P ( ∣ s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ∣ ≥ ε s l ∥ α ( l − 1 ) ( x i ) ∥ 2 ) = P ( ∣ ( w j ( l ) ) T α ( l − 1 ) ( x i ) ∥ α ( l − 1 ) ( x i ) ∥ 2 ∣ ≥ ε ) = P ( ∣ ( w j ( l ) ) T α ( l − 1 ) ( x i ) ∥ α ( l − 1 ) ( x i ) ∥ 2 ∣ ≥ ε ) ≤ 2 exp ( − ε 2 a l 2 ν 0 2 ) , \begin{aligned}
\mathbb{P}\left(\left|s_l\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(l-1)}\left(x_i\right)\right|\ge\varepsilon s_l\left\|\boldsymbol{\alpha}^{(l-1)}\left(x_i\right)\right\|_2\right)
&=\mathbb{P}\left(\left|\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\frac{\boldsymbol{\alpha}^{(l-1)}\left(x_i\right)}{\left\|\boldsymbol{\alpha}^{(l-1)}\left(x_i\right)\right\|_2}\right|\ge\varepsilon\right)\\
&=\mathbb{P}\left(\left|\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\frac{\boldsymbol{\alpha}^{(l-1)}\left(x_i\right)}{\left\|\boldsymbol{\alpha}^{(l-1)}\left(x_i\right)\right\|_2}\right|\ge\varepsilon\right)\\
&\le2\exp\left(-\frac{\varepsilon^2a_l}{2\nu_0^2}\right),
\end{aligned} P ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ≥ ε s l α ( l − 1 ) ( x i ) 2 ) = P ( ( w j ( l ) ) T α ( l − 1 ) ( x i ) 2 α ( l − 1 ) ( x i ) ≥ ε ) = P ( ( w j ( l ) ) T α ( l − 1 ) ( x i ) 2 α ( l − 1 ) ( x i ) ≥ ε ) ≤ 2 exp ( − 2 ν 0 2 ε 2 a l ) , since ( w j ( l ) ) T α ( l − 1 ) ( x i ) ∥ α ( l − 1 ) ( x i ) ∥ 2 ∼ N ( 0 , ν 0 2 a l ) \left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\frac{\boldsymbol{\alpha}^{(l-1)}(x_i)}{\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2}\sim\mathcal{N}\left(0,\frac{\nu_0^2}{a_l}\right) ( w j ( l ) ) T ∥ α ( l − 1 ) ( x i ) ∥ 2 α ( l − 1 ) ( x i ) ∼ N ( 0 , a l ν 0 2 ) A ( l − 1 ) \mathbf{A}^{(l-1)} A ( l − 1 ) [111 Exercise 2.6] .
Denoting ε ~ = ε s l ∥ α ( l − 1 ) ( x i ) ∥ 2 \widetilde{\varepsilon}=\varepsilon s_l\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2 ε = ε s l α ( l − 1 ) ( x i ) 2
P ( ∣ s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ∣ ≥ ε ~ ) ≤ 2 exp ( − ε ~ 2 a l 2 ν 0 2 s l 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 2 ) . \mathbb{P}\left(\left|s_l\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(l-1)}(x_i)\right|\ge\widetilde{\varepsilon}\right)\le2\exp\left(-\frac{\widetilde{\varepsilon}^2a_l}{2\nu_0^2s_l^2\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2^2}\right). P ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ≥ ε ) ≤ 2 exp ( − 2 ν 0 2 s l 2 α ( l − 1 ) ( x i ) 2 2 ε 2 a l ) . Since ϕ \phi ϕ α ϕ \alpha_\phi α ϕ ϕ ( 0 ) = 0 \phi(0)=0 ϕ ( 0 ) = 0 ∣ α j ( l ) ( x i ) ∣ = ∣ ϕ ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ) ∣ ≤ α ϕ ∣ s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ∣ \left|\alpha_j^{(l)}(x_i)\right|=\left|\phi\left(s_l\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(l-1)}(x_i)\right)\right|\le\alpha_\phi\left|s_l\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(l-1)}(x_i)\right| α j ( l ) ( x i ) = ϕ ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ) ≤ α ϕ s l ( w j ( l ) ) T α ( l − 1 ) ( x i )
P ( ∣ ϕ ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ) ∣ ≥ ε ~ ) ≤ P ( ∣ s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ∣ ≥ ε ~ α ϕ ) ≤ 2 exp ( − ε ~ 2 a l 2 ν 0 2 α ϕ 2 s l 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 2 ) , \mathbb{P}\left(\left|\phi\left(s_l\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(l-1)}(x_i)\right)\right|\ge\widetilde{\varepsilon}\right)\le\mathbb{P}\left(\left|s_l\left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(l-1)}(x_i)\right|\ge\frac{\widetilde{\varepsilon}}{\alpha_\phi}\right)
\le2\exp\left(-\frac{\widetilde{\varepsilon}^2a_l}{2\nu_0^2\alpha_\phi^2s_l^2\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2^2}\right), P ( ϕ ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ) ≥ ε ) ≤ P ( s l ( w j ( l ) ) T α ( l − 1 ) ( x i ) ≥ α ϕ ε ) ≤ 2 exp ( − 2 ν 0 2 α ϕ 2 s l 2 α ( l − 1 ) ( x i ) 2 2 ε 2 a l ) , proving ∣ α j ( l ) ( x i ) ∣ \left|\alpha_j^{(l)}(x_i)\right| α j ( l ) ( x i ) ∥ α j ( l ) ( x i ) ∥ ψ 2 ≤ c 2 α ϕ ν 0 s l a l ∥ α ( l − 1 ) ( x i ) ∥ 2 \left\|\alpha_j^{(l)}(x_i)\right\|_{\psi_2}\le\frac{c\sqrt{2}\alpha_\phi\nu_0s_l}{\sqrt{a_l}}\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2 α j ( l ) ( x i ) ψ 2 ≤ a l c 2 α ϕ ν 0 s l α ( l − 1 ) ( x i ) 2 c > 0 c>0 c > 0
We now produce a tighter alternative to Lemma A.1 of Banerjee et al. [9 ] , to be used later.
Proposition D.2. For all l ∈ [ L ] l\in[L] l ∈ [ L ] j ∈ [ m l ] j\in[m_l] j ∈ [ m l ] ∥ A : , j ( l ) ∥ 2 \left\|\mathbf{A}_{:,j}^{(l)}\right\|_2 A : , j ( l ) 2 A ( l − 1 ) \mathbf{A}^{(l-1)} A ( l − 1 ) ∥ ∥ A : , j ( l ) ∥ 2 2 ∥ ψ 2 ≤ c 2 α ϕ ν 0 s l a l ∥ A ( l − 1 ) ∥ F \left\|\left\|\mathbf{A}_{:,j}^{(l)}\right\|_2^2\right\|_{\psi_2}\le c\frac{\sqrt{2}\alpha_\phi\nu_0s_l}{\sqrt{a_l}}\left\|\mathbf{A}^{(l-1)}\right\|_F A : , j ( l ) 2 2 ψ 2 ≤ c a l 2 α ϕ ν 0 s l A ( l − 1 ) F c > 0 c>0 c > 0
Proof. Observe that ∥ A : , j ( l ) ∥ 2 2 = ∑ i = 1 n ( α j ( l ) ( x i ) ) 2 = ∑ i = 1 n ∣ α j ( l ) ( x i ) ∣ 2 \left\|\mathbf{A}_{:,j}^{(l)}\right\|_2^2=\sum_{i=1}^{n}\left(\alpha_j^{(l)}(x_i)\right)^2=\sum_{i=1}^{n}\left|\alpha_j^{(l)}(x_i)\right|^2 A : , j ( l ) 2 2 = ∑ i = 1 n ( α j ( l ) ( x i ) ) 2 = ∑ i = 1 n α j ( l ) ( x i ) 2 ∀ i ∈ [ n ] \forall i\in[n] ∀ i ∈ [ n ] ∥ α j ( l ) ( x i ) ∥ ψ 2 ≤ c 2 α ϕ ν 0 s l a l ∥ α ( l − 1 ) ( x i ) ∥ 2 \left\|\alpha_j^{(l)}(x_i)\right\|_{\psi_2}\le c\frac{\sqrt{2}\alpha_\phi\nu_0s_l}{\sqrt{a_l}}\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2 α j ( l ) ( x i ) ψ 2 ≤ c a l 2 α ϕ ν 0 s l α ( l − 1 ) ( x i ) 2 ∥ ∥ α j ( l ) ( x i ) ∥ 2 ∥ ψ 2 ≤ ( c 2 α ϕ ν 0 s l a l ) 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 2 \left\|\left\|\alpha_j^{(l)}(x_i)\right\|^2\right\|_{\psi_2}\le\left(c\frac{\sqrt{2}\alpha_\phi\nu_0s_l}{\sqrt{a_l}}\right)^2\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2^2 α j ( l ) ( x i ) 2 ψ 2 ≤ ( c a l 2 α ϕ ν 0 s l ) 2 α ( l − 1 ) ( x i ) 2 2 [111 Lemma 2.8.5] . From the triangle inequality for the subexponential norm we have
∥ ∥ A : , j ( l ) ∥ 2 2 ∥ ψ 1 = ∥ ∑ i = 1 n ∣ α j ( l ) ( x i ) ∣ 2 ∥ ψ 1 ≤ ∑ i = 1 n ∥ ∣ α j ( l ) ( x i ) ∣ 2 ∥ ψ 1 ≤ ( c 2 α ϕ ν 0 s l a l ) 2 ∑ i = 1 n ∥ α ( l − 1 ) ( x i ) ∥ 2 2 = ( c 2 α ϕ ν 0 s l a l ) 2 ∥ A ( l − 1 ) ∥ F 2 \left\|\left\|\mathbf{A}_{:,j}^{(l)}\right\|_2^2\right\|_{\psi_1}
=\left\|\sum_{i=1}^{n}\left|\alpha_j^{(l)}(x_i)\right|^2\right\|_{\psi_1}
\le\sum_{i=1}^{n}\left\|\left|\alpha_j^{(l)}(x_i)\right|^2\right\|_{\psi_1}
\le\left(c\frac{\sqrt{2}\alpha_\phi\nu_0s_l}{\sqrt{a_l}}\right)^2\sum_{i=1}^{n}\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|_2^2
=\left(c\frac{\sqrt{2}\alpha_\phi\nu_0s_l}{\sqrt{a_l}}\right)^2\left\|\mathbf{A}^{(l-1)}\right\|_F^2 A : , j ( l ) 2 2 ψ 1 = i = 1 ∑ n α j ( l ) ( x i ) 2 ψ 1 ≤ i = 1 ∑ n α j ( l ) ( x i ) 2 ψ 1 ≤ ( c a l 2 α ϕ ν 0 s l ) 2 i = 1 ∑ n α ( l − 1 ) ( x i ) 2 2 = ( c a l 2 α ϕ ν 0 s l ) 2 A ( l − 1 ) F 2 ⇒ ∥ ∥ A : , j ( l ) ∥ 2 ∥ ψ 2 ≤ c 2 α ϕ ν 0 s l a l ∥ A ( l − 1 ) ∥ F \Rightarrow\left\|\left\|\mathbf{A}_{:,j}^{(l)}\right\|_2\right\|_{\psi_2}\le c\frac{\sqrt{2}\alpha_\phi\nu_0s_l}{\sqrt{a_l}}\left\|\mathbf{A}^{(l-1)}\right\|_F ⇒ A : , j ( l ) 2 ψ 2 ≤ c a l 2 α ϕ ν 0 s l A ( l − 1 ) F
We proceed to prove a version of Lemma 4.1 in Banerjee et al. [9 ] , adjusted to general parameterizations and fixing a minor issue that we found in the original proof.
Proposition D.3. Let λ l ≜ λ min ( E g ∼ N ( 0 m l − 1 , v 0 2 a l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] ) \lambda_l \triangleq\lambda_{\min}\left(\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\boldsymbol{0}_{m_{l-1}},\frac{v_0^2}{a_l^2}\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\right]\right) λ l ≜ λ m i n E g ∼ N ( 0 m l − 1 , a l 2 v 0 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] l ∈ [ L ] l\in[L] l ∈ [ L ] m l = Ω ( max [ n , ln ( n L δ ) 1 λ l s l 2 a l 2 ∥ A ( l − 1 ) ∥ F 2 max [ 1 , ln ( 1 λ l s l 2 a l 2 ∥ A ( l − 1 ) ∥ F 2 ) ] ] ) m_l=\Omega\left(\max\left[n,\ln\left(\frac{nL}{\delta}\right)\frac{1}{\lambda_l}\frac{s_l^2}{a_l^2}\left\|\mathbf{A}^{(l-1)}\right\|_{\mathrm{F}}^2\max\left[1,\ln\left(\frac{1}{\lambda_l}\frac{s_l^2}{a_l^2}\left\|\mathbf{A}^{(l-1)}\right\|_{\mathrm{F}}^2\right)\right]\right]\right) m l = Ω ( max [ n , ln ( δ n L ) λ l 1 a l 2 s l 2 A ( l − 1 ) F 2 max [ 1 , ln ( λ l 1 a l 2 s l 2 A ( l − 1 ) F 2 ) ] ] ) 1 − δ L 1-\frac{\delta}{L} 1 − L δ W ( l ) \mathbf{W}^{(l)} W ( l )
λ min ( A ( l ) ( A ( l ) ) T ) ≥ m l λ l 4 . \lambda_{\min}\left(\mathbf{A}^{(l)}\left(\mathbf{A}^{(l)}\right)^{\mathsf{T}}\right)\ge\frac{m_l\lambda_l}{4}. λ m i n ( A ( l ) ( A ( l ) ) T ) ≥ 4 m l λ l . Proof. The proof follows theirs very closely for the most part. A necessary condition for A ( l ) ( A ( l ) ) T \mathbf{A}^{(l)}\left(\mathbf{A}^{(l)}\right)^{\mathsf{T}} A ( l ) ( A ( l ) ) T m l ≥ n m_l\ge n m l ≥ n m l ≥ n m_l\ge n m l ≥ n
We define the matrix A ^ ( l ) ∈ R n × m l \widehat{\mathbf{A}}^{(l)}\in\mathbb{R}^{n\times m_l} A ( l ) ∈ R n × m l j ∈ [ m l ] j\in[m_l] j ∈ [ m l ] j j j A ^ ( l ) \widehat{\mathbf{A}}^{(l)} A ( l )
A ^ : , j ( l ) = ϕ ( s l A ( l − 1 ) w j ( l ) ) 1 [ ∥ ϕ ( s l A ( l − 1 ) w j ( l ) ) ∥ 2 ≤ t ] . \widehat{\mathbf{A}}_{:,j}^{(l)}=\phi\left(s_l\mathbf{A}^{(l-1)}\mathbf{w}_j^{(l)}\right)\mathbf{1}_{\left[\left\|\phi\left(s_l\mathbf{A}^{(l-1)}\mathbf{w}_j^{(l)}\right)\right\|_2\le t\right]}. A : , j ( l ) = ϕ ( s l A ( l − 1 ) w j ( l ) ) 1 [ ϕ ( s l A ( l − 1 ) w j ( l ) ) 2 ≤ t ] . where ( w j ( l ) ) T \left(\mathbf{w}_j^{(l)}\right)^{\mathsf{T}} ( w j ( l ) ) T j j j W ( l ) \mathbf{W}^{(l)} W ( l ) 1 [ B ] \mathbf{1}_{[B]} 1 [ B ] B B B t > 0 t>0 t > 0
λ min ( A ^ ( l ) ( A ^ ( l ) ) T ) ≤ λ min ( A ( l ) ( A ( l ) ) T ) . (13) \lambda_{\min}\left(\widehat{\mathbf{A}}^{(l)}\left(\widehat{\mathbf{A}}^{(l)}\right)^{\mathsf{T}}\right)\le\lambda_{\min}\left(\mathbf{A}^{(l)}\left(\mathbf{A}^{(l)}\right)^{\mathsf{T}}\right). \tag*{(13)} λ m i n ( A ( l ) ( A ( l ) ) T ) ≤ λ m i n ( A ( l ) ( A ( l ) ) T ) . (13) Since A ^ : , j ( l ) ∈ R n \widehat{\mathbf{A}}_{:,j}^{(l)}\in\mathbb{R}^n A : , j ( l ) ∈ R n λ max ( A ^ : , j ( l ) ( A ^ : , j ( l ) ) T ) = ∥ A ^ : , j ( l ) ∥ 2 ≤ t 2 \lambda_{\max}\left(\widehat{\mathbf{A}}_{:,j}^{(l)}\left(\widehat{\mathbf{A}}_{:,j}^{(l)}\right)^{\mathsf{T}}\right)=\left\|\widehat{\mathbf{A}}_{:,j}^{(l)}\right\|^2\le t^2 λ m a x ( A : , j ( l ) ( A : , j ( l ) ) T ) = A : , j ( l ) 2 ≤ t 2 A ^ : , j ( l ) \widehat{\mathbf{A}}_{:,j}^{(l)} A : , j ( l ) A ^ ( l ) ( A ^ ( l ) ) T = ∑ j ∈ [ m l ] A ^ : , j ( l ) ( A ^ : , j ( l ) ) T \widehat{\mathbf{A}}^{(l)}\left(\widehat{\mathbf{A}}^{(l)}\right)^{\mathsf{T}}=\sum_{j\in[m_l]}\widehat{\mathbf{A}}_{:,j}^{(l)}\left(\widehat{\mathbf{A}}_{:,j}^{(l)}\right)^{\mathsf{T}} A ( l ) ( A ( l ) ) T = ∑ j ∈ [ m l ] A : , j ( l ) ( A : , j ( l ) ) T ϵ ∈ [ 0 , 1 ) \epsilon\in[0,1) ϵ ∈ [ 0 , 1 )
P ( λ min ( A ^ ( l ) ( A ^ ( l ) ) T ) ≤ ( 1 − ϵ ) λ min ( E W ( l ) [ A ^ ( l ) ( A ^ ( l ) ) T ] ) ) ≤ n ( e − ϵ ( 1 − ϵ ) 1 − ϵ ) λ min ( E W ( l ) [ A ^ ( l ) ( A ^ ( l ) ) T ] ) / t 2 . \begin{aligned}
\mathbb{P}\left(\lambda_{\min}\left(\widehat{\mathbf{A}}^{(l)}\left(\widehat{\mathbf{A}}^{(l)}\right)^{\mathsf{T}}\right)\le(1-\epsilon)\lambda_{\min}\left(\mathbb{E}_{\mathbf{W}^{(l)}}\left[\widehat{\mathbf{A}}^{(l)}\left(\widehat{\mathbf{A}}^{(l)}\right)^{\mathsf{T}}\right]\right)\right) \\
&\le n\left(\frac{e^{-\epsilon}}{(1-\epsilon)^{1-\epsilon}}\right)^{\lambda_{\min}\left(\mathbb{E}_{\mathbf{W}^{(l)}}\left[\widehat{\mathbf{A}}^{(l)}\left(\widehat{\mathbf{A}}^{(l)}\right)^{\mathsf{T}}\right]\right)/t^2}.
\end{aligned} P ( λ m i n ( A ( l ) ( A ( l ) ) T ) ≤ ( 1 − ϵ ) λ m i n ( E W ( l ) [ A ( l ) ( A ( l ) ) T ] ) ) ≤ n ( ( 1 − ϵ ) 1 − ϵ e − ϵ ) λ m i n ( E W ( l ) [ A ( l ) ( A ( l ) ) T ] ) / t 2 . Recall that ∀ j ∈ [ m l ] \forall j\in[m_l] ∀ j ∈ [ m l ] w j ( l ) ∼ N ( 0 m l − 1 , v 0 2 a l 2 I m l − 1 ) \mathbf{w}_j^{(l)}\sim\mathcal{N}\left(\mathbf{0}_{m_{l-1}},\frac{v_0^2}{a_l^2}\mathbf{I}_{m_{l-1}}\right) w j ( l ) ∼ N ( 0 m l − 1 , a l 2 v 0 2 I m l − 1 )
E W ( l ) [ A ^ ( l ) ( A ^ ( l ) ) T ] = E W ( l ) [ ∑ j ∈ [ m l ] A ^ : , j ( l ) ( A ^ : , j ( l ) ) T ] = ∑ j ∈ [ m l ] E w j ( l ) [ A ^ : , j ( l ) ( A ^ : , j ( l ) ) T ] = m l E g ∼ N ( 0 m l − 1 , v 0 2 a l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 ≤ t ] ] . \begin{aligned}
\mathbb{E}_{\mathbf{W}^{(l)}}\left[\widehat{\mathbf{A}}^{(l)}\left(\widehat{\mathbf{A}}^{(l)}\right)^{\mathsf{T}}\right] \\
&=\mathbb{E}_{\mathbf{W}^{(l)}}\left[\sum_{j\in[m_l]}\widehat{\mathbf{A}}_{:,j}^{(l)}\left(\widehat{\mathbf{A}}_{:,j}^{(l)}\right)^{\mathsf{T}}\right] \\
&=\sum_{j\in[m_l]}\mathbb{E}_{\mathbf{w}_j^{(l)}}\left[\widehat{\mathbf{A}}_{:,j}^{(l)}\left(\widehat{\mathbf{A}}_{:,j}^{(l)}\right)^{\mathsf{T}}\right] \\
&=m_l\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\mathbf{0}_{m_{l-1}},\frac{v_0^2}{a_l^2}\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\mathbf{1}_{\left[\left\|\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)\right\|_2\le t\right]}\right].
\end{aligned} E W ( l ) [ A ( l ) ( A ( l ) ) T ] = E W ( l ) j ∈ [ m l ] ∑ A : , j ( l ) ( A : , j ( l ) ) T = j ∈ [ m l ] ∑ E w j ( l ) [ A : , j ( l ) ( A : , j ( l ) ) T ] = m l E g ∼ N ( 0 m l − 1 , a l 2 v 0 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 ≤ t ] ] . Denoting
G = E g ∼ N ( 0 m l − 1 , v 0 2 a l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] , \mathbf{G}=\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\mathbf{0}_{m_{l-1}},\frac{v_0^2}{a_l^2}\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\right], G = E g ∼ N ( 0 m l − 1 , a l 2 v 0 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] , G ^ = E g ∼ N ( 0 m l − 1 , v 0 2 a l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 ≤ t ] ] , \widehat{\mathbf{G}}=\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\mathbf{0}_{m_{l-1}},\frac{v_0^2}{a_l^2}\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\mathbf{1}_{\left[\left\|\phi\left(s_l\mathbf{A}^{(l-1)}\boldsymbol{g}\right)\right\|_2\le t\right]}\right], G = E g ∼ N ( 0 m l − 1 , a l 2 v 0 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 ≤ t ] ] , we have for ϵ = 1 2 \epsilon= \frac{1}{2} ϵ = 2 1
P ( λ min ( A ^ ( l ) ( A ^ ( l ) ) T ) ≤ m l 2 λ min ( G ^ ) ) ≤ n ( e − 1 / 2 ( 1 / 2 ) 1 / 2 ) m l λ min ( G ^ ) / t 2 = n exp ( ln ( 2 e ) m l λ min ( G ^ ) t 2 ) = exp ( − ( 1 − ln 2 ) m l λ min ( G ^ ) 2 t 2 + ln n ) . \begin{aligned}
\mathbb{P}\left(\lambda_{\min}\left(\widehat{A}^{(l)}\left(\widehat{A}^{(l)}\right)^{\mathsf{T}}\right) \le\frac{m_l}{2}\lambda_{\min}\left(\widehat{G}\right)\right) \le n\left(\frac{e^{-1/2}}{(1/2)^{1/2}}\right)^{m_l\lambda_{\min}(\widehat{G})/t^2} \\
&= n\exp\left(\ln\left(\sqrt{\frac{2}{e}}\right)\frac{m_l\lambda_{\min}\left(\widehat{G}\right)}{t^2}\right) \\
&= \exp\left(-\frac{(1-\ln2)m_l\lambda_{\min}\left(\widehat{G}\right)}{2t^2}+\ln n\right).
\end{aligned} P ( λ m i n ( A ( l ) ( A ( l ) ) T ) ≤ 2 m l λ m i n ( G ) ) ≤ n ( ( 1/2 ) 1/2 e − 1/2 ) m l λ m i n ( G ) / t 2 = n exp ln ( e 2 ) t 2 m l λ m i n ( G ) = exp − 2 t 2 ( 1 − ln 2 ) m l λ m i n ( G ) + ln n . We denote δ ~ ≜ exp ( − ( 1 − ln 2 ) m l λ min ( G ^ ) 2 t 2 + ln n ) \tilde{\delta} \triangleq\exp\left(-\frac{(1-\ln2)m_l\lambda_{\min}\left(\widehat{G}\right)}{2t^2}+\ln n\right) δ ~ ≜ exp ( − 2 t 2 ( 1 − l n 2 ) m l λ m i n ( G ) + ln n ) 1 − δ ~ 1-\tilde{\delta} 1 − δ ~
λ min ( A ^ ( l ) ( A ^ ( l ) ) T ) ≥ m l 2 λ min ( G ^ ) . (14) \lambda_{\min}\left(\widehat{A}^{(l)}\left(\widehat{A}^{(l)}\right)^{\mathsf{T}}\right) \ge\frac{m_l}{2}\lambda_{\min}\left(\widehat{G}\right). \tag*{(14)} λ m i n ( A ( l ) ( A ( l ) ) T ) ≥ 2 m l λ m i n ( G ) . (14) Now we observe that
∥ G ^ − G ∥ 2 = ∥ E g ∼ N ( 0 m l − 1 , ν l 2 α l I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ⋅ ( 1 − 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 ≤ t ] ) ] ∥ 2 = ∥ E g ∼ N ( 0 m l − 1 , ν l 2 α l I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] ] ∥ 2 \begin{aligned}
\left\|\widehat{G}-G\right\|_2 = \left\|\mathbb{E}_{g\sim\mathcal{N}\left(0_{m_{l-1}},\frac{\nu_l^2}{\alpha_l}I_{m_{l-1}}\right)}\left[\phi\left(s_lA^{(l-1)}g\right)\phi\left(s_lA^{(l-1)}g\right)^{\mathsf{T}}\cdot\left(1-\mathbb{1}_{\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2\le t\right]}\right)\right]\right\|_2 \\
&= \left\|\mathbb{E}_{g\sim\mathcal{N}\left(0_{m_{l-1}},\frac{\nu_l^2}{\alpha_l}I_{m_{l-1}}\right)}\left[\phi\left(s_lA^{(l-1)}g\right)\phi\left(s_lA^{(l-1)}g\right)^{\mathsf{T}}\mathbb{1}_{\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>t\right]}\right]\right\|_2
\end{aligned} G − G 2 = E g ∼ N ( 0 m l − 1 , α l ν l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ⋅ ( 1 − 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 ≤ t ] ) ] 2 = E g ∼ N ( 0 m l − 1 , α l ν l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] ] 2 [ Jensen ] ≤ E g ∼ N ( 0 m l − 1 , ν l 2 α l I m l − 1 ) [ ∥ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] ∥ 2 ] [\text{Jensen}] \le\mathbb{E}_{g\sim\mathcal{N}\left(0_{m_{l-1}},\frac{\nu_l^2}{\alpha_l}I_{m_{l-1}}\right)}\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\phi\left(s_lA^{(l-1)}g\right)^{\mathsf{T}}\mathbb{1}_{\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>t\right]}\right\|_2\right] [ Jensen ] ≤ E g ∼ N ( 0 m l − 1 , α l ν l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] 2 ]
= E g ∼ N ( 0 m l − 1 , ν l 2 α l I m l − 1 ) [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] ] = ∫ s = 0 ∞ P ( ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] > s ) d s = ∫ s = 0 ∞ P ( ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] > s ) d s , \begin{aligned}
= \mathbb{E}_{g\sim\mathcal{N}\left(0_{m_{l-1}},\frac{\nu_l^2}{\alpha_l}I_{m_{l-1}}\right)}\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2^2\mathbb{1}_{\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>t\right]}\right] \\
&= \int_{s=0}^{\infty}\mathbb{P}\left(\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2^2\mathbb{1}_{\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>t\right]}>s\right)\,ds \\
&= \int_{s=0}^{\infty}\mathbb{P}\left(\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2\mathbb{1}_{\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>t\right]}>\sqrt{s}\right)\,ds,
\end{aligned} = E g ∼ N ( 0 m l − 1 , α l ν l 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) 2 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] ] = ∫ s = 0 ∞ P ( ϕ ( s l A ( l − 1 ) g ) 2 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] > s ) d s = ∫ s = 0 ∞ P ( ϕ ( s l A ( l − 1 ) g ) 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] > s ) d s , where Jensen could be used since the operator norm is convex. This is the first place our proof deviates from that of Banerjee et al. [9 ] , by correcting the first transition in the following:
∥ G ^ − G ∥ 2 ≤ ∫ s = 0 ∞ P ( ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] > s ) d s = ∫ 0 t 2 P ( ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ) d s + ∫ t 2 ∞ P ( ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > s ) d s , \begin{aligned}
\left\|\widehat{G}-G\right\|_2 \le\int_{s=0}^{\infty}\mathbb{P}\left(\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2\mathbb{1}_{\left[\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>t\right]}>\sqrt{s}\right)\,ds \\
&= \int_0^{t^2}\mathbb{P}\left(\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>t\right)\,ds+\int_{t^2}^{\infty}\mathbb{P}\left(\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2>\sqrt{s}\right)\,ds,
\end{aligned} G − G 2 ≤ ∫ s = 0 ∞ P ( ϕ ( s l A ( l − 1 ) g ) 2 1 [ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 > t ] > s ) d s = ∫ 0 t 2 P ( ϕ ( s l A ( l − 1 ) g ) 2 > t ) d s + ∫ t 2 ∞ P ( ϕ ( s l A ( l − 1 ) g ) 2 > s ) d s , since when 0 < s < t 0<\sqrt{s}<t 0 < s < t x 1 [ x > t ] > s x\mathbb{1}_{[x>t]}>\sqrt{s} x 1 [ x > t ] > s x > t x>t x > t s ≥ t \sqrt{s}\ge t s ≥ t x 1 [ x > t ] > s x\mathbb{1}_{[x>t]}>\sqrt{s} x 1 [ x > t ] > s x > s x>\sqrt{s} x > s
From Proposition D.2 we have that ∥ A : , j ( l ) ∥ 2 = ∥ ϕ ( s l A ( l − 1 ) w j ( l ) ) ∥ 2 \left\|A_{:,j}^{(l)}\right\|_2=\left\|\phi\left(s_lA^{(l-1)}w_j^{(l)}\right)\right\|_2 A : , j ( l ) 2 = ϕ ( s l A ( l − 1 ) w j ( l ) ) 2 ∥ ∥ A : , j ( l ) ∥ 2 ∥ ψ 2 ≤ c 2 α ϕ ν 0 s l α l ∥ A ( l − 1 ) ∥ F \left\|\left\|A_{:,j}^{(l)}\right\|_2\right\|_{\psi_2}\le c\frac{\sqrt{2\alpha_\phi\nu_0s_l}}{\sqrt{\alpha_l}}\left\|A^{(l-1)}\right\|_{\mathrm{F}} A : , j ( l ) 2 ψ 2 ≤ c α l 2 α ϕ ν 0 s l A ( l − 1 ) F g g g w j ( l ) w_j^{(l)} w j ( l ) ∥ ∥ ϕ ( s l A ( l − 1 ) g ) ∥ 2 ∥ ψ 2 ≤ c 2 α ϕ ν 0 s l α l ∥ A ( l − 1 ) ∥ F \left\|\left\|\phi\left(s_lA^{(l-1)}g\right)\right\|_2\right\|_{\psi_2}\le c\frac{\sqrt{2\alpha_\phi\nu_0s_l}}{\sqrt{\alpha_l}}\left\|A^{(l-1)}\right\|_{\mathrm{F}} ϕ ( s l A ( l − 1 ) g ) 2 ψ 2 ≤ c α l 2 α ϕ ν 0 s l A ( l − 1 ) F b ≜ c 2 α ϕ ν 0 s l α l ∥ A ( l − 1 ) ∥ F b\triangleq c\frac{\sqrt{2\alpha_\phi\nu_0s_l}}{\sqrt{\alpha_l}}\left\|A^{(l-1)}\right\|_{\mathrm{F}} b ≜ c α l 2 α ϕ ν 0 s l A ( l − 1 ) F
∥ G ^ − G ∥ 2 ≤ ∫ 0 t 2 P ( ∥ ϕ ( s A ( l − 1 ) g ) ∥ 2 > t ) d s + ∫ t 2 ∞ P ( ∥ ϕ ( s A ( l − 1 ) g ) ∥ 2 > s ) d s \left\|\widehat{G}-G\right\|_{2}\leq\int_{0}^{t^{2}}\mathbb{P}\left(\left\|\phi\left(sA^{(l-1)}g\right)\right\|_{2}>t\right)\,ds+\int_{t^{2}}^{\infty}\mathbb{P}\left(\left\|\phi\left(sA^{(l-1)}g\right)\right\|_{2}>\sqrt{s}\right)\,ds G − G 2 ≤ ∫ 0 t 2 P ( ϕ ( s A ( l − 1 ) g ) 2 > t ) d s + ∫ t 2 ∞ P ( ϕ ( s A ( l − 1 ) g ) 2 > s ) d s ≤ 2 t 2 exp ( − c 1 b 2 t 2 ) + ∫ t 2 ∞ 2 exp ( − c 1 b 2 s ) d s = 2 t 2 exp ( − c 1 b 2 t 2 ) + 2 b 2 c 1 exp ( − c 1 b 2 t 2 ) = 2 ( t 2 + b 2 c 1 ) exp ( − c 1 b 2 t 2 ) . \begin{aligned}
\leq2t^{2}\exp\left(-\frac{c_{1}}{b^{2}}t^{2}\right)+\int_{t^{2}}^{\infty}2\exp\left(-\frac{c_{1}}{b^{2}}s\right)\,ds \\
&=2t^{2}\exp\left(-\frac{c_{1}}{b^{2}}t^{2}\right)+2\frac{b^{2}}{c_{1}}\exp\left(-\frac{c_{1}}{b^{2}}t^{2}\right) \\
&=2\left(t^{2}+\frac{b^{2}}{c_{1}}\right)\exp\left(-\frac{c_{1}}{b^{2}}t^{2}\right).
\end{aligned} ≤ 2 t 2 exp ( − b 2 c 1 t 2 ) + ∫ t 2 ∞ 2 exp ( − b 2 c 1 s ) d s = 2 t 2 exp ( − b 2 c 1 t 2 ) + 2 c 1 b 2 exp ( − b 2 c 1 t 2 ) = 2 ( t 2 + c 1 b 2 ) exp ( − b 2 c 1 t 2 ) . Let
R l ≜ max { e , 4 b 2 λ l c 1 } , t ≜ b c 1 2 ln ( 8 R l ) . R_l\triangleq\max\left\{e,\frac{4b^{2}}{\lambda_l c_{1}}\right\},\qquad t\triangleq\frac{b}{\sqrt{c_{1}}}\sqrt{2\ln(8R_l)}. R l ≜ max { e , λ l c 1 4 b 2 } , t ≜ c 1 b 2 ln ( 8 R l ) . Then all logarithms and roots are well-defined. Using the bound above, and writing
u ≜ c 1 t 2 b 2 = 2 ln ( 8 R l ) , u\triangleq\frac{c_{1}t^{2}}{b^{2}}=2\ln(8R_l), u ≜ b 2 c 1 t 2 = 2 ln ( 8 R l ) , we get
∥ G ^ − G ∥ 2 ≤ 2 ( t 2 + b 2 c 1 ) exp ( − c 1 b 2 t 2 ) = 2 b 2 c 1 ( u + 1 ) e − u = 2 b 2 c 1 ( 2 ln ( 8 R l ) + 1 ) 1 ( 8 R l ) 2 . \begin{aligned}
\left\|\widehat{G}-G\right\|_{2}\leq2\left(t^{2}+\frac{b^{2}}{c_{1}}\right)\exp\left(-\frac{c_{1}}{b^{2}}t^{2}\right) \\
&=2\frac{b^{2}}{c_{1}}(u+1)e^{-u} \\
&=2\frac{b^{2}}{c_{1}}\left(2\ln(8R_l)+1\right)\frac{1}{(8R_l)^{2}}.
\end{aligned} G − G 2 ≤ 2 ( t 2 + c 1 b 2 ) exp ( − b 2 c 1 t 2 ) = 2 c 1 b 2 ( u + 1 ) e − u = 2 c 1 b 2 ( 2 ln ( 8 R l ) + 1 ) ( 8 R l ) 2 1 . Since R l ≥ e R_l\geq e R l ≥ e
2 ln ( 8 R l ) + 1 ≤ 8 R l . 2\ln(8R_l)+1\leq8R_l. 2 ln ( 8 R l ) + 1 ≤ 8 R l . Hence,
∥ G ^ − G ∥ 2 ≤ 2 b 2 c 1 8 R l 64 R l 2 = 1 4 b 2 c 1 R l . \left\|\widehat{G}-G\right\|_{2}\leq2\frac{b^{2}}{c_{1}}\frac{8R_l}{64R_l^{2}}=\frac{1}{4}\frac{b^{2}}{c_{1}R_l}. G − G 2 ≤ 2 c 1 b 2 64 R l 2 8 R l = 4 1 c 1 R l b 2 . Moreover, by definition of R l R_l R l
R l ≥ 4 b 2 λ l c 1 , ⟹ b 2 c 1 R l ≤ λ l 4 . R_l\geq\frac{4b^{2}}{\lambda_l c_{1}},\qquad\Longrightarrow\qquad\frac{b^{2}}{c_{1}R_l}\leq\frac{\lambda_l}{4}. R l ≥ λ l c 1 4 b 2 , ⟹ c 1 R l b 2 ≤ 4 λ l . Therefore,
∥ G ^ − G ∥ 2 ≤ λ l 16 ≤ λ l 2 . \left\|\widehat{G}-G\right\|_{2}\leq\frac{\lambda_l}{16}\leq\frac{\lambda_l}{2}. G − G 2 ≤ 16 λ l ≤ 2 λ l . From the definition of G G G
λ l = λ min ( G ) , \lambda_l=\lambda_{\min}(G), λ l = λ m i n ( G ) , so Weyl’s inequality gives
λ min ( G ^ ) ≥ λ min ( G ) − ∥ G ^ − G ∥ 2 ≥ λ l 2 . \lambda_{\min}\left(\widehat{G}\right)\geq\lambda_{\min}(G)-\left\|\widehat{G}-G\right\|_{2}\geq\frac{\lambda_l}{2}. λ m i n ( G ) ≥ λ m i n ( G ) − G − G 2 ≥ 2 λ l . From (13) , (14) , we then have with probability at least 1 − δ ~ 1-\widetilde{\delta} 1 − δ
λ min ( A ( l ) ( A ( l ) ) T ) ≥ λ min ( A ^ ( l ) ( A ^ ( l ) ) T ) ≥ m l 2 λ min ( G ^ ) ≥ m l λ l 4 . \lambda_{\min}\left(A^{(l)}\left(A^{(l)}\right)^{\mathsf{T}}\right)\geq\lambda_{\min}\left(\widehat{A}^{(l)}\left(\widehat{A}^{(l)}\right)^{\mathsf{T}}\right)\geq\frac{m_l}{2}\lambda_{\min}\left(\widehat{G}\right)\geq\frac{m_l\lambda_l}{4}. λ m i n ( A ( l ) ( A ( l ) ) T ) ≥ λ m i n ( A ( l ) ( A ( l ) ) T ) ≥ 2 m l λ m i n ( G ) ≥ 4 m l λ l . where
δ ~ = exp ( − ( 1 − ln 2 ) m l λ min ( G ^ ) 2 t 2 + ln n ) ≤ exp ( − ( 1 − ln 2 ) m l λ l 4 t 2 + ln n ) = exp ( − ( 1 − ln 2 ) c 1 8 m l λ l b 2 ln ( 8 R l ) + ln n ) ≤ exp ( − c 2 m l λ l b 2 max [ 1 , ln ( 4 b 2 λ l c 1 ) ] + ln n ) , \begin{aligned}
\widetilde{\delta}
&= \exp\left(-\frac{(1-\ln2)m_l\lambda_{\min}(\widehat{G})}{2t^2}+\ln n\right)\\
&\le\exp\left(-\frac{(1-\ln2)m_l\lambda_l}{4t^2}+\ln n\right)\\
&= \exp\left(-\frac{(1-\ln2)c_1}{8}\frac{m_l\lambda_l}{b^2\ln(8R_l)}+\ln n\right)\\
&\le\exp\left(-c_2\frac{m_l\lambda_l}{b^2\max\left[1,\ln\left(\frac{4b^2}{\lambda_l c_1}\right)\right]}+\ln n\right),
\end{aligned} δ = exp ( − 2 t 2 ( 1 − ln 2 ) m l λ m i n ( G ) + ln n ) ≤ exp ( − 4 t 2 ( 1 − ln 2 ) m l λ l + ln n ) = exp ( − 8 ( 1 − ln 2 ) c 1 b 2 ln ( 8 R l ) m l λ l + ln n ) ≤ exp − c 2 b 2 max [ 1 , ln ( λ l c 1 4 b 2 ) ] m l λ l + ln n , for an absolute constant c 2 > 0 c_2>0 c 2 > 0
ln ( 8 R l ) = O ( max [ 1 , ln ( 4 b 2 λ l c 1 ) ] ) . \ln(8R_l)=O\left(\max\left[1,\ln\left(\frac{4b^2}{\lambda_l c_1}\right)\right]\right). ln ( 8 R l ) = O ( max [ 1 , ln ( λ l c 1 4 b 2 ) ] ) . We now demand δ ~ ≤ δ L \widetilde{\delta}\le\frac{\delta}{L} δ ≤ L δ
m l ≥ 1 c 2 ln ( n L δ ) b 2 λ l max [ 1 , ln ( 4 b 2 λ l c 1 ) ] . m_l\ge\frac{1}{c_2}\ln\left(\frac{nL}{\delta}\right)\frac{b^2}{\lambda_l}\max\left[1,\ln\left(\frac{4b^2}{\lambda_l c_1}\right)\right]. m l ≥ c 2 1 ln ( δ n L ) λ l b 2 max [ 1 , ln ( λ l c 1 4 b 2 ) ] . Substituting
b = c 2 α ϕ ϑ 0 s l a l ∥ A ( l − 1 ) ∥ F , b=c\frac{\sqrt{2\alpha_\phi\vartheta_0s_l}}{\sqrt{a_l}}\left\|\mathbf{A}^{(l-1)}\right\|_{\mathrm{F}}, b = c a l 2 α ϕ ϑ 0 s l A ( l − 1 ) F , we obtain the sufficient condition
m l ≳ ln ( n L δ ) 1 λ l s l 2 a l ∥ A ( l − 1 ) ∥ F 2 max [ 1 , ln ( 1 λ l s l 2 a l ∥ A ( l − 1 ) ∥ F 2 ) ] , m_l\gtrsim\ln\left(\frac{nL}{\delta}\right)\frac{1}{\lambda_l}\frac{s_l^2}{a_l}\left\|\mathbf{A}^{(l-1)}\right\|_{\mathrm{F}}^2\max\left[1,\ln\left(\frac{1}{\lambda_l}\frac{s_l^2}{a_l}\left\|\mathbf{A}^{(l-1)}\right\|_{\mathrm{F}}^2\right)\right], m l ≳ ln ( δ n L ) λ l 1 a l s l 2 A ( l − 1 ) F 2 max [ 1 , ln ( λ l 1 a l s l 2 A ( l − 1 ) F 2 ) ] , after absorbing constants depending only on α ϕ , ϑ 0 , c , c 1 \alpha_\phi,\vartheta_0,c,c_1 α ϕ , ϑ 0 , c , c 1 m l ≥ n m_l\ge n m l ≥ n m l m_l m l δ ~ ≤ δ / L \widetilde{\delta}\le\delta/L δ ≤ δ / L 1 − δ / L 1-\delta/L 1 − δ / L
We now proceed to prove an analogue of their [9 ] Lemma 4.2, again adjusting for general parameterization and, importantly, improving the bound’s probability.
First, we need to restate their Proposition A.2, and adjust the proof to our modified settings. Specifically, in this proof we differ from theirs by the generalized gradient bound α ϕ \alpha_\phi α ϕ
Proposition D.4. With c ϕ , σ 0 ≜ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z ) ] c_{\phi,\sigma_0}\triangleq\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(z)] c ϕ , σ 0 ≜ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z )] ϕ ( 0 ) = 0 \phi(0)=0 ϕ ( 0 ) = 0 ∣ ϕ ˙ ∣ ≤ α ϕ \lvert\dot{\phi}\rvert\le\alpha_\phi ∣ ϕ ˙ ∣ ≤ α ϕ ζ > 0 \zeta>0 ζ > 0
c ϕ , σ 0 − α ϕ 2 σ 0 2 ∣ ζ 2 − 1 ∣ ≤ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( ζ z ) ] ≤ c ϕ , σ 0 + α ϕ 2 σ 0 2 ∣ ζ 2 − 1 ∣ . c_{\phi,\sigma_0}-\alpha_\phi^2\sigma_0^2\lvert\zeta^2-1\rvert
\le\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(\zeta z)]
\le c_{\phi,\sigma_0}+\alpha_\phi^2\sigma_0^2\lvert\zeta^2-1\rvert. c ϕ , σ 0 − α ϕ 2 σ 0 2 ∣ ζ 2 − 1 ∣ ≤ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( ζ z )] ≤ c ϕ , σ 0 + α ϕ 2 σ 0 2 ∣ ζ 2 − 1 ∣ . Proof. Observe that the bounded gradient implies ϕ \phi ϕ α ϕ \alpha_\phi α ϕ ϕ ( 0 ) = 0 \phi(0)=0 ϕ ( 0 ) = 0 ∣ ϕ ( ⋅ ) ∣ ≤ α ϕ ∣ ⋅ ∣ \lvert\phi(\cdot)\rvert\le\alpha_\phi\lvert\cdot\rvert ∣ ϕ ( ⋅ )∣ ≤ α ϕ ∣ ⋅ ∣
∣ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( ζ z ) ] − E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z ) ] ∣ = ∣ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( ζ z ) − ϕ 2 ( z ) ] ∣ ≤ [Jensen] E z ∼ N ( 0 , σ 0 2 ) [ ∣ ϕ 2 ( ζ z ) − ϕ 2 ( z ) ∣ ] = E z ∼ N ( 0 , σ 0 2 ) [ ∣ ϕ ( ζ z ) − ϕ ( z ) ∣ ∣ ϕ ( ζ z ) + ϕ ( z ) ∣ ] ≤ [ ϕ is α ϕ -Lipschitz ] E z ∼ N ( 0 , σ 0 2 ) [ α ϕ ∣ z ∣ ∣ ζ − 1 ∣ ∣ ϕ ( ζ z ) + ϕ ( z ) ∣ ] ≤ [Triangle Inequality] α ϕ ∣ ζ − 1 ∣ E z ∼ N ( 0 , σ 0 2 ) [ ∣ z ∣ ( ∣ ϕ ( ζ z ) ∣ + ∣ ϕ ( z ) ∣ ) ] ≤ α ϕ ∣ ζ − 1 ∣ E z ∼ N ( 0 , σ 0 2 ) [ ∣ z ∣ α ϕ ( ∣ ζ z ∣ + ∣ z ∣ ) ] = α ϕ 2 ∣ ζ − 1 ∣ ( ∣ ζ ∣ + 1 ) E z ∼ N ( 0 , σ 0 2 ) [ ∣ z ∣ 2 ] = [ ζ > 0 ] α ϕ 2 σ 0 2 ∣ ζ 2 − 1 ∣ . \begin{aligned}
\left|\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(\zeta z)]-\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(z)]\right|
&=\left|\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(\zeta z)-\phi^2(z)]\right|\\
&\mathrel{\underset{\text{[Jensen]}}{\le}}\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\left|\phi^2(\zeta z)-\phi^2(z)\right|\right]\\
&=\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\left|\phi(\zeta z)-\phi(z)\right|\left|\phi(\zeta z)+\phi(z)\right|\right]\\
&\mathrel{\underset{\left[\phi\text{ is }\alpha_\phi\text{-Lipschitz}\right]}{\le}}\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\alpha_\phi\lvert z\rvert\lvert\zeta-1\rvert\lvert\phi(\zeta z)+\phi(z)\rvert\right]\\
&\mathrel{\underset{\text{[Triangle Inequality]}}{\le}}\alpha_\phi\lvert\zeta-1\rvert\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\lvert z\rvert\left(\lvert\phi(\zeta z)\rvert+\lvert\phi(z)\rvert\right)\right]\\
&\le\alpha_\phi\lvert\zeta-1\rvert\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\lvert z\rvert\alpha_\phi\left(\lvert\zeta z\rvert+\lvert z\rvert\right)\right]\\
&=\alpha_\phi^2\lvert\zeta-1\rvert(\lvert\zeta\rvert+1)\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\lvert z\rvert^2\right]\\
&\mathrel{\underset{[\zeta>0]}{=}}\alpha_\phi^2\sigma_0^2\lvert\zeta^2-1\rvert.
\end{aligned} E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( ζ z )] − E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z )] = E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( ζ z ) − ϕ 2 ( z )] [Jensen] ≤ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( ζ z ) − ϕ 2 ( z ) ] = E z ∼ N ( 0 , σ 0 2 ) [ ∣ ϕ ( ζ z ) − ϕ ( z ) ∣ ∣ ϕ ( ζ z ) + ϕ ( z ) ∣ ] [ ϕ is α ϕ -Lipschitz ] ≤ E z ∼ N ( 0 , σ 0 2 ) [ α ϕ ∣ z ∣ ∣ ζ − 1 ∣ ∣ ϕ ( ζ z ) + ϕ ( z )∣ ] [Triangle Inequality] ≤ α ϕ ∣ ζ − 1 ∣ E z ∼ N ( 0 , σ 0 2 ) [ ∣ z ∣ ( ∣ ϕ ( ζ z )∣ + ∣ ϕ ( z )∣ ) ] ≤ α ϕ ∣ ζ − 1 ∣ E z ∼ N ( 0 , σ 0 2 ) [ ∣ z ∣ α ϕ ( ∣ ζ z ∣ + ∣ z ∣ ) ] = α ϕ 2 ∣ ζ − 1 ∣ (∣ ζ ∣ + 1 ) E z ∼ N ( 0 , σ 0 2 ) [ ∣ z ∣ 2 ] [ ζ > 0 ] = α ϕ 2 σ 0 2 ∣ ζ 2 − 1 ∣ . Since E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z ) ] = c ϕ , σ 0 \mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}[\phi^2(z)]=c_{\phi,\sigma_0} E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( z )] = c ϕ , σ 0
Proposition D.5. Let h C ( l ) ≜ ∑ i = 0 l ( α l v 0 ) 2 h_C(l)\triangleq\sum_{i=0}^{l}(\alpha_l v_0)^2 h C ( l ) ≜ ∑ i = 0 l ( α l v 0 ) 2 ∀ l ∈ [ L ] \forall l\in[L] ∀ l ∈ [ L ] a l s l 2 = m l − 1 \frac{a_l}{s_l^2}=m_{l-1} s l 2 a l = m l − 1 16 m ≜ min l ∈ [ L ] m l = Ω ( h C ( L ) max [ h C ( L ) , 1 ] ln ( n L ) ) m\triangleq\min_{l\in[L]}m_l=\Omega\left(h_C(L)\max[h_C(L),1]\ln(nL)\right) m ≜ min l ∈ [ L ] m l = Ω ( h C ( L ) max [ h C ( L ) , 1 ] ln ( n L ) ) 1 − 2 n L exp ( − c 3 m c ϕ , σ 0 max [ h C ( L ) , 1 ] ) 1-2nL\exp\left(-c_3\frac{m}{c_{\phi,\sigma_0}\max[h_C(L),1]}\right) 1 − 2 n L exp ( − c 3 c ϕ , σ 0 m a x [ h C ( L ) , 1 ] m ) { W ( l ) } l ∈ [ L ] \{\mathbf{W}^{(l)}\}_{l\in[L]} { W ( l ) } l ∈ [ L ] i ∈ [ n ] , l ∈ [ L ] i\in[n],l\in[L] i ∈ [ n ] , l ∈ [ L ]
c ϕ , σ 0 2 m l ≤ ∥ α ( l ) ( x i ) ∥ 2 2 ≤ 3 c ϕ , σ 0 2 m l . \frac{c_{\phi,\sigma_0}}{2}m_l\le\left\|\boldsymbol{\alpha}^{(l)}(\mathbf{x}_i)\right\|_2^2\le\frac{3c_{\phi,\sigma_0}}{2}m_l. 2 c ϕ , σ 0 m l ≤ α ( l ) ( x i ) 2 2 ≤ 2 3 c ϕ , σ 0 m l . Proof. We start by following their proof exactly, using induction over the layers to prove the following holds for all i ∈ [ n ] , l ∈ [ L ] i\in[n],l\in[L] i ∈ [ n ] , l ∈ [ L ]
c ϕ , σ 0 m l ( 1 − h C ( l ) 2 h C ( L ) ) ≤ ∥ α ( l ) ( x i ) ∥ 2 2 ≤ c ϕ , σ 0 m l ( 1 + h C ( l ) 2 h C ( L ) ) . c_{\phi,\sigma_0}m_l\left(1-\frac{h_C(l)}{2h_C(L)}\right)\le\left\|\boldsymbol{\alpha}^{(l)}(\mathbf{x}_i)\right\|_2^2\le c_{\phi,\sigma_0}m_l\left(1+\frac{h_C(l)}{2h_C(L)}\right). c ϕ , σ 0 m l ( 1 − 2 h C ( L ) h C ( l ) ) ≤ α ( l ) ( x i ) 2 2 ≤ c ϕ , σ 0 m l ( 1 + 2 h C ( L ) h C ( l ) ) . For l = 0 l=0 l = 0 ∥ α ( 0 ) ( x i ) ∥ 2 2 = ∥ x i ∥ 2 2 = c ϕ , σ 0 d = c ϕ , σ 0 m 0 \left\|\boldsymbol{\alpha}^{(0)}(\mathbf{x}_i)\right\|_2^2=\left\|\mathbf{x}_i\right\|_2^2=c_{\phi,\sigma_0}d=c_{\phi,\sigma_0}m_0 α ( 0 ) ( x i ) 2 2 = ∥ x i ∥ 2 2 = c ϕ , σ 0 d = c ϕ , σ 0 m 0 l ∈ [ L ] l\in[L] l ∈ [ L ] l − 1 l-1 l − 1 i ∈ [ n ] i\in[n] i ∈ [ n ]
c ϕ , σ 0 m l − 1 ( 1 − h C ( l − 1 ) 2 h C ( L ) ) ≤ ∥ α ( l − 1 ) ( x i ) ∥ 2 2 ≤ c ϕ , σ 0 m l − 1 ( 1 + h C ( l − 1 ) 2 h C ( L ) ) c_{\phi,\sigma_0}m_{l-1}\left(1-\frac{h_C(l-1)}{2h_C(L)}\right)\le\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2^2\le c_{\phi,\sigma_0}m_{l-1}\left(1+\frac{h_C(l-1)}{2h_C(L)}\right) c ϕ , σ 0 m l − 1 ( 1 − 2 h C ( L ) h C ( l − 1 ) ) ≤ α ( l − 1 ) ( x i ) 2 2 ≤ c ϕ , σ 0 m l − 1 ( 1 + 2 h C ( L ) h C ( l − 1 ) ) ∣ ∥ α ( l − 1 ) ( x i ) ∥ 2 2 c ϕ , σ 0 m l − 1 − 1 ∣ ≤ h C ( l − 1 ) 2 h C ( L ) . (15) \left|\frac{\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2^2}{c_{\phi,\sigma_0}m_{l-1}}-1\right|\le\frac{h_C(l-1)}{2h_C(L)}. \tag*{(15)} c ϕ , σ 0 m l − 1 α ( l − 1 ) ( x i ) 2 2 − 1 ≤ 2 h C ( L ) h C ( l − 1 ) . (15) Now we observe for all i ∈ [ n ] i\in[n] i ∈ [ n ]
E W ( l ) [ ∥ α ( l ) ( x i ) ∥ 2 2 ] \mathbb{E}_{\mathbf{W}^{(l)}}\left[\left\|\boldsymbol{\alpha}^{(l)}(\mathbf{x}_i)\right\|_2^2\right] E W ( l ) [ α ( l ) ( x i ) 2 2 ] = E W ( l ) [ ∑ j ∈ [ m l ] ( α j ( l ) ( x i ) ) 2 ] = ∑ j ∈ [ m l ] E w j ( l ) [ ( α j ( l ) ( x i ) ) 2 ] =\mathbb{E}_{\mathbf{W}^{(l)}}\left[\sum_{j\in[m_l]}\left(\alpha_j^{(l)}(\mathbf{x}_i)\right)^2\right]
=\sum_{j\in[m_l]}\mathbb{E}_{w_j^{(l)}}\left[\left(\alpha_j^{(l)}(\mathbf{x}_i)\right)^2\right] = E W ( l ) j ∈ [ m l ] ∑ ( α j ( l ) ( x i ) ) 2 = j ∈ [ m l ] ∑ E w j ( l ) [ ( α j ( l ) ( x i ) ) 2 ] = m l E g ∼ N ( 0 , σ 0 2 a l I m l − 1 ) [ ( ϕ ( s l g T α ( l − 1 ) ( x i ) ) ) 2 ] =m_l\mathbb{E}_{g\sim\mathcal{N}\left(0,\frac{\sigma_0^2}{a_l}\mathbf{I}_{m_{l-1}}\right)}\left[\left(\phi\left(s_l\mathbf{g}^{\mathsf{T}}\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right)\right)^2\right] = m l E g ∼ N ( 0 , a l σ 0 2 I m l − 1 ) [ ( ϕ ( s l g T α ( l − 1 ) ( x i ) ) ) 2 ] = m l E g ∼ N ( 0 , σ 0 2 m l − 1 c ϕ , σ 0 a l I m l − 1 ) [ ( ϕ ( s l ∥ α ( l − 1 ) ( x i ) ∥ 2 ⋅ g T α ( l − 1 ) ( x i ) ∥ α ( l − 1 ) ( x i ) ∥ 2 ) ) 2 ] =m_l\mathbb{E}_{g\sim\mathcal{N}\left(0,\frac{\sigma_0^2}{m_{l-1}c_{\phi,\sigma_0}a_l}\mathbf{I}_{m_{l-1}}\right)}\left[\left(\phi\left(s_l\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2\cdot\mathbf{g}^{\mathsf{T}}\frac{\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)}{\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2}\right)\right)^2\right] = m l E g ∼ N ( 0 , m l − 1 c ϕ , σ 0 a l σ 0 2 I m l − 1 ) ( ϕ ( s l α ( l − 1 ) ( x i ) 2 ⋅ g T α ( l − 1 ) ( x i ) 2 α ( l − 1 ) ( x i ) ) ) 2 = m l E z ∼ N ( 0 , σ 0 2 c ϕ , σ 0 a l ) [ ( ϕ ( s l ∥ α ( l − 1 ) ( x i ) ∥ 2 z ) ) 2 ] =m_l\mathbb{E}_{z\sim\mathcal{N}\left(0,\frac{\sigma_0^2}{c_{\phi,\sigma_0}a_l}\right)}\left[\left(\phi\left(s_l\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2z\right)\right)^2\right] = m l E z ∼ N ( 0 , c ϕ , σ 0 a l σ 0 2 ) [ ( ϕ ( s l α ( l − 1 ) ( x i ) 2 z ) ) 2 ] = m l E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( s l ∥ α ( l − 1 ) ( x i ) ∥ 2 c ϕ , σ 0 a l z ) ] . =m_l\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\phi^2\left(\frac{s_l\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2}{\sqrt{c_{\phi,\sigma_0}a_l}}z\right)\right]. = m l E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( c ϕ , σ 0 a l s l α ( l − 1 ) ( x i ) 2 z ) ] . Then, from Proposition D.4, we have
∣ E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( s l ∥ α ( l − 1 ) ( x i ) ∥ 2 c ϕ , σ 0 a l z ) ] − c ϕ , σ 0 ∣ ≤ α ϕ 2 σ 0 2 ( s l ∥ α ( l − 1 ) ( x i ) ∥ 2 c ϕ , σ 0 a l ) 2 − 1 \left|\mathbb{E}_{z\sim\mathcal{N}(0,\sigma_0^2)}\left[\phi^2\left(\frac{s_l\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2}{\sqrt{c_{\phi,\sigma_0}a_l}}z\right)\right]-c_{\phi,\sigma_0}\right|
\le\alpha_\phi^2\sigma_0^2\left(\frac{s_l\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2}{\sqrt{c_{\phi,\sigma_0}a_l}}\right)^2-1 E z ∼ N ( 0 , σ 0 2 ) [ ϕ 2 ( c ϕ , σ 0 a l s l α ( l − 1 ) ( x i ) 2 z ) ] − c ϕ , σ 0 ≤ α ϕ 2 σ 0 2 ( c ϕ , σ 0 a l s l α ( l − 1 ) ( x i ) 2 ) 2 − 1 ∣ E W ( l ) [ ∥ α ( l ) ( x i ) ∥ 2 2 ] − m l c ϕ , σ 0 ∣ ≤ m l α ϕ 2 σ 0 2 ∣ s l 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 2 c ϕ , σ 0 a l − 1 ∣ . \left|\mathbb{E}_{\mathbf{W}^{(l)}}\left[\left\|\boldsymbol{\alpha}^{(l)}(\mathbf{x}_i)\right\|_2^2\right]-m_lc_{\phi,\sigma_0}\right|
\le m_l\alpha_\phi^2\sigma_0^2\left|\frac{s_l^2\left\|\boldsymbol{\alpha}^{(l-1)}(\mathbf{x}_i)\right\|_2^2}{c_{\phi,\sigma_0}a_l}-1\right|. E W ( l ) [ α ( l ) ( x i ) 2 2 ] − m l c ϕ , σ 0 ≤ m l α ϕ 2 σ 0 2 c ϕ , σ 0 a l s l 2 α ( l − 1 ) ( x i ) 2 2 − 1 . From a l s l 2 = m l − 1 \frac{a_l}{s_l^2}=m_{l-1} s l 2 a l = m l − 1 (15) , and get:
m l ( c ϕ , σ 0 − α ϕ 2 σ 0 2 h C ( l − 1 ) 2 h C ( L ) ) ≤ E w ( l ) [ ∥ α ( l ) ( x i ) ∥ 2 2 ] ≤ m l ( c ϕ , σ 0 + α ϕ 2 σ 0 2 h C ( l − 1 ) 2 h C ( L ) ) . (16) m_l\left(c_{\phi,\sigma_0}-\alpha_\phi^2\sigma_0^2\frac{h_C(l-1)}{2h_C(L)}\right)\leq\mathbb{E}_{\mathbf{w}^{(l)}}\left[\left\|\boldsymbol{\alpha}^{(l)}(x_i)\right\|_2^2\right]\leq m_l\left(c_{\phi,\sigma_0}+\alpha_\phi^2\sigma_0^2\frac{h_C(l-1)}{2h_C(L)}\right). \tag*{(16)} m l ( c ϕ , σ 0 − α ϕ 2 σ 0 2 2 h C ( L ) h C ( l − 1 ) ) ≤ E w ( l ) [ α ( l ) ( x i ) 2 2 ] ≤ m l ( c ϕ , σ 0 + α ϕ 2 σ 0 2 2 h C ( L ) h C ( l − 1 ) ) . (16) By Proposition D.1, for all l ∈ [ L ] l\in[L] l ∈ [ L ] j ∈ [ m l ] j\in[m_l] j ∈ [ m l ] i ∈ [ n ] i\in[n] i ∈ [ n ] ∣ α j ( l ) ( x i ) ∣ \left|\alpha_j^{(l)}(x_i)\right| α j ( l ) ( x i ) A ( l − 1 ) \mathbf{A}^{(l-1)} A ( l − 1 ) ∣ α j ( l ) ( x i ) ∣ 2 \left|\alpha_j^{(l)}(x_i)\right|^2 α j ( l ) ( x i ) 2 K ≜ ∥ α j ( l ) ( x i ) ∥ ψ 1 2 = ∥ ∥ α j ( l ) ( x i ) ∥ ψ 2 2 ∥ ψ 2 K\triangleq\left\|\alpha_j^{(l)}(x_i)\right\|_{\psi_1}^2=\left\|\left\|\alpha_j^{(l)}(x_i)\right\|_{\psi_2}^2\right\|_{\psi_2} K ≜ α j ( l ) ( x i ) ψ 1 2 = α j ( l ) ( x i ) ψ 2 2 ψ 2 [111 Lemma 2.8.5] . By applying the induction hypothesis and using a l s l 2 = m l − 1 \frac{a_l}{s_l^2}=m_{l-1} s l 2 a l = m l − 1
K = ∥ ∣ α j ( l ) ( x i ) ∣ 2 ∥ ψ 2 ≤ 2 c ϕ 2 σ 0 2 s l 2 a l ∥ α ( l − 1 ) ( x i ) ∥ 2 2 ≤ 2 c ϕ 2 σ 1 2 a l σ 0 2 c ϕ , σ 0 ⋅ 3 2 c ϕ , σ 0 a l s l 2 = 3 c ϕ 2 σ 0 2 . \begin{aligned}
K=\left\|\left|\alpha_j^{(l)}(x_i)\right|^2\right\|_{\psi_2}\leq\frac{2c_{\phi}^2\sigma_0^2s_l^2}{a_l}\left\|\alpha^{(l-1)}(x_i)\right\|_2^2 \\
&\leq\frac{2c_{\phi}^2\sigma_1^2}{a_l}\frac{\sigma_0^2}{c_{\phi,\sigma_0}}\cdot\frac{3}{2}c_{\phi,\sigma_0}\frac{a_l}{s_l^2}=3c_{\phi}^2\sigma_0^2.
\end{aligned} K = α j ( l ) ( x i ) 2 ψ 2 ≤ a l 2 c ϕ 2 σ 0 2 s l 2 α ( l − 1 ) ( x i ) 2 2 ≤ a l 2 c ϕ 2 σ 1 2 c ϕ , σ 0 σ 0 2 ⋅ 2 3 c ϕ , σ 0 s l 2 a l = 3 c ϕ 2 σ 0 2 . Conditioned on A ( l − 1 ) \mathbf{A}^{(l-1)} A ( l − 1 ) ∣ α j ( l ) ( x i ) ∣ 2 \left|\alpha_j^{(l)}(x_i)\right|^2 α j ( l ) ( x i ) 2 j ∈ [ m l ] j\in[m_l] j ∈ [ m l ] i ∈ [ n ] i\in[n] i ∈ [ n ] [111 Theorem 2.9.1] :
P ( ∣ ∑ j = 1 m l ( ∣ α j ( l ) ( x i ) ∣ 2 − E [ ∣ α j ( l ) ( x i ) ∣ 2 ] ) ∣ ≥ t ) ≤ 2 exp ( − c 2 min ( t 2 m l K 2 , t K ) ) \mathbb{P}\left(\left|\sum_{j=1}^{m_l}\left(\left|\alpha_j^{(l)}(x_i)\right|^2-\mathbb{E}\left[\left|\alpha_j^{(l)}(x_i)\right|^2\right]\right)\right|\geq t\right)\leq2\exp\left(-c_2\min\left(\frac{t^2}{m_lK^2},\frac{t}{K}\right)\right) P ( j = 1 ∑ m l ( α j ( l ) ( x i ) 2 − E [ α j ( l ) ( x i ) 2 ] ) ≥ t ) ≤ 2 exp ( − c 2 min ( m l K 2 t 2 , K t ) ) P ( ∣ ∑ j = 1 m l ( ∣ α j ( l ) ( x i ) ∣ 2 ) − E [ ∑ j = 1 m l ∣ α j ( l ) ( x i ) ∣ 2 ] ∣ ≥ t ) ≤ 2 exp ( − c 2 min ( t 2 m l K 2 , t K ) ) \mathbb{P}\left(\left|\sum_{j=1}^{m_l}\left(\left|\alpha_j^{(l)}(x_i)\right|^2\right)-\mathbb{E}\left[\sum_{j=1}^{m_l}\left|\alpha_j^{(l)}(x_i)\right|^2\right]\right|\geq t\right)\leq2\exp\left(-c_2\min\left(\frac{t^2}{m_lK^2},\frac{t}{K}\right)\right) P ( j = 1 ∑ m l ( α j ( l ) ( x i ) 2 ) − E [ j = 1 ∑ m l α j ( l ) ( x i ) 2 ] ≥ t ) ≤ 2 exp ( − c 2 min ( m l K 2 t 2 , K t ) ) P ( ∣ ∥ α ( l ) ( x i ) ∥ 2 2 − E ∥ α ( l ) ( x i ) ∥ 2 2 ∣ ≥ t ) ≤ 2 exp ( − c 2 min ( t 2 m l K 2 , t K ) ) , \mathbb{P}\left(\left|\left\|\alpha^{(l)}(x_i)\right\|_2^2-\mathbb{E}\left\|\alpha^{(l)}(x_i)\right\|_2^2\right|\geq t\right)\leq2\exp\left(-c_2\min\left(\frac{t^2}{m_lK^2},\frac{t}{K}\right)\right), P ( α ( l ) ( x i ) 2 2 − E α ( l ) ( x i ) 2 2 ≥ t ) ≤ 2 exp ( − c 2 min ( m l K 2 t 2 , K t ) ) , Combining this with Eq. (16) , we get for all i ∈ [ n ] i\in[n] i ∈ [ n ] 1 − 2 exp ( − c 2 min ( t 2 m l K 2 , t K ) ) 1-2\exp\left(-c_2\min\left(\frac{t^2}{m_lK^2},\frac{t}{K}\right)\right) 1 − 2 exp ( − c 2 min ( m l K 2 t 2 , K t ) )
∣ ∥ α ( l ) ( x i ) ∥ 2 2 − m l c ϕ , σ 0 ∣ ≤ m l α ϕ 2 σ 0 2 h C ( l − 1 ) 2 h C ( L ) + t . \left|\left\|\alpha^{(l)}(x_i)\right\|_2^2-m_lc_{\phi,\sigma_0}\right|\leq m_l\alpha_\phi^2\sigma_0^2\frac{h_C(l-1)}{2h_C(L)}+t. α ( l ) ( x i ) 2 2 − m l c ϕ , σ 0 ≤ m l α ϕ 2 σ 0 2 2 h C ( L ) h C ( l − 1 ) + t . This is where our proof diverges from the original proof. By selecting t = c ϕ , σ 0 m l 2 h C ( L ) t=\frac{c_{\phi,\sigma_0}m_l}{2h_C(L)} t = 2 h C ( L ) c ϕ , σ 0 m l ν 0 2 = σ 0 2 c ϕ , σ 0 \nu_0^2=\frac{\sigma_0^2}{c_{\phi,\sigma_0}} ν 0 2 = c ϕ , σ 0 σ 0 2
∣ ∥ α ( l ) ( x i ) ∥ 2 2 − m l c ϕ , σ 0 ∣ ≤ m l c ϕ , σ 0 ( α ϕ 2 ν 0 2 h C ( l − 1 ) 2 h C ( L ) + 1 2 h C ( L ) ) . \left|\left\|\alpha^{(l)}(x_i)\right\|_2^2-m_lc_{\phi,\sigma_0}\right|\leq m_lc_{\phi,\sigma_0}\left(\alpha_\phi^2\nu_0^2\frac{h_C(l-1)}{2h_C(L)}+\frac{1}{2h_C(L)}\right). α ( l ) ( x i ) 2 2 − m l c ϕ , σ 0 ≤ m l c ϕ , σ 0 ( α ϕ 2 ν 0 2 2 h C ( L ) h C ( l − 1 ) + 2 h C ( L ) 1 ) . Now applying h C ( l ) = 1 + α ϕ 2 ν 0 2 h C ( l − 1 ) h_C(l)=1+\alpha_\phi^2\nu_0^2h_C(l-1) h C ( l ) = 1 + α ϕ 2 ν 0 2 h C ( l − 1 )
m l c ϕ , σ 0 ( 1 − h C ( l ) 2 h C ( L ) ) ≤ ∥ α ( l ) ( x i ) ∥ 2 2 ≤ m l c ϕ , σ 0 ( 1 + h C ( l ) 2 h C ( L ) ) . m_lc_{\phi,\sigma_0}\left(1-\frac{h_C(l)}{2h_C(L)}\right)\leq\left\|\alpha^{(l)}(x_i)\right\|_2^2\leq m_lc_{\phi,\sigma_0}\left(1+\frac{h_C(l)}{2h_C(L)}\right). m l c ϕ , σ 0 ( 1 − 2 h C ( L ) h C ( l ) ) ≤ α ( l ) ( x i ) 2 2 ≤ m l c ϕ , σ 0 ( 1 + 2 h C ( L ) h C ( l ) ) . The probability of this event is at least
1 − 2 exp ( − c 2 min ( t 2 m l K 2 , t K ) ) = 1 − 2 exp ( − c 2 min ( c ϕ , σ 0 2 m l 2 4 h C ( L ) 2 m l ⋅ 9 c ϕ 4 σ 0 4 , c ϕ , σ 0 m l 2 h C ( L ) ⋅ 3 c ϕ 2 σ 0 2 ) ) = 1 − 2 exp ( − c 2 m l 6 c ϕ 2 h C ( L ) ν 0 2 min ( 1 6 c ϕ 2 h C ( L ) ν 0 2 , 1 ) ) \begin{aligned}
1-2\exp\left(-c_2\min\left(\frac{t^2}{m_lK^2},\frac{t}{K}\right)\right) \\
&=1-2\exp\left(-c_2\min\left(\frac{c_{\phi,\sigma_0}^2m_l^2}{4h_C(L)^2m_l\cdot9c_\phi^4\sigma_0^4},\frac{c_{\phi,\sigma_0}m_l}{2h_C(L)\cdot3c_\phi^2\sigma_0^2}\right)\right) \\
&=1-2\exp\left(-c_2\frac{m_l}{6c_\phi^2h_C(L)\nu_0^2}\min\left(\frac{1}{6c_\phi^2h_C(L)\nu_0^2},1\right)\right)
\end{aligned} 1 − 2 exp ( − c 2 min ( m l K 2 t 2 , K t ) ) = 1 − 2 exp ( − c 2 min ( 4 h C ( L ) 2 m l ⋅ 9 c ϕ 4 σ 0 4 c ϕ , σ 0 2 m l 2 , 2 h C ( L ) ⋅ 3 c ϕ 2 σ 0 2 c ϕ , σ 0 m l ) ) = 1 − 2 exp ( − c 2 6 c ϕ 2 h C ( L ) ν 0 2 m l min ( 6 c ϕ 2 h C ( L ) ν 0 2 1 , 1 ) ) ≥ 1 − 2 exp ( − c 3 m l h C ( L ) max [ h C ( L ) , 1 ] ) , \geq1-2\exp\left(-c_3\frac{m_l}{h_C(L)\max\left[h_C(L),1\right]}\right), ≥ 1 − 2 exp ( − c 3 h C ( L ) max [ h C ( L ) , 1 ] m l ) , for some constant c 3 > 0 c_3 > 0 c 3 > 0 ∀ l ∈ [ L ] \forall l \in[L] ∀ l ∈ [ L ] m l ≥ m m_l \ge m m l ≥ m i ∈ [ n ] i \in[n] i ∈ [ n ]
1 − 2 n L exp ( − c 3 m h C ( L ) max [ h C ( L ) , 1 ] ) . 1 - 2nL \exp\left(-c_3 \frac{m}{h_C(L)\max[h_C(L),1]}\right). 1 − 2 n L exp ( − c 3 h C ( L ) max [ h C ( L ) , 1 ] m ) . Similarly to Banerjee et al. [9 ] , this result leads to an upper bound on ∥ A ( l ) ∥ F 2 \left\|A^{(l)}\right\|_{\mathrm{F}}^2 A ( l ) F 2 l ∈ [ L ] l \in[L] l ∈ [ L ]
∥ A ( l ) ∥ F 2 = ∑ i ∈ [ n ] ∥ α ( l ) ( x i ) ∥ 2 2 ≤ 3 2 c ϕ , σ 0 n m l , \left\|A^{(l)}\right\|_{\mathrm{F}}^2 = \sum_{i\in[n]}\left\|\alpha^{(l)}(x_i)\right\|_2^2 \le\frac{3}{2}c_{\phi,\sigma_0}nm_l, A ( l ) F 2 = i ∈ [ n ] ∑ α ( l ) ( x i ) 2 2 ≤ 2 3 c ϕ , σ 0 n m l , with probability at least 1 − 2 n L exp ( − c 3 m h C ( L ) max [ h C ( L ) , 1 ] ) 1 - 2nL \exp\left(-c_3 \frac{m}{h_C(L)\max[h_C(L),1]}\right) 1 − 2 n L exp ( − c 3 h C ( L ) m a x [ h C ( L ) , 1 ] m ) m l m_l m l l ∈ [ L ] l \in[L] l ∈ [ L ]
m l = Ω ( max [ n , ln ( n L δ ) 1 λ l s l 2 a l m l − 1 max [ 1 , ln ( 1 λ l s l 2 a l m l − 1 ) ] ] ) = Ω ( max [ n , ln ( n L δ ) n λ l max [ 1 , ln ( n λ l ) ] ] ) . \begin{aligned}
m_l = \Omega\left(\max\left[n,\ln\left(\frac{nL}{\delta}\right)\frac{1}{\lambda_l}\frac{s_l^2}{a_l}m_{l-1}\max\left[1,\ln\left(\frac{1}{\lambda_l}\frac{s_l^2}{a_l}m_{l-1}\right)\right]\right]\right) \\
&= \Omega\left(\max\left[n,\ln\left(\frac{nL}{\delta}\right)\frac{n}{\lambda_l}\max\left[1,\ln\left(\frac{n}{\lambda_l}\right)\right]\right]\right).
\end{aligned} m l = Ω ( max [ n , ln ( δ n L ) λ l 1 a l s l 2 m l − 1 max [ 1 , ln ( λ l 1 a l s l 2 m l − 1 ) ] ] ) = Ω ( max [ n , ln ( δ n L ) λ l n max [ 1 , ln ( λ l n ) ] ] ) . Since we assumed L = O ( 1 ) L = O(1) L = O ( 1 ) h C ( L ) = O ( 1 ) h_C(L) = O(1) h C ( L ) = O ( 1 ) m ≜ min l ∈ [ L ] m l m \triangleq\min_{l\in[L]}m_l m ≜ min l ∈ [ L ] m l 1 − exp ( − c 4 m ) 1 - \exp(-c_4m) 1 − exp ( − c 4 m )
Their Lemma 4.3 can be used without proving again, since by assuming ∀ l ∈ [ L ] \forall l \in[L] ∀ l ∈ [ L ] a l s l 2 = m l − 1 \frac{a_l}{s_l^2}=m_{l-1} s l 2 a l = m l − 1 λ l \lambda_l λ l l ∈ [ L ] l \in[L] l ∈ [ L ]
λ l ≜ λ min ( E g ∼ N ( 0 m l − 1 , v 0 2 a l I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] ) = λ min ( E g ∼ N ( 0 m l − 1 , v 0 2 I m l − 1 ) [ ϕ ( s l a l A ( l − 1 ) g ) ϕ ( s l a l A ( l − 1 ) g ) T ] ) = λ min ( E g ∼ N ( 0 m l − 1 , v 0 2 I m l − 1 ) [ ϕ ( 1 m l − 1 A ( l − 1 ) g ) ϕ ( 1 m l − 1 A ( l − 1 ) g ) T ] ) . \begin{aligned}
\lambda_l \triangleq\lambda_{\min}\left(\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\boldsymbol{0}_{m_{l-1}},\frac{v_0^2}{a_l}\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(s_lA^{(l-1)}\boldsymbol{g}\right)\phi\left(s_lA^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\right]\right) \\
&= \lambda_{\min}\left(\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\boldsymbol{0}_{m_{l-1}},v_0^2\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(\frac{s_l}{\sqrt{a_l}}A^{(l-1)}\boldsymbol{g}\right)\phi\left(\frac{s_l}{\sqrt{a_l}}A^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\right]\right) \\
&= \lambda_{\min}\left(\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\boldsymbol{0}_{m_{l-1}},v_0^2\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(\frac{1}{\sqrt{m_{l-1}}}A^{(l-1)}\boldsymbol{g}\right)\phi\left(\frac{1}{\sqrt{m_{l-1}}}A^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\right]\right).
\end{aligned} λ l ≜ λ m i n E g ∼ N ( 0 m l − 1 , a l v 0 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] = λ m i n ( E g ∼ N ( 0 m l − 1 , v 0 2 I m l − 1 ) [ ϕ ( a l s l A ( l − 1 ) g ) ϕ ( a l s l A ( l − 1 ) g ) T ] ) = λ m i n ( E g ∼ N ( 0 m l − 1 , v 0 2 I m l − 1 ) [ ϕ ( m l − 1 1 A ( l − 1 ) g ) ϕ ( m l − 1 1 A ( l − 1 ) g ) T ] ) . We also note that their proof of their Lemma 4.3 relies on having the events of Proposition D.3, Proposition D.5 occur for every layer l ∈ [ L ] l \in[L] l ∈ [ L ]
Proposition D.6. Assume m l = Ω ( max [ n , ln ( n L δ ) n λ l max [ 1 , ln ( n λ l ) ] ] ) m_l = \Omega\left(\max\left[n,\ln\left(\frac{nL}{\delta}\right)\frac{n}{\lambda_l}\max\left[1,\ln\left(\frac{n}{\lambda_l}\right)\right]\right]\right) m l = Ω ( max [ n , ln ( δ n L ) λ l n max [ 1 , ln ( λ l n ) ] ] ) ∀ l ∈ [ L ] \forall l \in[L] ∀ l ∈ [ L ] l ∈ [ L ] l \in[L] l ∈ [ L ]
λ l ≜ λ min ( E g ∼ N ( 0 m l − 1 , v 0 2 a l I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] ) ≥ c ~ λ 1 , \lambda_l \triangleq\lambda_{\min}\left(\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\boldsymbol{0}_{m_{l-1}},\frac{v_0^2}{a_l}\mathbf{I}_{m_{l-1}}\right)}\left[\phi\left(s_lA^{(l-1)}\boldsymbol{g}\right)\phi\left(s_lA^{(l-1)}\boldsymbol{g}\right)^{\mathsf{T}}\right]\right) \ge\tilde{c}\lambda_1, λ l ≜ λ m i n E g ∼ N ( 0 m l − 1 , a l v 0 2 I m l − 1 ) [ ϕ ( s l A ( l − 1 ) g ) ϕ ( s l A ( l − 1 ) g ) T ] ≥ c ~ λ 1 , with probability at least 1 − δ − exp ( − c 4 m ) 1-\delta-\exp(-c_4m) 1 − δ − exp ( − c 4 m ) c ~ l \tilde{c}_l c ~ l c ~ l = Θ ( 1 ) \tilde{c}_l = \Theta(1) c ~ l = Θ ( 1 )
λ 1 ≜ λ min ( E g ∼ N ( 0 d , v 0 2 a 1 I d ) [ ϕ ( s 1 X g ) ϕ ( s 1 X g ) T ] ) . \lambda_1 \triangleq\lambda_{\min}\left(\mathbb{E}_{\boldsymbol{g}\sim\mathcal{N}\left(\boldsymbol{0}_d,\frac{v_0^2}{a_1}\mathbf{I}_d\right)}\left[\phi\left(s_1\mathbf{X}\boldsymbol{g}\right)\phi\left(s_1\mathbf{X}\boldsymbol{g}\right)^{\mathsf{T}}\right]\right). λ 1 ≜ λ m i n E g ∼ N ( 0 d , a 1 v 0 2 I d ) [ ϕ ( s 1 X g ) ϕ ( s 1 X g ) T ] . Conclusion
We are now ready to prove the main result of this section:
Lemma D.7. For m l = Ω ( max [ n , ln ( n L δ ) n λ 1 ln ( n λ 1 ) ] ) m_l = \Omega\left(\max\left[n,\ln\left(\frac{nL}{\delta}\right)\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\right]\right) m l = Ω ( max [ n , ln ( δ n L ) λ 1 n ln ( λ 1 n ) ] ) ∀ l ∈ [ L ] \forall l \in[L] ∀ l ∈ [ L ] 1 − δ − exp ( − c 1 m ) 1-\delta-\exp(-c_1m) 1 − δ − exp ( − c 1 m ) { W ( l ) } l ∈ [ L ′ ] \{\mathbf{W}^{(l)}\}_{l\in[L']} { W ( l ) } l ∈ [ L ′ ]
λ min ( K N T K ( θ ) ) ≥ c 0 s L 2 m λ 1 . \lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta)\right) \ge c_0s_L^2m\lambda_1. λ m i n ( K NTK ( θ ) ) ≥ c 0 s L 2 m λ 1 . Proof. First, we note that λ 1 \lambda_1 λ 1 [85 ] ):
λ 1 = λ min ( E g ∼ N ( 0 , v 0 2 a 1 2 I d ) [ ϕ ( s 1 X g ) ϕ ( s 1 X g ) T ] ) ≤ 1 n Tr ( E g ∼ N ( 0 , v 0 2 a 1 2 I d ) [ ϕ ( s 1 X g ) ϕ ( s 1 X g ) T ] ) = 1 n E g ∼ N ( 0 , v 0 2 a 1 2 I d ) ∥ ϕ ( s 1 X g ) ∥ 2 ≤ 1 n s 1 2 α ϕ 2 E g ∼ N ( 0 , v 0 2 a 1 2 I d ) ∥ X g ∥ 2 = 1 n s 1 2 α ϕ 2 n c ϕ , σ 0 E z ∼ N ( 0 , v 0 2 a 1 2 ) ∣ z ∣ 2 = α ϕ 2 v 0 2 c ϕ , σ 0 s 1 2 d a 1 = α ϕ 2 v 0 2 c ϕ , σ 0 = O ( 1 ) . (17) \begin{aligned}
\lambda_1 &= \lambda_{\min}\left(\mathbb{E}_{\mathbf{g}\sim\mathcal{N}\left(\mathbf{0},\frac{v_0^2}{a_1^2}\mathbf{I}_d\right)}\left[\phi(s_1\mathbf{Xg})\phi(s_1\mathbf{Xg})^{\mathsf{T}}\right]\right)\\
&\le\frac{1}{n}\operatorname{Tr}\left(\mathbb{E}_{\mathbf{g}\sim\mathcal{N}\left(\mathbf{0},\frac{v_0^2}{a_1^2}\mathbf{I}_d\right)}\left[\phi(s_1\mathbf{Xg})\phi(s_1\mathbf{Xg})^{\mathsf{T}}\right]\right)\\
&= \frac{1}{n}\mathbb{E}_{\mathbf{g}\sim\mathcal{N}\left(\mathbf{0},\frac{v_0^2}{a_1^2}\mathbf{I}_d\right)}\left\|\phi(s_1\mathbf{Xg})\right\|^2\\
&\le\frac{1}{n}s_1^2\alpha_\phi^2\mathbb{E}_{\mathbf{g}\sim\mathcal{N}\left(\mathbf{0},\frac{v_0^2}{a_1^2}\mathbf{I}_d\right)}\|\mathbf{Xg}\|^2
= \frac{1}{n}s_1^2\alpha_\phi^2 n c_{\phi,\sigma_0}\mathbb{E}_{z\sim\mathcal{N}\left(0,\frac{v_0^2}{a_1^2}\right)}|z|^2\\
&= \alpha_\phi^2v_0^2c_{\phi,\sigma_0}\frac{s_1^2d}{a_1}
= \alpha_\phi^2v_0^2c_{\phi,\sigma_0}=O(1). \tag*{(17)}
\end{aligned} λ 1 = λ m i n E g ∼ N ( 0 , a 1 2 v 0 2 I d ) [ ϕ ( s 1 Xg ) ϕ ( s 1 Xg ) T ] ≤ n 1 Tr E g ∼ N ( 0 , a 1 2 v 0 2 I d ) [ ϕ ( s 1 Xg ) ϕ ( s 1 Xg ) T ] = n 1 E g ∼ N ( 0 , a 1 2 v 0 2 I d ) ∥ ϕ ( s 1 Xg ) ∥ 2 ≤ n 1 s 1 2 α ϕ 2 E g ∼ N ( 0 , a 1 2 v 0 2 I d ) ∥ Xg ∥ 2 = n 1 s 1 2 α ϕ 2 n c ϕ , σ 0 E z ∼ N ( 0 , a 1 2 v 0 2 ) ∣ z ∣ 2 = α ϕ 2 v 0 2 c ϕ , σ 0 a 1 s 1 2 d = α ϕ 2 v 0 2 c ϕ , σ 0 = O ( 1 ) . (17) Then, ln ( n λ 1 ) ≥ 1 \ln\left(\frac{n}{\lambda_1}\right)\ge1 ln ( λ 1 n ) ≥ 1 n n n
By plugging Proposition D.6 into Proposition D.3, we have for all l ∈ [ L ] l\in[L] l ∈ [ L ]
λ min ( A ( l ) ( A ( l ) ) T ) ≥ m λ l 4 ≥ c 0 m λ 1 , \lambda_{\min}\left(A^{(l)}\left(A^{(l)}\right)^{\mathsf{T}}\right)\ge\frac{m\lambda_l}{4}\ge c_0m\lambda_1, λ m i n ( A ( l ) ( A ( l ) ) T ) ≥ 4 m λ l ≥ c 0 m λ 1 , with probability at least 1 − δ − exp ( − c 1 m ) 1-\delta-\exp(-c_1m) 1 − δ − exp ( − c 1 m ) [9 ] , we observe that
K N T K ( θ ) = ∇ f ( θ ) ∇ f ( θ ) T = ∑ l ∈ [ L ] [ ∂ f ( θ ) ∂ vec ( W ( l ) ) ] [ ∂ f ( θ ) ∂ vec ( W ( l ) ) ] T + [ ∂ f ( θ ) ∂ v ] [ ∂ f ( θ ) ∂ v ] T \mathbf{K}_{\mathrm{NTK}}(\theta)=\nabla f(\theta)\nabla f(\theta)^{\mathsf{T}}=\sum_{l\in[L]}\left[\frac{\partial f(\theta)}{\partial\operatorname{vec}\left(\mathbf{W}^{(l)}\right)}\right]\left[\frac{\partial f(\theta)}{\partial\operatorname{vec}\left(\mathbf{W}^{(l)}\right)}\right]^{\mathsf{T}}+\left[\frac{\partial f(\theta)}{\partial\mathbf{v}}\right]\left[\frac{\partial f(\theta)}{\partial\mathbf{v}}\right]^{\mathsf{T}} K NTK ( θ ) = ∇ f ( θ ) ∇ f ( θ ) T = l ∈ [ L ] ∑ [ ∂ vec ( W ( l ) ) ∂ f ( θ ) ] [ ∂ vec ( W ( l ) ) ∂ f ( θ ) ] T + [ ∂ v ∂ f ( θ ) ] [ ∂ v ∂ f ( θ ) ] T ≽ [ ∂ f ( θ ) ∂ v ] [ ∂ f ( θ ) ∂ v ] T = s L + 1 2 A ( L ) ( A ( L ) ) T , \succcurlyeq\left[\frac{\partial f(\theta)}{\partial\mathbf{v}}\right]\left[\frac{\partial f(\theta)}{\partial\mathbf{v}}\right]^{\mathsf{T}}=s_{L+1}^2A^{(L)}\left(A^{(L)}\right)^{\mathsf{T}}, ≽ [ ∂ v ∂ f ( θ ) ] [ ∂ v ∂ f ( θ ) ] T = s L + 1 2 A ( L ) ( A ( L ) ) T , so, by the Rayleigh-Ritz characterization of the minimum eigenvalue of a symmetric matrix,
λ min ( K N T K ( θ ) ) = min ∥ u ∥ = 1 u T K N T K ( θ ) u = u ⋆ T K N T K ( θ ) u ⋆ ≥ u ⋆ T s L + 1 2 A ( L ) ( A ( L ) ) T u ⋆ ≥ s L + 1 2 λ min ( A ( L ) ( A ( L ) ) T ) ≥ c 0 s L + 1 2 m λ 1 . \begin{aligned}
\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta)\right)
&=\min_{\|\mathbf{u}\|=1}\mathbf{u}^{\mathsf{T}}\mathbf{K}_{\mathrm{NTK}}(\theta)\mathbf{u}
=\mathbf{u}_\star^{\mathsf{T}}\mathbf{K}_{\mathrm{NTK}}(\theta)\mathbf{u}_\star\\
&\ge\mathbf{u}_\star^{\mathsf{T}}s_{L+1}^2A^{(L)}\left(A^{(L)}\right)^{\mathsf{T}}\mathbf{u}_\star
\ge s_{L+1}^2\lambda_{\min}\left(A^{(L)}\left(A^{(L)}\right)^{\mathsf{T}}\right)\\
&\ge c_0s_{L+1}^2m\lambda_1.
\end{aligned} λ m i n ( K NTK ( θ ) ) = ∥ u ∥ = 1 min u T K NTK ( θ ) u = u ⋆ T K NTK ( θ ) u ⋆ ≥ u ⋆ T s L + 1 2 A ( L ) ( A ( L ) ) T u ⋆ ≥ s L + 1 2 λ m i n ( A ( L ) ( A ( L ) ) T ) ≥ c 0 s L + 1 2 m λ 1 . By choosing δ = exp ( − c ~ m λ 1 ln n λ 1 ) \delta=\exp\left(-\widetilde{c}\frac{m}{\lambda_1}\ln\frac{n}{\lambda_1}\right) δ = exp ( − c λ 1 m ln λ 1 n ) c ~ ∈ ( 0 , 1 ) \widetilde{c}\in(0,1) c ∈ ( 0 , 1 ) L = O ( 1 ) L=O(1) L = O ( 1 ) m = Ω ( n λ 1 ln n λ 1 ln n ) m=\Omega\left(\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\ln n\right) m = Ω ( λ 1 n ln λ 1 n ln n ) 1 − exp ( − c ~ m λ 1 ln n λ 1 ) − exp ( − c 1 m ) ≥ 1 − exp ( − c ~ 1 m λ 1 ln n λ 1 ) 1-\exp\left(-\widetilde{c}\frac{m}{\lambda_1}\ln\frac{n}{\lambda_1}\right)-\exp(-c_1m)\ge1-\exp\left(-\widetilde{c}_1\frac{m}{\lambda_1}\ln\frac{n}{\lambda_1}\right) 1 − exp ( − c λ 1 m ln λ 1 n ) − exp ( − c 1 m ) ≥ 1 − exp ( − c 1 λ 1 m ln λ 1 n ) c ~ 1 \widetilde{c}_1 c 1
Auxiliary claims
Claim E.1. Let A ∈ R n × m A \in\mathbb{R}^{n \times m} A ∈ R n × m B ∈ R n × m B \in\mathbb{R}^{n \times m} B ∈ R n × m B B B A A A
λ min ( B B T ) ≤ λ min ( A A T ) . \lambda_{\min}(BB^{\mathsf{T}}) \leq\lambda_{\min}(AA^{\mathsf{T}}). λ m i n ( B B T ) ≤ λ m i n ( A A T ) . Proof. Since each column of B B B A A A
S = diag ( s 1 , … , s n ) , s i ∈ { 0 , 1 } , S = \operatorname{diag}(s_1,\ldots,s_n), \qquad s_i \in\{0,1\}, S = diag ( s 1 , … , s n ) , s i ∈ { 0 , 1 } , such that
Then
B B T = ( A S ) ( A S ) T = A S S T A T = A S 2 A T = A S A T , BB^{\mathsf{T}} = (AS)(AS)^{\mathsf{T}} = ASS^{\mathsf{T}}A^{\mathsf{T}} = AS^2A^{\mathsf{T}} = ASA^{\mathsf{T}}, B B T = ( A S ) ( A S ) T = A S S T A T = A S 2 A T = A S A T , where we used the facts that S T = S S^{\mathsf{T}} = S S T = S S 2 = S S^2 = S S 2 = S
A A T − B B T = A A T − A S A T = A ( I − S ) A T . AA^{\mathsf{T}} - BB^{\mathsf{T}} = AA^{\mathsf{T}} - ASA^{\mathsf{T}} = A(I-S)A^{\mathsf{T}}. A A T − B B T = A A T − A S A T = A ( I − S ) A T . The matrix I − S I-S I − S 1 − s i ∈ { 0 , 1 } 1-s_i \in\{0,1\} 1 − s i ∈ { 0 , 1 } x ∈ R m x \in\mathbb{R}^{m} x ∈ R m
x T A ( I − S ) A T x = ( A T x ) T ( I − S ) ( A T x ) ≥ 0 , x^{\mathsf{T}}A(I-S)A^{\mathsf{T}}x = (A^{\mathsf{T}}x)^{\mathsf{T}}(I-S)(A^{\mathsf{T}}x) \geq0, x T A ( I − S ) A T x = ( A T x ) T ( I − S ) ( A T x ) ≥ 0 , leading to
A A T ⪰ B B T . AA^{\mathsf{T}} \succeq BB^{\mathsf{T}}. A A T ⪰ B B T . Finally, by the Rayleigh-Ritz characterization of the minimum eigenvalue of a symmetric matrix,
λ min ( M ) = min ∥ x ∥ = 1 x T M x . \lambda_{\min}(M) = \min_{\lVert x\rVert=1} x^{\mathsf{T}}Mx. λ m i n ( M ) = ∥ x ∥ = 1 min x T M x . Since B B T ⪯ A A T BB^{\mathsf{T}} \preceq AA^{\mathsf{T}} B B T ⪯ A A T x T B B T x ≤ x T A A T x x^{\mathsf{T}}BB^{\mathsf{T}}x \leq x^{\mathsf{T}}AA^{\mathsf{T}}x x T B B T x ≤ x T A A T x x x x
λ min ( B B T ) = min ∥ x ∥ = 1 x T B B T x ≤ min ∥ x ∥ = 1 x T A A T x = λ min ( A A T ) . \lambda_{\min}(BB^{\mathsf{T}}) = \min_{\lVert x\rVert=1} x^{\mathsf{T}}BB^{\mathsf{T}}x \leq\min_{\lVert x\rVert=1} x^{\mathsf{T}}AA^{\mathsf{T}}x = \lambda_{\min}(AA^{\mathsf{T}}). λ m i n ( B B T ) = ∥ x ∥ = 1 min x T B B T x ≤ ∥ x ∥ = 1 min x T A A T x = λ m i n ( A A T ) . Claim E.2. Let z ∼ N ( 0 m , I m ) z \sim\mathcal{N}(0_m,I_m) z ∼ N ( 0 m , I m ) p ≥ 1 p \geq1 p ≥ 1
∥ ∥ z ∥ 2 ∥ L p ≜ ( E [ ∥ z ∥ 2 p ] ) 1 / p ≤ C p + m , \lVert\lVert z\rVert_2\rVert_{L_p} \triangleq\left(\mathbb{E}\left[\lVert z\rVert_2^p\right]\right)^{1/p} \leq C\sqrt{p}+\sqrt{m}, ∥∥ z ∥ 2 ∥ L p ≜ ( E [ ∥ z ∥ 2 p ] ) 1/ p ≤ C p + m , where C > 0 C>0 C > 0
Proof. From Vershynin [111 Theorem 3.1.1 and Exercise 2.24(d)] we have ∥ ∥ z ∥ 2 − m ∥ ψ 2 ≤ 8 3 C 1 \lVert\lVert z\rVert_2-\sqrt{m}\rVert_{\psi_2}\leq\frac{8}{3}C_1 ∥∥ z ∥ 2 − m ∥ ψ 2 ≤ 3 8 C 1
∥ ∥ z ∥ 2 ∥ L p ≤ ∥ ∥ z ∥ 2 − m ∥ L p + ∥ m ∥ L p \lVert\lVert z\rVert_2\rVert_{L_p} \leq\lVert\lVert z\rVert_2-\sqrt{m}\rVert_{L_p}+\lVert\sqrt{m}\rVert_{L_p} ∥∥ z ∥ 2 ∥ L p ≤ ∥∥ z ∥ 2 − m ∥ L p + ∥ m ∥ L p ≤ ∥ ∥ z ∥ 2 − m ∥ L p + ∥ m ∥ L p [Minkowski’s inequality] \leq\lVert\lVert z\rVert_2-\sqrt{m}\rVert_{L_p}+\lVert\sqrt{m}\rVert_{L_p}
\qquad\text{[Minkowski’s inequality]} ≤ ∥∥ z ∥ 2 − m ∥ L p + ∥ m ∥ L p [Minkowski’s inequality] ≤ C 2 8 3 C 1 p + m citep‘[111, Proposition 2.6.6]‘ \leq C_2\frac{8}{3}C_1\sqrt{p}+\sqrt{m}
\qquad\text{{citep}`[111, Proposition 2.6.6]`} ≤ C 2 3 8 C 1 p + m citep ‘[111, Proposition 2.6.6]‘ = C p + m . = C\sqrt{p}+\sqrt{m}. = C p + m . Claim E.3. Let G ∈ R m × n G \in\mathbb{R}^{m \times n} G ∈ R m × n c > 0 c>0 c > 0 p ≥ 1 p \geq1 p ≥ 1
∥ ∥ G ∥ ∥ L p ≤ c ( m + n + p ) . \lVert\lVert G\rVert\rVert_{L_p} \leq c\left(\sqrt{m}+\sqrt{n}+\sqrt{p}\right). ∥∥ G ∥ ∥ L p ≤ c ( m + n + p ) . As a result, there exists a constant c ′ > 0 c'>0 c ′ > 0
∥ ∥ G ∥ 2 ∥ L p ≤ c ′ ( m + n + p ) . \lVert\lVert G\rVert^2\rVert_{L_p} \leq c'(m+n+p). ∥∥ G ∥ 2 ∥ L p ≤ c ′ ( m + n + p ) . Proof. From Vershynin [111, Theorem 4.4.3], there exists a constant c 1 c_1 c 1 t > 0 t>0 t > 0
P ( ∥ G ∥ ≥ c 1 ( m + n + t ) ) ≤ 2 exp ( − t 2 ) . \mathbb{P}\left(\lVert G\rVert\ge c_1\left(\sqrt{m}+\sqrt{n}+t\right)\right)\le2\exp\left(-t^2\right). P ( ∥ G ∥ ≥ c 1 ( m + n + t ) ) ≤ 2 exp ( − t 2 ) . Let Y ≜ max { 0 , ∥ G ∥ − c 1 ( m + n ) } Y\triangleq\max\{0,\lVert G\rVert-c_1(\sqrt{m}+\sqrt{n})\} Y ≜ max { 0 , ∥ G ∥ − c 1 ( m + n )}
P ( ∣ Y ∣ ≥ t ) ≤ 2 exp ( − t 2 / c 1 2 ) . \mathbb{P}(|Y|\ge t)\le2\exp\left(-t^2/c_1^2\right). P ( ∣ Y ∣ ≥ t ) ≤ 2 exp ( − t 2 / c 1 2 ) . Therefore, by Vershynin [111, Proposition 2.6.1], ∣ Y ∣ |Y| ∣ Y ∣ ∥ ∣ Y ∣ ∥ L p ≤ c 2 p \lVert|Y|\rVert_{L_p}\le c_2\sqrt{p} ∥ ∣ Y ∣ ∥ L p ≤ c 2 p c 2 > 0 c_2>0 c 2 > 0 p ≥ 1 p\ge1 p ≥ 1
∥ ∥ G ∥ ∥ L p ≤ ∥ ∣ Y ∣ + c 1 ( m + n ) ∥ L p ≤ c 1 ( m + n ) + ∥ ∣ Y ∣ ∥ L p ≤ max { c 1 , c 2 } ( m + n + p ) . \begin{aligned}
\lVert\lVert G\rVert\rVert_{L_p}\le\lVert|Y|+c_1(\sqrt{m}+\sqrt{n})\rVert_{L_p} \\
&\le c_1(\sqrt{m}+\sqrt{n})+\lVert|Y|\rVert_{L_p}\le\max\{c_1,c_2\}(\sqrt{m}+\sqrt{n}+\sqrt{p}).
\end{aligned} ∥∥ G ∥ ∥ L p ≤ ∥ ∣ Y ∣ + c 1 ( m + n ) ∥ L p ≤ c 1 ( m + n ) + ∥ ∣ Y ∣ ∥ L p ≤ max { c 1 , c 2 } ( m + n + p ) . Using Claim E.6 and applying the inequality ( a + b ) 2 ≤ 2 a 2 + 2 b 2 (a+b)^2\le2a^2+2b^2 ( a + b ) 2 ≤ 2 a 2 + 2 b 2
∥ ∥ G ∥ 2 ∥ L p ≤ ∥ ∥ G ∥ ∥ L 2 p 2 ≤ ( max { c 1 , c 2 } ) 2 ( m + n + 2 p ) 2 ≤ ( max { c 1 , c 2 } ) 2 ( 2 m + 4 n + 8 p ) . \begin{aligned}
\lVert\lVert G\rVert^2\rVert_{L_p}\le\lVert\lVert G\rVert\rVert_{L_{2p}}^2\le\left(\max\{c_1,c_2\}\right)^2\left(\sqrt{m}+\sqrt{n}+\sqrt{2p}\right)^2 \\
&\le\left(\max\{c_1,c_2\}\right)^2(2m+4n+8p).
\end{aligned} ∥∥ G ∥ 2 ∥ L p ≤ ∥∥ G ∥ ∥ L 2 p 2 ≤ ( max { c 1 , c 2 } ) 2 ( m + n + 2 p ) 2 ≤ ( max { c 1 , c 2 } ) 2 ( 2 m + 4 n + 8 p ) . Claim E.4 (Tower property for L p L^p L p X X X σ \sigma σ F \mathcal{F} F p > 0 p>0 p > 0
∥ X ∥ L p = ∥ ∥ X ∣ F ∥ L p ∥ L p . \lVert X\rVert_{L_p}=\lVert\lVert X\mid\mathcal{F}\rVert_{L_p}\rVert_{L_p}. ∥ X ∥ L p = ∥∥ X ∣ F ∥ L p ∥ L p . Proof. Observe
∥ ∥ X ∣ F ∥ L p ∥ L p = ( E [ ( ∥ X ∣ F ∥ L p ) p ] ) 1 / p = ( E [ ( ( E [ ∣ X ∣ p ∣ F ] ) 1 / p ) p ] ) 1 / p \lVert\lVert X\mid\mathcal{F}\rVert_{L_p}\rVert_{L_p}
=\left(\mathbb{E}\left[\left(\lVert X\mid\mathcal{F}\rVert_{L_p}\right)^p\right]\right)^{1/p}
=\left(\mathbb{E}\left[\left(\left(\mathbb{E}\left[|X|^p\mid\mathcal{F}\right]\right)^{1/p}\right)^p\right]\right)^{1/p} ∥∥ X ∣ F ∥ L p ∥ L p = ( E [ ( ∥ X ∣ F ∥ L p ) p ] ) 1/ p = ( E [ ( ( E [ ∣ X ∣ p ∣ F ] ) 1/ p ) p ] ) 1/ p = ( E [ E [ ∣ X ∣ p ∣ F ] ] ) 1 / p = ( E [ ∣ X ∣ p ] ) 1 / p = ∥ X ∥ L p . =\left(\mathbb{E}\left[\mathbb{E}\left[|X|^p\mid\mathcal{F}\right]\right]\right)^{1/p}
=\left(\mathbb{E}\left[|X|^p\right]\right)^{1/p}
=\lVert X\rVert_{L_p}. = ( E [ E [ ∣ X ∣ p ∣ F ] ] ) 1/ p = ( E [ ∣ X ∣ p ] ) 1/ p = ∥ X ∥ L p . Claim E.5. Suppose that, for some deterministic K > 0 K>0 K > 0 X X X σ \sigma σ F \mathcal{F} F
∥ X ∣ F ∥ ψ 2 ≤ K \lVert X\mid\mathcal{F}\rVert_{\psi_2}\le K ∥ X ∣ F ∥ ψ 2 ≤ K almost surely. Then,
∥ X ∥ ψ 2 ≤ c K . \lVert X\rVert_{\psi_2}\le cK. ∥ X ∥ ψ 2 ≤ cK . Proof. By Vershynin [111, Proposition 2.6.6], for all p ≥ 1 p\ge1 p ≥ 1
∥ X ∣ F ∥ L p ≤ C K p . \lVert X\mid\mathcal{F}\rVert_{L_p}\le CK\sqrt{p}. ∥ X ∣ F ∥ L p ≤ C K p . Then,
∥ ∥ X ∣ F ∥ L p ∥ L p ≤ ∥ C K p ∥ L p = C K p . \lVert\lVert X\mid\mathcal{F}\rVert_{L_p}\rVert_{L_p}\le\lVert CK\sqrt{p}\rVert_{L_p}=CK\sqrt{p}. ∥∥ X ∣ F ∥ L p ∥ L p ≤ ∥ C K p ∥ L p = C K p . By Claim E.4,
∥ X ∥ L p = ∥ ∥ X ∣ F ∥ L p ∥ L p ≤ C K p , \lVert X\rVert_{L_p}=\lVert\lVert X\mid\mathcal{F}\rVert_{L_p}\rVert_{L_p}\le CK\sqrt{p}, ∥ X ∥ L p = ∥∥ X ∣ F ∥ L p ∥ L p ≤ C K p , for all p ≥ 1 p\ge1 p ≥ 1
∥ X ∥ ψ 2 ≤ c K . \lVert X\rVert_{\psi_2}\le cK. ∥ X ∥ ψ 2 ≤ cK . Claim E.6. Given some RV X X X p > 0 p>0 p > 0 k > 0 k>0 k > 0
∥ X k ∥ L p = ∥ X ∥ L k p k . \lVert X^k\rVert_{L_p}=\lVert X\rVert_{L_{kp}}^k. ∥ X k ∥ L p = ∥ X ∥ L k p k . Proof. Observe
∥ X k ∥ L p = ( E [ ( X k ) p ] ) 1 / p = ( ( E [ X k p ] ) 1 / ( k p ) ) k = ∥ X ∥ L k p k . \lVert X^k\rVert_{L_p}
=\left(\mathbb{E}\left[(X^k)^p\right]\right)^{1/p}
=\left(\left(\mathbb{E}\left[X^{kp}\right]\right)^{1/(kp)}\right)^k
=\lVert X\rVert_{L_{kp}}^k. ∥ X k ∥ L p = ( E [ ( X k ) p ] ) 1/ p = ( ( E [ X k p ] ) 1/ ( k p ) ) k = ∥ X ∥ L k p k . Claim E.7 (Hölder’s inequality for L p L^p L p For random variables X X X Y Y Y p > 0 p>0 p > 0 r , s ≥ 1 r,s\geq1 r , s ≥ 1 1 r + 1 s = 1 \frac{1}{r}+\frac{1}{s}=1 r 1 + s 1 = 1
∥ X Y ∥ L p ≤ ∥ X ∥ L p r ∥ Y ∥ L p s . \|XY\|_{L^p}\leq\|X\|_{L^{pr}}\|Y\|_{L^{ps}}. ∥ X Y ∥ L p ≤ ∥ X ∥ L p r ∥ Y ∥ L p s . Proof. Observe
∥ X Y ∥ L p = ( E [ ( ∣ X Y ∣ ) p ] ) 1 / p = ( E [ ∣ X ∣ p ∣ Y ∣ p ] ) 1 / p \|XY\|_{L^p}=\left(\mathbb{E}\left[(|XY|)^p\right]\right)^{1/p}=\left(\mathbb{E}\left[|X|^p|Y|^p\right]\right)^{1/p} ∥ X Y ∥ L p = ( E [ ( ∣ X Y ∣ ) p ] ) 1/ p = ( E [ ∣ X ∣ p ∣ Y ∣ p ] ) 1/ p ≤ [H o ¨ lder] ( ( E [ ∣ X ∣ p r ] ) 1 / r ( E [ ∣ Y ∣ p s ] ) 1 / s ) 1 / p = ∥ X ∥ L p r ∥ Y ∥ L p s . \mathrel{\stackrel{\text{[Hölder]}}{\leq}}\left(\left(\mathbb{E}\left[|X|^{pr}\right]\right)^{1/r}\left(\mathbb{E}\left[|Y|^{ps}\right]\right)^{1/s}\right)^{1/p}=\|X\|_{L^{pr}}\|Y\|_{L^{ps}}. ≤ [H o ¨ lder] ( ( E [ ∣ X ∣ p r ] ) 1/ r ( E [ ∣ Y ∣ p s ] ) 1/ s ) 1/ p = ∥ X ∥ L p r ∥ Y ∥ L p s . Proposition E.8. For the neural network defined in § 4 , under the weight initialization and assumptions therein, we have that the loss at initialization satisfies
L ( θ 0 ) ≤ R y 2 + 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) , \mathcal{L}(\theta_0)\leq R_y^2+\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0), L ( θ 0 ) ≤ R y 2 + n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) , and
∥ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ ψ 1 ≤ c M ϕ 2 s L + 1 2 m L a L + 1 , \left\|\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right\|_{\psi_1}\leq cM_\phi^2\frac{s_{L+1}^2m_L}{a_{L+1}}, n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) ψ 1 ≤ c M ϕ 2 a L + 1 s L + 1 2 m L , where c > 0 c>0 c > 0 L ψ ≜ c M ϕ 2 s L + 1 2 m L a L + 1 L_\psi\triangleq cM_\phi^2\frac{s_{L+1}^2m_L}{a_{L+1}} L ψ ≜ c M ϕ 2 a L + 1 s L + 1 2 m L ξ ≤ L ψ − 1 \xi\leq L_\psi^{-1} ξ ≤ L ψ − 1
E [ e ξ L ( θ 0 ) ] ≤ 2 e ξ R y 2 . \mathbb{E}\left[e^{\xi\mathcal{L}(\theta_0)}\right]\leq2e^{\xi R_y^2}. E [ e ξ L ( θ 0 ) ] ≤ 2 e ξ R y 2 . Proof.
L ( θ 0 ) = 1 2 n ∑ i ∈ [ n ] ( f i ( θ 0 ) − y i ) 2 ≤ 1 2 n ∑ i ∈ [ n ] ( 2 f i 2 ( θ 0 ) + 2 y i 2 ) ≤ R y 2 + 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) , \begin{aligned}
\mathcal{L}(\theta_0)=\frac{1}{2n}\sum_{i\in[n]}(f_i(\theta_0)-y_i)^2\leq\frac{1}{2n}\sum_{i\in[n]}\left(2f_i^2(\theta_0)+2y_i^2\right) \\
&\leq R_y^2+\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0),
\end{aligned} L ( θ 0 ) = 2 n 1 i ∈ [ n ] ∑ ( f i ( θ 0 ) − y i ) 2 ≤ 2 n 1 i ∈ [ n ] ∑ ( 2 f i 2 ( θ 0 ) + 2 y i 2 ) ≤ R y 2 + n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) , where for all i ∈ [ n ] i\in[n] i ∈ [ n ]
f i 2 ( θ 0 ) = ( s L + 1 v T ϕ ( s L W ( L ) α ( L − 1 ) ( x i ) ) ) 2 = s L + 1 2 ( ∑ j ∈ [ m L ] v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ) 2 . \begin{aligned}
f_i^2(\theta_0)=\left(s_{L+1}v^\mathsf{T}\phi\left(s_L\mathbf{W}^{(L)}\boldsymbol{\alpha}^{(L-1)}(x_i)\right)\right)^2 \\
&=s_{L+1}^2\left(\sum_{j\in[m_L]}v_j\phi\left(s_L(\mathbf{w}^{(L)}_j)^\mathsf{T}\boldsymbol{\alpha}^{(L-1)}(x_i)\right)\right)^2.
\end{aligned} f i 2 ( θ 0 ) = ( s L + 1 v T ϕ ( s L W ( L ) α ( L − 1 ) ( x i ) ) ) 2 = s L + 1 2 j ∈ [ m L ] ∑ v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) 2 . Conditioning on all weights except v v v
ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) , j ∈ [ m L ] \phi\left(s_L(\mathbf{w}^{(L)}_j)^\mathsf{T}\boldsymbol{\alpha}^{(L-1)}(x_i)\right),\ j\in[m_L] ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) , j ∈ [ m L ] are fixed. Moreover, v j v_j v j j ∈ [ m L ] j\in[m_L] j ∈ [ m L ]
∥ v j ∥ ψ 2 ≤ C 1 a L + 1 . \|v_j\|_{\psi_2}\leq C\frac{1}{\sqrt{a_{L+1}}}. ∥ v j ∥ ψ 2 ≤ C a L + 1 1 . Then, conditioned on all weights except v \mathbf{v} v [111 Proposition 2.7.1]
∥ s L + 1 ( ∑ j ∈ [ m L ] v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ) | { W ( l ) } l ∈ [ L ] ∥ ψ 2 ≤ s L + 1 C ∑ j ∈ [ m L ] ∥ v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) | { W ( l ) } l ∈ [ L ] ∥ ψ 2 2 ≤ C 1 s L + 1 ∑ j ∈ [ m L ] ϕ 2 ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ∥ v j ∥ ψ 2 2 ≤ C 1 s L + 1 ∑ j ∈ [ m L ] M ϕ 2 C 2 a L + 1 = C 3 M ϕ s L + 1 m L a L + 1 . \begin{aligned}
\left\|s_{L+1}\left(\sum_{j\in[m_L]}v_j\phi\left(s_L\left(\mathbf{w}_j^{(L)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(L-1)}(\mathbf{x}_i)\right)\right)\middle|\left\{\mathbf{W}^{(l)}\right\}_{l\in[L]}\right\|_{\psi_2}
\le s_{L+1}\sqrt{C\sum_{j\in[m_L]}\left\|v_j\phi\left(s_L\left(\mathbf{w}_j^{(L)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(L-1)}(\mathbf{x}_i)\right)\middle|\left\{\mathbf{W}^{(l)}\right\}_{l\in[L]}\right\|_{\psi_2}^{2}} \\
&\le C_1s_{L+1}\sqrt{\sum_{j\in[m_L]}\phi^2\left(s_L\left(\mathbf{w}_j^{(L)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(L-1)}(\mathbf{x}_i)\right)\|\mathbf{v}_j\|_{\psi_2}^{2}} \\
&\le C_1s_{L+1}\sqrt{\sum_{j\in[m_L]}M_\phi^2\frac{C_2}{a_{L+1}}}=C_3M_\phi s_{L+1}\sqrt{\frac{m_L}{a_{L+1}}}.
\end{aligned} s L + 1 j ∈ [ m L ] ∑ v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) { W ( l ) } l ∈ [ L ] ψ 2 ≤ s L + 1 C j ∈ [ m L ] ∑ v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) { W ( l ) } l ∈ [ L ] ψ 2 2 ≤ C 1 s L + 1 j ∈ [ m L ] ∑ ϕ 2 ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ∥ v j ∥ ψ 2 2 ≤ C 1 s L + 1 j ∈ [ m L ] ∑ M ϕ 2 a L + 1 C 2 = C 3 M ϕ s L + 1 a L + 1 m L . Applying Claim E.5, we can remove the conditioning and get
∥ s L + 1 ( ∑ j ∈ [ m L ] v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ) ∥ ψ 2 ≤ C 4 M ϕ s L + 1 m L a L + 1 . \left\|s_{L+1}\left(\sum_{j\in[m_L]}v_j\phi\left(s_L\left(\mathbf{w}_j^{(L)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(L-1)}(\mathbf{x}_i)\right)\right)\right\|_{\psi_2}
\le C_4M_\phi s_{L+1}\sqrt{\frac{m_L}{a_{L+1}}}. s L + 1 j ∈ [ m L ] ∑ v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ψ 2 ≤ C 4 M ϕ s L + 1 a L + 1 m L . Now, applying Vershynin [111 Proposition 2.8.1]
∥ ( s L + 1 v T ϕ ( s L W ( L ) α ( L − 1 ) ( x i ) ) ) 2 ∥ ψ 1 = ∥ s L + 1 ( ∑ j ∈ [ m L ] v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ) ∥ ψ 2 2 ≤ C 5 M ϕ 2 s L + 1 2 m L a L + 1 . \begin{aligned}
\left\|\left(s_{L+1}\mathbf{v}^{\mathsf{T}}\phi\left(s_L\mathbf{W}^{(L)}\boldsymbol{\alpha}^{(L-1)}(\mathbf{x}_i)\right)\right)^2\right\|_{\psi_1} \\
&=\left\|s_{L+1}\left(\sum_{j\in[m_L]}v_j\phi\left(s_L\left(\mathbf{w}_j^{(L)}\right)^{\mathsf{T}}\boldsymbol{\alpha}^{(L-1)}(\mathbf{x}_i)\right)\right)\right\|_{\psi_2}^{2}
\le C_5M_\phi^2\frac{s_{L+1}^2m_L}{a_{L+1}}.
\end{aligned} ( s L + 1 v T ϕ ( s L W ( L ) α ( L − 1 ) ( x i ) ) ) 2 ψ 1 = s L + 1 j ∈ [ m L ] ∑ v j ϕ ( s L ( w j ( L ) ) T α ( L − 1 ) ( x i ) ) ψ 2 2 ≤ C 5 M ϕ 2 a L + 1 s L + 1 2 m L . We conclude by triangle inequality that
∥ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ ψ 1 ≤ C 5 M ϕ 2 s L + 1 2 m L a L + 1 . \left\|\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right\|_{\psi_1}
\le C_5M_\phi^2\frac{s_{L+1}^2m_L}{a_{L+1}}. n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) ψ 1 ≤ C 5 M ϕ 2 a L + 1 s L + 1 2 m L . Now note that for all ξ > 0 \xi>0 ξ > 0
E [ e ξ L ( θ 0 ) ] ≤ e ξ R y 2 E [ exp ( ξ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ) ] . \mathbb{E}\left[e^{\xi\mathcal{L}(\theta_0)}\right]\le e^{\xi R_y^2}\mathbb{E}\left[\exp\left(\xi\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right)\right]. E [ e ξ L ( θ 0 ) ] ≤ e ξ R y 2 E exp ξ n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) . By choosing ξ ≤ 1 ∥ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ ψ 1 \xi\le\frac{1}{\left\|\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right\|_{\psi_1}} ξ ≤ ∥ n 1 ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ ψ 1 1 [111 Proposition 2.8.1] that
E [ exp ( ξ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ) ] ≤ 2. \mathbb{E}\left[\exp\left(\xi\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right)\right]\le2. E exp ξ n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) ≤ 2. By choosing ξ ≤ 1 C 5 M ϕ 2 s L + 1 2 m L a L + 1 ≤ 1 ∥ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ ψ 1 , \xi\le\frac{1}{C_5M_\phi^2\frac{s_{L+1}^2m_L}{a_{L+1}}}\le\frac{1}{\left\|\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right\|_{\psi_1}}, ξ ≤ C 5 M ϕ 2 a L + 1 s L + 1 2 m L 1 ≤ ∥ n 1 ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ ψ 1 1 , □ \square □
Claim E.9. Given the ordinary differential inequality for y ( t ) y(t) y ( t )
y ˙ ≤ − b y + a , \dot{y}\le-by+a, y ˙ ≤ − b y + a , for some a , b > 0 a,b>0 a , b > 0 y ( 0 ) y(0) y ( 0 ) t ≥ 0 t\ge0 t ≥ 0
y ( t ) ≤ ( y ( 0 ) − y ∞ ) e − b t + y ∞ , y(t)\le(y(0)-y_\infty)e^{-bt}+y_\infty, y ( t ) ≤ ( y ( 0 ) − y ∞ ) e − b t + y ∞ , where y ∞ ≜ a b y_\infty\triangleq\frac{a}{b} y ∞ ≜ b a
Proof. Define z ( t ) = y ( t ) − y ∞ z(t)=y(t)-y_\infty z ( t ) = y ( t ) − y ∞ y ∞ = a b y_\infty=\frac{a}{b} y ∞ = b a
z ˙ = y ˙ ≤ − b y + a = − b ( z + a b ) + a = − b z . \dot{z}=\dot{y}\le-by+a=-b\left(z+\frac{a}{b}\right)+a=-bz. z ˙ = y ˙ ≤ − b y + a = − b ( z + b a ) + a = − b z . The claim follows from Grönwall’s inequality; we include the elementary derivation for completeness. Observe that
d d t ( e b t z ( t ) ) = b e b t z ( t ) + e b t z ˙ ( t ) ≤ b e b t z ( t ) − b e b t z ( t ) = 0 , \frac{\mathrm{d}}{\mathrm{d}t}\left(e^{bt}z(t)\right)=be^{bt}z(t)+e^{bt}\dot{z}(t)\le be^{bt}z(t)-be^{bt}z(t)=0, d t d ( e b t z ( t ) ) = b e b t z ( t ) + e b t z ˙ ( t ) ≤ b e b t z ( t ) − b e b t z ( t ) = 0 , hence e b t z ( t ) e^{bt}z(t) e b t z ( t ) t ≥ 0 t\ge0 t ≥ 0
e b t z ( t ) ≤ z ( 0 ) e^{bt}z(t)\le z(0) e b t z ( t ) ≤ z ( 0 ) ⟹ e b t ( y ( t ) − y ∞ ) ≤ y ( 0 ) − y ∞ \Longrightarrow e^{bt}\left(y(t)-y_\infty\right)\le y(0)-y_\infty ⟹ e b t ( y ( t ) − y ∞ ) ≤ y ( 0 ) − y ∞ y ( t ) ≤ ( y ( 0 ) − y ∞ ) e − b t + y ∞ . y(t)\le\left(y(0)-y_\infty\right)e^{-bt}+y_\infty. y ( t ) ≤ ( y ( 0 ) − y ∞ ) e − b t + y ∞ . Well-posedness of the dynamics and Itô’s formula
Global well-posedness and stationarity of regularized CLD
Two places in the paper use the fact that Eq. (1) has a solution defined for all times whose law is absolutely continuous with respect to the Gibbs measure, and that this measure is stationary for the process: Corollary 3.1, which applies Lemma 2.2 ; and the verification of the Itô hypotheses below, which needs the path to be bounded on compact time intervals. The vanishing of the stationary probability flux does not by itself supply any of this, as Remark F.2 explains. We therefore record the required statement and derive it from the analytic theory of Itô SDEs with non-smooth coefficients developed by Lee et al. [66 ] . In this subsection, we refer to Lee et al. [66 ] whenever we use "they" or "their". We note that this proof was LLM-assisted (Claude), and was verified by the authors, both mathematically and for proper citations.
Proposition F.1 (Non-explosion, stationarity, and absolute continuity). Suppose N p ≥ 2 N_p \ge2 N p ≥ 2 L ∈ C 1 ( R N p ) \mathcal{L} \in C^1(\mathbb{R}^{N_p}) L ∈ C 1 ( R N p ) L ≥ 0 \mathcal{L} \ge0 L ≥ 0 ∇ L \nabla\mathcal{L} ∇ L Λ ≻ 0 \Lambda\succ0 Λ ≻ 0 θ 0 ∼ ν ≜ N ( 0 N p , Λ − 1 ) \theta_0 \sim\nu\triangleq\mathcal{N}(0_{N_p},\Lambda^{-1}) θ 0 ∼ ν ≜ N ( 0 N p , Λ − 1 ) (1) has a pathwise-unique global strong solution, which is a time-homogeneous Markov process. The Gibbs measure
p ∞ ( d θ ) ∝ e − β L ( θ ) φ Λ ( θ ) d θ (18) p_\infty(\mathrm{d}\theta) \propto e^{-\beta\mathcal{L}(\theta)}\varphi_\Lambda(\theta)\,\mathrm{d}\theta \tag*{(18)} p ∞ ( d θ ) ∝ e − β L ( θ ) φ Λ ( θ ) d θ (18) is stationary for it, and its law p t p_t p t p t ≪ p ∞ ∼ ν p_t \ll p_\infty\sim\nu p t ≪ p ∞ ∼ ν
The hypotheses are weaker than those imposed in the body: they ask only that L \mathcal{L} L C 2 C^2 C 2 C 1 C^1 C 1 ∇ L \nabla\mathcal{L} ∇ L L \mathcal{L} L L ≥ 0 \mathcal{L} \ge0 L ≥ 0 L 2 L_2 L 2 Z Z Z
Proof. The cited framework. Lee et al. [66 ] study, on R d \mathbb{R}^d R d d ≥ 2 d \ge2 d ≥ 2
X t = x 0 + ∫ 0 t S ( X s ) d w s + ∫ 0 t b ( X s ) d s , 0 ≤ t < ζ , (19) X_t=x_0+\int_0^t S(X_s)\,\mathrm{d}w_s+\int_0^t b(X_s)\,\mathrm{d}s,\qquad0\le t<\zeta, \tag*{(19)} X t = x 0 + ∫ 0 t S ( X s ) d w s + ∫ 0 t b ( X s ) d s , 0 ≤ t < ζ , (19) with ζ \zeta ζ X X X C 2 ( R d ) C^2(\mathbb{R}^d) C 2 ( R d )
G f = 1 2 Tr ( A ∇ 2 f ) + b ⊤ ∇ f , A ≜ S S ⊤ , (20) \mathcal{G}f=\frac{1}{2}\operatorname{Tr}(A\nabla^2 f)+b^\top\nabla f,\qquad A\triangleq SS^\top, \tag*{(20)} G f = 2 1 Tr ( A ∇ 2 f ) + b ⊤ ∇ f , A ≜ S S ⊤ , (20) together with a measure μ \mu μ ρ \rho ρ infinitesimally invariant , meaning ∫ G f d μ = 0 \int\mathcal{G}f\,\mathrm{d}\mu=0 ∫ G f d μ = 0 f ∈ C c ∞ ( R d ) f\in C^\infty_c(\mathbb{R}^d) f ∈ C c ∞ ( R d ) C c ∞ ( R d ) C^\infty_c(\mathbb{R}^d) C c ∞ ( R d ) R d \mathbb{R}^d R d
β ρ , A ≜ 1 2 ∇ A + A ∇ ρ 2 ρ , ( ∇ A ) i ≜ ∑ j = 1 d ∂ j A i j , (21) \beta^{\rho,A}\triangleq\frac{1}{2}\nabla A+\frac{A\nabla\rho}{2\rho},\qquad(\nabla A)_i\triangleq\sum_{j=1}^d\partial_jA_{ij}, \tag*{(21)} β ρ , A ≜ 2 1 ∇ A + 2 ρ A ∇ ρ , ( ∇ A ) i ≜ j = 1 ∑ d ∂ j A ij , (21) where ∇ A \nabla A ∇ A A A A b = β ρ , A b=\beta^{\rho,A} b = β ρ , A symmetric case . Their condition (c) [66 Section 3.3] asks, for some p ∈ ( d , ∞ ) p\in(d,\infty) p ∈ ( d , ∞ ) S i j ∈ H l o c 1 , p ( R d ) ∩ C ( R d ) S_{ij}\in H_{\mathrm{loc}}^{1,p}(\mathbb{R}^d)\cap C(\mathbb{R}^d) S ij ∈ H loc 1 , p ( R d ) ∩ C ( R d ) i , j i,j i , j A A A b ∈ L l o c p ( R d , R d ) b\in L_{\mathrm{loc}}^{p}(\mathbb{R}^d,\mathbb{R}^d) b ∈ L loc p ( R d , R d ) L l o c p ( R d ) L_{\mathrm{loc}}^{p}(\mathbb{R}^d) L loc p ( R d ) p p p R d \mathbb{R}^d R d H l o c 1 , p H_{\mathrm{loc}}^{1,p} H loc 1 , p L l o c p L_{\mathrm{loc}}^p L loc p C ( R d ) C(\mathbb{R}^d) C ( R d ) locally strictly elliptic means that on every relatively compact open V ⊂ R d V\subset\mathbb{R}^d V ⊂ R d ξ ⊤ A ξ \xi^\top A\xi ξ ⊤ A ξ ∥ ξ ∥ 2 \lVert\xi\rVert^2 ∥ ξ ∥ 2 V V V ( T t ) t > 0 (T_t)_{t>0} ( T t ) t > 0 G \mathcal{G} G L s ( R d , μ ) L^s(\mathbb{R}^d,\mu) L s ( R d , μ ) s s s μ \mu μ s ∈ [ 1 , ∞ ] s\in[1,\infty] s ∈ [ 1 , ∞ ] [66 Definition 2.7] ; it is sub-Markovian, so T t 1 ≤ 1 T_t1\le1 T t 1 ≤ 1 M \mathsf{M} M i.e. , a strong Markov process, here with continuous sample paths, with lifetime ζ \zeta ζ ( P x 0 ) x 0 ∈ R d (\mathbb{P}_{x_0})_{x_0\in\mathbb{R}^d} ( P x 0 ) x 0 ∈ R d P x 0 \mathbb{P}_{x_0} P x 0 M \mathsf{M} M x 0 x_0 x 0 ( P t ) t ≥ 0 (P_t)_{t\ge0} ( P t ) t ≥ 0 ( T t ) t > 0 (T_t)_{t>0} ( T t ) t > 0 P t ( x 0 , E ) = P x 0 ( X t ∈ E ) P_t(x_0,E)=\mathbb{P}_{x_0}(X_t\in E) P t ( x 0 , E ) = P x 0 ( X t ∈ E ) x 0 ∈ R d x_0\in\mathbb{R}^d x 0 ∈ R d t ≥ 0 t\ge0 t ≥ 0 E ⊆ R d E\subseteq\mathbb{R}^d E ⊆ R d [66 Theorems 2.31 and 3.11] .
Their Sections 2.3 and 3.1–3.3 assume in addition that the density of μ \mu μ ρ ∈ H l o c 1 , p ( R d ) ∩ C ( R d ) \rho\in H^{1,p}_{\mathrm{loc}}(\mathbb{R}^{d}) \cap C(\mathbb{R}^{d}) ρ ∈ H loc 1 , p ( R d ) ∩ C ( R d ) ( A , 0 , 0 ) (A,0,0) ( A , 0 , 0 ) β ρ , A \beta^{\rho,A} β ρ , A [66 eqs. (2.10)–(2.14)] . Steps 1 and 2 below exhibit our equation as exactly this member, and verify (2.10)–(2.14) directly, so that nothing here rests on an unchecked hypothesis.
We import four results, each stated under the assumptions just described:
(R1) Conservativeness [66 Corollary 2.16(i)] : if A i j ∈ L 1 ( μ ) A_{ij} \in L^{1}(\mu) A ij ∈ L 1 ( μ ) b i − β i ρ , A ∈ L 1 ( μ ) b_i-\beta_i^{\rho,A} \in L^{1}(\mu) b i − β i ρ , A ∈ L 1 ( μ ) i , j i,j i , j T t 1 = 1 T_t1=1 T t 1 = 1
(R2) Non-explosion [66 Corollary 3.23] : ( T t ) t > 0 (T_t)_{t>0} ( T t ) t > 0 P x 0 ( ζ = ∞ ) = 1 \mathbb{P}_{x_0}(\zeta=\infty)=1 P x 0 ( ζ = ∞ ) = 1 x 0 ∈ R d x_0 \in\mathbb{R}^{d} x 0 ∈ R d
(R3) Strong well-posedness [66 Theorem 3.52] : under (c) and non-explosion, (19) has, for every x 0 ∈ R d x_0 \in\mathbb{R}^{d} x 0 ∈ R d M M M x 0 x_0 x 0
(R4) Absolute continuity and invariance [66 Proposition 3.1(i)] : P t ( x 0 , ⋅ ) P_t(x_0,\cdot) P t ( x 0 , ⋅ ) μ \mu μ x 0 ∈ R d x_0 \in\mathbb{R}^{d} x 0 ∈ R d t > 0 t>0 t > 0 μ \mu μ [66 Remark 2.13(i)] and hence for the process [66 Remark 3.44] .
Step 1: the dictionary, and condition (c). We read (1) as (19) with d = N p ≥ 2 d=N_p \ge2 d = N p ≥ 2
S = 2 β − 1 I N p , A = S S T = 2 β − 1 I N p , b = − ∇ L − β − 1 Λ θ , (22) S=\sqrt{2\beta^{-1}}I_{N_p}, \qquad A=SS^{\mathsf{T}}=2\beta^{-1}I_{N_p}, \qquad b=-\nabla\mathcal{L}-\beta^{-1}\Lambda\theta, \tag*{(22)} S = 2 β − 1 I N p , A = S S T = 2 β − 1 I N p , b = − ∇ L − β − 1 Λ θ , (22) and with ρ \rho ρ p ∞ p_\infty p ∞ (18) , so that μ = ρ d θ = p ∞ \mu=\rho\,\mathrm{d}\theta=p_\infty μ = ρ d θ = p ∞ L ≥ 0 \mathcal{L}\ge0 L ≥ 0
Z = ∫ R N p e − β L ( θ ) − 1 2 θ T Λ θ d θ ≤ ∫ R N p e − 1 2 θ T Λ θ d θ = ( 2 π ) N p / 2 det Λ < ∞ , (23) Z=\int_{\mathbb{R}^{N_p}}e^{-\beta\mathcal{L}(\theta)-\frac{1}{2}\theta^{\mathsf{T}}\Lambda\theta}\,\mathrm{d}\theta
\leq\int_{\mathbb{R}^{N_p}}e^{-\frac{1}{2}\theta^{\mathsf{T}}\Lambda\theta}\,\mathrm{d}\theta
=\frac{(2\pi)^{N_p/2}}{\sqrt{\det\Lambda}}<\infty, \tag*{(23)} Z = ∫ R N p e − β L ( θ ) − 2 1 θ T Λ θ d θ ≤ ∫ R N p e − 2 1 θ T Λ θ d θ = det Λ ( 2 π ) N p /2 < ∞ , (23) so the normalizing constant of (18) is finite and strictly positive and p ∞ p_\infty p ∞ ν \nu ν e − β L e^{-\beta\mathcal{L}} e − β L
Condition (c) is immediate here, because both S S S A A A constant : their entries lie in H l o c 1 , p ∩ C H^{1,p}_{\mathrm{loc}}\cap C H loc 1 , p ∩ C p p p ξ T A ξ = 2 β − 1 ∥ ξ ∥ 2 \xi^{\mathsf{T}}A\xi=2\beta^{-1}\|\xi\|^{2} ξ T A ξ = 2 β − 1 ∥ ξ ∥ 2 A A A b b b L ∈ C 1 \mathcal{L}\in C^{1} L ∈ C 1 L l o c p L^{p}_{\mathrm{loc}} L loc p p p p (19) satisfies (c) for every p ∈ ( N p , ∞ ) p\in(N_p,\infty) p ∈ ( N p , ∞ )
Step 2: the standing assumptions, and symmetry. The density ρ ∝ e − β L − 1 2 θ T Λ θ \rho\propto e^{-\beta\mathcal{L}-\frac{1}{2}\theta^{\mathsf{T}}\Lambda\theta} ρ ∝ e − β L − 2 1 θ T Λ θ C 1 C^{1} C 1 L \mathcal{L} L H l o c 1 , p ∩ C H^{1,p}_{\mathrm{loc}}\cap C H loc 1 , p ∩ C p p p [66 Remark 2.28] requires; consequently ρ ∈ H l o c 1 , 2 \sqrt{\rho}\in H^{1,2}_{\mathrm{loc}} ρ ∈ H loc 1 , 2 p ∞ p_\infty p ∞ A A A p ∞ p_\infty p ∞ b b b p ∞ p_\infty p ∞
For the symmetry, A A A (21) , leaving
β ρ , A = A ∇ ρ 2 ρ = β − 1 ∇ log ρ = β − 1 ∇ ( − β L − 1 2 θ T Λ θ ) = − ∇ L − β − 1 Λ θ = b , (24) \beta^{\rho,A}=\frac{A\nabla\rho}{2\rho}=\beta^{-1}\nabla\log\rho=\beta^{-1}\nabla\left(-\beta\mathcal{L}-\frac{1}{2}\theta^{\mathsf{T}}\Lambda\theta\right)=-\nabla\mathcal{L}-\beta^{-1}\Lambda\theta=b, \tag*{(24)} β ρ , A = 2 ρ A ∇ ρ = β − 1 ∇ log ρ = β − 1 ∇ ( − β L − 2 1 θ T Λ θ ) = − ∇ L − β − 1 Λ θ = b , (24) by (22) . This is the symmetric case. It remains to check their eq. (2.14). Let f ∈ C c ∞ ( R N p ) f\in C^\infty_c(\mathbb{R}^{N_p}) f ∈ C c ∞ ( R N p ) A = 2 β − 1 I N p A=2\beta^{-1}I_{N_p} A = 2 β − 1 I N p (20) is β − 1 Δ f \beta^{-1}\Delta f β − 1 Δ f
β − 1 ∫ R N p Δ f ρ d θ = − β − 1 ∫ R N p ∇ f T ∇ ρ d θ = − ∫ R N p ( β − 1 ∇ ρ ρ ) T ∇ f d p ∞ = − ∫ R N p ( β ρ , A ) T ∇ f d p ∞ , (25) \beta^{-1}\int_{\mathbb{R}^{N_p}}\Delta f\,\rho\,\mathrm{d}\theta
=-\beta^{-1}\int_{\mathbb{R}^{N_p}}\nabla f^{\mathsf{T}}\nabla\rho\,\mathrm{d}\theta
=-\int_{\mathbb{R}^{N_p}}\left(\beta^{-1}\frac{\nabla\rho}{\rho}\right)^{\mathsf{T}}\nabla f\,\mathrm{d}p_\infty
=-\int_{\mathbb{R}^{N_p}}\left(\beta^{\rho,A}\right)^{\mathsf{T}}\nabla f\,\mathrm{d}p_\infty, \tag*{(25)} β − 1 ∫ R N p Δ f ρ d θ = − β − 1 ∫ R N p ∇ f T ∇ ρ d θ = − ∫ R N p ( β − 1 ρ ∇ ρ ) T ∇ f d p ∞ = − ∫ R N p ( β ρ , A ) T ∇ f d p ∞ , (25) the last equality by Eq. (24) . Restoring the first-order part b T ∇ f \mathbf{b}^{\mathsf{T}}\nabla f b T ∇ f (20) ,
∫ R N p G f d p ∞ = ∫ R N p ( b − β ρ , A ) T ∇ f d p ∞ = 0 , (26) \int_{\mathbb{R}^{N_p}} \mathcal{G}f\,\mathrm{d}p_{\infty}=\int_{\mathbb{R}^{N_p}}\left(\mathbf{b}-\boldsymbol{\beta}^{\rho,\mathbf{A}}\right)^{\mathsf{T}}\nabla f\,\mathrm{d}p_{\infty}=0, \tag*{(26)} ∫ R N p G f d p ∞ = ∫ R N p ( b − β ρ , A ) T ∇ f d p ∞ = 0 , (26) again by Eq. (24) . So p ∞ p_{\infty} p ∞ [66 ] consumes.
Step 3: conservativeness and non-explosion. Both integrability conditions of (R1) are immediate. The drift discrepancy b − β ρ , A \mathbf{b}-\boldsymbol{\beta}^{\rho,\mathbf{A}} b − β ρ , A (24) , and A \mathbf{A} A p ∞ p_{\infty} p ∞ (23) , so
∫ R N p ∣ A i j ∣ d p ∞ = 2 β − 1 δ i j < ∞ , ∀ i , j ∈ [ N p ] . (27) \int_{\mathbb{R}^{N_p}} |A_{ij}|\,\mathrm{d}p_{\infty}=2\beta^{-1}\delta_{ij}<\infty,\qquad\forall i,j\in[N_p]. \tag*{(27)} ∫ R N p ∣ A ij ∣ d p ∞ = 2 β − 1 δ ij < ∞ , ∀ i , j ∈ [ N p ] . (27) Hence (R1) gives T t 1 = 1 T_t1=1 T t 1 = 1 P θ ( ζ = ∞ ) = 1 \mathbb{P}_{\theta}(\zeta=\infty)=1 P θ ( ζ = ∞ ) = 1 θ ∈ R N p \theta\in\mathbb{R}^{N_p} θ ∈ R N p
Step 4: the conclusions. Step 1 supplies (c) and Step 3 non-explosion, so (R3) gives, for every deterministic starting point, a pathwise-unique strong solution of Eq. (1) defined for all t ≥ 0 t\geq0 t ≥ 0 M \mathsf{M} M ( P t ) t ≥ 0 (P_t)_{t\geq0} ( P t ) t ≥ 0 S \mathbf{S} S b = − ∇ L − β − 1 Λ θ \mathbf{b}=-\nabla\mathcal{L}-\beta^{-1}\boldsymbol{\Lambda}\theta b = − ∇ L − β − 1 Λ θ ∇ L \nabla\mathcal{L} ∇ L [58 Chapter 5, Theorems 2.5 and 2.9] , applied after localization. By that independence, conditioning on θ 0 = θ \theta_0=\theta θ 0 = θ θ \theta θ P θ \mathbb{P}_{\theta} P θ
For the stationarity, Eq. (24) puts us in the symmetric case and Step 3 gives conservativeness, so p ∞ p_{\infty} p ∞ E E E p ∞ ( E ) = 0 p_{\infty}(E)=0 p ∞ ( E ) = 0 θ 0 = θ \theta_0=\theta θ 0 = θ P θ \mathbb{P}_{\theta} P θ E E E t t t P t ( θ , E ) P_t(\theta,E) P t ( θ , E ) θ 0 ∼ ν \theta_0\sim\nu θ 0 ∼ ν
p t ( E ) = ∫ R N p P t ( θ , E ) ν ( d θ ) , (28) p_t(E)=\int_{\mathbb{R}^{N_p}}P_t(\theta,E)\,\nu(\mathrm{d}\theta), \tag*{(28)} p t ( E ) = ∫ R N p P t ( θ , E ) ν ( d θ ) , (28) and by (R4) the integrand vanishes for every θ \theta θ t > 0 t>0 t > 0 p t ( E ) = 0 p_t(E)=0 p t ( E ) = 0 p 0 = ν ∼ p ∞ p_0=\nu\sim p_{\infty} p 0 = ν ∼ p ∞ t = 0 t=0 t = 0 p t ≪ p ∞ ∼ ν p_t\ll p_{\infty}\sim\nu p t ≪ p ∞ ∼ ν
Remark F.2 (Why the flux computation is not enough). The verification cited in Footnote 6 shows that the Gibbs density solves the stationary Fokker–Planck equation , its probability flux vanishing identically. That is strictly weaker than invariance under the transition semigroup, which is what Proposition F.1 asserts and what Lemma 2.2 consumes. If the equation explodes with positive probability, the minimal semigroup is sub-Markov and, for a symmetric diffusion, the reference measure is only sub-invariant , with equality if and only if the process is conservative. The flux computation is local and cannot see this: it is insensitive to whether mass escapes to infinity in finite time. This is why the stationarity claim waits until non-explosion has been established, rather than being read off the flux alone.
Introduction and smooth activations
We first recall the standard multidimensional Itô formula, as stated in Mao [75 Chapter 1, Theorem 6.4] , adapted to our notation.
Theorem F.3 (Multidimensional Itô formula). Let ( w t ) t ≥ 0 (\mathbf{w}_t)_{t\geq0} ( w t ) t ≥ 0 R N p \mathbb{R}^{N_p} R N p u t ∈ R N p \mathbf{u}_t\in\mathbb{R}^{N_p} u t ∈ R N p V t ∈ R N p × N p \mathbf{V}_t\in\mathbb{R}^{N_p\times N_p} V t ∈ R N p × N p
∫ 0 T ∥ u t ∥ d t < ∞ , ∫ 0 T ∥ V t ∥ F 2 d t < ∞ , \int_0^T\|\mathbf{u}_t\|\,\mathrm{d}t<\infty,\qquad\int_0^T\|\mathbf{V}_t\|_{\mathrm{F}}^2\,\mathrm{d}t<\infty, ∫ 0 T ∥ u t ∥ d t < ∞ , ∫ 0 T ∥ V t ∥ F 2 d t < ∞ , a.s. for every T > 0 T>0 T > 0 θ t ∈ R N p \theta_t\in\mathbb{R}^{N_p} θ t ∈ R N p
d θ t = u t d t + V t d w t . \mathrm{d}\theta_t=\mathbf{u}_t\mathrm{d}t+\mathbf{V}_t\mathrm{d}\mathbf{w}_t. d θ t = u t d t + V t d w t . If L ∈ C 2 , 1 ( R N p × R + ; R ) \mathcal{L} \in C^{2,1}(\mathbb{R}^{N_p} \times\mathbb{R}_{+}; \mathbb{R}) L ∈ C 2 , 1 ( R N p × R + ; R ) 17 L ( θ t , t ) \mathcal{L}(\theta_t,t) L ( θ t , t )
d L ( θ t , t ) = ∂ L ∂ t ( θ t , t ) d t + ∇ θ L ( θ t , t ) ⊤ u t d t + 1 2 Tr ( V t ⊤ ∇ θ 2 L ( θ t , t ) V t ) d t + ∇ θ L ( θ t , t ) ⊤ V t d w t . \mathrm{d}\mathcal{L}(\theta_t,t)=\frac{\partial\mathcal{L}}{\partial t}(\theta_t,t)\,\mathrm{d}t+\nabla_\theta\mathcal{L}(\theta_t,t)^{\top}\mathbf{u}_t\mathrm{d}t
+\frac{1}{2}\operatorname{Tr}\left(\mathbf{V}_t^{\top}\nabla_\theta^2\mathcal{L}(\theta_t,t)\mathbf{V}_t\right)\mathrm{d}t+\nabla_\theta\mathcal{L}(\theta_t,t)^{\top}\mathbf{V}_t\mathrm{d}\mathbf{w}_t. d L ( θ t , t ) = ∂ t ∂ L ( θ t , t ) d t + ∇ θ L ( θ t , t ) ⊤ u t d t + 2 1 Tr ( V t ⊤ ∇ θ 2 L ( θ t , t ) V t ) d t + ∇ θ L ( θ t , t ) ⊤ V t d w t . In our setting, as presented in § 4 ,
L ( θ ) = 1 2 n ∥ f ( θ ) − y ∥ 2 , \mathcal{L}(\theta)=\frac{1}{2n}\lVert f(\theta)-\mathbf{y}\rVert^2, L ( θ ) = 2 n 1 ∥ f ( θ ) − y ∥ 2 , u t = − ∇ θ L ( θ t ) − β − 1 Λ θ t , \mathbf{u}_t=-\nabla_\theta\mathcal{L}(\theta_t)-\beta^{-1}\Lambda\theta_t, u t = − ∇ θ L ( θ t ) − β − 1 Λ θ t , V t = 2 β − 1 I N p . \mathbf{V}_t=\sqrt{2\beta^{-1}}I_{N_p}. V t = 2 β − 1 I N p . The diffusion condition is immediate, since
∫ 0 T ∥ V t ∥ F 2 d t = 2 β − 1 N p T < ∞ . \int_0^T\lVert\mathbf{V}_t\rVert_{\mathrm{F}}^2\mathrm{d}t=2\beta^{-1}N_pT<\infty. ∫ 0 T ∥ V t ∥ F 2 d t = 2 β − 1 N p T < ∞. For the drift condition, Proposition F.1 gives that the solution θ t \theta_t θ t [ 0 , T ] [0,T] [ 0 , T ] ( θ t ) 0 ≤ t ≤ T (\theta_t)_{0\leq t\leq T} ( θ t ) 0 ≤ t ≤ T f f f ∇ θ L \nabla_\theta\mathcal{L} ∇ θ L u t \mathbf{u}_t u t [ 0 , T ] [0,T] [ 0 , T ]
∫ 0 T ∥ u t ∥ d t < ∞ \int_0^T\lVert\mathbf{u}_t\rVert\mathrm{d}t<\infty ∫ 0 T ∥ u t ∥ d t < ∞ a.s. for every T > 0 T>0 T > 0
When the activation function is twice continuously differentiable, the loss L \mathcal{L} L θ \theta θ Theorem F.3 hold.
Extension to once continuously differentiable activations
We next describe how to extend the argument to activations that are only once continuously differentiable, such as softsign and ELU [33 ] . For this purpose, we use the following one-dimensional generalized Itô formula from Ondrejat and Seidler [90 Proposition 1.1] , adapted to our notation.
Theorem F.4 (Generalized one-dimensional Itô formula). Let X t X_t X t F ∈ C 1 ( R ) F\in C^1(\mathbb{R}) F ∈ C 1 ( R ) 18 F ′ F' F ′ g : R → R g:\mathbb{R}\to\mathbb{R} g : R → R g = F ′ ′ g=F'' g = F ′′ R \mathbb{R} R t ≥ 0 t\geq0 t ≥ 0
F ( X t ) − F ( X 0 ) = ∫ 0 t F ′ ( X s ) d X s + 1 2 ∫ 0 t g ( X s ) d ⟨ X ⟩ s , F(X_t)-F(X_0)=\int_0^t F'(X_s)\,\mathrm{d}X_s+\frac{1}{2}\int_0^t g(X_s)\,\mathrm{d}\langle X\rangle_s, F ( X t ) − F ( X 0 ) = ∫ 0 t F ′ ( X s ) d X s + 2 1 ∫ 0 t g ( X s ) d ⟨ X ⟩ s , a.s.
Remark F.5 (A practical sufficient condition). The absolute-continuity assumption in Theorem F.4 is implied by a simpler and more familiar condition: it is enough that F ′ F' F ′ [ a , b ] [a,b] [ a , b ] L a , b L_{a,b} L a , b
∣ F ′ ( x ) − F ′ ( y ) ∣ ≤ L a , b ∣ x − y ∣ ∀ x , y ∈ [ a , b ] . |F'(x)-F'(y)|\leq L_{a,b}|x-y|\qquad\forall x,y\in[a,b]. ∣ F ′ ( x ) − F ′ ( y ) ∣ ≤ L a , b ∣ x − y ∣ ∀ x , y ∈ [ a , b ] . Under this condition, F ′ F' F ′ F ′ ′ F'' F ′′ Theorem F.4 applies with any Borel representative g : R → R g:\mathbb{R}\to\mathbb{R} g : R → R g = F ′ ′ g=F'' g = F ′′
We will use the following specialization to scalar Itô processes driven by the same Brownian motion w t \mathbf{w}_t w t
Corollary F.6. Let u t ∈ R u_t \in\mathbb{R} u t ∈ R v t ∈ R N p v_t \in\mathbb{R}^{N_p} v t ∈ R N p
∫ 0 T ∣ u t ∣ d t < ∞ , ∫ 0 T ∥ v t ∥ 2 d t < ∞ \int_0^T |u_t|\,\mathrm{d}t < \infty,\qquad\int_0^T \lVert v_t\rVert^2\,\mathrm{d}t < \infty ∫ 0 T ∣ u t ∣ d t < ∞ , ∫ 0 T ∥ v t ∥ 2 d t < ∞ a.s. for every T > 0 T > 0 T > 0 X t ∈ R X_t \in\mathbb{R} X t ∈ R
d X t = u t d t + v t T d w t . \mathrm{d}X_t = u_t\mathrm{d}t + v_t^{\mathsf{T}}\,\mathrm{d}w_t. d X t = u t d t + v t T d w t . Let F ∈ C 1 ( R ) F \in C^1(\mathbb{R}) F ∈ C 1 ( R ) F ′ F' F ′ g : R → R g : \mathbb{R} \to\mathbb{R} g : R → R g = F ′ ′ g = F'' g = F ′′ R \mathbb{R} R
F ( X t ) − F ( X 0 ) = ∫ 0 t [ F ′ ( X s ) u s + 1 2 g ( X s ) ∥ v s ∥ 2 ] d s + ∫ 0 t F ′ ( X s ) v s T d w s . F(X_t) - F(X_0) = \int_0^t \left[F'(X_s)u_s + \frac{1}{2}g(X_s)\lVert v_s\rVert^2\right]\mathrm{d}s + \int_0^t F'(X_s)v_s^{\mathsf{T}}\mathrm{d}w_s. F ( X t ) − F ( X 0 ) = ∫ 0 t [ F ′ ( X s ) u s + 2 1 g ( X s ) ∥ v s ∥ 2 ] d s + ∫ 0 t F ′ ( X s ) v s T d w s . a.s. for every t ≥ 0 t \ge0 t ≥ 0
d F ( X t ) = F ′ ( X t ) u t d t + 1 2 g ( X t ) ∥ v t ∥ 2 d t + F ′ ( X t ) v t T d w t . \mathrm{d}F(X_t) = F'(X_t)u_t\mathrm{d}t + \frac{1}{2}g(X_t)\lVert v_t\rVert^2\mathrm{d}t + F'(X_t)v_t^{\mathsf{T}}\mathrm{d}w_t. d F ( X t ) = F ′ ( X t ) u t d t + 2 1 g ( X t ) ∥ v t ∥ 2 d t + F ′ ( X t ) v t T d w t . Proof. The martingale part of X t X_t X t
∫ 0 t v s T d w s . \int_0^t v_s^{\mathsf{T}}\mathrm{d}w_s. ∫ 0 t v s T d w s . Hence its quadratic variation is
⟨ X ⟩ t = ∫ 0 t ∥ v s ∥ 2 d s . \langle X\rangle_t = \int_0^t \lVert v_s\rVert^2\mathrm{d}s. ⟨ X ⟩ t = ∫ 0 t ∥ v s ∥ 2 d s . Substituting this identity into Theorem F.4 gives the result.
We now explain why this scalar generalized formula is sufficient for neural networks with pointwise activations. Although θ t ∈ R N p \theta_t \in\mathbb{R}^{N_p} θ t ∈ R N p
Assume that the activation function satisfies
ϕ ∈ C 1 ( R ) \phi\in C^1(\mathbb{R}) ϕ ∈ C 1 ( R ) and that ϕ ′ \phi' ϕ ′ ϕ ′ \phi' ϕ ′ ϕ ′ ′ \phi'' ϕ ′′ ϕ ( X t ) \phi(X_t) ϕ ( X t ) X t X_t X t
Write the layer pre-activations and post-activations as
A ~ ℓ , t = s ℓ A ℓ − 1 , t W ℓ , t T , A ℓ , t = ϕ ( A ~ ℓ , t ) , \widetilde{A}_{\ell,t} = s_\ell A_{\ell-1,t}W_{\ell,t}^{\mathsf{T}},\qquad A_{\ell,t} = \phi\left(\widetilde{A}_{\ell,t}\right), A ℓ , t = s ℓ A ℓ − 1 , t W ℓ , t T , A ℓ , t = ϕ ( A ℓ , t ) , where ϕ \phi ϕ A 0 , t = X A_{0,t} = X A 0 , t = X W ℓ , t W_{\ell,t} W ℓ , t θ t \theta_t θ t A ℓ − 1 , t A_{\ell-1,t} A ℓ − 1 , t W ℓ , t W_{\ell,t} W ℓ , t A ℓ − 1 , t A_{\ell-1,t} A ℓ − 1 , t A ~ ℓ , t \widetilde{A}_{\ell,t} A ℓ , t A ~ ℓ , t i , j \widetilde{A}_{\ell,t}^{i,j} A ℓ , t i , j
Applying Corollary F.6 to the scalar process A ~ ℓ , t i , j \widetilde{A}_{\ell,t}^{i,j} A ℓ , t i , j
A ℓ , t i , j = ϕ ( A ~ ℓ , t i , j ) A_{\ell,t}^{i,j} = \phi\left(\widetilde{A}_{\ell,t}^{i,j}\right) A ℓ , t i , j = ϕ ( A ℓ , t i , j ) is again an Itô process. This completes the induction over layers. Consequently, all network output coordinates are Itô processes. Since the squared loss
L ( z ) = 1 2 n ∥ z − y ∥ 2 \mathcal{L}(z) = \frac{1}{2n}\lVert z-y\rVert^2 L ( z ) = 2 n 1 ∥ z − y ∥ 2 is a smooth function of the output vector z z z L ( θ t ) \mathcal{L}(\theta_t) L ( θ t )
Proof of global convergence of regularized CLD for shallow networks
Introduction and goal
Recall Theorem 5.1 (Shallow networks). Consider the network defined in § 4 , under the assumptions therein, with L = 1 L=1 L = 1 (1) , with Λ \Lambda Λ (6) . Denoting κ ˉ ≜ max ( 1 , κ ) \bar{\kappa}\triangleq\max(1,\kappa) κ ˉ ≜ max ( 1 , κ ) m = Ω ( κ ˉ 2 β n λ 1 ln n λ 1 ln ( κ ˉ 4 β ) ) m=\Omega\left(\bar{\kappa}^{2}\beta\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\ln\left(\bar{\kappa}^{4}\beta\right)\right) m = Ω ( κ ˉ 2 β λ 1 n ln λ 1 n ln ( κ ˉ 4 β ) ) β = Ω ( n λ 1 ) \beta=\Omega\left(\frac{n}{\lambda_{1}}\right) β = Ω ( λ 1 n ) t ≥ 0 t\ge0 t ≥ 0
E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c 1 s L + 1 2 λ 1 m n t + L ∞ , (29) \mathbb{E}\mathcal{L}(\theta_t)\leq\left(\mathbb{E}\mathcal{L}(\theta_0)-L_{\infty}\right)e^{-c_1s_{L+1}^{2}\frac{\lambda_1^m}{n}t}+L_{\infty}, \tag*{(29)} E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c 1 s L + 1 2 n λ 1 m t + L ∞ , (29) where c 1 c_1 c 1
L ∞ ≲ ( 1 + κ ˉ 4 β ) exp ( − c 2 m κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 2 n 2 λ 1 2 ( 1 κ ˉ 4 + s 1 4 d 2 s 2 4 m 2 ) + β − 1 n 2 λ 1 2 κ ˉ 2 κ 2 s 1 4 d 2 s 2 4 m 2 + β − 1 n λ 1 ( 1 + κ 2 s 1 2 d s 2 2 m ) . (30) \begin{aligned}
L_{\infty}\lesssim(1+\bar{\kappa}^{4}\beta)\exp\left(-c_2\frac{m}{\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-2}\frac{n^2}{\lambda_1^2}\left(\frac{1}{\bar{\kappa}^{4}}+\frac{s_1^4d^2}{s_2^4m^2}\right) \\
&+\beta^{-1}\frac{n^2}{\lambda_1^2}\bar{\kappa}^{2}\kappa^2\frac{s_1^4d^2}{s_2^4m^2}+\beta^{-1}\frac{n}{\lambda_1}\left(1+\kappa^2\frac{s_1^2d}{s_2^2m}\right). \tag*{(30)}
\end{aligned} L ∞ ≲ ( 1 + κ ˉ 4 β ) exp ( − c 2 κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 2 λ 1 2 n 2 ( κ ˉ 4 1 + s 2 4 m 2 s 1 4 d 2 ) + β − 1 λ 1 2 n 2 κ ˉ 2 κ 2 s 2 4 m 2 s 1 4 d 2 + β − 1 λ 1 n ( 1 + κ 2 s 2 2 m s 1 2 d ) . (30) Proof. We begin by noting that for the MSE loss, as defined in Eq. (4) , we have
∇ L ( θ t ) = 1 n ∇ f ( θ t ) T r ( θ t ) , ∇ 2 L ( θ t ) = 1 n ( ∇ f ( θ t ) T ∇ f ( θ t ) + ∑ i = 1 n r i ( θ t ) ∇ 2 f i ( θ t ) ) . \nabla\mathcal{L}(\theta_t)=\frac{1}{n}\nabla f(\theta_t)^{\mathsf{T}}r(\theta_t),\qquad
\nabla^2\mathcal{L}(\theta_t)=\frac{1}{n}\left(\nabla f(\theta_t)^{\mathsf{T}}\nabla f(\theta_t)+\sum_{i=1}^{n}r_i(\theta_t)\nabla^2f_i(\theta_t)\right). ∇ L ( θ t ) = n 1 ∇ f ( θ t ) T r ( θ t ) , ∇ 2 L ( θ t ) = n 1 ( ∇ f ( θ t ) T ∇ f ( θ t ) + i = 1 ∑ n r i ( θ t ) ∇ 2 f i ( θ t ) ) . Then, substituting these into Eq. (1) , we get
d d t E L ( θ t ) = 1 n ( − n − 1 E [ r ( θ t ) T K N T K ( θ t ) r ( θ t ) ] + β − 1 E [ Tr ( ∇ f ( θ t ) T ∇ f ( θ t ) ) ] + β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] − β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Λ θ t ] ) . (31) \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)
=\frac{1}{n}\left(-n^{-1}\mathbb{E}\left[r(\theta_t)^{\mathsf{T}}K_{\mathrm{NTK}}(\theta_t)r(\theta_t)\right]
+\beta^{-1}\mathbb{E}\left[\operatorname{Tr}\left(\nabla f(\theta_t)^{\mathsf{T}}\nabla f(\theta_t)\right)\right]
+\beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]
-\beta^{-1}\mathbb{E}\left[r(\theta_t)^{\mathsf{T}}\nabla f(\theta_t)\Lambda\theta_t\right]\right). \tag*{(31)} d t d E L ( θ t ) = n 1 ( − n − 1 E [ r ( θ t ) T K NTK ( θ t ) r ( θ t ) ] + β − 1 E [ Tr ( ∇ f ( θ t ) T ∇ f ( θ t ) ) ] + β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] − β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Λ θ t ] ) . (31) We wish to bound every term in the right-hand side of Eq. (31) , using Corollary 3.1. In order to use Corollary 3.1, we must bound ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α \left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}} ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1
E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ≤ 2 e 2 ( α − 1 ) β R y 2 , \mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\leq2e^{2(\alpha-1)\beta R_y^2}, E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ≤ 2 e 2 ( α − 1 ) β R y 2 , provided that 2 ( α − 1 ) β ≤ L ψ − 1 2(\alpha-1)\beta\leq L_{\psi}^{-1} 2 ( α − 1 ) β ≤ L ψ − 1 L ψ ≜ c M ϕ 2 s L + 1 2 L + 1 m L a L + 1 = c M ϕ 2 κ 2 L_{\psi}\triangleq cM_{\phi}^{2}\frac{s_{L+1}^{2L+1}m^L}{a_{L+1}}=cM_{\phi}^{2}\kappa^2 L ψ ≜ c M ϕ 2 a L + 1 s L + 1 2 L + 1 m L = c M ϕ 2 κ 2 α ≤ 1 + 1 2 c M ϕ 2 κ 2 β \alpha\leq1+\frac{1}{2cM_{\phi}^{2}\kappa^2\beta} α ≤ 1 + 2 c M ϕ 2 κ 2 β 1 2 e 2 ( α − 1 ) β R y 2 = O ( 1 ) 2e^{2(\alpha-1)\beta R_y^2}=O(1) 2 e 2 ( α − 1 ) β R y 2 = O ( 1 ) α ≤ 1 + 1 2 c M ϕ 2 β \alpha\leq1+\frac{1}{2cM_{\phi}^{2}\beta} α ≤ 1 + 2 c M ϕ 2 β 1 α = 1 + 1 2 c M ϕ 2 max ( κ 2 , 1 ) β \alpha=1+\frac{1}{2cM_{\phi}^{2}\max(\kappa^2,1)\beta} α = 1 + 2 c M ϕ 2 m a x ( κ 2 , 1 ) β 1 ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α \left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}} ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1 O ( 1 ) O(1) O ( 1 ) κ ˉ ≜ max ( κ , 1 ) \bar{\kappa}\triangleq\max(\kappa,1) κ ˉ ≜ max ( κ , 1 )
α = 1 + 1 2 c M ϕ 2 κ ˉ 2 β , \alpha=1+\frac{1}{2cM_{\phi}^{2}\bar{\kappa}^{2}\beta}, α = 1 + 2 c M ϕ 2 κ ˉ 2 β 1 , and
( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α ≤ ( 2 e 2 ( α − 1 ) β R y 2 ) 1 2 α ≤ 2 e R y 2 2 c M ϕ 2 κ ˉ 2 = 2 e R y 2 2 c M ϕ 2 κ ˉ 2 = O ( 1 ) , \left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}}
\leq\left(2e^{2(\alpha-1)\beta R_y^2}\right)^{\frac{1}{2\alpha}}
\leq\sqrt{2}e^{\frac{R_y^2}{2cM_{\phi}^{2}\bar{\kappa}^{2}}}
=\sqrt{2}e^{\frac{R_y^2}{2cM_{\phi}^{2}\bar{\kappa}^{2}}}=O(1), ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1 ≤ ( 2 e 2 ( α − 1 ) β R y 2 ) 2 α 1 ≤ 2 e 2 c M ϕ 2 κ ˉ 2 R y 2 = 2 e 2 c M ϕ 2 κ ˉ 2 R y 2 = O ( 1 ) , since κ ˉ ≥ 1 \bar{\kappa}\geq1 κ ˉ ≥ 1 η ≜ 2 e R y 2 2 c M ϕ 2 κ ˉ 2 = O ( 1 ) \eta\triangleq\sqrt{2}e^{\frac{R_y^2}{2cM_{\phi}^{2}\bar{\kappa}^{2}}}=O(1) η ≜ 2 e 2 c M ϕ 2 κ ˉ 2 R y 2 = O ( 1 )
Given this choice of α \alpha α
α α − 1 = 1 + 2 c M ϕ 2 κ ˉ 2 β = c β κ ˉ 2 β , \frac{\alpha}{\alpha-1}=1+2cM_{\phi}^{2}\bar{\kappa}^{2}\beta=c_{\beta}\bar{\kappa}^{2}\beta, α − 1 α = 1 + 2 c M ϕ 2 κ ˉ 2 β = c β κ ˉ 2 β , where we defined c β ≜ 2 c M ϕ 2 + 1 κ ˉ 2 β c_{\beta}\triangleq2cM_{\phi}^{2}+\frac{1}{\bar{\kappa}^{2}\beta} c β ≜ 2 c M ϕ 2 + κ ˉ 2 β 1 κ ˉ 2 β ≫ 1 \bar{\kappa}^{2}\beta\gg1 κ ˉ 2 β ≫ 1 c β = O ( 1 ) c_{\beta}=O(1) c β = O ( 1 ) α \alpha α non-negative measurable function φ : R N p → R \varphi:\mathbb{R}^{N_{p}}\to\mathbb{R} φ : R N p → R
E p t φ ( θ t ) ≤ η ∥ φ ( θ 0 ) ∥ L c β κ ˉ 2 β . (32) \mathbb{E}_{p_{t}}\varphi(\theta_{t})\leq\eta\|\varphi(\theta_{0})\|_{L^{c_{\beta}\bar{\kappa}^{2}\beta}}. \tag*{(32)} E p t φ ( θ t ) ≤ η ∥ φ ( θ 0 ) ∥ L c β κ ˉ 2 β . (32) We proceed to bound each of the four terms in Eq. (31) using the change-of-measure inequality Eq. (32) .
First term
We have
− n − 1 E [ r ( θ t ) T K N T K ( θ t ) r ( θ t ) ] ≤ − n − 1 E [ ∥ r ( θ t ) ∥ 2 2 λ min ( K N T K ( θ t ) ) ] = − n − 1 E [ 2 n L ( θ t ) λ min ( K N T K ( θ t ) ) ] = − 2 E [ L ( θ t ) λ min ( K N T K ( θ t ) ) ] . \begin{aligned}
-n^{-1}\mathbb{E}\left[\boldsymbol{r}(\theta_{t})^{\mathsf{T}}\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\boldsymbol{r}(\theta_{t})\right]\leq-n^{-1}\mathbb{E}\left[\|\boldsymbol{r}(\theta_{t})\|_{2}^{2}\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right] \\
&=-n^{-1}\mathbb{E}\left[2n\mathcal{L}(\theta_{t})\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right] \\
&=-2\mathbb{E}\left[\mathcal{L}(\theta_{t})\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right].
\end{aligned} − n − 1 E [ r ( θ t ) T K NTK ( θ t ) r ( θ t ) ] ≤ − n − 1 E [ ∥ r ( θ t ) ∥ 2 2 λ m i n ( K NTK ( θ t ) ) ] = − n − 1 E [ 2 n L ( θ t ) λ m i n ( K NTK ( θ t ) ) ] = − 2 E [ L ( θ t ) λ m i n ( K NTK ( θ t ) ) ] . By Lemma 5.3, if m = Ω ( n λ 1 ln n λ 1 ln n ) m=\Omega\left(\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\ln n\right) m = Ω ( λ 1 n ln λ 1 n ln n ) λ min ( K N T K ( θ 0 ) ) ≥ c 0 λ 1 s 2 m \lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{0})\right)\geq c_{0}\lambda_{1}s^{2}m λ m i n ( K NTK ( θ 0 ) ) ≥ c 0 λ 1 s 2 m 1 − exp ( − c ~ 1 n λ 1 m ln n ) 1-\exp\left(-\widetilde{c}_{1}\frac{n}{\lambda_{1}}\frac{m}{\ln n}\right) 1 − exp ( − c 1 λ 1 n l n n m ) λ ⋆ ≜ c 0 λ 1 s 2 m \lambda_{\star}\triangleq c_{0}\lambda_{1}s^{2}m λ ⋆ ≜ c 0 λ 1 s 2 m δ ≜ exp ( − c ~ 1 n λ 1 m ln n ) \delta\triangleq\exp\left(-\widetilde{c}_{1}\frac{n}{\lambda_{1}}\frac{m}{\ln n}\right) δ ≜ exp ( − c 1 λ 1 n l n n m )
− n − 1 E [ r ( θ t ) T K N T K ( θ t ) r ( θ t ) ] ≤ − 2 E [ L ( θ t ) λ min ( K N T K ( θ t ) ) ] = − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) ( λ ⋆ − λ min ( K N T K ( θ t ) ) ) ] ≤ − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) max ( 0 , λ ⋆ − λ min ( K N T K ( θ t ) ) ) ] . \begin{aligned}
-n^{-1}\mathbb{E}\left[\boldsymbol{r}(\theta_{t})^{\mathsf{T}}\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\boldsymbol{r}(\theta_{t})\right] \\
&\leq-2\mathbb{E}\left[\mathcal{L}(\theta_{t})\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right] \\
&=-2\lambda_{\star}\mathbb{E}\left[\mathcal{L}(\theta_{t})\right]+2\mathbb{E}\left[\mathcal{L}(\theta_{t})\left(\lambda_{\star}-\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right)\right] \\
&\leq-2\lambda_{\star}\mathbb{E}\left[\mathcal{L}(\theta_{t})\right]+2\mathbb{E}\left[\mathcal{L}(\theta_{t})\max\left(0,\lambda_{\star}-\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right)\right].
\end{aligned} − n − 1 E [ r ( θ t ) T K NTK ( θ t ) r ( θ t ) ] ≤ − 2 E [ L ( θ t ) λ m i n ( K NTK ( θ t ) ) ] = − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) ( λ ⋆ − λ m i n ( K NTK ( θ t ) ) ) ] ≤ − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ t ) ) ) ] . We upper bound E [ L ( θ t ) max ( 0 , λ ⋆ − λ min ( K N T K ( θ t ) ) ) ] \mathbb{E}\left[\mathcal{L}(\theta_{t})\max\left(0,\lambda_{\star}-\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right)\right] E [ L ( θ t ) max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ t ) ) ) ] (32) ):
E [ L ( θ t ) max ( 0 , λ ⋆ − λ min ( K N T K ( θ t ) ) ) ] ≤ η ∥ L ( θ 0 ) max ( 0 , λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) ∥ L c β κ ˉ 2 β \begin{aligned}
\mathbb{E}\left[\mathcal{L}(\theta_{t})\max\left(0,\lambda_{\star}-\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{t})\right)\right)\right] \\
&\leq\eta\left\|\mathcal{L}(\theta_{0})\max\left(0,\lambda_{\star}-\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{0})\right)\right)\right\|_{L^{c_{\beta}\bar{\kappa}^{2}\beta}}
\end{aligned} E [ L ( θ t ) max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ t ) ) ) ] ≤ η ∥ L ( θ 0 ) max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ 0 ) ) ) ∥ L c β κ ˉ 2 β [H o ¨ lder, Claim E.7] ≤ η ∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ∥ max ( 0 , λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) ∥ L 2 c β κ ˉ 2 β . \text{[Hölder, Claim E.7]}\leq\eta\|\mathcal{L}(\theta_{0})\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}}\left\|\max\left(0,\lambda_{\star}-\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_{0})\right)\right)\right\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}}. [H o ¨ lder, Claim E.7] ≤ η ∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ∥ max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ 0 ) ) ) ∥ L 2 c β κ ˉ 2 β . From Proposition E.8,
∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ≤ ∥ R y 2 + 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ L 2 c β κ ˉ 2 β \| \mathcal{L}(\theta_0) \|_{L^{2c_\beta\bar{\kappa}^{2\beta}}}
\le\left\| R_y^2+\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right\|_{L^{2c_\beta\bar{\kappa}^{2\beta}}} ∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ≤ R y 2 + n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) L 2 c β κ ˉ 2 β ≤ [ Minkowski’s inequality ] ≤ ∥ R y 2 ∥ L 2 c β κ ˉ 2 β + ∥ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ≤ R y 2 + C ⋅ 2 c β κ ˉ 2 β ∥ 1 n ∑ i ∈ [ n ] f i 2 ( θ 0 ) ∥ ψ 1 ≤ R y 2 + C 2 c β κ ˉ − 4 β . \begin{aligned}
\mathrel{\phantom{\le}}[\text{Minkowski's inequality}]
\le\left\|R_y^2\right\|_{L^{2c_\beta\bar{\kappa}^{2\beta}}}
+\left\|\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right\|_{L^{2c_\beta\bar{\kappa}^{2\beta}}} \\
&\le R_y^2+C\cdot2c_\beta\bar{\kappa}^{2\beta}
\left\|\frac{1}{n}\sum_{i\in[n]}f_i^2(\theta_0)\right\|_{\psi_1} \\
&\le R_y^2+C_2c_\beta\bar{\kappa}^{-4}\beta.
\end{aligned} ≤ [ Minkowski’s inequality ] ≤ R y 2 L 2 c β κ ˉ 2 β + n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) L 2 c β κ ˉ 2 β ≤ R y 2 + C ⋅ 2 c β κ ˉ 2 β n 1 i ∈ [ n ] ∑ f i 2 ( θ 0 ) ψ 1 ≤ R y 2 + C 2 c β κ ˉ − 4 β . Splitting according to the high-probability event
E = { λ ⋆ ≤ λ min ( K N T K ( θ 0 ) ) } , \mathcal{E}=\{\lambda_\star\le\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0))\}, E = { λ ⋆ ≤ λ m i n ( K NTK ( θ 0 ))} , we have for the second term,
∥ max ( 0 , λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) ∥ L 2 c β κ ˉ 2 β = ( E [ max ( 0 , λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) 2 c β κ ˉ 2 β ] ) 1 / ( 2 c β κ ˉ 2 β ) ≤ ( 0 ⋅ P ( λ ⋆ ≤ λ min ( K N T K ( θ 0 ) ) ) + λ ⋆ 2 c β κ ˉ 2 β P ( λ ⋆ > λ min ( K N T K ( θ 0 ) ) ) ) 1 / ( 2 c β κ ˉ 2 β ) ≤ λ ⋆ δ 1 / ( 2 c β κ ˉ 2 β ) = c 0 λ 1 s 2 2 m exp ( − c ~ 2 m c β κ ˉ 2 β λ 1 2 ln n λ 1 ln n λ 1 ) , \begin{aligned}
\left\|\max\left(0,\lambda_\star-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0))\right)\right\|_{L^{2c_\beta\bar{\kappa}^{2\beta}}} \\
&=\left(\mathbb{E}\left[\max\left(0,\lambda_\star-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0))\right)^{2c_\beta\bar{\kappa}^{2\beta}}\right]\right)^{1/(2c_\beta\bar{\kappa}^{2\beta})} \\
&\le\left(0\cdot\mathbb{P}\left(\lambda_\star\le\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0))\right)+\lambda_\star^{2c_\beta\bar{\kappa}^{2\beta}}\mathbb{P}\left(\lambda_\star>\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0))\right)\right)^{1/(2c_\beta\bar{\kappa}^{2\beta})} \\
&\le\lambda_\star\delta^{1/(2c_\beta\bar{\kappa}^{2\beta})}
=c_0\lambda_1s_2^{2m}\exp\left(-\widetilde{c}_2\frac{m}{c_\beta\bar{\kappa}^{2\beta}\lambda_1^2}\ln\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right),
\end{aligned} ∥ max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ 0 )) ) ∥ L 2 c β κ ˉ 2 β = ( E [ max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ 0 )) ) 2 c β κ ˉ 2 β ] ) 1/ ( 2 c β κ ˉ 2 β ) ≤ ( 0 ⋅ P ( λ ⋆ ≤ λ m i n ( K NTK ( θ 0 )) ) + λ ⋆ 2 c β κ ˉ 2 β P ( λ ⋆ > λ m i n ( K NTK ( θ 0 )) ) ) 1/ ( 2 c β κ ˉ 2 β ) ≤ λ ⋆ δ 1/ ( 2 c β κ ˉ 2 β ) = c 0 λ 1 s 2 2 m exp ( − c 2 c β κ ˉ 2 β λ 1 2 m ln λ 1 n ln λ 1 n ) , where in the first transition we used the fact that the empirical NTK matrix is PSD, hence λ min ( K N T K ) ≥ 0 \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}})\ge0 λ m i n ( K NTK ) ≥ 0
− n − 1 E [ r ( θ t ) ⊤ K N T K ( θ t ) r ( θ t ) ] ≤ − 2 c 0 λ 1 s 2 2 m E [ L ( θ t ) ] + 2 η c 0 λ 1 s 2 2 m ( R y 2 + C 2 c β κ ˉ − 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β λ 1 2 ln n λ 1 ln n λ 1 ) -n^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^\top\mathbf{K}_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)\right]
\le-2c_0\lambda_1s_2^{2m}\mathbb{E}[\mathcal{L}(\theta_t)]
+2\eta c_0\lambda_1s_2^{2m}\left(R_y^2+C_2c_\beta\bar{\kappa}^{-4}\beta\right)
\exp\left(-\widetilde{c}_2\frac{m}{c_\beta\bar{\kappa}^{2\beta}\lambda_1^2}\ln\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right) − n − 1 E [ r ( θ t ) ⊤ K NTK ( θ t ) r ( θ t ) ] ≤ − 2 c 0 λ 1 s 2 2 m E [ L ( θ t )] + 2 η c 0 λ 1 s 2 2 m ( R y 2 + C 2 c β κ ˉ − 4 β ) exp ( − c 2 c β κ ˉ 2 β λ 1 2 m ln λ 1 n ln λ 1 n ) ≲ − λ 1 s 2 2 m E [ L ( θ t ) ] + λ 1 s 2 2 m ( 1 + κ ˉ − 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β λ 1 2 ln n λ 1 ln n λ 1 ) . \lesssim-\lambda_1s_2^{2m}\mathbb{E}[\mathcal{L}(\theta_t)]
+\lambda_1s_2^{2m}\left(1+\bar{\kappa}^{-4}\beta\right)
\exp\left(-\widetilde{c}_2\frac{m}{c_\beta\bar{\kappa}^{2\beta}\lambda_1^2}\ln\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right). ≲ − λ 1 s 2 2 m E [ L ( θ t )] + λ 1 s 2 2 m ( 1 + κ ˉ − 4 β ) exp ( − c 2 c β κ ˉ 2 β λ 1 2 m ln λ 1 n ln λ 1 n ) . Second term
Observe that,
β − 1 E Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) = β − 1 E [ ∑ j = 1 N p ∑ i = 1 n ( ( ∇ f i ( θ t ) ) j ) 2 ] = β − 1 ∑ i = 1 n E [ ∥ ∇ f i ( θ t ) ∥ 2 ] . \begin{aligned}
\beta^{-1}\mathbb{E}\operatorname{Tr}\left(\nabla f(\theta_t)^\top\nabla f(\theta_t)\right)
=\beta^{-1}\mathbb{E}\left[\sum_{j=1}^{N_p}\sum_{i=1}^{n}\left((\nabla f_i(\theta_t))_j\right)^2\right] \\
&=\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\|\nabla f_i(\theta_t)\|^2\right].
\end{aligned} β − 1 E Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) = β − 1 E j = 1 ∑ N p i = 1 ∑ n ( ( ∇ f i ( θ t ) ) j ) 2 = β − 1 i = 1 ∑ n E [ ∥∇ f i ( θ t ) ∥ 2 ] . Now, for all i ∈ [ n ] i \in[n] i ∈ [ n ] α ( x i ) ≜ ϕ ( s 1 W x i ) \alpha(x_i) \triangleq\phi(s_1W x_i) α ( x i ) ≜ ϕ ( s 1 W x i )
f i ( θ ) = s 2 v T α ( x i ) ⇒ ∂ f i ( θ ) ∂ v = s 2 α ( x i ) = s 2 ϕ ( s 1 W x i ) , \begin{aligned}
f_i(\theta)=s_2\mathbf{v}^{\mathsf{T}}\alpha(x_i) \\
&\Rightarrow\frac{\partial f_i(\theta)}{\partial\mathbf{v}}=s_2\alpha(x_i)=s_2\phi(s_1W x_i),
\end{aligned} f i ( θ ) = s 2 v T α ( x i ) ⇒ ∂ v ∂ f i ( θ ) = s 2 α ( x i ) = s 2 ϕ ( s 1 W x i ) , ∀ j ∈ [ m ] , ∂ f i ( θ ) ∂ w j = s 2 v j ∂ ( α ( x i ) ) j ∂ w j , \forall j\in[m],\qquad\frac{\partial f_i(\theta)}{\partial\mathbf{w}_j}=s_2v_j\frac{\partial(\alpha(x_i))_j}{\partial\mathbf{w}_j}, ∀ j ∈ [ m ] , ∂ w j ∂ f i ( θ ) = s 2 v j ∂ w j ∂ ( α ( x i ) ) j , ∂ ( α ( x i ) ) j ∂ w j = ϕ ′ ( s 1 w j T x i ) s 1 x i ⇒ ∀ j ∈ [ m ] , ∂ f i ( θ ) ∂ w j = s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i . \begin{aligned}
\frac{\partial(\alpha(x_i))_j}{\partial\mathbf{w}_j}=\phi'(s_1\mathbf{w}_j^{\mathsf{T}}x_i)s_1x_i \\
&\Rightarrow\forall j\in[m],\qquad\frac{\partial f_i(\theta)}{\partial\mathbf{w}_j}=s_2v_j\phi'(s_1\mathbf{w}_j^{\mathsf{T}}x_i)s_1x_i.
\end{aligned} ∂ w j ∂ ( α ( x i ) ) j = ϕ ′ ( s 1 w j T x i ) s 1 x i ⇒ ∀ j ∈ [ m ] , ∂ w j ∂ f i ( θ ) = s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i . Then we have:
∥ ∇ f i ( θ ) ∥ 2 = ∥ s 2 ϕ ( s 1 W x i ) ∥ 2 + ∑ j = 1 m ∥ s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i ∥ 2 \|\nabla f_i(\theta)\|^2=\|s_2\phi(s_1W x_i)\|^2+\sum_{j=1}^{m}\|s_2v_j\phi'(s_1\mathbf{w}_j^{\mathsf{T}}x_i)s_1x_i\|^2 ∥∇ f i ( θ ) ∥ 2 = ∥ s 2 ϕ ( s 1 W x i ) ∥ 2 + j = 1 ∑ m ∥ s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i ∥ 2 [ ∣ ϕ ( ⋅ ) ∣ ≤ α ϕ ∣ ⋅ ∣ , ∣ ϕ ′ ∣ ≤ α ϕ ] ≤ α ϕ 2 s 2 2 s 1 2 ( ∥ W x i ∥ 2 + ∑ j = 1 m v j 2 ∥ x i ∥ 2 ) = α ϕ 2 s 2 2 s 1 2 c ϕ , σ 0 d ( ν 0 2 a 1 ∥ z ( 1 ) ∥ 2 + 1 a 2 ∥ z ( 2 ) ∥ 2 ) , \begin{aligned}
\left[|\phi(\cdot)|\leq\alpha_\phi|\cdot|,\ |\phi'|\leq\alpha_\phi\right]\leq\alpha_\phi^2s_2^2s_1^2\left(\|W x_i\|^2+\sum_{j=1}^{m}v_j^2\|x_i\|^2\right) \\
&=\alpha_\phi^2s_2^2s_1^2c_{\phi,\sigma_0}d\left(\frac{\nu_0^2}{a_1}\|z^{(1)}\|^2+\frac{1}{a_2}\|z^{(2)}\|^2\right),
\end{aligned} [ ∣ ϕ ( ⋅ ) ∣ ≤ α ϕ ∣ ⋅ ∣ , ∣ ϕ ′ ∣ ≤ α ϕ ] ≤ α ϕ 2 s 2 2 s 1 2 ( ∥ W x i ∥ 2 + j = 1 ∑ m v j 2 ∥ x i ∥ 2 ) = α ϕ 2 s 2 2 s 1 2 c ϕ , σ 0 d ( a 1 ν 0 2 ∥ z ( 1 ) ∥ 2 + a 2 1 ∥ z ( 2 ) ∥ 2 ) , where z ( 1 ) , z ( 2 ) ∼ N ( 0 m , I m ) z^{(1)},z^{(2)}\sim\mathcal{N}(0_m,I_m) z ( 1 ) , z ( 2 ) ∼ N ( 0 m , I m ) p 0 p_0 p 0
Observe that for all p ≥ 1 p\geq1 p ≥ 1
∥ ∥ z ∥ 2 ∥ L p = ∥ ∥ z ∥ 2 ∥ L 2 p 2 ≤ ( C 2 p + m ) 2 ≤ 4 C p 2 + 2 m . \left\|\|z\|^2\right\|_{L_p}=\left\|\|z\|^2\right\|_{L_{2p}}^2\leq\left(C\sqrt{2p}+\sqrt{m}\right)^2\leq4C_p^2+2m. ∥ z ∥ 2 L p = ∥ z ∥ 2 L 2 p 2 ≤ ( C 2 p + m ) 2 ≤ 4 C p 2 + 2 m . So, at initialization, using Minkowski’s inequality,
∥ ∥ ∇ f i ( θ ) ∥ 2 ∥ L p ≤ α ϕ 2 s 2 2 s 1 2 d ( ν 0 2 a 1 ∥ ∥ z ( 1 ) ∥ 2 ∥ L p + 1 a 2 ∥ ∥ z ( 2 ) ∥ 2 ∥ L p ) ≤ α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 d ( ν 0 2 a 1 + 1 a 2 ) ( 4 C p 2 + 2 m ) . \begin{aligned}
\left\|\|\nabla f_i(\theta)\|^2\right\|_{L_p}\leq\alpha_\phi^2s_2^2s_1^2d\left(\frac{\nu_0^2}{a_1}\left\|\|z^{(1)}\|^2\right\|_{L_p}+\frac{1}{a_2}\left\|\|z^{(2)}\|^2\right\|_{L_p}\right) \\
&\leq\alpha_\phi^2c_{\phi,\sigma_0}s_2^2s_1^2d\left(\frac{\nu_0^2}{a_1}+\frac{1}{a_2}\right)(4C_p^2+2m).
\end{aligned} ∥∇ f i ( θ ) ∥ 2 L p ≤ α ϕ 2 s 2 2 s 1 2 d ( a 1 ν 0 2 ∥ z ( 1 ) ∥ 2 L p + a 2 1 ∥ z ( 2 ) ∥ 2 L p ) ≤ α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 d ( a 1 ν 0 2 + a 2 1 ) ( 4 C p 2 + 2 m ) . Using the second law ((32) ), we get that at time t ≥ 0 t\geq0 t ≥ 0 i ∈ [ n ] i\in[n] i ∈ [ n ]
E [ ∥ ∇ f i ( θ t ) ∥ 2 ] ≤ η ∥ ∇ f i ( θ 0 ) ∥ L c β κ ˉ 2 β 2 ≤ η α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 d ( ν 0 2 a 1 + 1 a 2 ) ( 4 C 2 c β κ ˉ 2 β + 2 m ) . \mathbb{E}\left[\|\nabla f_i(\theta_t)\|^2\right]\leq\eta\|\nabla f_i(\theta_0)\|_{L_{c_\beta\bar{\kappa}^2\beta}}^2\leq\eta\alpha_\phi^2c_{\phi,\sigma_0}s_2^2s_1^2d\left(\frac{\nu_0^2}{a_1}+\frac{1}{a_2}\right)\left(4C^2c_\beta\bar{\kappa}^2\beta+2m\right). E [ ∥∇ f i ( θ t ) ∥ 2 ] ≤ η ∥∇ f i ( θ 0 ) ∥ L c β κ ˉ 2 β 2 ≤ η α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 d ( a 1 ν 0 2 + a 2 1 ) ( 4 C 2 c β κ ˉ 2 β + 2 m ) . Summing over i i i
β − 1 E Tr ( ∇ f ( θ t ) T ∇ f ( θ t ) ) = β − 1 ∑ i = 1 n E [ ∥ ∇ f i ( θ t ) ∥ 2 ] ≤ η α ϕ 2 c ϕ , σ 0 β − 1 n s 2 2 s 1 2 d ( ν 0 2 a 1 + 1 a 2 ) ( C 3 c β κ ˉ 2 β + 2 m ) . \begin{aligned}
\beta^{-1}\mathbb{E}\operatorname{Tr}\left(\nabla f(\theta_t)^{\mathsf{T}}\nabla f(\theta_t)\right)=\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\|\nabla f_i(\theta_t)\|^2\right] \\
&\leq\eta\alpha_\phi^2c_{\phi,\sigma_0}\beta^{-1}ns_2^2s_1^2d\left(\frac{\nu_0^2}{a_1}+\frac{1}{a_2}\right)\left(C_3c_\beta\bar{\kappa}^2\beta+2m\right).
\end{aligned} β − 1 E Tr ( ∇ f ( θ t ) T ∇ f ( θ t ) ) = β − 1 i = 1 ∑ n E [ ∥∇ f i ( θ t ) ∥ 2 ] ≤ η α ϕ 2 c ϕ , σ 0 β − 1 n s 2 2 s 1 2 d ( a 1 ν 0 2 + a 2 1 ) ( C 3 c β κ ˉ 2 β + 2 m ) . ≲ β − 1 n s 2 2 s 1 2 d ( 1 a 1 + 1 a 2 ) ( κ ˉ 2 β + m ) . \lesssim\beta^{-1}ns_2^2s_1^2d\left(\frac{1}{a_1}+\frac{1}{a_2}\right)(\bar{\kappa}^2\beta+m). ≲ β − 1 n s 2 2 s 1 2 d ( a 1 1 + a 2 1 ) ( κ ˉ 2 β + m ) . Third term
We have
β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] = β − 1 ∑ i = 1 n E [ r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) Δ f i ( θ t ) ∣ ] \beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]=\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)\Delta f_i(\theta_t)\right]
\leq\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\left|r_i(\theta_t)\Delta f_i(\theta_t)\right|\right] β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] = β − 1 i = 1 ∑ n E [ r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 i = 1 ∑ n E [ ∣ r i ( θ t ) Δ f i ( θ t ) ∣ ] [Cauchy-Schwarz] ≤ β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] \leq\beta^{-1}\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[|r_i(\theta_t)|^2\right]\mathbb{E}\left[|\Delta f_i(\theta_t)|^2\right]} ≤ β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] E [ ∣Δ f i ( θ t ) ∣ 2 ]
= β − 1 ∑ i = 1 n 2 E [ ℓ i ( θ t ) ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] , =\beta^{-1}\sum_{i=1}^{n}\sqrt{2\mathbb{E}\left[\ell_i(\theta_t)\right]\mathbb{E}\left[|\Delta f_i(\theta_t)|^2\right]}, = β − 1 i = 1 ∑ n 2 E [ ℓ i ( θ t ) ] E [ ∣Δ f i ( θ t ) ∣ 2 ] , where ℓ i ( θ ) = 1 2 ∣ r i ( θ ) ∣ 2 \ell_i(\theta)=\frac{1}{2}|r_i(\theta)|^2 ℓ i ( θ ) = 2 1 ∣ r i ( θ ) ∣ 2
Note that, for all i ∈ [ n ] i\in[n] i ∈ [ n ]
Δ f i ( θ ) = ∑ j ∈ [ m ] ∂ 2 f i ( θ ) ∂ v j 2 + ∑ j ∈ [ m ] ∑ k ∈ [ d ] ∂ 2 f i ( θ ) ∂ W j k 2 . \Delta f_i(\theta)=\sum_{j\in[m]}\frac{\partial^2f_i(\theta)}{\partial v_j^2}+\sum_{j\in[m]}\sum_{k\in[d]}\frac{\partial^2f_i(\theta)}{\partial W_{jk}^2}. Δ f i ( θ ) = j ∈ [ m ] ∑ ∂ v j 2 ∂ 2 f i ( θ ) + j ∈ [ m ] ∑ k ∈ [ d ] ∑ ∂ W j k 2 ∂ 2 f i ( θ ) . Observe that for all i ∈ [ n ] i\in[n] i ∈ [ n ] f i ( θ ) = s 2 v T α ( x i ) f_i(\theta)=s_2\mathbf{v}^{\mathsf{T}}\boldsymbol{\alpha}(x_i) f i ( θ ) = s 2 v T α ( x i ) α ( x i ) \boldsymbol{\alpha}(x_i) α ( x i ) v \mathbf{v} v j ∈ [ m ] j\in[m] j ∈ [ m ]
∂ f i ( θ ) ∂ v j = s 2 α j ( x i ) , \frac{\partial f_i(\theta)}{\partial v_j}=s_2\alpha_j(x_i), ∂ v j ∂ f i ( θ ) = s 2 α j ( x i ) , ∂ 2 f i ( θ ) ∂ v j 2 = 0. \frac{\partial^2 f_i(\theta)}{\partial v_j^2}=0. ∂ v j 2 ∂ 2 f i ( θ ) = 0. In terms of the derivatives w.r.t. W j k W_{jk} W j k ∀ j ∈ [ m ] , k ∈ [ d ] \forall j\in[m],\,k\in[d] ∀ j ∈ [ m ] , k ∈ [ d ]
∂ f i ( θ ) ∂ W j k = s 2 v T ∂ α ( x i ) ∂ W j k = s 2 ∑ j ′ ∈ [ m ] v j ′ ∂ α j ′ ( x i ) ∂ W j k , \frac{\partial f_i(\theta)}{\partial W_{jk}}=s_2\mathbf{v}^{\mathsf{T}}\frac{\partial\boldsymbol{\alpha}(x_i)}{\partial W_{jk}}=s_2\sum_{j'\in[m]}v_{j'}\frac{\partial\alpha_{j'}(x_i)}{\partial W_{jk}}, ∂ W j k ∂ f i ( θ ) = s 2 v T ∂ W j k ∂ α ( x i ) = s 2 j ′ ∈ [ m ] ∑ v j ′ ∂ W j k ∂ α j ′ ( x i ) , where
α j ′ ( x i ) = ϕ ( s 1 w j ′ T x i ) , \alpha_{j'}(x_i)=\phi\left(s_1\mathbf{w}_{j'}^{\mathsf{T}}x_i\right), α j ′ ( x i ) = ϕ ( s 1 w j ′ T x i ) , denoting w j ′ ∈ R d \mathbf{w}_{j'}\in\mathbb{R}^{d} w j ′ ∈ R d j ′ j' j ′ W \mathbf{W} W
∀ j ′ ≠ j , ∂ α j ′ ( x i ) ∂ W j k = 0 , \forall j'\neq j,\quad\frac{\partial\alpha_{j'}(x_i)}{\partial W_{jk}}=0, ∀ j ′ = j , ∂ W j k ∂ α j ′ ( x i ) = 0 , ∂ α j ( x i ) ∂ W j k = ϕ ′ ( s 1 w j T x i ) s 1 x i , k . \frac{\partial\alpha_j(x_i)}{\partial W_{jk}}=\phi'\left(s_1\mathbf{w}_j^{\mathsf{T}}x_i\right)s_1x_{i,k}. ∂ W j k ∂ α j ( x i ) = ϕ ′ ( s 1 w j T x i ) s 1 x i , k . We then have
∂ f i ( θ ) ∂ W j k = s 2 s 1 v j x i , k ϕ ′ ( s 1 w j T x i ) (33) \frac{\partial f_i(\theta)}{\partial W_{jk}}=s_2s_1v_jx_{i,k}\phi'\left(s_1\mathbf{w}_j^{\mathsf{T}}x_i\right) \tag*{(33)} ∂ W j k ∂ f i ( θ ) = s 2 s 1 v j x i , k ϕ ′ ( s 1 w j T x i ) (33) ⟹ ∂ 2 f i ( θ ) ∂ W j k 2 = s 2 s 1 v j x i , k ϕ ′ ′ ( s 1 w j T x i ) s 1 x i , k = s 2 s 1 2 v j x i , k 2 ϕ ′ ′ ( s 1 w j T x i ) . \Longrightarrow\frac{\partial^2 f_i(\theta)}{\partial W_{jk}^2}=s_2s_1v_jx_{i,k}\phi''\left(s_1\mathbf{w}_j^{\mathsf{T}}x_i\right)s_1x_{i,k}
=s_2s_1^2v_jx_{i,k}^2\phi''\left(s_1\mathbf{w}_j^{\mathsf{T}}x_i\right). ⟹ ∂ W j k 2 ∂ 2 f i ( θ ) = s 2 s 1 v j x i , k ϕ ′′ ( s 1 w j T x i ) s 1 x i , k = s 2 s 1 2 v j x i , k 2 ϕ ′′ ( s 1 w j T x i ) . Calculating the Laplacian, we get
Δ f i ( θ ) = ∑ j ∈ [ m ] ∂ 2 f i ( θ ) ∂ v j 2 + ∑ j ∈ [ m ] ∑ k ∈ [ d ] ∂ 2 f i ( θ ) ∂ W j k 2 = ∑ j ∈ [ m ] ∑ k ∈ [ d ] s 2 s 1 2 v j x i , k 2 ϕ ′ ′ ( s 1 w j T x i ) = s 2 s 1 2 ( ∑ k ∈ [ d ] x i , k 2 ) ( ∑ j ∈ [ m ] v j ϕ ′ ′ ( s 1 w j T x i ) ) = [ ∥ x i ∥ 2 = c ϕ , σ 0 d ] = s 2 s 1 2 c ϕ , σ 0 d v T ϕ ′ ′ ( s 1 W x i ) . \begin{aligned}
\Delta f_i(\boldsymbol{\theta}) &= \sum_{j\in[m]} \frac{\partial^2 f_i(\boldsymbol{\theta})}{\partial v_j^2}+\sum_{j\in[m]}\sum_{k\in[d]}\frac{\partial^2 f_i(\boldsymbol{\theta})}{\partial W_{jk}^2}\\
&=\sum_{j\in[m]}\sum_{k\in[d]}s_2s_1^2v_jx_{i,k}^2\phi^{\prime\prime}\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\\
&=s_2s_1^2\left(\sum_{k\in[d]}x_{i,k}^2\right)\left(\sum_{j\in[m]}v_j\phi^{\prime\prime}\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\right)\\
&=\left[\left\|\mathbf{x}_i\right\|^2=c_{\phi,\sigma_0}d\right]=s_2s_1^2c_{\phi,\sigma_0}d\mathbf{v}^{\mathsf{T}}\phi^{\prime\prime}\left(s_1\mathbf{W}\mathbf{x}_i\right).
\end{aligned} Δ f i ( θ ) = j ∈ [ m ] ∑ ∂ v j 2 ∂ 2 f i ( θ ) + j ∈ [ m ] ∑ k ∈ [ d ] ∑ ∂ W j k 2 ∂ 2 f i ( θ ) = j ∈ [ m ] ∑ k ∈ [ d ] ∑ s 2 s 1 2 v j x i , k 2 ϕ ′′ ( s 1 w j T x i ) = s 2 s 1 2 k ∈ [ d ] ∑ x i , k 2 j ∈ [ m ] ∑ v j ϕ ′′ ( s 1 w j T x i ) = [ ∥ x i ∥ 2 = c ϕ , σ 0 d ] = s 2 s 1 2 c ϕ , σ 0 d v T ϕ ′′ ( s 1 W x i ) . In order to bound E [ ∣ Δ f i ( θ ) ∣ 2 ] \mathbb{E}\left[\left|\Delta f_i(\boldsymbol{\theta})\right|^2\right] E [ ∣ Δ f i ( θ ) ∣ 2 ] ∥ X ∥ L p , ρ ≜ ( E q [ ∣ X ∣ p ] ) 1 / p \left\|\mathbf{X}\right\|_{L^p,\rho}\triangleq\left(\mathbb{E}_q\left[|\mathbf{X}|^p\right]\right)^{1/p} ∥ X ∥ L p , ρ ≜ ( E q [ ∣ X ∣ p ] ) 1/ p
E p t [ ∣ Δ f i ( θ t ) ∣ 2 ] = E p t [ ∣ s 2 s 1 2 c ϕ , σ 0 d v T ϕ ′ ′ ( s 1 W x i ) ∣ 2 ] ≤ η ( s 2 s 1 2 c ϕ , σ 0 d ) 2 ∥ v T ϕ ′ ′ ( s 1 W x i ) ∥ L c β κ ˉ 2 β , ρ 0 2 . \begin{aligned}
\mathbb{E}_{p_t}\left[\left|\Delta f_i(\boldsymbol{\theta}_t)\right|^2\right]
&=\mathbb{E}_{p_t}\left[\left|s_2s_1^2c_{\phi,\sigma_0}d\mathbf{v}^{\mathsf{T}}\phi^{\prime\prime}\left(s_1\mathbf{W}\mathbf{x}_i\right)\right|^2\right]\\
&\leq\eta\left(s_2s_1^2c_{\phi,\sigma_0}d\right)^2\left\|\mathbf{v}^{\mathsf{T}}\phi^{\prime\prime}\left(s_1\mathbf{W}\mathbf{x}_i\right)\right\|_{L^{c_\beta\bar{\kappa}^2\beta,\rho_0}}^2.
\end{aligned} E p t [ ∣ Δ f i ( θ t ) ∣ 2 ] = E p t [ s 2 s 1 2 c ϕ , σ 0 d v T ϕ ′′ ( s 1 W x i ) 2 ] ≤ η ( s 2 s 1 2 c ϕ , σ 0 d ) 2 v T ϕ ′′ ( s 1 W x i ) L c β κ ˉ 2 β , ρ 0 2 . We have
v T ϕ ′ ′ ( s 1 W x i ) = ∑ j ∈ [ m ] v j ϕ ′ ′ ( s 1 w j T x i ) . \mathbf{v}^{\mathsf{T}}\phi^{\prime\prime}\left(s_1\mathbf{W}\mathbf{x}_i\right)=\sum_{j\in[m]}v_j\phi^{\prime\prime}\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right). v T ϕ ′′ ( s 1 W x i ) = j ∈ [ m ] ∑ v j ϕ ′′ ( s 1 w j T x i ) . This is a sum of zero-mean, independent and subgaussian random variables, each of them satisfying
∣ v j ϕ ′ ′ ( s 1 w j T x i ) ∣ ≤ β ϕ ∣ v j ∣ \left|v_j\phi^{\prime\prime}\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\right|\leq\beta_\phi|v_j| v j ϕ ′′ ( s 1 w j T x i ) ≤ β ϕ ∣ v j ∣ ∥ ∣ v j ϕ ′ ′ ( s 1 w j T x i ) ∣ ∥ ψ 2 ≤ β ϕ ∥ v j ∥ ψ 2 ≤ C β ϕ 1 α 2 . \left\|\left|v_j\phi^{\prime\prime}\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\right|\right\|_{\psi_2}\leq\beta_\phi\left\|v_j\right\|_{\psi_2}\leq C\beta_\phi\frac{1}{\sqrt{\alpha_2}}. v j ϕ ′′ ( s 1 w j T x i ) ψ 2 ≤ β ϕ ∥ v j ∥ ψ 2 ≤ C β ϕ α 2 1 . Then, by Vershynin [111 Proposition 2.7.1] ,
∥ ∣ v T ϕ ′ ′ ( s 1 W x i ) ∣ 2 ∥ ψ 1 = ∥ ( v T ϕ ′ ′ ( s 1 W x i ) ) ∥ ψ 2 2 ≤ C 1 m ( C β ϕ 1 α 2 ) 2 = C 2 β ϕ 2 m α 2 . \begin{aligned}
\left\|\left|\mathbf{v}^{\mathsf{T}}\phi^{\prime\prime}\left(s_1\mathbf{W}\mathbf{x}_i\right)\right|^2\right\|_{\psi_1}
&=\left\|\left(\mathbf{v}^{\mathsf{T}}\phi^{\prime\prime}\left(s_1\mathbf{W}\mathbf{x}_i\right)\right)\right\|_{\psi_2}^2\\
&\leq C_1m\left(C\beta_\phi\frac{1}{\sqrt{\alpha_2}}\right)^2=C_2\beta_\phi^2\frac{m}{\alpha_2}.
\end{aligned} v T ϕ ′′ ( s 1 W x i ) 2 ψ 1 = ( v T ϕ ′′ ( s 1 W x i ) ) ψ 2 2 ≤ C 1 m ( C β ϕ α 2 1 ) 2 = C 2 β ϕ 2 α 2 m . Thus, by Vershynin [111 Proposition 2.8.1] ,
∥ ∣ v T ϕ ′ ′ ( s 1 W x i ) ∣ 2 ∥ L c β κ ˉ 2 β ≤ C 3 C 2 β ϕ 2 m α 2 c β κ ˉ 2 β . \left\|\left|\mathbf{v}^{\mathsf{T}}\phi^{\prime\prime}\left(s_1\mathbf{W}\mathbf{x}_i\right)\right|^2\right\|_{L^{c_\beta\bar{\kappa}^2\beta}}\leq C_3C_2\beta_\phi^2\frac{m}{\alpha_2}c_\beta\bar{\kappa}^2\beta. v T ϕ ′′ ( s 1 W x i ) 2 L c β κ ˉ 2 β ≤ C 3 C 2 β ϕ 2 α 2 m c β κ ˉ 2 β . Consequently,
β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 ∑ i = 1 n 2 E [ ℓ i ( θ t ) ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ β − 1 s 2 s 1 2 c ϕ , σ 0 d 2 η C 3 C 2 β ϕ 2 m α 2 c β κ ˉ 2 β ∑ i = 1 n E [ ℓ i ( θ t ) ] ≤ [Jensen] β − 1 / 2 s 2 s 1 2 c ϕ , σ 0 d 2 η C 3 C 2 β ϕ 2 m α 2 c β κ ˉ 2 β n 1 n ∑ i = 1 n E [ ℓ i ( θ t ) ] = β − 1 / 2 n s 2 s 1 2 c ϕ , σ 0 d 2 η C 3 C 2 β ϕ 2 m α 2 c β κ ˉ 2 β E L ( θ t ) ≲ β − 1 / 2 n s 2 s 1 2 d m α 2 κ ˉ 2 E L ( θ t ) . \begin{aligned}
\beta^{-1}\mathbb{E}\left[\sum_{i=1}^n r_i(\boldsymbol{\theta}_t)\Delta f_i(\boldsymbol{\theta}_t)\right]
&\leq\beta^{-1}\sum_{i=1}^n\sqrt{2\mathbb{E}\left[\ell_i(\boldsymbol{\theta}_t)\right]\mathbb{E}\left[\left|\Delta f_i(\boldsymbol{\theta}_t)\right|^2\right]}\\
&\leq\beta^{-1}s_2s_1^2c_{\phi,\sigma_0}d\sqrt{2\eta C_3C_2\beta_\phi^2\frac{m}{\alpha_2}c_\beta\bar{\kappa}^2\beta}\sum_{i=1}^n\sqrt{\mathbb{E}\left[\ell_i(\boldsymbol{\theta}_t)\right]}\\
&\underset{\text{[Jensen]}}{\leq}\beta^{-1/2}s_2s_1^2c_{\phi,\sigma_0}d\sqrt{2\eta C_3C_2\beta_\phi^2\frac{m}{\alpha_2}c_\beta\bar{\kappa}^2\beta n}\sqrt{\frac{1}{n}\sum_{i=1}^n\mathbb{E}\left[\ell_i(\boldsymbol{\theta}_t)\right]}\\
&=\beta^{-1/2}ns_2s_1^2c_{\phi,\sigma_0}d\sqrt{2\eta C_3C_2\beta_\phi^2\frac{m}{\alpha_2}c_\beta\bar{\kappa}^2\beta}\sqrt{\mathbb{E}\mathcal{L}(\boldsymbol{\theta}_t)}\\
&\lesssim\beta^{-1/2}ns_2s_1^2d\sqrt{\frac{m}{\alpha_2}\bar{\kappa}^2\sqrt{\mathbb{E}\mathcal{L}(\boldsymbol{\theta}_t)}}.
\end{aligned} β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 i = 1 ∑ n 2 E [ ℓ i ( θ t ) ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ β − 1 s 2 s 1 2 c ϕ , σ 0 d 2 η C 3 C 2 β ϕ 2 α 2 m c β κ ˉ 2 β i = 1 ∑ n E [ ℓ i ( θ t ) ] [Jensen] ≤ β − 1/2 s 2 s 1 2 c ϕ , σ 0 d 2 η C 3 C 2 β ϕ 2 α 2 m c β κ ˉ 2 β n n 1 i = 1 ∑ n E [ ℓ i ( θ t ) ] = β − 1/2 n s 2 s 1 2 c ϕ , σ 0 d 2 η C 3 C 2 β ϕ 2 α 2 m c β κ ˉ 2 β E L ( θ t ) ≲ β − 1/2 n s 2 s 1 2 d α 2 m κ ˉ 2 E L ( θ t ) . Fourth term
Note that under ϕ ( 0 ) = 0 \phi(0)=0 ϕ ( 0 ) = 0 ∣ ϕ ′ ∣ ≤ α ϕ |\phi'|\le\alpha_\phi ∣ ϕ ′ ∣ ≤ α ϕ ∣ ϕ ( ⋅ ) ∣ ≤ α ϕ ∣ ⋅ ∣ |\phi(\cdot)|\le\alpha_\phi|\cdot| ∣ ϕ ( ⋅ ) ∣ ≤ α ϕ ∣ ⋅ ∣ ∀ x ∈ R \forall x\in\mathbb{R} ∀ x ∈ R
∣ ϕ ( x ) − x ϕ ′ ( x ) ∣ ≤ ∣ ϕ ( x ) ∣ + ∣ x ϕ ′ ( x ) ∣ ≤ 2 α ϕ ∣ x ∣ , |\phi(x)-x\phi'(x)|\le|\phi(x)|+|x\phi'(x)|\le2\alpha_\phi|x|, ∣ ϕ ( x ) − x ϕ ′ ( x ) ∣ ≤ ∣ ϕ ( x ) ∣ + ∣ x ϕ ′ ( x ) ∣ ≤ 2 α ϕ ∣ x ∣ , Thus, setting C ϕ = 2 α ϕ C_\phi=2\alpha_\phi C ϕ = 2 α ϕ
Claim G.2. There exists a constant C ϕ > 0 C_\phi>0 C ϕ > 0 x ∈ R x\in\mathbb{R} x ∈ R ∣ ϕ ( x ) − x ϕ ′ ( x ) ∣ ≤ C ϕ ∣ x ∣ |\phi(x)-x\phi'(x)|\le C_\phi|x| ∣ ϕ ( x ) − x ϕ ′ ( x ) ∣ ≤ C ϕ ∣ x ∣ 19
We now have:
− β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] = − β − 1 ∑ i = 1 n E [ r i ( θ t ) ∇ f i ( θ t ) ⊤ Λ θ t ] . -\beta^{-1}\mathbb{E}\left[\boldsymbol{r}(\theta_t)^\top\nabla f(\theta_t)\Lambda\theta_t\right]
=-\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)\nabla f_i(\theta_t)^\top\Lambda\theta_t\right]. − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] = − β − 1 i = 1 ∑ n E [ r i ( θ t ) ∇ f i ( θ t ) ⊤ Λ θ t ] . Observe for each i ∈ [ n ] i\in[n] i ∈ [ n ]
f i ( θ ) = s 2 v ⊤ α ( x i ) f_i(\theta)=s_2\boldsymbol{v}^\top\boldsymbol{\alpha}(x_i) f i ( θ ) = s 2 v ⊤ α ( x i ) α ( x i ) = ϕ ( s 1 W x i ) \boldsymbol{\alpha}(x_i)=\phi(s_1W x_i) α ( x i ) = ϕ ( s 1 W x i ) ∂ f i ( θ ) ∂ v = s 2 α ( x i ) , \frac{\partial f_i(\theta)}{\partial\boldsymbol{v}}=s_2\boldsymbol{\alpha}(x_i), ∂ v ∂ f i ( θ ) = s 2 α ( x i ) , Recall from Eq. (33) that
∂ f i ( θ ) ∂ W j k = s 2 s 1 v j x i , k ϕ ′ ( s 1 w j ⊤ x i ) \frac{\partial f_i(\theta)}{\partial W_{jk}}
=s_2s_1v_jx_{i,k}\phi'(s_1\boldsymbol{w}_j^\top x_i) ∂ W j k ∂ f i ( θ ) = s 2 s 1 v j x i , k ϕ ′ ( s 1 w j ⊤ x i ) ⟹ ∂ f i ( θ ) ∂ w j = s 2 s 1 v j ϕ ′ ( s 1 w j ⊤ x i ) x i . \Longrightarrow\frac{\partial f_i(\theta)}{\partial\boldsymbol{w}_j}
=s_2s_1v_j\phi'(s_1\boldsymbol{w}_j^\top x_i)x_i. ⟹ ∂ w j ∂ f i ( θ ) = s 2 s 1 v j ϕ ′ ( s 1 w j ⊤ x i ) x i . We then have,
∇ f i ( θ ) ⊤ Λ θ = s 2 ϕ ( s 1 W x i ) ⊤ a 2 I m v + ∑ j = 1 m s 2 v j ϕ ′ ( s 1 w j ⊤ x i ) s 1 x i ⊤ a 1 v 0 2 I d w j \nabla f_i(\theta)^\top\Lambda\theta
=s_2\phi(s_1Wx_i)^\top a_2I_m\boldsymbol{v}
+\sum_{j=1}^{m}s_2v_j\phi'(s_1\boldsymbol{w}_j^\top x_i)s_1x_i^\top\frac{a_1}{v_0^2}I_d\boldsymbol{w}_j ∇ f i ( θ ) ⊤ Λ θ = s 2 ϕ ( s 1 W x i ) ⊤ a 2 I m v + j = 1 ∑ m s 2 v j ϕ ′ ( s 1 w j ⊤ x i ) s 1 x i ⊤ v 0 2 a 1 I d w j [ define g ( x ) = ϕ ′ ( x ) x ] = s 2 ( a 2 v ⊤ ϕ ( s 1 W x i ) + a 1 v 0 2 ∑ j = 1 m v j g ( s 1 w j ⊤ x i ) ) = s 2 ( a 2 v ⊤ ϕ ( s 1 W x i ) + a 1 v 0 2 v ⊤ g ( s 1 W x i ) ) = s 2 v ⊤ ( a 2 ϕ ( s 1 W x i ) + a 1 v 0 2 g ( s 1 W x i ) ) . \begin{aligned}
\left[\text{define }g(x)=\phi'(x)x\right]
=s_2\left(a_2\boldsymbol{v}^\top\phi(s_1Wx_i)
+\frac{a_1}{v_0^2}\sum_{j=1}^{m}v_jg(s_1\boldsymbol{w}_j^\top x_i)\right) \\
&=s_2\left(a_2\boldsymbol{v}^\top\phi(s_1Wx_i)
+\frac{a_1}{v_0^2}\boldsymbol{v}^\top g(s_1Wx_i)\right) \\
&=s_2\boldsymbol{v}^\top\left(a_2\phi(s_1Wx_i)+\frac{a_1}{v_0^2}g(s_1Wx_i)\right).
\end{aligned} [ define g ( x ) = ϕ ′ ( x ) x ] = s 2 ( a 2 v ⊤ ϕ ( s 1 W x i ) + v 0 2 a 1 j = 1 ∑ m v j g ( s 1 w j ⊤ x i ) ) = s 2 ( a 2 v ⊤ ϕ ( s 1 W x i ) + v 0 2 a 1 v ⊤ g ( s 1 W x i ) ) = s 2 v ⊤ ( a 2 ϕ ( s 1 W x i ) + v 0 2 a 1 g ( s 1 W x i ) ) . Defining γ ( x ) ≜ ϕ ( x ) − g ( x ) \gamma(x)\triangleq\phi(x)-g(x) γ ( x ) ≜ ϕ ( x ) − g ( x )
r i ( θ ) ∇ f i ( θ ) ⊤ Λ θ = s 2 ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) v ⊤ ( a 2 ϕ ( s 1 W x i ) + a 1 v 0 2 g ( s 1 W x i ) ) = s 2 ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) ⋅ v ⊤ ( a 2 ϕ ( s 1 W x i ) + a 1 v 0 2 ( g ( s 1 W x i ) − ϕ ( s 1 W x i ) + ϕ ( s 1 W x i ) ) ) = s 2 ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) ⋅ v ⊤ ( ( a 2 + a 1 v 0 2 ) ϕ ( s 1 W x i ) − a 1 v 0 2 γ ( s 1 W x i ) ) = ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) ⋅ [ ( a 2 + a 1 v 0 2 ) ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i + y i ) − s 2 a 1 v 0 2 v ⊤ γ ( s 1 W x i ) ] \begin{aligned}
r_i(\theta)\nabla f_i(\theta)^\top\Lambda\theta
=s_2\left(s_2\boldsymbol{v}^\top\phi(s_1Wx_i)-y_i\right)\boldsymbol{v}^\top\left(a_2\phi(s_1Wx_i)+\frac{a_1}{v_0^2}g(s_1Wx_i)\right) \\
&=s_2\left(s_2\boldsymbol{v}^\top\phi(s_1Wx_i)-y_i\right)
\cdot\boldsymbol{v}^\top\left(a_2\phi(s_1Wx_i)+\frac{a_1}{v_0^2}\left(g(s_1Wx_i)-\phi(s_1Wx_i)+\phi(s_1Wx_i)\right)\right) \\
&=s_2\left(s_2\boldsymbol{v}^\top\phi(s_1Wx_i)-y_i\right)
\cdot\boldsymbol{v}^\top\left(\left(a_2+\frac{a_1}{v_0^2}\right)\phi(s_1Wx_i)-\frac{a_1}{v_0^2}\gamma(s_1Wx_i)\right) \\
&=\left(s_2\boldsymbol{v}^\top\phi(s_1Wx_i)-y_i\right)
\cdot\left[\left(a_2+\frac{a_1}{v_0^2}\right)\left(s_2\boldsymbol{v}^\top\phi(s_1Wx_i)-y_i+y_i\right)-s_2\frac{a_1}{v_0^2}\boldsymbol{v}^\top\gamma(s_1Wx_i)\right]
\end{aligned} r i ( θ ) ∇ f i ( θ ) ⊤ Λ θ = s 2 ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) v ⊤ ( a 2 ϕ ( s 1 W x i ) + v 0 2 a 1 g ( s 1 W x i ) ) = s 2 ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) ⋅ v ⊤ ( a 2 ϕ ( s 1 W x i ) + v 0 2 a 1 ( g ( s 1 W x i ) − ϕ ( s 1 W x i ) + ϕ ( s 1 W x i ) ) ) = s 2 ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) ⋅ v ⊤ ( ( a 2 + v 0 2 a 1 ) ϕ ( s 1 W x i ) − v 0 2 a 1 γ ( s 1 W x i ) ) = ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i ) ⋅ [ ( a 2 + v 0 2 a 1 ) ( s 2 v ⊤ ϕ ( s 1 W x i ) − y i + y i ) − s 2 v 0 2 a 1 v ⊤ γ ( s 1 W x i ) ] \begin{aligned}
&= \left(\left(a_2+\frac{a_1}{v_0^2}\right)\left(s_2\mathbf{v}^{\mathsf T}\phi\left(s_1\mathbf{W}\mathbf{x}_i\right)-y_i\right)^2 \\
&\quad+\left(s_2\mathbf{v}^{\mathsf T}\phi\left(s_1\mathbf{W}\mathbf{x}_i\right)-y_i\right)\left(\left(a_2+\frac{a_1}{v_0^2}\right)y_i-s_2\frac{a_1}{v_0^2}\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right)\right) \\
&= \left(a_2+\frac{a_1}{v_0^2}\right)2\ell_i(\theta)+r_i(\theta)\left(\left(a_2+\frac{a_1}{v_0^2}\right)y_i-s_2\frac{a_1}{v_0^2}\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right),
\end{aligned}
where ℓ i ( θ ) = 1 2 ∣ r i ( θ ) ∣ 2 \ell_i(\theta)=\frac{1}{2}\lvert r_i(\theta)\rvert^2 ℓ i ( θ ) = 2 1 ∣ r i ( θ ) ∣ 2
⟹ − β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Δ θ t ] = − β − 1 ∑ i = 1 n E [ r i ( θ t ) ∇ f i ( θ t ) T Δ θ t ] = − β − 1 [ 2 ( a 2 + a 1 v 0 2 ) n E L ( θ t ) + ∑ i = 1 n E [ r i ( θ t ) ( ( a 2 + a 1 v 0 2 ) y i − s 2 a 1 v 0 2 v T γ ( s 1 W x i ) ) ] ] ≤ − 2 β − 1 n ( a 2 + a 1 v 0 2 ) E L ( θ t ) + β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ ∣ ( a 2 + a 1 v 0 2 ) y i − s 2 a 1 v 0 2 v T γ ( s 1 W x i ) ∣ ] \begin{aligned}
&\Longrightarrow-\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\mathsf T}\nabla f(\theta_t)\Delta\theta_t\right] \\
&=-\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)\nabla f_i(\theta_t)^{\mathsf T}\Delta\theta_t\right] \\
&=-\beta^{-1}\left[2\left(a_2+\frac{a_1}{v_0^2}\right)n\mathbb{E}\mathcal{L}(\theta_t)+\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)\left(\left(a_2+\frac{a_1}{v_0^2}\right)y_i-s_2\frac{a_1}{v_0^2}\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right)\right]\right] \\
&\leq-2\beta^{-1}n\left(a_2+\frac{a_1}{v_0^2}\right)\mathbb{E}\mathcal{L}(\theta_t) \\
&\quad+\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\left\lvert r_i(\theta_t)\right\rvert\left\lvert\left(a_2+\frac{a_1}{v_0^2}\right)y_i-s_2\frac{a_1}{v_0^2}\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right\rvert\right]
\end{aligned} ⟹ − β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Δ θ t ] = − β − 1 i = 1 ∑ n E [ r i ( θ t ) ∇ f i ( θ t ) T Δ θ t ] = − β − 1 [ 2 ( a 2 + v 0 2 a 1 ) n E L ( θ t ) + i = 1 ∑ n E [ r i ( θ t ) ( ( a 2 + v 0 2 a 1 ) y i − s 2 v 0 2 a 1 v T γ ( s 1 W x i ) ) ] ] ≤ − 2 β − 1 n ( a 2 + v 0 2 a 1 ) E L ( θ t ) + β − 1 i = 1 ∑ n E [ ∣ r i ( θ t ) ∣ ( a 2 + v 0 2 a 1 ) y i − s 2 v 0 2 a 1 v T γ ( s 1 W x i ) ] [ ( 1 ) ] ≤ − 2 β − 1 n ( a 2 + a 1 v 0 2 ) E L ( θ t ) + β − 1 ∑ i = 1 n E [ r i ( θ t ) 2 ] E [ ( ( a 2 + a 1 v 0 2 ) y i − s 2 a 1 v 0 2 v T γ ( s 1 W x i ) ) 2 ] \begin{aligned}
[(1)]\leq{}&-2\beta^{-1}n\left(a_2+\frac{a_1}{v_0^2}\right)\mathbb{E}\mathcal{L}(\theta_t)\\
&+\beta^{-1}\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[r_i(\theta_t)^2\right]\mathbb{E}\left[\left(\left(a_2+\frac{a_1}{v_0^2}\right)y_i-s_2\frac{a_1}{v_0^2}\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right)^2\right]}
\end{aligned} [( 1 )] ≤ − 2 β − 1 n ( a 2 + v 0 2 a 1 ) E L ( θ t ) + β − 1 i = 1 ∑ n E [ r i ( θ t ) 2 ] E [ ( ( a 2 + v 0 2 a 1 ) y i − s 2 v 0 2 a 1 v T γ ( s 1 W x i ) ) 2 ] [ ( 2 ) ] ≤ − 2 β − 1 n ( a 2 + a 1 v 0 2 ) E L ( θ t ) + β − 1 ∑ i = 1 n 2 E [ ℓ i ( θ t ) ] E [ 2 ( a 2 + a 1 v 0 2 ) 2 y i 2 + 2 s 2 2 a 1 2 v 0 4 ∣ v T γ ( s 1 W x i ) ∣ 2 ] \begin{aligned}
[(2)]\leq{}&-2\beta^{-1}n\left(a_2+\frac{a_1}{v_0^2}\right)\mathbb{E}\mathcal{L}(\theta_t)\\
&+\beta^{-1}\sum_{i=1}^{n}\sqrt{2\mathbb{E}\left[\ell_i(\theta_t)\right]\mathbb{E}\left[2\left(a_2+\frac{a_1}{v_0^2}\right)^2y_i^2+2s_2^2\frac{a_1^2}{v_0^4}\left\lvert\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right\rvert^2\right]}
\end{aligned} [( 2 )] ≤ − 2 β − 1 n ( a 2 + v 0 2 a 1 ) E L ( θ t ) + β − 1 i = 1 ∑ n 2 E [ ℓ i ( θ t ) ] E [ 2 ( a 2 + v 0 2 a 1 ) 2 y i 2 + 2 s 2 2 v 0 4 a 1 2 ∣ v T γ ( s 1 W x i ) ∣ 2 ] [ ( 3 ) ] ≤ − 2 β − 1 n ( a 2 + a 1 v 0 2 ) E L ( θ t ) + 2 β − 1 ∑ i = 1 n E [ ℓ i ( θ t ) ] ( E y [ ( a 2 + a 1 v 0 2 ) y i ] + s 2 a 1 v 0 2 E p t [ ∣ v T γ ( s 1 W x i ) ∣ 2 ] ) , \begin{aligned}
[(3)]\leq{}&-2\beta^{-1}n\left(a_2+\frac{a_1}{v_0^2}\right)\mathbb{E}\mathcal{L}(\theta_t)\\
&+2\beta^{-1}\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[\ell_i(\theta_t)\right]}\left(\mathbb{E}_{y}\left[\left(a_2+\frac{a_1}{v_0^2}\right)y_i\right]+s_2\frac{a_1}{v_0^2}\sqrt{\mathbb{E}_{p_t}\left[\left\lvert\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right\rvert^2\right]}\right),
\end{aligned} [( 3 )] ≤ − 2 β − 1 n ( a 2 + v 0 2 a 1 ) E L ( θ t ) + 2 β − 1 i = 1 ∑ n E [ ℓ i ( θ t ) ] ( E y [ ( a 2 + v 0 2 a 1 ) y i ] + s 2 v 0 2 a 1 E p t [ ∣ v T γ ( s 1 W x i ) ∣ 2 ] ) , where in transition (1) we used Cauchy–Schwarz, in transition (2) we used ( a − b ) 2 ≤ 2 a 2 + 2 b 2 (a-b)^2\leq2a^2+2b^2 ( a − b ) 2 ≤ 2 a 2 + 2 b 2 a + b ≤ a + b \sqrt{a+b}\leq\sqrt{a}+\sqrt{b} a + b ≤ a + b
Bounding E p t [ ∣ v T γ ( s 1 W x i ) ∣ 2 ] \mathbb{E}_{p_t}\left[\left\lvert\mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right\rvert^2\right] E p t [ v T γ ( s 1 W x i ) 2 ]
Conditioned on W \mathbf{W} W γ ( s 1 W x i ) \gamma\left(s_1\mathbf{W}\mathbf{x}_i\right) γ ( s 1 W x i ) v T γ ( s 1 W x i ) \mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right) v T γ ( s 1 W x i ) v p 0 ∼ N ( 0 m , 1 a 2 I m ) \mathbf{v}_{p_0}\sim\mathcal{N}\left(\mathbf{0}_m,\frac{1}{a_2}\mathbf{I}_m\right) v p 0 ∼ N ( 0 m , a 2 1 I m )
v T γ ( s 1 W x i ) = d 1 a 2 ∥ γ ( s 1 W x i ) ∥ 2 G , \mathbf{v}^{\mathsf T}\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\mathrel{\overset{d}{=}}\frac{1}{\sqrt{a_2}}\left\|\gamma\left(s_1\mathbf{W}\mathbf{x}_i\right)\right\|_2G, v T γ ( s 1 W x i ) = d a 2 1 ∥ γ ( s 1 W x i ) ∥ 2 G , where G ∼ N ( 0 , 1 ) G \sim\mathcal{N}(0,1) G ∼ N ( 0 , 1 ) W \mathbf{W} W = d \overset{d}{=} = d p ≥ 1 p \ge1 p ≥ 1
∥ v ⊤ γ ( s 1 W x i ) ∥ L p ∣ W = 1 a 2 ∥ γ ( s 1 W x i ) ∥ 2 ∥ G ∥ L p . \left\|\mathbf{v}^{\top}\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{L_{p}\mid\mathbf{W}} = \frac{1}{\sqrt{a_{2}}}\left\|\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{2}\|G\|_{L_{p}}. v ⊤ γ ( s 1 W x i ) L p ∣ W = a 2 1 ∥ γ ( s 1 W x i ) ∥ 2 ∥ G ∥ L p . \text{\[[111, Proposition 2.6.6]\]} \le\frac{1}{\sqrt{a_{2}}}\left\|\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{2}C\sqrt{p}.
By using the tower property for p p p
∥ v ⊤ γ ( s 1 W x i ) ∥ L p = ∥ ∥ v ⊤ γ ( s 1 W x i ) ∣ W ∥ L p ∥ L p ≤ 1 a 2 C p ∥ ∥ γ ( s 1 W x i ) ∥ 2 ∥ L p . \left\|\mathbf{v}^{\top}\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{L_{p}} = \left\|\left\|\mathbf{v}^{\top}\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\mid\mathbf{W}\right\|_{L_{p}}\right\|_{L_{p}} \le\frac{1}{\sqrt{a_{2}}}C\sqrt{p}\left\|\left\|\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{2}\right\|_{L_{p}}. v ⊤ γ ( s 1 W x i ) L p = v ⊤ γ ( s 1 W x i ) ∣ W L p L p ≤ a 2 1 C p ∥ ∥ γ ( s 1 W x i ) ∥ 2 ∥ L p . Note that
∥ γ ( s 1 W x i ) ∥ 2 ≤ C ϕ ∥ s 1 W x i ∥ 2 = C ϕ s 1 ∑ j = 1 m ( w j ⊤ x i ) 2 = C ϕ s 1 c ϕ , σ 0 d ν 0 2 a 1 ∥ z ∥ , \left\|\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{2} \le C_{\phi}\left\|s_{1}\mathbf{W}\mathbf{x}_{i}\right\|_{2} = C_{\phi}s_{1}\sqrt{\sum_{j=1}^{m}\left(\mathbf{w}_{j}^{\top}\mathbf{x}_{i}\right)^{2}} = C_{\phi}s_{1}\sqrt{c_{\phi,\sigma_{0}}d\frac{\nu_{0}^{2}}{a_{1}}}\|\mathbf{z}\|, ∥ γ ( s 1 W x i ) ∥ 2 ≤ C ϕ ∥ s 1 W x i ∥ 2 = C ϕ s 1 j = 1 ∑ m ( w j ⊤ x i ) 2 = C ϕ s 1 c ϕ , σ 0 d a 1 ν 0 2 ∥ z ∥ , where
z ≜ 1 c ϕ , σ 0 d ν 0 2 a 1 W x i , \mathbf{z} \triangleq\frac{1}{\sqrt{c_{\phi,\sigma_{0}}d\frac{\nu_{0}^{2}}{a_{1}}}}\mathbf{W}\mathbf{x}_{i}, z ≜ c ϕ , σ 0 d a 1 ν 0 2 1 W x i , so that z ∼ N ( 0 m , I m ) \mathbf{z}\sim\mathcal{N}(\mathbf{0}_{m},\mathbf{I}_{m}) z ∼ N ( 0 m , I m ) Claim E.2 gives:
∥ ∥ γ ( s 1 W x i ) ∥ 2 ∥ L p ≤ C ϕ s 1 c ϕ , σ 0 d ν 0 2 a 1 ∥ ∥ z ∥ ∥ L p ≤ C ϕ s 1 c ϕ , σ 0 d ν 0 2 a 1 ( C p + m ) . \left\|\left\|\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{2}\right\|_{L_{p}} \le C_{\phi}s_{1}\sqrt{c_{\phi,\sigma_{0}}d\frac{\nu_{0}^{2}}{a_{1}}}\|\|\mathbf{z}\|\|_{L_{p}} \le C_{\phi}s_{1}\sqrt{c_{\phi,\sigma_{0}}d\frac{\nu_{0}^{2}}{a_{1}}}\left(C\sqrt{p}+\sqrt{m}\right). ∥ ∥ γ ( s 1 W x i ) ∥ 2 ∥ L p ≤ C ϕ s 1 c ϕ , σ 0 d a 1 ν 0 2 ∥∥ z ∥ ∥ L p ≤ C ϕ s 1 c ϕ , σ 0 d a 1 ν 0 2 ( C p + m ) . Applying Eq. (32) to this nonnegative quantity gives, for all t ≥ 0 t\ge0 t ≥ 0
E p , t [ ∣ v ⊤ γ ( s 1 W x i ) ∣ 2 ] ≤ η ∥ v ⊤ γ ( s 1 W x i ) ∥ L p 0 c β κ ˉ 2 β 2 \mathbb{E}_{p,t}\left[\left|\mathbf{v}^{\top}\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right|^{2}\right] \le\eta\left\|\mathbf{v}^{\top}\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{L^{c_{\beta}\bar{\kappa}^{2}\beta}_{p_{0}}}^{2} E p , t [ v ⊤ γ ( s 1 W x i ) 2 ] ≤ η v ⊤ γ ( s 1 W x i ) L p 0 c β κ ˉ 2 β 2 \begin{aligned}
\text{\[{xref}`Claim E.6`\]} = \eta\left\|\mathbf{v}^{\top}\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}_{p_{0}}}^{2} \\
&\le\eta\left(\frac{1}{\sqrt{a_{2}}}C_{1}\sqrt{2c_{\beta}\bar{\kappa}^{2}\beta}\left\|\left\|\gamma\left(s_{1}\mathbf{W}\mathbf{x}_{i}\right)\right\|_{2}\right\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}_{p_{0}}}\right)^{2} \\
&\le\eta\left(\frac{1}{\sqrt{a_{2}}}C_{1}\sqrt{2c_{\beta}\bar{\kappa}^{2}\beta}C_{\phi}s_{1}\sqrt{c_{\phi,\sigma_{0}}d\frac{\nu_{0}^{2}}{a_{1}}}\left(C_{2}\sqrt{2c_{\beta}\bar{\kappa}^{2}\beta}+\sqrt{m}\right)\right)^{2} \\
&\le\eta C_{3}c_{\phi,\sigma_{0}}\nu_{0}^{2}s_{1}^{2}\frac{d}{a_{2}a_{1}}c_{\beta}\bar{\kappa}^{2}\beta\left(C_{4}c_{\beta}\bar{\kappa}^{2}\beta+m\right)
\end{aligned}
≲ s 1 2 d a 2 a 1 κ ˉ 2 β ( κ ˉ 2 β + m ) . \lesssim s_{1}^{2}\frac{d}{a_{2}a_{1}}\bar{\kappa}^{2}\beta\left(\bar{\kappa}^{2}\beta+m\right). ≲ s 1 2 a 2 a 1 d κ ˉ 2 β ( κ ˉ 2 β + m ) . Summary of fourth term
− β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ≤ − 2 β − 1 n ( a 2 + a 1 ν 0 2 ) E L ( θ t ) + 2 β − 1 ∑ i = 1 n E [ ℓ i ( θ t ) ] ( R y ( a 2 + a 1 ν 0 2 ) + s 2 a 1 ν 0 2 E p t [ ∣ v ⊤ γ ( s 1 W x i ) ∣ 2 ] ) ≲ − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 ∑ i = 1 n E [ ℓ i ( θ t ) ] ( a 2 + a 1 + s 2 a 1 s 1 2 d a 2 a 1 κ ˉ 2 β ( κ ˉ 2 β + m ) ) ≲ [Jensen] − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 + s 2 a 1 s 1 2 d a 2 a 1 κ ˉ 2 β ( κ ˉ 2 β + m ) ) 1 n ∑ i = 1 n E [ ℓ i ( θ t ) ] ≲ − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 + s 2 s 1 d a 1 a 2 κ ˉ 2 β ( κ ˉ 2 β + m ) ) E L ( θ t ) . -\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}\nabla f(\theta_t)\Lambda\theta_t\right]
\leq-2\beta^{-1}n\left(a_2+\frac{a_1}{\nu_0^2}\right)\mathcal{E}\mathcal{L}(\theta_t)
+2\beta^{-1}\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[\ell_i(\theta_t)\right]}\left(R_y\left(a_2+\frac{a_1}{\nu_0^2}\right)+\frac{s_2a_1}{\nu_0^2}\sqrt{\mathbb{E}_{p_t}\left[\left|\mathbf{v}^{\top}\gamma(s_1\mathbf{W}x_i)\right|^2\right]}\right)
\lesssim-\beta^{-1}n(a_2+a_1)\mathcal{E}\mathcal{L}(\theta_t)
+\beta^{-1}\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[\ell_i(\theta_t)\right]}\left(a_2+a_1+s_2a_1\sqrt{s_1^2\frac{d}{a_2a_1}\bar{\kappa}^2\beta\left(\bar{\kappa}^2\beta+m\right)}\right)
\overset{\text{[Jensen]}}{\lesssim} -\beta^{-1}n(a_2+a_1)\mathcal{E}\mathcal{L}(\theta_t)
+\beta^{-1}n\left(a_2+a_1+s_2a_1\sqrt{s_1^2\frac{d}{a_2a_1}\bar{\kappa}^2\beta\left(\bar{\kappa}^2\beta+m\right)}\right)\sqrt{\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[\ell_i(\theta_t)\right]}
\lesssim-\beta^{-1}n(a_2+a_1)\mathcal{E}\mathcal{L}(\theta_t)
+\beta^{-1}n\left(a_2+a_1+s_2s_1\sqrt{\frac{da_1}{a_2}\bar{\kappa}^2\beta\left(\bar{\kappa}^2\beta+m\right)}\right)\sqrt{\mathcal{E}\mathcal{L}(\theta_t)}. − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ≤ − 2 β − 1 n ( a 2 + ν 0 2 a 1 ) E L ( θ t ) + 2 β − 1 i = 1 ∑ n E [ ℓ i ( θ t ) ] ( R y ( a 2 + ν 0 2 a 1 ) + ν 0 2 s 2 a 1 E p t [ ∣ v ⊤ γ ( s 1 W x i ) ∣ 2 ] ) ≲ − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 i = 1 ∑ n E [ ℓ i ( θ t ) ] ( a 2 + a 1 + s 2 a 1 s 1 2 a 2 a 1 d κ ˉ 2 β ( κ ˉ 2 β + m ) ) ≲ [Jensen] − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 + s 2 a 1 s 1 2 a 2 a 1 d κ ˉ 2 β ( κ ˉ 2 β + m ) ) n 1 i = 1 ∑ n E [ ℓ i ( θ t ) ] ≲ − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 + s 2 s 1 a 2 d a 1 κ ˉ 2 β ( κ ˉ 2 β + m ) ) E L ( θ t ) . Importantly, observe that for E L ( θ t ) ≤ 1 \mathcal{E}\mathcal{L}(\theta_t)\leq1 E L ( θ t ) ≤ 1
− β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 ) E L ( θ t ) = β − 1 n ( a 2 + a 1 ) E L ( θ t ) ( 1 − E L ( θ t ) ) ≥ 0 , -\beta^{-1}n(a_2+a_1)\mathcal{E}\mathcal{L}(\theta_t)+\beta^{-1}n(a_2+a_1)\sqrt{\mathcal{E}\mathcal{L}(\theta_t)}
=\beta^{-1}n(a_2+a_1)\sqrt{\mathcal{E}\mathcal{L}(\theta_t)}\left(1-\sqrt{\mathcal{E}\mathcal{L}(\theta_t)}\right)\geq0, − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 ) E L ( θ t ) = β − 1 n ( a 2 + a 1 ) E L ( θ t ) ( 1 − E L ( θ t ) ) ≥ 0 , thus this upper bound does not yield a descent contribution from the regularization term in the small-loss regime.
Bounding the entire right-hand side
Recall:
d d t E L ( θ t ) = 1 n ( − n − 1 E [ r ( θ t ) ⊤ K N T K ( θ t ) r ( θ t ) ] + β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] + β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ) . \frac{\mathrm{d}}{\mathrm{d}t}\mathcal{E}\mathcal{L}(\theta_t)
=\frac{1}{n}\left(-n^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}K_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)\right]
+\beta^{-1}\mathbb{E}\left[\operatorname{Tr}\left(\nabla f(\theta_t)^{\top}\nabla f(\theta_t)\right)\right]
+\beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]
-\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}\nabla f(\theta_t)\Lambda\theta_t\right]\right). d t d E L ( θ t ) = n 1 ( − n − 1 E [ r ( θ t ) ⊤ K NTK ( θ t ) r ( θ t ) ] + β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] + β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ) . Bounding each term, we obtained:
− n − 1 E [ r ( θ t ) ⊤ K N T K ( θ t ) r ( θ t ) ] ≲ − λ 1 s 2 2 m E [ L ( θ t ) ] + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) , + β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] ≲ β − 1 n s 2 2 s 1 2 d ( 1 a 1 + 1 a 2 ) ( κ ˉ 2 β + m ) , + β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] ≲ β − 1 / 2 n s 2 s 1 2 d m a 2 κ ˉ 2 E L ( θ t ) , \begin{aligned}
-n^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}K_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)\right]\lesssim-\lambda_1s_2^2m\mathbb{E}\left[\mathcal{L}(\theta_t)\right]
+\lambda_1s_2^2m\left(1+\bar{\kappa}^4\beta\right)\exp\left(-\widetilde{c}_2\frac{m}{c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right), \\
&+\beta^{-1}\mathbb{E}\left[\operatorname{Tr}\left(\nabla f(\theta_t)^{\top}\nabla f(\theta_t)\right)\right]\lesssim\beta^{-1}ns_2^2s_1^2d\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\left(\bar{\kappa}^2\beta+m\right), \\
&+\beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]\lesssim\beta^{-1/2}ns_2s_1^2d\sqrt{\frac{m}{a_2\bar{\kappa}^2}}\sqrt{\mathbb{E}\mathcal{L}(\theta_t)},
\end{aligned} − n − 1 E [ r ( θ t ) ⊤ K NTK ( θ t ) r ( θ t ) ] ≲ − λ 1 s 2 2 m E [ L ( θ t ) ] + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) , + β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] ≲ β − 1 n s 2 2 s 1 2 d ( a 1 1 + a 2 1 ) ( κ ˉ 2 β + m ) , + β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] ≲ β − 1/2 n s 2 s 1 2 d a 2 κ ˉ 2 m E L ( θ t ) , − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ≲ − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 + s 2 s 1 d a 1 a 2 κ ˉ 2 β ( κ ˉ 2 β + m ) ) E L ( θ t ) . -\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}\nabla f(\theta_t)\Lambda\theta_t\right]\lesssim-\beta^{-1}n(a_2+a_1)\mathbb{E}\mathcal{L}(\theta_t)+\beta^{-1}n
\left(a_2+a_1+s_2s_1\sqrt{\frac{da_1}{a_2}}\sqrt{\bar{\kappa}^2\beta\left(\bar{\kappa}^2\beta+m\right)}\right)\sqrt{\mathbb{E}\mathcal{L}(\theta_t)}. − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ≲ − β − 1 n ( a 2 + a 1 ) E L ( θ t ) + β − 1 n ( a 2 + a 1 + s 2 s 1 a 2 d a 1 κ ˉ 2 β ( κ ˉ 2 β + m ) ) E L ( θ t ) . Together, we can bound d d t E L ( θ t ) \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t) d t d E L ( θ t ) t ≥ 0 t\ge0 t ≥ 0
d d t E L ( θ t ) ≲ 1 n ( − ( λ 1 s 2 2 m + β − 1 n ( a 2 + a 1 ) ) E L ( θ t ) + [ β − 1 n ( a 2 + a 1 + s 2 s 1 d a 1 a 2 κ ˉ 2 β ( κ ˉ 2 β + m ) ) + β − 1 / 2 n s 2 s 1 2 d m a 2 κ ˉ 2 ] E L ( θ t ) + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 1 n s 2 2 s 1 2 d ( 1 a 1 + 1 a 2 ) ( κ ˉ 2 β + m ) ) . \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)
\lesssim\frac{1}{n}\left(-\left(\lambda_1s_2^2m+\beta^{-1}n(a_2+a_1)\right)\mathbb{E}\mathcal{L}(\theta_t)
+\left[\beta^{-1}n\left(a_2+a_1+s_2s_1\sqrt{\frac{da_1}{a_2}}\sqrt{\bar{\kappa}^2\beta\left(\bar{\kappa}^2\beta+m\right)}\right)+\beta^{-1/2}ns_2s_1^2d\sqrt{\frac{m}{a_2\bar{\kappa}^2}}\right]\sqrt{\mathbb{E}\mathcal{L}(\theta_t)}
+\lambda_1s_2^2m\left(1+\bar{\kappa}^4\beta\right)\exp\left(-\widetilde{c}_2\frac{m}{c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-1}ns_2^2s_1^2d\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\left(\bar{\kappa}^2\beta+m\right)\right). d t d E L ( θ t ) ≲ n 1 ( − ( λ 1 s 2 2 m + β − 1 n ( a 2 + a 1 ) ) E L ( θ t ) + [ β − 1 n ( a 2 + a 1 + s 2 s 1 a 2 d a 1 κ ˉ 2 β ( κ ˉ 2 β + m ) ) + β − 1/2 n s 2 s 1 2 d a 2 κ ˉ 2 m ] E L ( θ t ) + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 1 n s 2 2 s 1 2 d ( a 1 1 + a 2 1 ) ( κ ˉ 2 β + m ) ) . The regularization contribution does not provide a useful descent term in the small-loss regime: for u = E L ( θ t ) ≤ 1 u=\mathbb{E}\mathcal{L}(\theta_t)\le1 u = E L ( θ t ) ≤ 1
− β − 1 n ( a 2 + a 1 ) u + β − 1 n ( a 2 + a 1 ) u = β − 1 n ( a 2 + a 1 ) u ( 1 − u ) ≥ 0. -\beta^{-1}n(a_2+a_1)u+\beta^{-1}n(a_2+a_1)\sqrt{u}=\beta^{-1}n(a_2+a_1)\sqrt{u}(1-\sqrt{u})\ge0. − β − 1 n ( a 2 + a 1 ) u + β − 1 n ( a 2 + a 1 ) u = β − 1 n ( a 2 + a 1 ) u ( 1 − u ) ≥ 0. Therefore, in closing the differential inequality, we rely only on the NTK term for the linear descent and treat all u \sqrt{u} u
d d t E L ( θ t ) ≲ 1 n ( − λ 1 s 2 2 m E L ( θ t ) + [ β − 1 n ( a 2 + a 1 + s 2 s 1 d a 1 a 2 κ ˉ 2 β ( κ ˉ 2 β + m ) ) + β − 1 / 2 n s 2 s 1 2 d m a 2 κ ˉ 2 ] E L ( θ t ) + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 1 n s 2 2 s 1 2 d ( 1 a 1 + 1 a 2 ) ( κ ˉ 2 β + m ) ) . \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)
\lesssim\frac{1}{n}\left(-\lambda_1s_2^2m\mathbb{E}\mathcal{L}(\theta_t)
+\left[\beta^{-1}n\left(a_2+a_1+s_2s_1\sqrt{\frac{da_1}{a_2}}\sqrt{\bar{\kappa}^2\beta\left(\bar{\kappa}^2\beta+m\right)}\right)+\beta^{-1/2}ns_2s_1^2d\sqrt{\frac{m}{a_2\bar{\kappa}^2}}\right]\sqrt{\mathbb{E}\mathcal{L}(\theta_t)}
+\lambda_1s_2^2m\left(1+\bar{\kappa}^4\beta\right)\exp\left(-\widetilde{c}_2\frac{m}{c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-1}ns_2^2s_1^2d\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\left(\bar{\kappa}^2\beta+m\right)\right). d t d E L ( θ t ) ≲ n 1 ( − λ 1 s 2 2 m E L ( θ t ) + [ β − 1 n ( a 2 + a 1 + s 2 s 1 a 2 d a 1 κ ˉ 2 β ( κ ˉ 2 β + m ) ) + β − 1/2 n s 2 s 1 2 d a 2 κ ˉ 2 m ] E L ( θ t ) + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 1 n s 2 2 s 1 2 d ( a 1 1 + a 2 1 ) ( κ ˉ 2 β + m ) ) . Next, in order to have
( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) ≪ 1 , \left(1+\bar{\kappa}^4\beta\right)\exp\left(-\widetilde{c}_2\frac{m}{c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)\ll1, ( 1 + κ ˉ 4 β ) exp ( − c 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) ≪ 1 , it is enough to assume,
m = Ω ( κ ˉ 2 β n λ 1 ln n λ 1 ln ( κ ˉ 4 β ) ) , m=\Omega\left(\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\ln\left(\bar{\kappa}^4\beta\right)\right), m = Ω ( κ ˉ 2 β λ 1 n ln λ 1 n ln ( κ ˉ 4 β ) ) , since β = Ω ( n λ 1 ) ≫ 1 \beta=\Omega\left(\frac{n}{\lambda_1}\right)\gg1 β = Ω ( λ 1 n ) ≫ 1 κ ˉ ≥ 1 \bar{\kappa}\ge1 κ ˉ ≥ 1
In particular, this implies m ≫ κ ˉ 2 β m \gg\bar{\kappa}^{2}\beta m ≫ κ ˉ 2 β
d d t E L ( θ t ) ≲ 1 n ( − λ 1 s 2 2 m E L ( θ t ) + [ β − 1 n ( a 2 + a 1 + s 2 s 1 d a 1 a 2 κ ˉ 2 β m ) + β − 1 / 2 n s 2 s 1 2 d m a 2 κ ˉ 2 ] E L ( θ t ) + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 1 n s 2 2 s 1 2 m d ( 1 a 1 + 1 a 2 ) ) . \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)
\lesssim\frac{1}{n}\left(-\lambda_1s_2^2m\mathbb{E}\mathcal{L}(\theta_t)
+\left[\beta^{-1}n\left(a_2+a_1+s_2s_1\sqrt{\frac{d a_1}{a_2}}\sqrt{\bar{\kappa}^{2}\beta m}\right)+\beta^{-1/2}ns_2s_1^2d\sqrt{\frac{m}{a_2\bar{\kappa}^{2}}}\right]\sqrt{\mathbb{E}\mathcal{L}(\theta_t)}
+\lambda_1s_2^2m\left(1+\bar{\kappa}^{4}\beta\right)\exp\left(-\tilde{c}_2\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-1}ns_2^2s_1^2md\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\right). d t d E L ( θ t ) ≲ n 1 ( − λ 1 s 2 2 m E L ( θ t ) + [ β − 1 n ( a 2 + a 1 + s 2 s 1 a 2 d a 1 κ ˉ 2 β m ) + β − 1/2 n s 2 s 1 2 d a 2 κ ˉ 2 m ] E L ( θ t ) + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 1 n s 2 2 s 1 2 m d ( a 1 1 + a 2 1 ) ) . By Young’s inequality, we have for any ε > 0 \varepsilon>0 ε > 0
b E L ( θ t ) ≤ b 2 2 ε + ε 2 E L ( θ t ) . b\sqrt{\mathbb{E}\mathcal{L}(\theta_t)}\leq\frac{b^2}{2\varepsilon}+\frac{\varepsilon}{2}\mathbb{E}\mathcal{L}(\theta_t). b E L ( θ t ) ≤ 2 ε b 2 + 2 ε E L ( θ t ) . Choosing ε = λ 1 s 2 2 m \varepsilon=\lambda_1s_2^2m ε = λ 1 s 2 2 m
d d t E L ( θ t ) ≲ 1 n ( − λ 1 2 s 2 2 m E L ( θ t ) + [ β − 1 n ( a 2 + a 1 + s 2 s 1 d a 1 a 2 κ ˉ 2 β m ) + β − 1 / 2 n s 2 s 1 2 d m a 2 κ ˉ 2 ] 2 2 λ 1 s 2 2 m + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 1 n s 2 2 s 1 2 m d ( 1 a 1 + 1 a 2 ) ) . \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}\mathcal{L}(\theta_t)
\lesssim\frac{1}{n}\left(-\frac{\lambda_1}{2}s_2^2m\mathbb{E}\mathcal{L}(\theta_t)
+\frac{\left[\beta^{-1}n\left(a_2+a_1+s_2s_1\sqrt{\frac{d a_1}{a_2}}\sqrt{\bar{\kappa}^{2}\beta m}\right)+\beta^{-1/2}ns_2s_1^2d\sqrt{\frac{m}{a_2\bar{\kappa}^{2}}}\right]^2}{2\lambda_1s_2^2m}
+\lambda_1s_2^2m\left(1+\bar{\kappa}^{4}\beta\right)\exp\left(-\tilde{c}_2\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-1}ns_2^2s_1^2md\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\right). d t d E L ( θ t ) ≲ n 1 − 2 λ 1 s 2 2 m E L ( θ t ) + 2 λ 1 s 2 2 m [ β − 1 n ( a 2 + a 1 + s 2 s 1 a 2 d a 1 κ ˉ 2 β m ) + β − 1/2 n s 2 s 1 2 d a 2 κ ˉ 2 m ] 2 + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 1 n s 2 2 s 1 2 m d ( a 1 1 + a 2 1 ) . ≲ 1 n ( − λ 1 s 2 2 m E L ( θ t ) + β − 2 n 2 ( a 2 2 + a 1 2 + s 2 2 s 1 2 d a 1 a 2 κ ˉ 2 β m ) + β − 1 n 2 s 2 2 s 1 4 d 2 m a 2 κ ˉ 2 λ 1 s 2 2 m + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 1 n s 2 2 s 1 2 m d ( 1 a 1 + 1 a 2 ) ) . \lesssim\frac{1}{n}\left(-\lambda_1s_2^2m\mathbb{E}\mathcal{L}(\theta_t)
+\frac{\beta^{-2}n^2\left(a_2^2+a_1^2+s_2^2s_1^2\frac{d a_1}{a_2}\bar{\kappa}^{2}\beta m\right)+\beta^{-1}n^2s_2^2s_1^4d^2\frac{m}{a_2}\bar{\kappa}^{2}}{\lambda_1s_2^2m}
+\lambda_1s_2^2m\left(1+\bar{\kappa}^{4}\beta\right)\exp\left(-\tilde{c}_2\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-1}ns_2^2s_1^2md\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\right). ≲ n 1 − λ 1 s 2 2 m E L ( θ t ) + λ 1 s 2 2 m β − 2 n 2 ( a 2 2 + a 1 2 + s 2 2 s 1 2 a 2 d a 1 κ ˉ 2 β m ) + β − 1 n 2 s 2 2 s 1 4 d 2 a 2 m κ ˉ 2 + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 1 n s 2 2 s 1 2 m d ( a 1 1 + a 2 1 ) . Then, from Claim E.9 we have
E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c s 2 2 λ 1 n t + L ∞ , \mathbb{E}\mathcal{L}(\theta_t)\leq\left(\mathbb{E}\mathcal{L}(\theta_0)-\mathcal{L}_{\infty}\right)e^{-cs_2^2\frac{\lambda_1}{n}t}+\mathcal{L}_{\infty}, E L ( θ t ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c s 2 2 n λ 1 t + L ∞ , where the loss floor may be chosen, up to universal multiplicative constants, as
L ∞ ≲ 1 λ 1 2 s 2 2 m ( β − 2 n 2 ( a 2 2 + a 1 2 + s 2 2 s 1 2 d a 1 a 2 κ ˉ 2 β m ) + β − 1 n 2 s 2 2 s 1 4 d 2 m a 2 κ ˉ 2 λ 1 s 2 2 m + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 1 n s 2 2 s 1 2 m d ( 1 a 1 + 1 a 2 ) ) \mathcal{L}_{\infty}\lesssim\frac{1}{\frac{\lambda_1}{2}s_2^2m}\left(\frac{\beta^{-2}n^2\left(a_2^2+a_1^2+s_2^2s_1^2\frac{d a_1}{a_2}\bar{\kappa}^{2}\beta m\right)+\beta^{-1}n^2s_2^2s_1^4d^2\frac{m}{a_2}\bar{\kappa}^{2}}{\lambda_1s_2^2m}
+\lambda_1s_2^2m\left(1+\bar{\kappa}^{4}\beta\right)\exp\left(-\tilde{c}_2\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-1}ns_2^2s_1^2md\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\right) L ∞ ≲ 2 λ 1 s 2 2 m 1 λ 1 s 2 2 m β − 2 n 2 ( a 2 2 + a 1 2 + s 2 2 s 1 2 a 2 d a 1 κ ˉ 2 β m ) + β − 1 n 2 s 2 2 s 1 4 d 2 a 2 m κ ˉ 2 + λ 1 s 2 2 m ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 1 n s 2 2 s 1 2 m d ( a 1 1 + a 2 1 ) ≲ β − 2 n 2 λ 1 2 ( a 2 2 s 2 4 m 2 + a 1 2 s 2 4 m 2 ) + β − 1 n 2 λ 1 2 s 1 2 d s 2 2 a 2 m κ ˉ 2 ( a 1 + s 1 2 d ) + ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 m c β κ ˉ 2 β n λ 1 ln n λ 1 ) + β − 1 n λ 1 s 1 2 d ( 1 a 1 + 1 a 2 ) . \lesssim\beta^{-2}\frac{n^2}{\lambda_1^2}\left(\frac{a_2^2}{s_2^4m^2}+\frac{a_1^2}{s_2^4m^2}\right)+\beta^{-1}\frac{n^2}{\lambda_1^2}\frac{s_1^2d}{s_2^2a_2m}\bar{\kappa}^{2}(a_1+s_1^2d)
+\left(1+\bar{\kappa}^{4}\beta\right)\exp\left(-\tilde{c}_2\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}}\right)+\beta^{-1}\frac{n}{\lambda_1}s_1^2d\left(\frac{1}{a_1}+\frac{1}{a_2}\right). ≲ β − 2 λ 1 2 n 2 ( s 2 4 m 2 a 2 2 + s 2 4 m 2 a 1 2 ) + β − 1 λ 1 2 n 2 s 2 2 a 2 m s 1 2 d κ ˉ 2 ( a 1 + s 1 2 d ) + ( 1 + κ ˉ 4 β ) exp ( − c ~ 2 c β κ ˉ 2 β λ 1 n ln λ 1 n m ) + β − 1 λ 1 n s 1 2 d ( a 1 1 + a 2 1 ) . Variance preservation and output scaling
In the proof of the lower bound for the minimal eigenvalue of the NTK (Lemma 5.3 ), we demanded a 1 2 s 1 2 = d \frac{a_{1}^{2}}{s_{1}^{2}}=d s 1 2 a 1 2 = d
\mathcal{L}_{\infty} \asymp\left(1+\bar{\kappa}^{4}\beta\right)\exp\left(-\frac{\tilde{c}_{2}m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)+\beta^{-2}\frac{n^{2}}{\lambda_{1}^{2}}\left(\frac{1}{\kappa^{4}}+\frac{s_{1}^{4}d^{2}}{s_{2}^{4}m^{2}}\right)
+ \beta^{-1}\frac{n^{2}}{\lambda_{1}^{2}}\bar{\kappa}^{2}\kappa^{2}\frac{s_{1}^{4}d^{2}}{s_{2}^{4}m^{2}}+\beta^{-1}\frac{n}{\lambda_{1}}\left(1+\kappa^{2}\frac{s_{1}^{2}d}{s_{2}^{2}m}\right).} \tag*{(34)}
Finally, since c β = O ( 1 ) c_{\beta}=O(1) c β = O ( 1 )
\mathcal{L}_{\infty} \asymp\left(1+\bar{\kappa}^{4}\beta\right)\exp\left(-\frac{c_{2}m}{\bar{\kappa}^{2}\beta\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)+\beta^{-2}\frac{n^{2}}{\lambda_{1}^{2}}\left(\frac{1}{\kappa^{4}}+\frac{s_{1}^{4}d^{2}}{s_{2}^{4}m^{2}}\right)
+ \beta^{-1}\frac{n^{2}}{\lambda_{1}^{2}}\bar{\kappa}^{2}\kappa^{2}\frac{s_{1}^{4}d^{2}}{s_{2}^{4}m^{2}}+\beta^{-1}\frac{n}{\lambda_{1}}\left(1+\kappa^{2}\frac{s_{1}^{2}d}{s_{2}^{2}m}\right).} \tag*{(35)}
Proof of global convergence of regularized CLD for deep networks
Setting
Network. We consider the neural network defined in § 4 . We denote m ≜ min l ∈ [ L ] m l m \triangleq\min_{l \in[L]} m_l m ≜ min l ∈ [ L ] m l
Intermediate quantities. Throughout the analysis, for every l ∈ [ L ] l \in[L] l ∈ [ L ] α ~ ( l ) ( x ) ∈ R m l \widetilde{\alpha}^{(l)}(x) \in\mathbb{R}^{m_l} α ( l ) ( x ) ∈ R m l
α ~ ( l ) ( x ) ≜ s l W ( l ) α ( l − 1 ) ( x ) \widetilde{\alpha}^{(l)}(x) \triangleq s_l W^{(l)} \alpha^{(l-1)}(x) α ( l ) ( x ) ≜ s l W ( l ) α ( l − 1 ) ( x ) and δ ( l ) ( x ) ∈ R m l \delta^{(l)}(x) \in\mathbb{R}^{m_l} δ ( l ) ( x ) ∈ R m l
δ ( l ) ( x ) ≜ ( ∂ f ( θ ; x ) ∂ α ~ ( l ) ( x ) ) T . \delta^{(l)}(x) \triangleq\left(\frac{\partial f(\theta;x)}{\partial\widetilde{\alpha}^{(l)}(x)}\right)^{\mathsf{T}}. δ ( l ) ( x ) ≜ ( ∂ α ( l ) ( x ) ∂ f ( θ ; x ) ) T . We can express δ ( l ) ( x ) \delta^{(l)}(x) δ ( l ) ( x )
δ ( L ) ( x ) = s L + 1 diag ( ϕ ′ ( α ~ ( L ) ( x ) ) ) v , \delta^{(L)}(x) = s_{L+1}\operatorname{diag}\left(\phi'\left(\widetilde{\alpha}^{(L)}(x)\right)\right)\mathbf{v}, δ ( L ) ( x ) = s L + 1 diag ( ϕ ′ ( α ( L ) ( x ) ) ) v , ∀ l ∈ [ L − 1 ] , δ ( l ) ( x ) = ( ∂ f ( θ ; x ) ∂ α ~ ( l + 1 ) ( x ) ∂ α ~ ( l + 1 ) ( x ) ∂ α ( l ) ( x ) ∂ α ( l ) ( x ) ∂ α ~ ( l ) ( x ) ) T = diag ( ϕ ′ ( α ~ ( l ) ( x ) ) ) s l + 1 ( W ( l + 1 ) ) T δ ( l + 1 ) ( x ) . (36) \begin{aligned}
\forall l \in[L-1],\quad\delta^{(l)}(x)
&= \left(\frac{\partial f(\theta;x)}{\partial\widetilde{\alpha}^{(l+1)}(x)}
\frac{\partial\widetilde{\alpha}^{(l+1)}(x)}{\partial\alpha^{(l)}(x)}
\frac{\partial\alpha^{(l)}(x)}{\partial\widetilde{\alpha}^{(l)}(x)}\right)^{\mathsf{T}} \\
&= \operatorname{diag}\left(\phi'\left(\widetilde{\alpha}^{(l)}(x)\right)\right)s_{l+1}\left(W^{(l+1)}\right)^{\mathsf{T}}\delta^{(l+1)}(x). \tag*{(36)}
\end{aligned} ∀ l ∈ [ L − 1 ] , δ ( l ) ( x ) = ( ∂ α ( l + 1 ) ( x ) ∂ f ( θ ; x ) ∂ α ( l ) ( x ) ∂ α ( l + 1 ) ( x ) ∂ α ( l ) ( x ) ∂ α ( l ) ( x ) ) T = diag ( ϕ ′ ( α ( l ) ( x ) ) ) s l + 1 ( W ( l + 1 ) ) T δ ( l + 1 ) ( x ) . (36) Note that for every l ∈ [ L ] l \in[L] l ∈ [ L ] ∂ f ( θ ; x ) ∂ W ( l ) ∈ R m l × m l − 1 \frac{\partial f(\theta;x)}{\partial W^{(l)}} \in\mathbb{R}^{m_l \times m_{l-1}} ∂ W ( l ) ∂ f ( θ ; x ) ∈ R m l × m l − 1
∂ f ( θ ; x ) ∂ W ( l ) = s l δ ( l ) ( x ) ( α ( l − 1 ) ( x ) ) T . (37) \frac{\partial f(\theta;x)}{\partial W^{(l)}} = s_l\delta^{(l)}(x)\left(\alpha^{(l-1)}(x)\right)^{\mathsf{T}}. \tag*{(37)} ∂ W ( l ) ∂ f ( θ ; x ) = s l δ ( l ) ( x ) ( α ( l − 1 ) ( x ) ) T . (37) Matrix norms. Throughout this section, unless stated otherwise, ∥ ⋅ ∥ \lVert\cdot\rVert ∥ ⋅ ∥
Sums and products. The analysis of the neural network’s gradient and Laplacian includes many sums and products over ranges of layers. To make the math as compact as possible, we sometimes combine terms in a way that makes a sum’s (or product’s) initial running index greater than its final index. We define such occurrences to be equal to 0 for sums and to 1 for products.
Use of ≲ \lesssim ≲ ≍ \asymp ≍ Throughout the proof, inequalities involving both descent and error terms are interpreted in the following one-sided asymptotic sense: an estimate of the form
X ≲ − A + B X \lesssim-A + B X ≲ − A + B means that there exist universal constants c , C > 0 c, C > 0 c , C > 0
X ≤ − c A + C B . X \leq-cA + CB. X ≤ − c A + C B . Thus, positive error terms may be enlarged and descent coefficients may be reduced by universal constant factors. These constants are absorbed into the final constants c 1 c_1 c 1 c 2 c_2 c 2 L ∞ \mathcal{L}_{\infty} L ∞
Similarly, when writing L ∞ ≍ F \mathcal{L}_{\infty} \asymp F L ∞ ≍ F L ∞ = C F \mathcal{L}_{\infty} = CF L ∞ = C F C > 0 C > 0 C > 0 ≍ \asymp ≍
For a ∈ R p a \in\mathbb{R}^{p} a ∈ R p b ∈ R q b \in\mathbb{R}^{q} b ∈ R q ∂ a ∂ b ∈ R p × q \frac{\partial a}{\partial b} \in\mathbb{R}^{p \times q} ∂ b ∂ a ∈ R p × q ( ∂ a ∂ b ) p q = ∂ a p ∂ b q \left(\frac{\partial a}{\partial b}\right)_{pq} = \frac{\partial a_p}{\partial b_q} ( ∂ b ∂ a ) pq = ∂ b q ∂ a p p = 1 p=1 p = 1 q = 1 q=1 q = 1
Supporting propositions
Proposition H.1. There exists a constant c > 0 c > 0 c > 0 p ≥ 1 p \ge1 p ≥ 1
∥ ∥ W ( l ) ∥ 2 ∥ L p ≤ c ν 0 2 a l ( m l − 1 + m l + p ) , ∀ l ∈ [ L ] ; \left\lVert\left\lVert\mathbf{W}^{(l)} \right\rVert^{2} \right\rVert_{L_p} \le c \frac{\nu_0^2}{a_l} (m_{l-1} + m_l + p), \quad\forall l \in[L]; W ( l ) 2 L p ≤ c a l ν 0 2 ( m l − 1 + m l + p ) , ∀ l ∈ [ L ] ; ∥ ∥ v ∥ 2 ∥ L p ≤ c 1 a L + 1 ( p + m L ) . \left\lVert\left\lVert\mathbf{v} \right\rVert^{2} \right\rVert_{L_p} \le c \frac{1}{a_{L+1}} (p + m_L). ∥ v ∥ 2 L p ≤ c a L + 1 1 ( p + m L ) . Proof. Let G ∈ R m × n \mathbf{G} \in\mathbb{R}^{m \times n} G ∈ R m × n c ′ > 0 c' > 0 c ′ > 0 l ∈ [ L ] l \in[L] l ∈ [ L ] p ≥ 1 p \ge1 p ≥ 1
∥ ∥ W ( l ) ∥ 2 ∥ L p ≤ ∥ ∥ W ( l ) ∥ 2 ∥ L 2 p = ν 0 2 a l ∥ ∥ G ( l ) ∥ 2 ∥ L 2 p ≤ c ′ ν 0 2 a l ( m l − 1 + m l + p ) . \left\lVert\left\lVert\mathbf{W}^{(l)} \right\rVert^{2} \right\rVert_{L_p} \le\left\lVert\left\lVert\mathbf{W}^{(l)} \right\rVert^{2} \right\rVert_{L_{2p}} = \frac{\nu_0^2}{a_l} \left\lVert\left\lVert\mathbf{G}^{(l)} \right\rVert^{2} \right\rVert_{L_{2p}} \le c' \frac{\nu_0^2}{a_l} (m_{l-1} + m_l + p). W ( l ) 2 L p ≤ W ( l ) 2 L 2 p = a l ν 0 2 G ( l ) 2 L 2 p ≤ c ′ a l ν 0 2 ( m l − 1 + m l + p ) . Additionally, denoting z ( L + 1 ) ∼ N ( 0 m L , I m L ) \mathbf{z}^{(L+1)} \sim\mathcal{N}(\mathbf{0}_{m_L}, \mathbf{I}_{m_L}) z ( L + 1 ) ∼ N ( 0 m L , I m L ) c ′ > 0 c' > 0 c ′ > 0
∥ ∥ v ∥ 2 ∥ L p ≤ ∥ ∥ v ∥ 2 ∥ L 2 p = 1 a L + 1 ∥ ∥ z ( L + 1 ) ∥ 2 ∥ L 2 p ≤ 1 a L + 1 ( c 2 p + m L ) 2 ≤ 4 c ′ 1 a L + 1 ( p + m L ) . \begin{aligned}
\left\lVert\left\lVert\mathbf{v} \right\rVert^{2} \right\rVert_{L_p} \le\left\lVert\left\lVert\mathbf{v} \right\rVert^{2} \right\rVert_{L_{2p}} = \frac{1}{a_{L+1}} \left\lVert\left\lVert\mathbf{z}^{(L+1)} \right\rVert^{2} \right\rVert_{L_{2p}} \\
&\le\frac{1}{a_{L+1}} \left(c\sqrt{2p}+\sqrt{m_L}\right)^2 \le4c' \frac{1}{a_{L+1}} (p+m_L).
\end{aligned} ∥ v ∥ 2 L p ≤ ∥ v ∥ 2 L 2 p = a L + 1 1 z ( L + 1 ) 2 L 2 p ≤ a L + 1 1 ( c 2 p + m L ) 2 ≤ 4 c ′ a L + 1 1 ( p + m L ) . :::
Proposition H.2. It holds that for all l ∈ [ L ] l \in[L] l ∈ [ L ]
∥ δ ( l ) ( x ) ∥ 2 ≤ α ϕ 2 ( L + 1 − l ) ( ∏ k = l + 1 L s k 2 ∥ W ( k ) ∥ 2 ) s L + 1 2 ∥ v ∥ 2 . \left\lVert\boldsymbol{\delta}^{(l)}(x) \right\rVert^2 \le\alpha_\phi^{2(L+1-l)} \left(\prod_{k=l+1}^{L} s_k^2 \left\lVert\mathbf{W}^{(k)} \right\rVert^2\right) s_{L+1}^2 \left\lVert\mathbf{v} \right\rVert^2. δ ( l ) ( x ) 2 ≤ α ϕ 2 ( L + 1 − l ) ( k = l + 1 ∏ L s k 2 W ( k ) 2 ) s L + 1 2 ∥ v ∥ 2 . In addition, at initialization, there exists a constant c > 0 c > 0 c > 0 p ≥ 1 p \ge1 p ≥ 1
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ L p ≤ ( c ν 0 α ϕ ) 2 ( L + 1 − l ) ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( p + m L ) . \left\lVert\left\lVert\boldsymbol{\delta}^{(l)}(x) \right\rVert^2 \right\rVert_{L_p} \le(c\nu_0\alpha_\phi)^{2(L+1-l)} \left(\prod_{k=l+1}^{L} \frac{s_k^2}{a_k}(m_{k-1}+m_k+p)\right)\frac{s_{L+1}^2}{a_{L+1}}(p+m_L). δ ( l ) ( x ) 2 L p ≤ ( c ν 0 α ϕ ) 2 ( L + 1 − l ) ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) a L + 1 s L + 1 2 ( p + m L ) . As a direct consequence, at initialization,
∥ ∥ δ ( l ) ( x ) ∥ ∥ L p ≤ ( c ν 0 α ϕ ) L + 1 − l ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( p + m L ) . \left\lVert\left\lVert\boldsymbol{\delta}^{(l)}(x) \right\rVert\right\rVert_{L_p} \le(c\nu_0\alpha_\phi)^{L+1-l} \left(\prod_{k=l+1}^{L} \sqrt{\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)}\right)\sqrt{\frac{s_{L+1}^2}{a_{L+1}}(p+m_L)}. δ ( l ) ( x ) L p ≤ ( c ν 0 α ϕ ) L + 1 − l k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) a L + 1 s L + 1 2 ( p + m L ) . Proof. By Eq. (36) , properties of the Euclidean norm and the spectral norm, and Assumption 4.2, we have that
∥ δ ( L ) ( x ) ∥ 2 ≤ s L + 1 2 ∥ diag ( ϕ ′ ( α ~ ( L ) ( x ) ) ) ∥ 2 ∥ v ∥ 2 ≤ α ϕ 2 s L + 1 2 ∥ v ∥ 2 . \left\lVert\boldsymbol{\delta}^{(L)}(x) \right\rVert^2 \le s_{L+1}^2 \left\lVert\operatorname{diag}\left(\phi'\left(\widetilde{\alpha}^{(L)}(x)\right)\right) \right\rVert^2 \left\lVert\mathbf{v} \right\rVert^2 \le\alpha_\phi^2 s_{L+1}^2 \left\lVert\mathbf{v} \right\rVert^2. δ ( L ) ( x ) 2 ≤ s L + 1 2 diag ( ϕ ′ ( α ( L ) ( x ) ) ) 2 ∥ v ∥ 2 ≤ α ϕ 2 s L + 1 2 ∥ v ∥ 2 . In addition, for all l ∈ [ L − 1 ] l \in[L-1] l ∈ [ L − 1 ]
∥ δ ( l ) ( x ) ∥ 2 ≤ s l + 1 2 ∥ diag ( ϕ ′ ( α ~ ( l ) ( x ) ) ) ∥ 2 ∥ W ( l + 1 ) ∥ 2 ∥ δ ( l + 1 ) ( x ) ∥ 2 \left\lVert\boldsymbol{\delta}^{(l)}(x) \right\rVert^2 \le s_{l+1}^2 \left\lVert\operatorname{diag}\left(\phi'\left(\widetilde{\alpha}^{(l)}(x)\right)\right) \right\rVert^2 \left\lVert\mathbf{W}^{(l+1)} \right\rVert^2 \left\lVert\boldsymbol{\delta}^{(l+1)}(x) \right\rVert^2 δ ( l ) ( x ) 2 ≤ s l + 1 2 diag ( ϕ ′ ( α ( l ) ( x ) ) ) 2 W ( l + 1 ) 2 δ ( l + 1 ) ( x ) 2 ≤ α ϕ 2 s l + 1 2 ∥ W ( l + 1 ) ∥ 2 ∥ δ ( l + 1 ) ( x ) ∥ 2 . \le\alpha_\phi^2s_{l+1}^2\left\lVert\mathbf{W}^{(l+1)} \right\rVert^2\left\lVert\boldsymbol{\delta}^{(l+1)}(x) \right\rVert^2. ≤ α ϕ 2 s l + 1 2 W ( l + 1 ) 2 δ ( l + 1 ) ( x ) 2 . Continuing recursively, we get
∥ δ ( l ) ( x ) ∥ 2 ≤ α ϕ 2 ( L + 1 − l ) ( ∏ k = l + 1 L s k 2 ∥ W ( k ) ∥ 2 ) s L + 1 2 ∥ v ∥ 2 . \left\lVert\boldsymbol{\delta}^{(l)}(x) \right\rVert^2 \le\alpha_\phi^{2(L+1-l)}\left(\prod_{k=l+1}^{L}s_k^2\left\lVert\mathbf{W}^{(k)} \right\rVert^2\right)s_{L+1}^2\left\lVert\mathbf{v} \right\rVert^2. δ ( l ) ( x ) 2 ≤ α ϕ 2 ( L + 1 − l ) ( k = l + 1 ∏ L s k 2 W ( k ) 2 ) s L + 1 2 ∥ v ∥ 2 . Together with Proposition H.1, we get
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ L p ≤ α ϕ 2 ( L + 1 − l ) ( ∏ k = l + 1 L c 3 ν 0 2 s k 2 a k ( m k − 1 + m k + p ) ) c 4 s L + 1 2 a L + 1 ( p + m L ) . \left\lVert\left\lVert\boldsymbol{\delta}^{(l)}(x) \right\rVert^2 \right\rVert_{L_p} \le\alpha_\phi^{2(L+1-l)}\left(\prod_{k=l+1}^{L}c_3\nu_0^2\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)\right)c_4\frac{s_{L+1}^2}{a_{L+1}}(p+m_L). δ ( l ) ( x ) 2 L p ≤ α ϕ 2 ( L + 1 − l ) ( k = l + 1 ∏ L c 3 ν 0 2 a k s k 2 ( m k − 1 + m k + p ) ) c 4 a L + 1 s L + 1 2 ( p + m L ) . Therefore, there exists a constant c 5 > 0 c_5 > 0 c 5 > 0
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ L p ≤ ( c 5 ν 0 2 α ϕ 2 ) L + 1 − l ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( p + m L ) . \left\|\left\|\delta^{(l)}(\mathbf{x})\right\|^2\right\|_{L^p} \le\left(c_5\nu_0^2\alpha_\phi^2\right)^{L+1-l}\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)\right)\frac{s_{L+1}^2}{a_{L+1}}(p+m_L). δ ( l ) ( x ) 2 L p ≤ ( c 5 ν 0 2 α ϕ 2 ) L + 1 − l ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) a L + 1 s L + 1 2 ( p + m L ) . Therefore,
∥ ∥ δ ( l ) ( x ) ∥ ∥ L p ≤ ( i ) ∥ ∥ δ ( l ) ( x ) ∥ ∥ L 2 p \left\|\left\|\delta^{(l)}(\mathbf{x})\right\|\right\|_{L^p}
\mathrel{\overset{(i)}{\le}} \left\|\left\|\delta^{(l)}(\mathbf{x})\right\|\right\|_{L^{2p}} δ ( l ) ( x ) L p ≤ ( i ) δ ( l ) ( x ) L 2 p = ( i i ) ∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ L p ≤ ( c 5 ν 0 α ϕ ) L + 1 − l ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( p + m L ) , \begin{aligned}
\mathrel{\overset{(ii)}{=}} \sqrt{\left\|\left\|\delta^{(l)}(\mathbf{x})\right\|^2\right\|_{L^p}} \\
&\le\left(\sqrt{c_5}\nu_0\alpha_\phi\right)^{L+1-l}\left(\prod_{k=l+1}^{L}\sqrt{\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)}\right)\sqrt{\frac{s_{L+1}^2}{a_{L+1}}(p+m_L)},
\end{aligned} = ( ii ) δ ( l ) ( x ) 2 L p ≤ ( c 5 ν 0 α ϕ ) L + 1 − l k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) a L + 1 s L + 1 2 ( p + m L ) , where ( i ) (i) ( i ) L p L^p L p ( i i ) (ii) ( ii )
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ L p = ( E [ ∥ δ ( l ) ( x ) ∥ 2 p ] ) 1 / p = ( E [ ∥ δ ( l ) ( x ) ∥ 2 p ] ) 1 / 2 p = ∥ ∥ δ ( l ) ( x ) ∥ ∥ L 2 p . \sqrt{\left\|\left\|\delta^{(l)}(\mathbf{x})\right\|^2\right\|_{L^p}}
= \sqrt{\left(\mathbb{E}\left[\left\|\delta^{(l)}(\mathbf{x})\right\|^{2p}\right]\right)^{1/p}}
= \left(\mathbb{E}\left[\left\|\delta^{(l)}(\mathbf{x})\right\|^{2p}\right]\right)^{1/2p}
= \left\|\left\|\delta^{(l)}(\mathbf{x})\right\|\right\|_{L^{2p}}. δ ( l ) ( x ) 2 L p = ( E [ δ ( l ) ( x ) 2 p ] ) 1/ p = ( E [ δ ( l ) ( x ) 2 p ] ) 1/2 p = δ ( l ) ( x ) L 2 p . □ \square □
Proposition H.3. For every l ∈ [ L ] l \in[L] l ∈ [ L ] c > 0 c > 0 c > 0 p ≥ 1 p \ge1 p ≥ 1
∥ ∥ α ( l ) ( x ) ∥ 2 ∥ L p ≤ ( c α ϕ 2 ν 0 2 ) l ∥ x ∥ 2 ∏ k = 1 l s k 2 a k ( p + m k ) . \left\|\left\|\alpha^{(l)}(\mathbf{x})\right\|^2\right\|_{L^p}
\le\left(c\alpha_\phi^2\nu_0^2\right)^l\|\mathbf{x}\|^2\prod_{k=1}^{l}\frac{s_k^2}{a_k}(p+m_k). α ( l ) ( x ) 2 L p ≤ ( c α ϕ 2 ν 0 2 ) l ∥ x ∥ 2 k = 1 ∏ l a k s k 2 ( p + m k ) . Proof. For any k ∈ [ L ] k \in[L] k ∈ [ L ] F k ≜ σ ( W ( 1 ) , … , W ( k ) ) \mathcal{F}_k \triangleq\sigma\left(\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(k)}\right) F k ≜ σ ( W ( 1 ) , … , W ( k ) ) σ \sigma σ ( W ( 1 ) , … , W ( k ) ) \left(\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(k)}\right) ( W ( 1 ) , … , W ( k ) ) F 0 \mathcal{F}_0 F 0 σ \sigma σ { z ( k ) ∈ R m k } k = 1 L \left\{\mathbf{z}^{(k)} \in\mathbb{R}^{m_k}\right\}_{k=1}^{L} { z ( k ) ∈ R m k } k = 1 L θ \theta θ z ( k ) ∼ N ( 0 m k , I m k ) \mathbf{z}^{(k)} \sim\mathcal{N}\left(0_{m_k},I_{m_k}\right) z ( k ) ∼ N ( 0 m k , I m k ) k ∈ [ L ] k \in[L] k ∈ [ L ]
Let l ∈ [ L ] l \in[L] l ∈ [ L ]
α ~ ( l ) ( x ) ∣ F l − 1 ∼ N ( 0 m l , ν 0 2 s l 2 a l ∥ α ( l − 1 ) ( x ) ∥ 2 I m l ) . \tilde{\alpha}^{(l)}(\mathbf{x}) \mid\mathcal{F}_{l-1}
\sim\mathcal{N}\left(0_{m_l},\nu_0^2\frac{s_l^2}{a_l}\left\|\alpha^{(l-1)}(\mathbf{x})\right\|^2 I_{m_l}\right). α ~ ( l ) ( x ) ∣ F l − 1 ∼ N ( 0 m l , ν 0 2 a l s l 2 α ( l − 1 ) ( x ) 2 I m l ) . Therefore,
α ~ ( l ) ( x ) ∣ F l − 1 = d ν 0 s l a l ∥ α ( l − 1 ) ( x ) ∥ z ( l ) , \tilde{\alpha}^{(l)}(\mathbf{x}) \mid\mathcal{F}_{l-1}
\mathrel{\overset{d}{=}} \nu_0\frac{s_l}{\sqrt{a_l}}\left\|\alpha^{(l-1)}(\mathbf{x})\right\|\mathbf{z}^{(l)}, α ~ ( l ) ( x ) ∣ F l − 1 = d ν 0 a l s l α ( l − 1 ) ( x ) z ( l ) , where = d \overset{d}{=} = d
∥ α ~ ( l ) ( x ) ∥ 2 ∣ F l − 1 = d ν 0 2 s l 2 a l ∥ z ( l ) ∥ 2 ∥ α ( l − 1 ) ( x ) ∥ 2 . \left\|\tilde{\alpha}^{(l)}(\mathbf{x})\right\|^2 \mid\mathcal{F}_{l-1}
\mathrel{\overset{d}{=}} \nu_0^2\frac{s_l^2}{a_l}\left\|\mathbf{z}^{(l)}\right\|^2\left\|\alpha^{(l-1)}(\mathbf{x})\right\|^2. α ~ ( l ) ( x ) 2 ∣ F l − 1 = d ν 0 2 a l s l 2 z ( l ) 2 α ( l − 1 ) ( x ) 2 . Together with Assumptions 4.2 and 4.3 and Claim E.7,
∥ ∥ α ( l ) ( x ) ∥ 2 ∥ L p ∣ F l − 1 ≤ ∥ α ϕ 2 ∥ α ~ ( l ) ( x ) ∥ 2 ∥ L p ∣ F l − 1 \left\|\left\|\alpha^{(l)}(\mathbf{x})\right\|^2\right\|_{L^p\mid\mathcal{F}_{l-1}}
\le\left\|\alpha_\phi^2\left\|\tilde{\alpha}^{(l)}(\mathbf{x})\right\|^2\right\|_{L^p\mid\mathcal{F}_{l-1}} α ( l ) ( x ) 2 L p ∣ F l − 1 ≤ α ϕ 2 α ~ ( l ) ( x ) 2 L p ∣ F l − 1 ≤ α ϕ 2 ν 0 2 s l 2 a l ∥ ∥ z ( l ) ∥ 2 ∥ L p ∥ ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p ∣ F l − 1 \le\alpha_\phi^2\nu_0^2\frac{s_l^2}{a_l}
\left\|\left\|\mathbf{z}^{(l)}\right\|^2\right\|_{L^p}
\left\|\left\|\alpha^{(l-1)}(\mathbf{x})\right\|^2\right\|_{L^p\mid\mathcal{F}_{l-1}} ≤ α ϕ 2 ν 0 2 a l s l 2 z ( l ) 2 L p α ( l − 1 ) ( x ) 2 L p ∣ F l − 1 = α ϕ 2 ν 0 2 s l 2 a l ∥ ∥ z ( l ) ∥ 2 ∥ L p ∥ α ( l − 1 ) ( x ) ∥ 2 . = \alpha_\phi^2\nu_0^2\frac{s_l^2}{a_l}
\left\|\left\|\mathbf{z}^{(l)}\right\|^2\right\|_{L^p}
\left\|\alpha^{(l-1)}(\mathbf{x})\right\|^2. = α ϕ 2 ν 0 2 a l s l 2 z ( l ) 2 L p α ( l − 1 ) ( x ) 2 . By Claim E.4,
∥ ∥ α ( l ) ( x ) ∥ 2 ∥ L p ≤ ∥ ∥ α ( l ) ( x ) ∥ 2 ∥ L p ∥ F l − 1 ∥ L p = ∥ α ϕ 2 ν 0 2 s l 2 a l ∥ z ( l ) ∥ 2 ∥ L p ∥ ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p = α ϕ 2 ν 0 2 s l 2 a l ∥ ∥ z ( l ) ∥ 2 ∥ L p ∥ ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p . \begin{aligned}
\left\|\left\|\alpha^{(l)}(x)\right\|^2\right\|_{L^p}
\leq
\left\|\left\|\alpha^{(l)}(x)\right\|^2\right\|_{L^p}
\left\|\mathcal{F}_{l-1}\right\|_{L^p} \\
&=
\left\|\alpha_{\phi}^{2}\nu_{0}^{2}\frac{s_l^2}{a_l}\left\|z^{(l)}\right\|^2\right\|_{L^p}
\left\|\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L^p} \\
&=
\alpha_{\phi}^{2}\nu_{0}^{2}\frac{s_l^2}{a_l}
\left\|\left\|z^{(l)}\right\|^2\right\|_{L^p}
\left\|\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L^p}.
\end{aligned} α ( l ) ( x ) 2 L p ≤ α ( l ) ( x ) 2 L p ∥ F l − 1 ∥ L p = α ϕ 2 ν 0 2 a l s l 2 z ( l ) 2 L p α ( l − 1 ) ( x ) 2 L p = α ϕ 2 ν 0 2 a l s l 2 z ( l ) 2 L p α ( l − 1 ) ( x ) 2 L p . Continuing recursively, recalling that α ( 0 ) ( x ) = x \alpha^{(0)}(x)=x α ( 0 ) ( x ) = x c 4 > 0 c_4>0 c 4 > 0
∥ ∥ α ( l ) ( x ) ∥ 2 ∥ L p ≤ ∥ x ∥ 2 ∏ k = 1 l α ϕ 2 ν 0 2 s k 2 a k ∥ ∥ z ( k ) ∥ 2 ∥ L p ≤ ∥ x ∥ 2 ∏ k = 1 l α ϕ 2 ν 0 2 s k 2 a k ( c 4 2 p + m k ) 2 ≤ ∥ x ∥ 2 ∏ k = 1 l 4 c 4 α ϕ 2 ν 0 2 s k 2 a k ( p + m k ) = ( 4 c 4 α ϕ 2 ν 0 2 ) l ∥ x ∥ 2 ∏ k = 1 l s k 2 a k ( p + m k ) . \begin{aligned}
\left\|\left\|\alpha^{(l)}(x)\right\|^2\right\|_{L^p}
\leq
\|x\|^2\prod_{k=1}^{l}\alpha_{\phi}^{2}\nu_{0}^{2}\frac{s_k^2}{a_k}
\left\|\left\|z^{(k)}\right\|^2\right\|_{L^p} \\
&\leq
\|x\|^2\prod_{k=1}^{l}\alpha_{\phi}^{2}\nu_{0}^{2}\frac{s_k^2}{a_k}
\left(c_4\sqrt{2p}+\sqrt{m_k}\right)^2 \\
&\leq
\|x\|^2\prod_{k=1}^{l}4c_4\alpha_{\phi}^{2}\nu_{0}^{2}\frac{s_k^2}{a_k}(p+m_k) \\
&=
\left(4c_4\alpha_{\phi}^{2}\nu_{0}^{2}\right)^l\|x\|^2
\prod_{k=1}^{l}\frac{s_k^2}{a_k}(p+m_k).
\end{aligned} α ( l ) ( x ) 2 L p ≤ ∥ x ∥ 2 k = 1 ∏ l α ϕ 2 ν 0 2 a k s k 2 z ( k ) 2 L p ≤ ∥ x ∥ 2 k = 1 ∏ l α ϕ 2 ν 0 2 a k s k 2 ( c 4 2 p + m k ) 2 ≤ ∥ x ∥ 2 k = 1 ∏ l 4 c 4 α ϕ 2 ν 0 2 a k s k 2 ( p + m k ) = ( 4 c 4 α ϕ 2 ν 0 2 ) l ∥ x ∥ 2 k = 1 ∏ l a k s k 2 ( p + m k ) . Proposition H.4. There exists a constant c c c l ∈ [ L ] l\in[L] l ∈ [ L ] p ≥ 1 p\geq1 p ≥ 1
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p ≤ ( c α ϕ 2 ν 0 2 ) 2 L ∥ x ∥ 2 ( ∏ k = 1 l − 1 s k 2 a k ( m k + p ) ) ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( m L + p ) . \left\|\left\|\delta^{(l)}(x)\right\|^2\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L^p}
\leq
\left(c\alpha_{\phi}^{2}\nu_{0}^{2}\right)^{2L}\|x\|^2
\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}(m_k+p)\right)
\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)\right)
\frac{s_{L+1}^2}{a_{L+1}}(m_L+p). δ ( l ) ( x ) 2 α ( l − 1 ) ( x ) 2 L p ≤ ( c α ϕ 2 ν 0 2 ) 2 L ∥ x ∥ 2 ( k = 1 ∏ l − 1 a k s k 2 ( m k + p ) ) ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) a L + 1 s L + 1 2 ( m L + p ) . Proof. By Proposition H.2, we have
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p ≤ α ϕ 2 ( L + 1 − l ) ( ∏ k = l + 1 L s k 2 ∥ ∥ W ( k ) ∥ 2 ∥ L p ) s L + 1 2 ∥ v ∥ L p 2 ∥ ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p . \left\|\left\|\delta^{(l)}(x)\right\|^2\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L^p}
\leq
\alpha_{\phi}^{2(L+1-l)}
\left(\prod_{k=l+1}^{L}s_k^2\left\|\left\|\mathbf{W}^{(k)}\right\|^2\right\|_{L^p}\right)
s_{L+1}^2\left\|\mathbf{v}\right\|_{L^p}^2
\left\|\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L^p}. δ ( l ) ( x ) 2 α ( l − 1 ) ( x ) 2 L p ≤ α ϕ 2 ( L + 1 − l ) ( k = l + 1 ∏ L s k 2 W ( k ) 2 L p ) s L + 1 2 ∥ v ∥ L p 2 α ( l − 1 ) ( x ) 2 L p . Note that ∥ α ( l − 1 ) ( x ) ∥ 2 \left\|\alpha^{(l-1)}(x)\right\|^2 α ( l − 1 ) ( x ) 2
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p ≤ α ϕ 2 ( L + 1 − l ) ( ∏ k = l + 1 L s k 2 ∥ ∥ W ( k ) ∥ 2 ∥ L p ) s L + 1 2 ∥ ∥ v ∥ 2 ∥ L p ∥ ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p . \left\|\left\|\delta^{(l)}(x)\right\|^2\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L^p}
\leq
\alpha_{\phi}^{2(L+1-l)}
\left(\prod_{k=l+1}^{L}s_k^2\left\|\left\|\mathbf{W}^{(k)}\right\|^2\right\|_{L^p}\right)
s_{L+1}^2\left\|\left\|\mathbf{v}\right\|^2\right\|_{L^p}
\left\|\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L^p}. δ ( l ) ( x ) 2 α ( l − 1 ) ( x ) 2 L p ≤ α ϕ 2 ( L + 1 − l ) ( k = l + 1 ∏ L s k 2 W ( k ) 2 L p ) s L + 1 2 ∥ v ∥ 2 L p α ( l − 1 ) ( x ) 2 L p . Using Propositions H.1 and H.3, we get that there exists a constant c > 0 c>0 c > 0
∥ ∥ δ ( l ) ( x ) ∥ 2 ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p ≤ α ϕ 2 ( L + 1 − l ) ( ∏ k = l + 1 L c γ 0 2 s k 2 a k ( m k − 1 + m k + p ) ) ⋅ c s L + 1 2 a L + 1 ( p + m L ) ∥ x ∥ 2 ( c α ϕ 2 ν 0 2 ) l − 1 ∏ k = 1 l − 1 s k 2 a k ( p + m k ) ≤ ( c α ϕ 2 ν 0 2 ) 2 L ∥ x ∥ 2 s L + 1 2 a L + 1 ( p + m L ) ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) ∏ k = 1 l − 1 s k 2 a k ( p + m k ) ≤ ( c α ϕ 2 ν 0 2 ) 2 L ∥ x ∥ 2 ( ∏ k = 1 l − 1 s k 2 a k ( m k + p ) ) ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( m L + p ) . \begin{aligned}
\left\|\left\|\delta^{(l)}(x)\right\|^{2}\left\|\alpha^{(l-1)}(x)\right\|^{2}\right\|_{L_p}
&\leq\alpha_{\phi}^{2(L+1-l)}
\left(\prod_{k=l+1}^{L} c\gamma_{0}^{2}\frac{s_k^{2}}{a_k}(m_{k-1}+m_k+p)\right)\\
&\quad\cdot c\frac{s_{L+1}^{2}}{a_{L+1}}(p+m_L)\|x\|^{2}(c\alpha_{\phi}^{2}\nu_{0}^{2})^{l-1}
\prod_{k=1}^{l-1}\frac{s_k^{2}}{a_k}(p+m_k)\\
&\leq(c\alpha_{\phi}^{2}\nu_{0}^{2})^{2L}\|x\|^{2}\frac{s_{L+1}^{2}}{a_{L+1}}(p+m_L)
\left(\prod_{k=l+1}^{L}\frac{s_k^{2}}{a_k}(m_{k-1}+m_k+p)\right)
\prod_{k=1}^{l-1}\frac{s_k^{2}}{a_k}(p+m_k)\\
&\leq(c\alpha_{\phi}^{2}\nu_{0}^{2})^{2L}\|x\|^{2}
\left(\prod_{k=1}^{l-1}\frac{s_k^{2}}{a_k}(m_k+p)\right)
\left(\prod_{k=l+1}^{L}\frac{s_k^{2}}{a_k}(m_{k-1}+m_k+p)\right)
\frac{s_{L+1}^{2}}{a_{L+1}}(m_L+p).
\end{aligned} δ ( l ) ( x ) 2 α ( l − 1 ) ( x ) 2 L p ≤ α ϕ 2 ( L + 1 − l ) ( k = l + 1 ∏ L c γ 0 2 a k s k 2 ( m k − 1 + m k + p ) ) ⋅ c a L + 1 s L + 1 2 ( p + m L ) ∥ x ∥ 2 ( c α ϕ 2 ν 0 2 ) l − 1 k = 1 ∏ l − 1 a k s k 2 ( p + m k ) ≤ ( c α ϕ 2 ν 0 2 ) 2 L ∥ x ∥ 2 a L + 1 s L + 1 2 ( p + m L ) ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) k = 1 ∏ l − 1 a k s k 2 ( p + m k ) ≤ ( c α ϕ 2 ν 0 2 ) 2 L ∥ x ∥ 2 ( k = 1 ∏ l − 1 a k s k 2 ( m k + p ) ) ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) a L + 1 s L + 1 2 ( m L + p ) . Introduction and goal
Recalling Eq. (11) from § 5.1 , we have that
d d t E [ L ( θ t ) ] = 1 n ( − n − 1 E [ r ( θ t ) T K N T K ( θ t ) r ( θ t ) ] + β − 1 E [ Tr ( ∇ f ( θ t ) T ∇ f ( θ t ) ) ] + β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] − β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Λ θ t ] ) . (38) \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}[\mathcal{L}(\theta_t)]
=\frac{1}{n}\left(-n^{-1}\mathbb{E}\left[r(\theta_t)^{\mathsf{T}}K_{\mathrm{NTK}}(\theta_t)r(\theta_t)\right]
+\beta^{-1}\mathbb{E}\left[\operatorname{Tr}\left(\nabla f(\theta_t)^{\mathsf{T}}\nabla f(\theta_t)\right)\right]
+\beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]
-\beta^{-1}\mathbb{E}\left[r(\theta_t)^{\mathsf{T}}\nabla f(\theta_t)\Lambda\theta_t\right]\right). \tag*{(38)} d t d E [ L ( θ t )] = n 1 ( − n − 1 E [ r ( θ t ) T K NTK ( θ t ) r ( θ t ) ] + β − 1 E [ Tr ( ∇ f ( θ t ) T ∇ f ( θ t ) ) ] + β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] − β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Λ θ t ] ) . (38) We wish to bound every term in the right-hand side of this equation using Corollary 3.1. In order to use Corollary 3.1, we must bound ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α \left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}} ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1
E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ≤ 2 e 2 ( α − 1 ) β R y 2 , \mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\leq2e^{2(\alpha-1)\beta R_y^2}, E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ≤ 2 e 2 ( α − 1 ) β R y 2 , provided that 2 ( α − 1 ) β ≤ L ψ − 1 2(\alpha-1)\beta\leq L_{\psi}^{-1} 2 ( α − 1 ) β ≤ L ψ − 1 L ψ ≜ c M ϕ 2 s L + 1 2 m L a L + 1 = c M ϕ 2 κ 2 L_{\psi}\triangleq cM_{\phi}^{2}\frac{s_{L+1}^{2}m_L}{a_{L+1}}=cM_{\phi}^{2}\kappa^{2} L ψ ≜ c M ϕ 2 a L + 1 s L + 1 2 m L = c M ϕ 2 κ 2 α ≤ 1 + 1 2 c M ϕ 2 κ 2 β \alpha\leq1+\frac{1}{2cM_{\phi}^{2}\kappa^{2}\beta} α ≤ 1 + 2 c M ϕ 2 κ 2 β 1 2 e 2 ( α − 1 ) β R y 2 = O ( 1 ) 2e^{2(\alpha-1)\beta R_y^2}=O(1) 2 e 2 ( α − 1 ) β R y 2 = O ( 1 ) α ≤ 1 + 1 2 c M ϕ 2 β \alpha\leq1+\frac{1}{2cM_{\phi}^{2}\beta} α ≤ 1 + 2 c M ϕ 2 β 1 α = 1 + 1 2 c M ϕ 2 max ( κ 2 , 1 ) β \alpha=1+\frac{1}{2cM_{\phi}^{2}\max(\kappa^{2},1)\beta} α = 1 + 2 c M ϕ 2 m a x ( κ 2 , 1 ) β 1 ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α \left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}} ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1 O ( 1 ) O(1) O ( 1 ) κ ˉ ≜ max ( κ , 1 ) \bar{\kappa}\triangleq\max(\kappa,1) κ ˉ ≜ max ( κ , 1 )
α = 1 + 1 2 c M ϕ 2 κ ˉ 2 β , \alpha=1+\frac{1}{2cM_{\phi}^{2}\bar{\kappa}^{2}\beta}, α = 1 + 2 c M ϕ 2 κ ˉ 2 β 1 , and
( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 1 2 α ≤ ( 2 e 2 ( α − 1 ) β R y 2 ) 1 2 α ≤ 2 e β R y 2 2 c M ϕ 2 κ ˉ 2 β = 2 e R y 2 2 c M ϕ 2 κ ˉ 2 = O ( 1 ) , \left(\mathbb{E}_{p_0}\left[e^{2(\alpha-1)\beta\mathcal{L}(\theta_0)}\right]\right)^{\frac{1}{2\alpha}}
\leq\left(2e^{2(\alpha-1)\beta R_y^2}\right)^{\frac{1}{2\alpha}}
\leq\sqrt{2}e^{\frac{\beta R_y^2}{2cM_{\phi}^{2}\bar{\kappa}^{2}\beta}}
=\sqrt{2}e^{\frac{R_y^2}{2cM_{\phi}^{2}\bar{\kappa}^{2}}}=O(1), ( E p 0 [ e 2 ( α − 1 ) β L ( θ 0 ) ] ) 2 α 1 ≤ ( 2 e 2 ( α − 1 ) β R y 2 ) 2 α 1 ≤ 2 e 2 c M ϕ 2 κ ˉ 2 β β R y 2 = 2 e 2 c M ϕ 2 κ ˉ 2 R y 2 = O ( 1 ) , since κ ˉ ≥ 1 \bar{\kappa}\geq1 κ ˉ ≥ 1 η ≜ 2 e R y 2 2 c M ϕ 2 κ ˉ 2 = O ( 1 ) \eta\triangleq\sqrt{2}e^{\frac{R_y^2}{2cM_{\phi}^{2}\bar{\kappa}^{2}}}=O(1) η ≜ 2 e 2 c M ϕ 2 κ ˉ 2 R y 2 = O ( 1 )
Given this choice of α \alpha α
α α − 1 = 1 + 2 c M ϕ 2 κ ˉ 2 β = c β κ ˉ 2 β , \frac{\alpha}{\alpha-1}=1+2cM_{\phi}^{2}\bar{\kappa}^{2}\beta=c_{\beta}\bar{\kappa}^{2}\beta, α − 1 α = 1 + 2 c M ϕ 2 κ ˉ 2 β = c β κ ˉ 2 β , where we defined c β ≜ 2 c M ϕ 2 + 1 κ ˉ 2 β c_{\beta}\triangleq2cM_{\phi}^{2}+\frac{1}{\bar{\kappa}^{2}\beta} c β ≜ 2 c M ϕ 2 + κ ˉ 2 β 1 κ ˉ 2 β ≫ 1 \bar{\kappa}^{2}\beta\gg1 κ ˉ 2 β ≫ 1 c β = O ( 1 ) c_{\beta}=O(1) c β = O ( 1 ) non-negative measurable function φ : R N p → R \varphi:\mathbb{R}^{N_{p}}\to\mathbb{R} φ : R N p → R
E p t [ φ ( θ t ) ] ≤ η ∥ φ ( θ 0 ) ∥ L c β κ ˉ 2 β . (39) \mathbb{E}_{p_t}[\varphi(\theta_t)]\leq\eta\|\varphi(\theta_0)\|_{L^{c_{\beta}\bar{\kappa}^{2}\beta}}. \tag*{(39)} E p t [ φ ( θ t )] ≤ η ∥ φ ( θ 0 ) ∥ L c β κ ˉ 2 β . (39) We proceed to bound each of the four terms in Eq. (11) using the change-of-measure inequality Eq. (39) .
First term
Observe that
− n − 1 E [ r ( θ t ) T K N T K ( θ t ) r ( θ t ) ] ≤ − n − 1 E [ ∥ r ( θ t ) ∥ 2 λ min ( K N T K ( θ t ) ) ] = − n − 1 E [ 2 n L ( θ t ) λ min ( K N T K ( θ t ) ) ] = − 2 E [ L ( θ t ) λ min ( K N T K ( θ t ) ) ] . \begin{aligned}
-n^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\mathsf{T}}\mathbf{K}_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)\right]\leq-n^{-1}\mathbb{E}\left[\|\mathbf{r}(\theta_t)\|^{2}\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t))\right] \\
&= -n^{-1}\mathbb{E}\left[2n\mathcal{L}(\theta_t)\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t))\right] \\
&= -2\mathbb{E}\left[\mathcal{L}(\theta_t)\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t))\right].
\end{aligned} − n − 1 E [ r ( θ t ) T K NTK ( θ t ) r ( θ t ) ] ≤ − n − 1 E [ ∥ r ( θ t ) ∥ 2 λ m i n ( K NTK ( θ t )) ] = − n − 1 E [ 2 n L ( θ t ) λ m i n ( K NTK ( θ t )) ] = − 2 E [ L ( θ t ) λ m i n ( K NTK ( θ t )) ] . From Lemma 5.3 we have that if m = Ω ( n λ 1 ln n λ 1 ln n ) m=\Omega\left(\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\ln n\right) m = Ω ( λ 1 n ln λ 1 n ln n ) λ min ( K N T K ( θ 0 ) ) ≥ c 0 λ 1 s L + 1 2 m \lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0))\geq c_{0}\lambda_{1}s_{L+1}^{2}m λ m i n ( K NTK ( θ 0 )) ≥ c 0 λ 1 s L + 1 2 m 1 − exp ( − c ~ 1 m ( n λ 1 ln n λ 1 ) − 1 ) 1-\exp\left(-\tilde{c}_{1}m\left(\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)^{-1}\right) 1 − exp ( − c ~ 1 m ( λ 1 n ln λ 1 n ) − 1 )
λ ⋆ = c 0 λ 1 s L + 1 2 m \lambda_{\star}=c_{0}\lambda_{1}s_{L+1}^{2}m λ ⋆ = c 0 λ 1 s L + 1 2 m and
δ = exp ( − c ~ 1 m ( n λ 1 ln n λ 1 ) − 1 ) . \delta=\exp\left(-\tilde{c}_{1}m\left(\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)^{-1}\right). δ = exp ( − c ~ 1 m ( λ 1 n ln λ 1 n ) − 1 ) . Observe that
− n − 1 E [ r ( θ t ) T K N T K ( θ t ) r ( θ t ) ] ≤ − 2 E [ L ( θ t ) λ min ( K N T K ( θ t ) ) ] = − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) ( λ ⋆ − λ min ( K N T K ( θ t ) ) ) ] ≤ − 2 λ ⋆ E [ L ( θ t ) ] + 2 E [ L ( θ t ) max ( 0 , λ ⋆ − λ min ( K N T K ( θ t ) ) ) ] . \begin{aligned}
-n^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\mathsf{T}}\mathbf{K}_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)\right] \\
&\leq-2\mathbb{E}\left[\mathcal{L}(\theta_t)\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t))\right] \\
&= -2\lambda_{\star}\mathbb{E}[\mathcal{L}(\theta_t)]+2\mathbb{E}\left[\mathcal{L}(\theta_t)(\lambda_{\star}-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t)))\right] \\
&\leq-2\lambda_{\star}\mathbb{E}[\mathcal{L}(\theta_t)]+2\mathbb{E}\left[\mathcal{L}(\theta_t)\max(0,\lambda_{\star}-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t)))\right].
\end{aligned} − n − 1 E [ r ( θ t ) T K NTK ( θ t ) r ( θ t ) ] ≤ − 2 E [ L ( θ t ) λ m i n ( K NTK ( θ t )) ] = − 2 λ ⋆ E [ L ( θ t )] + 2 E [ L ( θ t ) ( λ ⋆ − λ m i n ( K NTK ( θ t ))) ] ≤ − 2 λ ⋆ E [ L ( θ t )] + 2 E [ L ( θ t ) max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ t ))) ] . We upper bound E [ L ( θ t ) max ( 0 , ( λ ⋆ − λ min ( K N T K ( θ t ) ) ) ) ] \mathbb{E}[\mathcal{L}(\theta_t)\max(0,(\lambda_{\star}-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t))))] E [ L ( θ t ) max ( 0 , ( λ ⋆ − λ m i n ( K NTK ( θ t ))))] (39) ) and Claim E.7:
E [ L ( θ t ) max ( 0 , ( λ ⋆ − λ min ( K N T K ( θ t ) ) ) ) ] ≤ η ∥ L ( θ 0 ) max ( 0 , ( λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) ) ∥ L c β κ ˉ 2 β ≤ η ∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ∥ max ( 0 , λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) ∥ L 2 c β κ ˉ 2 β . \begin{aligned}
\mathbb{E}\left[\mathcal{L}(\theta_t)\max(0,(\lambda_{\star}-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_t))))\right] \\
&\leq\eta\|\mathcal{L}(\theta_0)\max(0,(\lambda_{\star}-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0))))\|_{L^{c_{\beta}\bar{\kappa}^{2}\beta}} \\
&\leq\eta\|\mathcal{L}(\theta_0)\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}}\|\max(0,\lambda_{\star}-\lambda_{\min}(\mathbf{K}_{\mathrm{NTK}}(\theta_0)))\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}}.
\end{aligned} E [ L ( θ t ) max ( 0 , ( λ ⋆ − λ m i n ( K NTK ( θ t )))) ] ≤ η ∥ L ( θ 0 ) max ( 0 , ( λ ⋆ − λ m i n ( K NTK ( θ 0 )))) ∥ L c β κ ˉ 2 β ≤ η ∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ∥ max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ 0 ))) ∥ L 2 c β κ ˉ 2 β . From Proposition E.8,
∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ≤ ∥ R y 2 + 1 n ∑ i = 1 n f i 2 ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ≤ ∥ R y 2 ∥ L 2 c β κ ˉ 2 β + ∥ 1 n ∑ i = 1 n f i 2 ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ( i ) ≤ R y 2 + C ⋅ 2 c β κ ˉ − 2 β ∥ 1 n ∑ i = 1 n f i 2 ( θ 0 ) ∥ ψ 1 ( i i ) ≤ R y 2 + C ⋅ 2 c β κ ˉ − 4 β , \begin{aligned}
\|\mathcal{L}(\theta_0)\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}}
\leq
\left\|R_y^2+\frac{1}{n}\sum_{i=1}^{n}f_i^2(\theta_0)\right\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}} \\
&\leq\|R_y^2\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}}+
\left\|\frac{1}{n}\sum_{i=1}^{n}f_i^2(\theta_0)\right\|_{L^{2c_{\beta}\bar{\kappa}^{2}\beta}}^{(i)} \\
&\leq R_y^2+C\cdot2c_{\beta}\bar{\kappa}^{-2}\beta
\left\|\frac{1}{n}\sum_{i=1}^{n}f_i^2(\theta_0)\right\|_{\psi_1}^{(ii)} \\
&\leq R_y^2+C\cdot2c_{\beta}\bar{\kappa}^{-4}\beta,
\end{aligned} ∥ L ( θ 0 ) ∥ L 2 c β κ ˉ 2 β ≤ R y 2 + n 1 i = 1 ∑ n f i 2 ( θ 0 ) L 2 c β κ ˉ 2 β ≤ ∥ R y 2 ∥ L 2 c β κ ˉ 2 β + n 1 i = 1 ∑ n f i 2 ( θ 0 ) L 2 c β κ ˉ 2 β ( i ) ≤ R y 2 + C ⋅ 2 c β κ ˉ − 2 β n 1 i = 1 ∑ n f i 2 ( θ 0 ) ψ 1 ( ii ) ≤ R y 2 + C ⋅ 2 c β κ ˉ − 4 β , where in ( i ) (i) ( i ) ( i i ) (ii) ( ii ) [111 Proposition 2.8.1] . For the other term, by splitting according to the high-probability NTK event, we get
∥ max ( 0 , λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) ∥ L 2 c β κ ˉ 2 β = E [ ( max ( 0 , λ ⋆ − λ min ( K N T K ( θ 0 ) ) ) ) 2 c β κ ˉ 2 β ] 1 / ( 2 c β κ ˉ 2 β ) ≤ ( 0 ⋅ P ( λ ⋆ ≤ λ min ( K N T K ( θ 0 ) ) ) + λ ⋆ 2 c β κ ˉ 2 β P ( λ ⋆ > λ min ( K N T K ( θ 0 ) ) ) ) 1 / ( 2 c β κ ˉ 2 β ) ≤ λ ⋆ δ 1 / ( 2 c β κ ˉ 2 β ) = c 0 λ 1 s L + 1 2 m exp ( − c ~ 1 m ( 2 c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) . \begin{aligned}
\left\|\max\left(0,\lambda_\star-\lambda_{\min}\left(K_{\mathrm{NTK}}(\theta_0)\right)\right)\right\|_{L^{2c_\beta\bar{\kappa}^{2}\beta}}
= \mathbb{E}\left[\left(\max\left(0,\lambda_\star-\lambda_{\min}\left(K_{\mathrm{NTK}}(\theta_0)\right)\right)\right)^{2c_\beta\bar{\kappa}^{2}\beta}\right]^{1/(2c_\beta\bar{\kappa}^{2}\beta)} \\
&\leq\left(0\cdot\mathbb{P}\left(\lambda_\star\leq\lambda_{\min}\left(K_{\mathrm{NTK}}(\theta_0)\right)\right)+\lambda_\star^{2c_\beta\bar{\kappa}^{2}\beta}\mathbb{P}\left(\lambda_\star>\lambda_{\min}\left(K_{\mathrm{NTK}}(\theta_0)\right)\right)\right)^{1/(2c_\beta\bar{\kappa}^{2}\beta)} \\
&\leq\lambda_\star\delta^{1/(2c_\beta\bar{\kappa}^{2}\beta)}
= c_0\lambda_1s_{L+1}^{2}m\exp\left(-\widetilde{c}_1m\left(2c_\beta\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right).
\end{aligned} ∥ max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ 0 ) ) ) ∥ L 2 c β κ ˉ 2 β = E [ ( max ( 0 , λ ⋆ − λ m i n ( K NTK ( θ 0 ) ) ) ) 2 c β κ ˉ 2 β ] 1/ ( 2 c β κ ˉ 2 β ) ≤ ( 0 ⋅ P ( λ ⋆ ≤ λ m i n ( K NTK ( θ 0 ) ) ) + λ ⋆ 2 c β κ ˉ 2 β P ( λ ⋆ > λ m i n ( K NTK ( θ 0 ) ) ) ) 1/ ( 2 c β κ ˉ 2 β ) ≤ λ ⋆ δ 1/ ( 2 c β κ ˉ 2 β ) = c 0 λ 1 s L + 1 2 m exp ( − c 1 m ( 2 c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) . where in the first transition we used the fact that the empirical NTK matrix is PSD, hence λ min ( K N T K ( θ 0 ) ) ≥ 0 \lambda_{\min}\left(K_{\mathrm{NTK}}(\theta_0)\right)\geq0 λ m i n ( K NTK ( θ 0 ) ) ≥ 0
− n − 1 E [ r ( θ t ) ⊤ K N T K ( θ t ) r ( θ t ) ] ≤ − 2 c 0 λ 1 s L + 1 2 m E [ L ( θ t ) ] + 2 η c 0 λ 1 s L + 1 2 m ( R y 2 + C ⋅ 2 c β κ ˉ 4 β ) exp ( − c ~ 1 m ( 2 c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) -n^{-1}\mathbb{E}\left[\boldsymbol{r}(\theta_t)^{\top}K_{\mathrm{NTK}}(\theta_t)\boldsymbol{r}(\theta_t)\right]
\leq-2c_0\lambda_1s_{L+1}^{2}m\mathbb{E}\left[\mathcal{L}(\theta_t)\right]
+2\eta c_0\lambda_1s_{L+1}^{2}m\left(R_y^2+C\cdot2c_\beta\bar{\kappa}^{4}\beta\right)\exp\left(-\widetilde{c}_1m\left(2c_\beta\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right) − n − 1 E [ r ( θ t ) ⊤ K NTK ( θ t ) r ( θ t ) ] ≤ − 2 c 0 λ 1 s L + 1 2 m E [ L ( θ t ) ] + 2 η c 0 λ 1 s L + 1 2 m ( R y 2 + C ⋅ 2 c β κ ˉ 4 β ) exp ( − c 1 m ( 2 c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) ≲ − λ 1 s L + 1 2 m E [ L ( θ t ) ] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( 2 c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) . \lesssim-\lambda_1s_{L+1}^{2}m\mathbb{E}\left[\mathcal{L}(\theta_t)\right]
+\eta\lambda_1s_{L+1}^{2}m\left(1+c_\beta\bar{\kappa}^{4}\beta\right)\exp\left(-\widetilde{c}_1m\left(2c_\beta\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right). ≲ − λ 1 s L + 1 2 m E [ L ( θ t ) ] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c 1 m ( 2 c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) . Second term
Observe that
β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] = β − 1 E [ ∑ j = 1 N p ∑ i = 1 n ( ∇ f i ( θ t ) j ) 2 ] = β − 1 ∑ i = 1 n E [ ∥ ∇ f i ( θ t ) ∥ 2 ] . \beta^{-1}\mathbb{E}\left[\operatorname{Tr}\left(\nabla f(\theta_t)^{\top}\nabla f(\theta_t)\right)\right]
=\beta^{-1}\mathbb{E}\left[\sum_{j=1}^{N_p}\sum_{i=1}^{n}\left(\nabla f_i(\theta_t)_j\right)^2\right]
=\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\nabla f_i(\theta_t)\right\|^2\right]. β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] = β − 1 E j = 1 ∑ N p i = 1 ∑ n ( ∇ f i ( θ t ) j ) 2 = β − 1 i = 1 ∑ n E [ ∥ ∇ f i ( θ t ) ∥ 2 ] . Let i ∈ [ n ] i\in[n] i ∈ [ n ] ∇ f i ( θ ) \nabla f_i(\theta) ∇ f i ( θ )
∥ ∇ f i ( θ ) ∥ 2 = ∥ ∂ f i ( θ ) ∂ v ∥ 2 + ∑ l = 1 L ∥ ∂ f i ( θ ) ∂ W ( l ) ∥ F 2 \left\|\nabla f_i(\theta)\right\|^2
=\left\|\frac{\partial f_i(\theta)}{\partial v}\right\|^2
+\sum_{l=1}^{L}\left\|\frac{\partial f_i(\theta)}{\partial W^{(l)}}\right\|_{\mathrm{F}}^2 ∥ ∇ f i ( θ ) ∥ 2 = ∂ v ∂ f i ( θ ) 2 + l = 1 ∑ L ∂ W ( l ) ∂ f i ( θ ) F 2 ≤ ( i ) s L + 1 2 ∥ α ( L ) ( x i ) ∥ 2 + ∑ l = 1 L s l 2 ∥ δ ( l ) ( x i ) ( α ( l − 1 ) ( x i ) ) ⊤ ∥ F 2 \overset{(i)}{\leq}s_{L+1}^{2}\left\|\alpha^{(L)}(x_i)\right\|^2
+\sum_{l=1}^{L}s_l^2\left\|\delta^{(l)}(x_i)\left(\alpha^{(l-1)}(x_i)\right)^{\top}\right\|_{\mathrm{F}}^2 ≤ ( i ) s L + 1 2 α ( L ) ( x i ) 2 + l = 1 ∑ L s l 2 δ ( l ) ( x i ) ( α ( l − 1 ) ( x i ) ) ⊤ F 2 = ( i i ) s L + 1 2 ∥ α ( L ) ( x i ) ∥ 2 + ∑ l = 1 L s l 2 ∥ δ ( l ) ( x i ) ∥ 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 , \overset{(ii)}{=}s_{L+1}^{2}\left\|\alpha^{(L)}(x_i)\right\|^2
+\sum_{l=1}^{L}s_l^2\left\|\delta^{(l)}(x_i)\right\|^2\left\|\alpha^{(l-1)}(x_i)\right\|^2, = ( ii ) s L + 1 2 α ( L ) ( x i ) 2 + l = 1 ∑ L s l 2 δ ( l ) ( x i ) 2 α ( l − 1 ) ( x i ) 2 , where in ( i ) (i) ( i ) (37) , and in ( i i ) (ii) ( ii ) ∥ u v ⊤ ∥ F = ∥ u ∥ ∥ v ∥ \left\|\boldsymbol{u}\boldsymbol{v}^{\top}\right\|_{\mathrm{F}}=\left\|\boldsymbol{u}\right\|\left\|\boldsymbol{v}\right\| u v ⊤ F = ∥ u ∥ ∥ v ∥
By Proposition H.3, there exists a constant c 4 > 0 c_4>0 c 4 > 0 p ≥ 1 p\geq1 p ≥ 1
∥ ∥ α ( L ) ( x i ) ∥ 2 ∥ L p ≤ ( c 4 α 0 2 ν 0 2 ) L ∥ x i ∥ 2 ∏ k = 1 L s k 2 a k ( m k + p ) . \left\|\left\|\alpha^{(L)}(x_i)\right\|^2\right\|_{L_p}
\leq\left(c_4\alpha_0^2\nu_0^2\right)^L\left\|x_i\right\|^2\prod_{k=1}^{L}\frac{s_k^2}{a_k}(m_k+p). α ( L ) ( x i ) 2 L p ≤ ( c 4 α 0 2 ν 0 2 ) L ∥ x i ∥ 2 k = 1 ∏ L a k s k 2 ( m k + p ) . By Proposition H.4, for every l ∈ [ L ] l\in[L] l ∈ [ L ] c 5 > 0 c_5>0 c 5 > 0
∥ ∥ δ ( l ) ( x i ) ∥ 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 ∥ L p ≤ ( c 5 α 0 2 ν 0 2 ) 2 L ∥ x i ∥ 2 ( ∏ k = 1 l − 1 s k 2 a k ( m k + p ) ) ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( m L + p ) . \left\|\left\|\delta^{(l)}(x_i)\right\|^2\left\|\alpha^{(l-1)}(x_i)\right\|^2\right\|_{L_p}
\leq\left(c_5\alpha_0^2\nu_0^2\right)^{2L}\left\|x_i\right\|^2
\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}(m_k+p)\right)
\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)\right)
\frac{s_{L+1}^2}{a_{L+1}}(m_L+p). δ ( l ) ( x i ) 2 α ( l − 1 ) ( x i ) 2 L p ≤ ( c 5 α 0 2 ν 0 2 ) 2 L ∥ x i ∥ 2 ( k = 1 ∏ l − 1 a k s k 2 ( m k + p ) ) ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) a L + 1 s L + 1 2 ( m L + p ) . Thus, using Minkowski’s inequality we get that at initialization,
∥ ∥ ∇ f i ( θ ) ∥ 2 ∥ L p ≤ s L + 1 2 ∥ ∥ α ( L ) ( x i ) ∥ 2 ∥ L p + ∑ l = 1 L s l 2 ∥ ∥ δ ( l ) ( x i ) ∥ 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 ∥ L p ≤ s L + 1 2 ( c 4 α ϕ 2 ν ϕ 2 ) L ∥ x i ∥ 2 ∏ k = 1 L s k 2 a k ( m k + p ) + ∑ l = 1 L s l 2 ( c 5 α ϕ 2 ν ϕ 2 ) 2 L ∥ x i ∥ 2 ( ∏ k = 1 l − 1 s k 2 a k ( m k + p ) ) ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) s L + 1 2 a L + 1 ( m L + p ) . \begin{aligned}
\left\| \left\| \nabla f_i(\theta) \right\|^2 \right\|_{L_p}
\le s_{L+1}^2 \left\| \left\| \alpha^{(L)}(x_i) \right\|^2 \right\|_{L_p}
+ \sum_{l=1}^{L} s_l^2 \left\| \left\| \delta^{(l)}(x_i) \right\|^2 \left\| \alpha^{(l-1)}(x_i) \right\|^2 \right\|_{L_p} \\
&\le s_{L+1}^2 \left(c_4\alpha_\phi^2\nu_\phi^2\right)^{L} \|x_i\|^2 \prod_{k=1}^{L} \frac{s_k^2}{a_k}(m_k+p)
+ \sum_{l=1}^{L} s_l^2 \left(c_5\alpha_\phi^2\nu_\phi^2\right)^{2L}\|x_i\|^2
\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}(m_k+p)\right)
\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)\right)
\frac{s_{L+1}^2}{a_{L+1}}(m_L+p).
\end{aligned} ∥ ∇ f i ( θ ) ∥ 2 L p ≤ s L + 1 2 α ( L ) ( x i ) 2 L p + l = 1 ∑ L s l 2 δ ( l ) ( x i ) 2 α ( l − 1 ) ( x i ) 2 L p ≤ s L + 1 2 ( c 4 α ϕ 2 ν ϕ 2 ) L ∥ x i ∥ 2 k = 1 ∏ L a k s k 2 ( m k + p ) + l = 1 ∑ L s l 2 ( c 5 α ϕ 2 ν ϕ 2 ) 2 L ∥ x i ∥ 2 ( k = 1 ∏ l − 1 a k s k 2 ( m k + p ) ) ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) a L + 1 s L + 1 2 ( m L + p ) . Denoting C = ( c 4 + c 5 ) α ϕ 2 ν ϕ 2 C=(c_4+c_5)\alpha_\phi^2\nu_\phi^2 C = ( c 4 + c 5 ) α ϕ 2 ν ϕ 2 ( m L + p ) (m_L+p) ( m L + p )
∥ ∥ ∇ f i ( θ 0 ) ∥ 2 ∥ L p ≤ c ϕ , σ 0 C 2 L d L + 1 s L + 1 2 ( m L + p ) ( s L 2 a L ∏ k = 1 L − 1 s k 2 a k ( m k + p ) + 1 a L + 1 ∑ l = 1 L ( ∏ k = 1 l − 1 s k 2 a k ( m k + p ) ) s l 2 ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + p ) ) ) . \left\| \left\| \nabla f_i(\theta_0) \right\|^2 \right\|_{L_p}
\le c_{\phi,\sigma_0} C^{2L}d_{L+1}s_{L+1}^2(m_L+p)
\left(
\frac{s_L^2}{a_L}\prod_{k=1}^{L-1}\frac{s_k^2}{a_k}(m_k+p)
+\frac{1}{a_{L+1}}\sum_{l=1}^{L}
\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}(m_k+p)\right)s_l^2
\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}(m_{k-1}+m_k+p)\right)
\right). ∥ ∇ f i ( θ 0 ) ∥ 2 L p ≤ c ϕ , σ 0 C 2 L d L + 1 s L + 1 2 ( m L + p ) ( a L s L 2 k = 1 ∏ L − 1 a k s k 2 ( m k + p ) + a L + 1 1 l = 1 ∑ L ( k = 1 ∏ l − 1 a k s k 2 ( m k + p ) ) s l 2 ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + p ) ) ) . Using this together with Eq. (39) , we eventually get
⋯ ≤ β − 1 ∑ i = 1 n E [ ∥ ∇ f i ( θ t ) ∥ 2 ] ≤ β − 1 ∑ i = 1 n η ∥ ∥ ∇ f i ( θ 0 ) ∥ 2 ∥ L c β κ ˉ 2 β ≤ c ϕ , σ 0 ( ( c 4 + c 5 ) α ϕ 2 ν ϕ 2 ) 2 L η β − 1 n d s L + 1 2 ( m L + c β κ ˉ 2 β ) ( s L 2 a L ∏ k = 1 L − 1 s k 2 a k ( m k + c β κ ˉ 2 β ) + 1 a L + 1 ∑ l = 1 L ( ∏ k = 1 l − 1 s k 2 a k ( m k + c β κ ˉ 2 β ) ) s l 2 ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + c β κ ˉ 2 β ) ) ) . \begin{aligned}
\cdots\le\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\left\|\nabla f_i(\theta_t)\right\|^2\right] \\
&\le\beta^{-1}\sum_{i=1}^{n}\eta\left\|\left\|\nabla f_i(\theta_0)\right\|^2\right\|_{L^{c_\beta\bar{\kappa}^2\beta}} \\
&\le c_{\phi,\sigma_0}\left((c_4+c_5)\alpha_\phi^2\nu_\phi^2\right)^{2L}\eta\beta^{-1}nds_{L+1}^2
\left(m_L+c_\beta\bar{\kappa}^2\beta\right)
\left(
\frac{s_L^2}{a_L}\prod_{k=1}^{L-1}\frac{s_k^2}{a_k}\left(m_k+c_\beta\bar{\kappa}^2\beta\right)
+\frac{1}{a_{L+1}}\sum_{l=1}^{L}
\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}\left(m_k+c_\beta\bar{\kappa}^2\beta\right)\right)s_l^2
\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}\left(m_{k-1}+m_k+c_\beta\bar{\kappa}^2\beta\right)\right)
\right).
\end{aligned} ⋯ ≤ β − 1 i = 1 ∑ n E [ ∥ ∇ f i ( θ t ) ∥ 2 ] ≤ β − 1 i = 1 ∑ n η ∥ ∇ f i ( θ 0 ) ∥ 2 L c β κ ˉ 2 β ≤ c ϕ , σ 0 ( ( c 4 + c 5 ) α ϕ 2 ν ϕ 2 ) 2 L η β − 1 n d s L + 1 2 ( m L + c β κ ˉ 2 β ) ( a L s L 2 k = 1 ∏ L − 1 a k s k 2 ( m k + c β κ ˉ 2 β ) + a L + 1 1 l = 1 ∑ L ( k = 1 ∏ l − 1 a k s k 2 ( m k + c β κ ˉ 2 β ) ) s l 2 ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + c β κ ˉ 2 β ) ) ) . Third term
A remark on the second derivative of the activations. In this section, when we work toward bounding the L p L^p L p ϕ ′ ′ ( α ~ ( l ) ( x ) ) \phi''\left(\tilde{\alpha}^{(l)}(x)\right) ϕ ′′ ( α ~ ( l ) ( x ) ) β ϕ \beta_\phi β ϕ ϕ ′ ′ \phi'' ϕ ′′ ∣ ϕ ′ ′ ∣ ≤ β ϕ |\phi''|\le\beta_\phi ∣ ϕ ′′ ∣ ≤ β ϕ L p L^p L p ϕ ′ ′ \phi'' ϕ ′′ 0 0 0 ∣ ϕ ′ ′ ∣ ≤ β ϕ |\phi''|\le\beta_\phi ∣ ϕ ′′ ∣ ≤ β ϕ ∣ ϕ ′ ′ ∣ ≤ β ϕ |\phi''|\le\beta_\phi ∣ ϕ ′′ ∣ ≤ β ϕ
Observe that
β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] = β − 1 ∑ i = 1 n E [ r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) Δ f i ( θ t ) ∣ ] ≤ β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ β − 1 ∑ i = 1 n 2 E [ ℓ i ( θ t ) ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] , \begin{aligned}
\beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n} r_i(\theta_t)\Delta f_i(\theta_t)\right]
&= \beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)\Delta f_i(\theta_t)\right] \\
&\le\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\left|r_i(\theta_t)\Delta f_i(\theta_t)\right|\right] \\
&\le\beta^{-1}\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[\left|r_i(\theta_t)\right|^2\right]\mathbb{E}\left[\left|\Delta f_i(\theta_t)\right|^2\right]} \\
&\le\beta^{-1}\sum_{i=1}^{n}\sqrt{2\mathbb{E}\left[\ell_i(\theta_t)\right]\mathbb{E}\left[\left|\Delta f_i(\theta_t)\right|^2\right]},
\end{aligned} β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] = β − 1 i = 1 ∑ n E [ r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 i = 1 ∑ n E [ ∣ r i ( θ t ) Δ f i ( θ t ) ∣ ] ≤ β − 1 i = 1 ∑ n E [ ∣ r i ( θ t ) ∣ 2 ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ β − 1 i = 1 ∑ n 2 E [ ℓ i ( θ t ) ] E [ ∣ Δ f i ( θ t ) ∣ 2 ] , where ℓ i ( θ ) = 1 2 ∣ r i ( θ ) ∣ 2 \ell_i(\theta)=\frac{1}{2}\left|r_i(\theta)\right|^2 ℓ i ( θ ) = 2 1 ∣ r i ( θ ) ∣ 2 Q > 0 Q>0 Q > 0
E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ η Q 2 , \mathbb{E}\left[\left|\Delta f_i(\theta_t)\right|^2\right]\le\eta Q^2, E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ η Q 2 , then we get
β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 η Q ∑ i = 1 n 2 E [ ℓ i ( θ t ) ] = β − 1 n η Q 1 n ∑ i = 1 n 2 E [ ℓ i ( θ t ) ] ≤ β − 1 n η Q 2 n ∑ i = 1 n E [ ℓ i ( θ t ) ] = β − 1 n η ⋅ 2 Q E [ L ( θ t ) ] . \begin{aligned}
\beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]
&\le\beta^{-1}\sqrt{\eta}Q\sum_{i=1}^{n}\sqrt{2\mathbb{E}\left[\ell_i(\theta_t)\right]} \\
&= \beta^{-1}n\sqrt{\eta}Q\frac{1}{n}\sum_{i=1}^{n}\sqrt{2\mathbb{E}\left[\ell_i(\theta_t)\right]} \\
&\le\beta^{-1}n\sqrt{\eta}Q\sqrt{\frac{2}{n}\sum_{i=1}^{n}\mathbb{E}\left[\ell_i(\theta_t)\right]}=\beta^{-1}n\sqrt{\eta}\cdot2Q\sqrt{\mathbb{E}\left[\mathcal{L}(\theta_t)\right]}.
\end{aligned} β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] ≤ β − 1 η Q i = 1 ∑ n 2 E [ ℓ i ( θ t ) ] = β − 1 n η Q n 1 i = 1 ∑ n 2 E [ ℓ i ( θ t ) ] ≤ β − 1 n η Q n 2 i = 1 ∑ n E [ ℓ i ( θ t ) ] = β − 1 n η ⋅ 2 Q E [ L ( θ t ) ] . We move on to obtain such Q Q Q ℓ ∈ [ L ] \ell\in[L] ℓ ∈ [ L ]
∂ f i ( θ ) ∂ W ( ℓ ) = s ℓ δ ( ℓ ) ( x i ) ( α ( ℓ − 1 ) ( x i ) ) T . \frac{\partial f_i(\theta)}{\partial W^{(\ell)}}=s_\ell\delta^{(\ell)}(x_i)\left(\alpha^{(\ell-1)}(x_i)\right)^{\mathsf{T}}. ∂ W ( ℓ ) ∂ f i ( θ ) = s ℓ δ ( ℓ ) ( x i ) ( α ( ℓ − 1 ) ( x i ) ) T . Therefore, for any j ∈ [ m ℓ ] j\in[m_\ell] j ∈ [ m ℓ ] k ∈ [ m ℓ − 1 ] k\in[m_{\ell-1}] k ∈ [ m ℓ − 1 ]
∂ f i ( θ ) ∂ W j k ( ℓ ) = s ℓ α k ( ℓ − 1 ) ( x i ) δ j ( ℓ ) ( x i ) . \frac{\partial f_i(\theta)}{\partial W_{jk}^{(\ell)}}=s_\ell\alpha_k^{(\ell-1)}(x_i)\delta_j^{(\ell)}(x_i). ∂ W j k ( ℓ ) ∂ f i ( θ ) = s ℓ α k ( ℓ − 1 ) ( x i ) δ j ( ℓ ) ( x i ) . Observing that ( α ( ℓ − 1 ) ( x i ) ) T \left(\alpha^{(\ell-1)}(x_i)\right)^{\mathsf{T}} ( α ( ℓ − 1 ) ( x i ) ) T W ( ℓ ) W^{(\ell)} W ( ℓ )
∂ 2 f i ( θ ) ∂ ( W j k ( ℓ ) ) 2 = s ℓ α k ( ℓ − 1 ) ( x i ) ∂ ∂ W j k ( ℓ ) δ j ( ℓ ) ( x i ) = s ℓ α k ( ℓ − 1 ) ( x i ) ∂ δ j ( ℓ ) ( x i ) ∂ α j ( ℓ ) ( x i ) ∂ α ~ j ( ℓ ) ( x i ) ∂ W j k ( ℓ ) = s ℓ 2 ( α k ( ℓ − 1 ) ( x i ) ) 2 ∂ δ j ( ℓ ) ( x i ) ∂ α ~ j ( ℓ ) ( x i ) . \begin{aligned}
\frac{\partial^2 f_i(\theta)}{\partial\left(W_{jk}^{(\ell)}\right)^2}
&=s_\ell\alpha_k^{(\ell-1)}(x_i)\frac{\partial}{\partial W_{jk}^{(\ell)}}\delta_j^{(\ell)}(x_i) \\
&=s_\ell\alpha_k^{(\ell-1)}(x_i)\frac{\partial\delta_j^{(\ell)}(x_i)}{\partial\alpha_j^{(\ell)}(x_i)}\frac{\partial\widetilde{\alpha}_j^{(\ell)}(x_i)}{\partial W_{jk}^{(\ell)}} \\
&=s_\ell^2\left(\alpha_k^{(\ell-1)}(x_i)\right)^2\frac{\partial\delta_j^{(\ell)}(x_i)}{\partial\widetilde{\alpha}_j^{(\ell)}(x_i)}.
\end{aligned} ∂ ( W j k ( ℓ ) ) 2 ∂ 2 f i ( θ ) = s ℓ α k ( ℓ − 1 ) ( x i ) ∂ W j k ( ℓ ) ∂ δ j ( ℓ ) ( x i ) = s ℓ α k ( ℓ − 1 ) ( x i ) ∂ α j ( ℓ ) ( x i ) ∂ δ j ( ℓ ) ( x i ) ∂ W j k ( ℓ ) ∂ α j ( ℓ ) ( x i ) = s ℓ 2 ( α k ( ℓ − 1 ) ( x i ) ) 2 ∂ α j ( ℓ ) ( x i ) ∂ δ j ( ℓ ) ( x i ) . Therefore, the contribution of layer l l l
∣ Δ W ( l ) f i ( θ ) ∣ = ∣ ∑ j = 1 m l ∑ k = 1 m l − 1 ∂ 2 f i ( θ ) ∂ ( W j k ( l ) ) 2 ∣ = s l 2 ∑ j = 1 m l ∂ δ j ( l ) ( x i ) ∂ α ~ j ( l ) ( x i ) ∑ k = 1 m l − 1 ( α k ( l − 1 ) ( x i ) ) 2 = s l 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 ∣ Tr ( ∂ δ ( l ) ( x i ) ∂ α ~ ( l ) ( x i ) ) ∣ = s l 2 ∥ α ( l − 1 ) ( x i ) ∥ 2 ∣ Tr ( ∂ 2 f i ( θ ) ∂ ( α ~ ( l ) ( x i ) ) 2 ) ∣ . \left|\Delta_{W^{(l)}} f_i(\theta)\right| =
\left|\sum_{j=1}^{m_l}\sum_{k=1}^{m_{l-1}}
\frac{\partial^2 f_i(\theta)}{\partial\left(W_{jk}^{(l)}\right)^2}\right|
=
s_l^2\sum_{j=1}^{m_l}
\frac{\partial\delta_j^{(l)}(x_i)}{\partial\widetilde{\alpha}_j^{(l)}(x_i)}
\sum_{k=1}^{m_{l-1}}\left(\alpha_k^{(l-1)}(x_i)\right)^2
=
s_l^2\left\|\alpha^{(l-1)}(x_i)\right\|^2
\left|\operatorname{Tr}\left(\frac{\partial\delta^{(l)}(x_i)}{\partial\widetilde{\alpha}^{(l)}(x_i)}\right)\right|
=
s_l^2\left\|\alpha^{(l-1)}(x_i)\right\|^2
\left|\operatorname{Tr}\left(\frac{\partial^2 f_i(\theta)}{\partial\left(\widetilde{\alpha}^{(l)}(x_i)\right)^2}\right)\right|. ∣ Δ W ( l ) f i ( θ ) ∣ = j = 1 ∑ m l k = 1 ∑ m l − 1 ∂ ( W j k ( l ) ) 2 ∂ 2 f i ( θ ) = s l 2 j = 1 ∑ m l ∂ α j ( l ) ( x i ) ∂ δ j ( l ) ( x i ) k = 1 ∑ m l − 1 ( α k ( l − 1 ) ( x i ) ) 2 = s l 2 α ( l − 1 ) ( x i ) 2 Tr ( ∂ α ( l ) ( x i ) ∂ δ ( l ) ( x i ) ) = s l 2 α ( l − 1 ) ( x i ) 2 Tr ( ∂ ( α ( l ) ( x i ) ) 2 ∂ 2 f i ( θ ) ) . Let p ≥ 1 p \ge1 p ≥ 1
∥ Δ W ( l ) f i ( θ ) ∥ L p ≤ s l 2 ∥ ∥ α ( l − 1 ) ( x i ) ∥ 2 ∣ Tr ( ∂ 2 f i ( θ ) ∂ ( α ~ ( l ) ( x i ) ) 2 ) ∣ ∥ L p ≤ s l 2 ∥ ∥ α ( l − 1 ) ( x i ) ∥ 2 ∥ L 2 p ∥ ∣ Tr ( ∂ 2 f i ( θ ) ∂ ( α ~ ( l ) ( x i ) ) 2 ) ∣ ∥ L 2 p . \left\|\Delta_{W^{(l)}} f_i(\theta)\right\|_{L_p}
\le s_l^2
\left\|
\left\|\alpha^{(l-1)}(x_i)\right\|^2
\left|\operatorname{Tr}\left(\frac{\partial^2 f_i(\theta)}{\partial\left(\widetilde{\alpha}^{(l)}(x_i)\right)^2}\right)\right|
\right\|_{L_p}
\le s_l^2
\left\|\left\|\alpha^{(l-1)}(x_i)\right\|^2\right\|_{L_{2p}}
\left\|
\left|\operatorname{Tr}\left(\frac{\partial^2 f_i(\theta)}{\partial\left(\widetilde{\alpha}^{(l)}(x_i)\right)^2}\right)\right|
\right\|_{L_{2p}}. ∥ Δ W ( l ) f i ( θ ) ∥ L p ≤ s l 2 α ( l − 1 ) ( x i ) 2 Tr ( ∂ ( α ( l ) ( x i ) ) 2 ∂ 2 f i ( θ ) ) L p ≤ s l 2 α ( l − 1 ) ( x i ) 2 L 2 p Tr ( ∂ ( α ( l ) ( x i ) ) 2 ∂ 2 f i ( θ ) ) L 2 p . Summing over all weight matrices,
∥ Δ f i ( θ 0 ) ∥ L p ≤ ∑ l = 1 L s l 2 ∥ ∥ α ( l − 1 ) ( x i ) ∥ 2 ∥ L 2 p ∥ ∣ Tr ( ∂ 2 f i ( θ ) ∂ ( α ~ ( l ) ( x i ) ) 2 ) ∣ ∥ L 2 p . \left\|\Delta f_i(\theta_0)\right\|_{L_p}
\le\sum_{l=1}^{L}s_l^2
\left\|\left\|\alpha^{(l-1)}(x_i)\right\|^2\right\|_{L_{2p}}
\left\|
\left|\operatorname{Tr}\left(\frac{\partial^2 f_i(\theta)}{\partial\left(\widetilde{\alpha}^{(l)}(x_i)\right)^2}\right)\right|
\right\|_{L_{2p}}. ∥ Δ f i ( θ 0 ) ∥ L p ≤ l = 1 ∑ L s l 2 α ( l − 1 ) ( x i ) 2 L 2 p Tr ( ∂ ( α ( l ) ( x i ) ) 2 ∂ 2 f i ( θ ) ) L 2 p . By Eq. (39) and Claim E.6,
E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ η ∥ Δ f i ( θ 0 ) ∥ L c β κ ˉ 2 β 2 ≤ η ∥ Δ f i ( θ 0 ) ∥ L 2 c β κ ˉ 2 β 2 ≤ η ( ∑ l = 1 L s l 2 ∥ ∥ α ( l − 1 ) ( x i ) ∥ 2 ∥ L 4 c β κ ˉ 2 β ∥ ∣ Tr ( ∂ 2 f i ( θ ) ∂ ( α ~ ( l ) ( x i ) ) 2 ) ∣ ∥ L 4 c β κ ˉ 2 β ) 2 . \mathbb{E}\left[\left|\Delta f_i(\theta_t)\right|^2\right]
\le\eta\left\|\Delta f_i(\theta_0)\right\|^2_{L^{c_\beta\bar{\kappa}^2\beta}}
\le\eta\left\|\Delta f_i(\theta_0)\right\|^2_{L^{2c_\beta\bar{\kappa}^2\beta}}
\le\eta\left(
\sum_{l=1}^{L}s_l^2
\left\|\left\|\alpha^{(l-1)}(x_i)\right\|^2\right\|_{L^{4c_\beta\bar{\kappa}^2\beta}}
\left\|
\left|\operatorname{Tr}\left(\frac{\partial^2 f_i(\theta)}{\partial\left(\widetilde{\alpha}^{(l)}(x_i)\right)^2}\right)\right|
\right\|_{L^{4c_\beta\bar{\kappa}^2\beta}}
\right)^2. E [ ∣ Δ f i ( θ t ) ∣ 2 ] ≤ η ∥ Δ f i ( θ 0 ) ∥ L c β κ ˉ 2 β 2 ≤ η ∥ Δ f i ( θ 0 ) ∥ L 2 c β κ ˉ 2 β 2 ≤ η ( l = 1 ∑ L s l 2 α ( l − 1 ) ( x i ) 2 L 4 c β κ ˉ 2 β Tr ( ∂ ( α ( l ) ( x i ) ) 2 ∂ 2 f i ( θ ) ) L 4 c β κ ˉ 2 β ) 2 . Therefore, a valid choice for Q Q Q
∑ l = 1 L s l 2 ∥ ∥ α ( l − 1 ) ( x i ) ∥ 2 ∥ L 4 c β κ ˉ 2 β ∥ ∣ Tr ( ∂ 2 f i ( θ ) ∂ ( α ~ ( l ) ( x i ) ) 2 ) ∣ ∥ L 4 c β κ ˉ 2 β . \sum_{l=1}^{L}s_l^2
\left\|\left\|\alpha^{(l-1)}(x_i)\right\|^2\right\|_{L^{4c_\beta\bar{\kappa}^2\beta}}
\left\|
\left|\operatorname{Tr}\left(\frac{\partial^2 f_i(\theta)}{\partial\left(\widetilde{\alpha}^{(l)}(x_i)\right)^2}\right)\right|
\right\|_{L^{4c_\beta\bar{\kappa}^2\beta}}. l = 1 ∑ L s l 2 α ( l − 1 ) ( x i ) 2 L 4 c β κ ˉ 2 β Tr ( ∂ ( α ( l ) ( x i ) ) 2 ∂ 2 f i ( θ ) ) L 4 c β κ ˉ 2 β . We can bound the first L 4 c β κ ˉ 2 β L^{4c_\beta\bar{\kappa}^2\beta} L 4 c β κ ˉ 2 β L 4 c β κ ˉ 2 β L^{4c_\beta\bar{\kappa}^2\beta} L 4 c β κ ˉ 2 β x \mathbf{x} x ∥ x ∥ 2 = c ϕ , σ 0 d \lVert\mathbf{x}\rVert^2=c_{\phi,\sigma_0}d ∥ x ∥ 2 = c ϕ , σ 0 d
Tr ( ( ∂ 2 f ( θ ; x ) ∂ ( α ~ ( l ) ( x ) ) 2 ) ) = ∑ j = 1 m l ∂ 2 f ( θ ; x ) ∂ ( α ~ j ( l ) ( x ) ) 2 = ∑ j = 1 m l s L + 1 v T ∂ 2 α ( L ) ( x ) ∂ ( α ~ j ( l ) ( x ) ) 2 = s L + 1 ∑ j = 1 m l ∑ r = 1 m L v r ∂ 2 α r ( L ) ( x ) ∂ ( α ~ j ( l ) ( x ) ) 2 = s L + 1 ∑ r = 1 m L v r ∑ j = 1 m l ∂ 2 α r ( L ) ( x ) ∂ ( α ~ j ( l ) ( x ) ) 2 . \begin{aligned}
\operatorname{Tr}\left(\left(\frac{\partial^2 f\left(\boldsymbol{\theta};\mathbf{x}\right)}{\partial\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)^2}\right)\right)
= \sum_{j=1}^{m_l}\frac{\partial^2 f\left(\boldsymbol{\theta};\mathbf{x}\right)}{\partial\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)^2} \\
&= \sum_{j=1}^{m_l}s_{L+1}\mathbf{v}^{\mathsf{T}}\frac{\partial^2\boldsymbol{\alpha}^{(L)}(\mathbf{x})}{\partial\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)^2} \\
&= s_{L+1}\sum_{j=1}^{m_l}\sum_{r=1}^{m_L}v_r\frac{\partial^2\alpha_r^{(L)}(\mathbf{x})}{\partial\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)^2}
= s_{L+1}\sum_{r=1}^{m_L}v_r\sum_{j=1}^{m_l}\frac{\partial^2\alpha_r^{(L)}(\mathbf{x})}{\partial\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)^2}.
\end{aligned} Tr ( ( ∂ ( α ( l ) ( x ) ) 2 ∂ 2 f ( θ ; x ) ) ) = j = 1 ∑ m l ∂ ( α j ( l ) ( x ) ) 2 ∂ 2 f ( θ ; x ) = j = 1 ∑ m l s L + 1 v T ∂ ( α j ( l ) ( x ) ) 2 ∂ 2 α ( L ) ( x ) = s L + 1 j = 1 ∑ m l r = 1 ∑ m L v r ∂ ( α j ( l ) ( x ) ) 2 ∂ 2 α r ( L ) ( x ) = s L + 1 r = 1 ∑ m L v r j = 1 ∑ m l ∂ ( α j ( l ) ( x ) ) 2 ∂ 2 α r ( L ) ( x ) . Let us denote, for any l ∈ [ L ] l\in[L] l ∈ [ L ] r ∈ [ m L ] r\in[m_L] r ∈ [ m L ] M ∈ R m l × m l \mathbf{M}\in\mathbb{R}^{m_l\times m_l} M ∈ R m l × m l
H r ( l ) ∈ R m l × m l , H r ( l ) ≜ ∂ 2 α r ( L ) ( x ) ∂ ( α ~ ( l ) ( x ) ) 2 ; \mathbf{H}^{(l)}_r\in\mathbb{R}^{m_l\times m_l},\mathbf{H}^{(l)}_r\triangleq\frac{\partial^2\alpha_r^{(L)}(\mathbf{x})}{\partial\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)^2}; H r ( l ) ∈ R m l × m l , H r ( l ) ≜ ∂ ( α ( l ) ( x ) ) 2 ∂ 2 α r ( L ) ( x ) ; T ( l ) ( M ) ∈ R m L , T ( l ) ( M ) ≜ ( Tr ( H 1 ( l ) M ) , … , Tr ( H m L ( l ) M ) ) T ; \mathbf{T}^{(l)}(\mathbf{M})\in\mathbb{R}^{m_L},\mathbf{T}^{(l)}(\mathbf{M})\triangleq\left(\operatorname{Tr}\left(\mathbf{H}^{(l)}_1\mathbf{M}\right),\ldots,\operatorname{Tr}\left(\mathbf{H}^{(l)}_{m_L}\mathbf{M}\right)\right)^{\mathsf{T}}; T ( l ) ( M ) ∈ R m L , T ( l ) ( M ) ≜ ( Tr ( H 1 ( l ) M ) , … , Tr ( H m L ( l ) M ) ) T ; D ( l ) ∈ R m l × m l , D ( l ) = diag ( ϕ ′ ( α ~ ( l ) ( x ) ) ) \mathbf{D}^{(l)}\in\mathbb{R}^{m_l\times m_l},\mathbf{D}^{(l)}=\operatorname{diag}\left(\phi'\left(\widetilde{\boldsymbol{\alpha}}^{(l)}(\mathbf{x})\right)\right) D ( l ) ∈ R m l × m l , D ( l ) = diag ( ϕ ′ ( α ( l ) ( x ) ) ) (these are functions of x \mathbf{x} x
Tr ( ∂ 2 f ( θ ; x ) ∂ ( α ~ ( l ) ( x ) ) 2 ) = s L + 1 ∑ r = 1 m L v r ∑ j = 1 m l ( H r ( l ) ) j j = s L + 1 ∑ r = 1 m L v r Tr ( H r ( l ) ) = s L + 1 v T T ( l ) ( I m l ) . \begin{aligned}
\operatorname{Tr}\left(\frac{\partial^2 f\left(\boldsymbol{\theta};\mathbf{x}\right)}{\partial\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)^2}\right)
= s_{L+1}\sum_{r=1}^{m_L}v_r\sum_{j=1}^{m_l}\left(\mathbf{H}^{(l)}_r\right)_{jj} \\
&= s_{L+1}\sum_{r=1}^{m_L}v_r\operatorname{Tr}\left(\mathbf{H}^{(l)}_r\right)
= s_{L+1}\mathbf{v}^{\mathsf{T}}\mathbf{T}^{(l)}(\mathbf{I}_{m_l}).
\end{aligned} Tr ( ∂ ( α ( l ) ( x ) ) 2 ∂ 2 f ( θ ; x ) ) = s L + 1 r = 1 ∑ m L v r j = 1 ∑ m l ( H r ( l ) ) j j = s L + 1 r = 1 ∑ m L v r Tr ( H r ( l ) ) = s L + 1 v T T ( l ) ( I m l ) . Let F ≜ σ ( W ( 1 ) , … , W ( L ) ) \mathcal{F}\triangleq\sigma\left(\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(L)}\right) F ≜ σ ( W ( 1 ) , … , W ( L ) ) σ \sigma σ σ ( W ( 1 ) , … , W ( L ) ) \sigma\left(\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(L)}\right) σ ( W ( 1 ) , … , W ( L ) )
Tr ( ∂ 2 f ( θ ; x ) ∂ ( α ~ ( l ) ( x ) ) 2 ∣ F ) ∼ N ( 0 , s L + 1 2 a L + 1 ∥ T ( l ) ( I m l ) ∥ 2 ) . \operatorname{Tr}\left(\left.\frac{\partial^2 f\left(\boldsymbol{\theta};\mathbf{x}\right)}{\partial\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)^2}\right|\mathcal{F}\right)
\sim\mathcal{N}\left(0,\frac{s_{L+1}^2}{a_{L+1}}\left\|\mathbf{T}^{(l)}(\mathbf{I}_{m_l})\right\|^2\right). Tr ( ∂ ( α ( l ) ( x ) ) 2 ∂ 2 f ( θ ; x ) F ) ∼ N ( 0 , a L + 1 s L + 1 2 T ( l ) ( I m l ) 2 ) . Therefore, for any p ≥ 1 p\geq1 p ≥ 1 c > 0 c>0 c > 0
∥ Tr ( ∂ 2 f ( θ ; x ) ∂ ( α ~ ( l ) ( x ) ) 2 ) ∥ L p ∣ F ≤ c ∥ T ( l ) ( I m l ) ∥ s L + 1 2 a L + 1 p , \left\|\operatorname{Tr}\left(\frac{\partial^2 f\left(\boldsymbol{\theta};\mathbf{x}\right)}{\partial\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)^2}\right)\right\|_{L_p\mid\mathcal{F}}
\leq c\left\|\mathbf{T}^{(l)}(\mathbf{I}_{m_l})\right\|\sqrt{\frac{s_{L+1}^2}{a_{L+1}}p}, Tr ( ∂ ( α ( l ) ( x ) ) 2 ∂ 2 f ( θ ; x ) ) L p ∣ F ≤ c T ( l ) ( I m l ) a L + 1 s L + 1 2 p , and together with Claim E.4,
∥ Tr ( ∂ 2 f ( θ ; x ) ∂ ( α ~ ( l ) ( x ) ) 2 ) ∥ L p ≤ c ∥ ∥ T ( l ) ( I m l ) ∥ ∥ L p s L + 1 2 a L + 1 p . \left\|\operatorname{Tr}\left(\frac{\partial^2 f\left(\boldsymbol{\theta};\mathbf{x}\right)}{\partial\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)^2}\right)\right\|_{L_p}
\leq c\left\|\left\|\mathbf{T}^{(l)}(\mathbf{I}_{m_l})\right\|\right\|_{L_p}\sqrt{\frac{s_{L+1}^2}{a_{L+1}}p}. Tr ( ∂ ( α ( l ) ( x ) ) 2 ∂ 2 f ( θ ; x ) ) L p ≤ c T ( l ) ( I m l ) L p a L + 1 s L + 1 2 p . Let us consider l = L l=L l = L
H r ( L ) = ϕ ′ ′ ( α ~ r ( L ) ( x ) ) e r e r T , \mathbf{H}^{(L)}_r=\phi''\left(\widetilde{\alpha}^{(L)}_r(\mathbf{x})\right)\mathbf{e}_r\mathbf{e}_r^{\mathsf{T}}, H r ( L ) = ϕ ′′ ( α r ( L ) ( x ) ) e r e r T , Thus, for every M \mathbf{M} M ∥ M ∥ ≤ 1 \lVert\mathbf{M}\rVert\leq1 ∥ M ∥ ≤ 1
∥ T ( L ) ( M ) ∥ = ∥ ( Tr ( H 1 ( L ) M ) , … , Tr ( H m L ( L ) M ) ) ∥ = ∥ ( ϕ ′ ′ ( α ~ 1 ( L ) ( x ) ) M 11 , … , ϕ ′ ′ ( α ~ m L ( L ) ( x ) ) M m L m L ) ∥ ≤ β ϕ m L . \begin{aligned}
\left\|\mathbf{T}^{(L)}(\mathbf{M})\right\|
=\left\|\left(\operatorname{Tr}\left(\mathbf{H}^{(L)}_1\mathbf{M}\right),\ldots,\operatorname{Tr}\left(\mathbf{H}^{(L)}_{m_L}\mathbf{M}\right)\right)\right\| \\
&=\left\|\left(\phi''\left(\widetilde{\alpha}^{(L)}_1(\mathbf{x})\right)M_{11},\ldots,\phi''\left(\widetilde{\alpha}^{(L)}_{m_L}(\mathbf{x})\right)M_{m_Lm_L}\right)\right\|
\leq\beta_{\phi}\sqrt{m_L}.
\end{aligned} T ( L ) ( M ) = ( Tr ( H 1 ( L ) M ) , … , Tr ( H m L ( L ) M ) ) = ( ϕ ′′ ( α 1 ( L ) ( x ) ) M 11 , … , ϕ ′′ ( α m L ( L ) ( x ) ) M m L m L ) ≤ β ϕ m L . Now consider l ∈ [ L − 1 ] l \in[L-1] l ∈ [ L − 1 ] T ( l ) ( M ) \mathbf{T}^{(l)}(\mathbf{M}) T ( l ) ( M ) r r r T ( l ) ( M ) \mathbf{T}^{(l)}(\mathbf{M}) T ( l ) ( M )
Tr ( H r ( l ) M ) = ∑ j = 1 m l ( H r ( l ) M ) j j = ∑ j = 1 m l ( ( H r ( l ) ) j : ( M ) : j ) . \operatorname{Tr}\left(\mathbf{H}^{(l)}_r\mathbf{M}\right)=\sum_{j=1}^{m_l}\left(\mathbf{H}^{(l)}_r\mathbf{M}\right)_{jj}=\sum_{j=1}^{m_l}\left(\left(\mathbf{H}^{(l)}_r\right)_{j:}\left(\mathbf{M}\right)_{:j}\right). Tr ( H r ( l ) M ) = j = 1 ∑ m l ( H r ( l ) M ) j j = j = 1 ∑ m l ( ( H r ( l ) ) j : ( M ) : j ) . Using Constantine and Savits [34 Theorem 2.1] with layer l + 1 l+1 l + 1
( H r ( l ) ) j k = ∂ 2 α r ( L ) ( x ) ∂ α ~ j ( l ) ( x ) ∂ α ~ k ( l ) ( x ) = ∂ α r ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) ∂ 2 α ~ ( l + 1 ) ( x ) ∂ α ~ j ( l ) ( x ) ∂ α ~ k ( l ) ( x ) + ( ∂ α ~ ( l + 1 ) ( x ) ∂ α ~ j ( l ) ( x ) ) T ∂ 2 α r ( L ) ( x ) ∂ ( α ~ ( l + 1 ) ( x ) ) 2 ∂ α ~ ( l + 1 ) ( x ) ∂ α ~ k ( l ) ( x ) = ∂ α r ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) s l + 1 W : j ( l + 1 ) ϕ ′ ′ ( α ~ j ( l ) ( x ) ) 1 { j = k } + s l + 1 2 ( ϕ ′ ′ ( α ~ j ( l ) ( x ) ) W : j ( l + 1 ) T ) H r ( l + 1 ) ( ϕ ′ ′ ( α ~ k ( l ) ( x ) ) W : k ( l + 1 ) ) = s l + 1 ϕ ′ ′ ( α ~ j ( l ) ( x ) ) ∂ α r ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) W : j ( l + 1 ) 1 { j = k } + s l + 1 2 [ D ( l ) ( W ( l + 1 ) ) T H r ( l + 1 ) W ( l + 1 ) D ( l ) ] j k . \begin{aligned}
\left(\mathbf{H}^{(l)}_r\right)_{jk}=\frac{\partial^2\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l)}_j(\mathbf{x})\partial\widetilde{\alpha}^{(l)}_k(\mathbf{x})} \\
&=\frac{\partial\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}\frac{\partial^2\widetilde{\alpha}^{(l+1)}(\mathbf{x})}{\partial\widetilde{\alpha}^{(l)}_j(\mathbf{x})\partial\widetilde{\alpha}^{(l)}_k(\mathbf{x})} \\
&+\left(\frac{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}{\partial\widetilde{\alpha}^{(l)}_j(\mathbf{x})}\right)^{\mathrm{T}}\frac{\partial^2\alpha^{(L)}_r(\mathbf{x})}{\partial\left(\widetilde{\alpha}^{(l+1)}(\mathbf{x})\right)^2}\frac{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}{\partial\widetilde{\alpha}^{(l)}_k(\mathbf{x})} \\
&=\frac{\partial\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}s_{l+1}\mathbf{W}^{(l+1)}_{:j}\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)\mathbf{1}_{\{j=k\}} \\
&+s_{l+1}^2\left(\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)\mathbf{W}^{(l+1)\mathrm{T}}_{:j}\right)\mathbf{H}^{(l+1)}_r\left(\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}_k(\mathbf{x})\right)\mathbf{W}^{(l+1)}_{:k}\right) \\
&=s_{l+1}\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)\frac{\partial\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}\mathbf{W}^{(l+1)}_{:j}\mathbf{1}_{\{j=k\}} \\
&+s_{l+1}^2\left[\mathbf{D}^{(l)}\left(\mathbf{W}^{(l+1)}\right)^{\mathrm{T}}\mathbf{H}^{(l+1)}_r\mathbf{W}^{(l+1)}\mathbf{D}^{(l)}\right]_{jk}.
\end{aligned} ( H r ( l ) ) j k = ∂ α j ( l ) ( x ) ∂ α k ( l ) ( x ) ∂ 2 α r ( L ) ( x ) = ∂ α ( l + 1 ) ( x ) ∂ α r ( L ) ( x ) ∂ α j ( l ) ( x ) ∂ α k ( l ) ( x ) ∂ 2 α ( l + 1 ) ( x ) + ( ∂ α j ( l ) ( x ) ∂ α ( l + 1 ) ( x ) ) T ∂ ( α ( l + 1 ) ( x ) ) 2 ∂ 2 α r ( L ) ( x ) ∂ α k ( l ) ( x ) ∂ α ( l + 1 ) ( x ) = ∂ α ( l + 1 ) ( x ) ∂ α r ( L ) ( x ) s l + 1 W : j ( l + 1 ) ϕ ′′ ( α j ( l ) ( x ) ) 1 { j = k } + s l + 1 2 ( ϕ ′′ ( α j ( l ) ( x ) ) W : j ( l + 1 ) T ) H r ( l + 1 ) ( ϕ ′′ ( α k ( l ) ( x ) ) W : k ( l + 1 ) ) = s l + 1 ϕ ′′ ( α j ( l ) ( x ) ) ∂ α ( l + 1 ) ( x ) ∂ α r ( L ) ( x ) W : j ( l + 1 ) 1 { j = k } + s l + 1 2 [ D ( l ) ( W ( l + 1 ) ) T H r ( l + 1 ) W ( l + 1 ) D ( l ) ] j k . Thus, denoting diag ( M ) ≜ ( M 11 , … , M m l m l ) T \operatorname{diag}(\mathbf{M})\triangleq(M_{11},\ldots,M_{m_lm_l})^{\mathrm{T}} diag ( M ) ≜ ( M 11 , … , M m l m l ) T
Tr ( H r ( l ) M ) = ∑ j = 1 m l s l + 1 ϕ ′ ′ ( α ~ j ( l ) ( x ) ) M j j ∂ α r ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) W : j ( l + 1 ) + ∑ j = 1 m l s l + 1 2 [ D ( l ) ( W ( l + 1 ) ) T H r ( l + 1 ) W ( l + 1 ) D ( l ) ] j : ( M ) : j = s l + 1 ∑ j = 1 m l ( ϕ ′ ′ ( α ~ ( l ) ( x ) ) ⊙ diag ( M ) ) j ( ∂ α r ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) W ( l + 1 ) ) j + s l + 1 2 Tr ( D ( l ) ( W ( l + 1 ) ) T H r ( l + 1 ) W ( l + 1 ) D ( l ) M ) = s l + 1 ∂ α r ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) W ( l + 1 ) ( ϕ ′ ′ ( α ~ ( l ) ( x ) ) ⊙ diag ( M ) ) + s l + 1 2 Tr ( H r ( l + 1 ) W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) = s l + 1 [ ∂ α r ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) W ( l + 1 ) ( ϕ ′ ′ ( α ~ ( l ) ( x ) ) ⊙ diag ( M ) ) ] r + s l + 1 2 Tr ( H r ( l + 1 ) W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) . \begin{aligned}
\operatorname{Tr}\left(\mathbf{H}^{(l)}_r\mathbf{M}\right)=\sum_{j=1}^{m_l}s_{l+1}\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}_j(\mathbf{x})\right)M_{jj}\frac{\partial\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}\mathbf{W}^{(l+1)}_{:j} \\
&+\sum_{j=1}^{m_l}s_{l+1}^2\left[\mathbf{D}^{(l)}\left(\mathbf{W}^{(l+1)}\right)^{\mathrm{T}}\mathbf{H}^{(l+1)}_r\mathbf{W}^{(l+1)}\mathbf{D}^{(l)}\right]_{j:}\left(\mathbf{M}\right)_{:j} \\
&=s_{l+1}\sum_{j=1}^{m_l}\left(\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)\odot\operatorname{diag}(\mathbf{M})\right)_j\left(\frac{\partial\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}\mathbf{W}^{(l+1)}\right)_j \\
&+s_{l+1}^2\operatorname{Tr}\left(\mathbf{D}^{(l)}\left(\mathbf{W}^{(l+1)}\right)^{\mathrm{T}}\mathbf{H}^{(l+1)}_r\mathbf{W}^{(l+1)}\mathbf{D}^{(l)}\mathbf{M}\right) \\
&=s_{l+1}\frac{\partial\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}\mathbf{W}^{(l+1)}\left(\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)\odot\operatorname{diag}(\mathbf{M})\right) \\
&+s_{l+1}^2\operatorname{Tr}\left(\mathbf{H}^{(l+1)}_r\mathbf{W}^{(l+1)}\mathbf{D}^{(l)}\mathbf{M}\mathbf{D}^{(l)}\left(\mathbf{W}^{(l+1)}\right)^{\mathrm{T}}\right) \\
&=s_{l+1}\left[\frac{\partial\alpha^{(L)}_r(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}\mathbf{W}^{(l+1)}\left(\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)\odot\operatorname{diag}(\mathbf{M})\right)\right]_r \\
&+s_{l+1}^2\operatorname{Tr}\left(\mathbf{H}^{(l+1)}_r\mathbf{W}^{(l+1)}\mathbf{D}^{(l)}\mathbf{M}\mathbf{D}^{(l)}\left(\mathbf{W}^{(l+1)}\right)^{\mathrm{T}}\right).
\end{aligned} Tr ( H r ( l ) M ) = j = 1 ∑ m l s l + 1 ϕ ′′ ( α j ( l ) ( x ) ) M j j ∂ α ( l + 1 ) ( x ) ∂ α r ( L ) ( x ) W : j ( l + 1 ) + j = 1 ∑ m l s l + 1 2 [ D ( l ) ( W ( l + 1 ) ) T H r ( l + 1 ) W ( l + 1 ) D ( l ) ] j : ( M ) : j = s l + 1 j = 1 ∑ m l ( ϕ ′′ ( α ( l ) ( x ) ) ⊙ diag ( M ) ) j ( ∂ α ( l + 1 ) ( x ) ∂ α r ( L ) ( x ) W ( l + 1 ) ) j + s l + 1 2 Tr ( D ( l ) ( W ( l + 1 ) ) T H r ( l + 1 ) W ( l + 1 ) D ( l ) M ) = s l + 1 ∂ α ( l + 1 ) ( x ) ∂ α r ( L ) ( x ) W ( l + 1 ) ( ϕ ′′ ( α ( l ) ( x ) ) ⊙ diag ( M ) ) + s l + 1 2 Tr ( H r ( l + 1 ) W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) = s l + 1 [ ∂ α ( l + 1 ) ( x ) ∂ α r ( L ) ( x ) W ( l + 1 ) ( ϕ ′′ ( α ( l ) ( x ) ) ⊙ diag ( M ) ) ] r + s l + 1 2 Tr ( H r ( l + 1 ) W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) . Recalling the definition of T ( l ) \mathbf{T}^{(l)} T ( l )
T ( l ) ( M ) = s l + 1 ∂ α ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) W ( l + 1 ) ( ϕ ′ ′ ( α ~ ( l ) ( x ) ) ⊙ diag ( M ) ) + s l + 1 2 T ( l + 1 ) ( W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) . \begin{aligned}
\mathbf{T}^{(l)}(\mathbf{M})=s_{l+1}\frac{\partial\alpha^{(L)}(\mathbf{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathbf{x})}\mathbf{W}^{(l+1)}\left(\phi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}(\mathbf{x})\right)\odot\operatorname{diag}(\mathbf{M})\right) \\
&+s_{l+1}^2\mathbf{T}^{(l+1)}\left(\mathbf{W}^{(l+1)}\mathbf{D}^{(l)}\mathbf{M}\mathbf{D}^{(l)}\left(\mathbf{W}^{(l+1)}\right)^{\mathrm{T}}\right).
\end{aligned} T ( l ) ( M ) = s l + 1 ∂ α ( l + 1 ) ( x ) ∂ α ( L ) ( x ) W ( l + 1 ) ( ϕ ′′ ( α ( l ) ( x ) ) ⊙ diag ( M ) ) + s l + 1 2 T ( l + 1 ) ( W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) . Hence, by the triangle inequality and other properties of the spectral norm,
∥ T ( l ) ( M ) ∥ ≤ s l + 1 ∥ ∂ α ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) ∥ ∥ W ( l + 1 ) ∥ ∥ φ ′ ′ ( α ~ ( l ) ( x ) ) ⊙ diag ( M ) ∥ + s l + 1 2 ∥ T ( l + 1 ) ( W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) ∥ . \left\|\mathrm{T}^{(l)}(\mathrm{M})\right\| \le s_{l+1}\left\|\frac{\partial\alpha^{(L)}(\mathrm{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathrm{x})}\right\|\left\|\mathrm{W}^{(l+1)}\right\|\left\|\varphi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}(\mathrm{x})\right)\odot\operatorname{diag}(\mathrm{M})\right\|
+ s_{l+1}^{2}\left\|\mathrm{T}^{(l+1)}\left(\mathrm{W}^{(l+1)}\mathrm{D}^{(l)}\mathrm{M}\mathrm{D}^{(l)}\left(\mathrm{W}^{(l+1)}\right)^{\mathsf{T}}\right)\right\|. T ( l ) ( M ) ≤ s l + 1 ∂ α ( l + 1 ) ( x ) ∂ α ( L ) ( x ) W ( l + 1 ) φ ′′ ( α ( l ) ( x ) ) ⊙ diag ( M ) + s l + 1 2 T ( l + 1 ) ( W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T ) . It holds that
∥ ∂ α ( L ) ( x ) ∂ α ~ ( l + 1 ) ( x ) ∥ = ∥ ∂ α ( L ) ( x ) ∂ α ~ ( L ) ( x ) ∂ α ~ ( L ) ( x ) ∂ α ( L − 1 ) ( x ) ⋯ ∂ α ( l + 2 ) ( x ) ∂ α ~ ( l + 2 ) ( x ) ∂ α ~ ( l + 2 ) ( x ) ∂ α ( l + 1 ) ( x ) ∂ α ( l + 1 ) ( x ) ∂ α ~ ( l + 1 ) ( x ) ∥ = ∥ D ( L ) s L W ( L ) ⋯ D ( l + 2 ) s l + 2 W ( l + 2 ) D ( l + 1 ) ∥ ≤ α ϕ L − l − 1 ∏ t = l + 2 L s t ∥ W ( t ) ∥ . \begin{aligned}
\left\|\frac{\partial\alpha^{(L)}(\mathrm{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathrm{x})}\right\|
=
\left\|\frac{\partial\alpha^{(L)}(\mathrm{x})}{\partial\widetilde{\alpha}^{(L)}(\mathrm{x})}\frac{\partial\widetilde{\alpha}^{(L)}(\mathrm{x})}{\partial\alpha^{(L-1)}(\mathrm{x})}\cdots\frac{\partial\alpha^{(l+2)}(\mathrm{x})}{\partial\widetilde{\alpha}^{(l+2)}(\mathrm{x})}\frac{\partial\widetilde{\alpha}^{(l+2)}(\mathrm{x})}{\partial\alpha^{(l+1)}(\mathrm{x})}\frac{\partial\alpha^{(l+1)}(\mathrm{x})}{\partial\widetilde{\alpha}^{(l+1)}(\mathrm{x})}\right\| \\
&= \left\|\mathrm{D}^{(L)}s_L\mathrm{W}^{(L)}\cdots\mathrm{D}^{(l+2)}s_{l+2}\mathrm{W}^{(l+2)}\mathrm{D}^{(l+1)}\right\| \\
&\le\alpha_{\phi}^{L-l-1}\prod_{t=l+2}^{L}s_t\left\|\mathrm{W}^{(t)}\right\|.
\end{aligned} ∂ α ( l + 1 ) ( x ) ∂ α ( L ) ( x ) = ∂ α ( L ) ( x ) ∂ α ( L ) ( x ) ∂ α ( L − 1 ) ( x ) ∂ α ( L ) ( x ) ⋯ ∂ α ( l + 2 ) ( x ) ∂ α ( l + 2 ) ( x ) ∂ α ( l + 1 ) ( x ) ∂ α ( l + 2 ) ( x ) ∂ α ( l + 1 ) ( x ) ∂ α ( l + 1 ) ( x ) = D ( L ) s L W ( L ) ⋯ D ( l + 2 ) s l + 2 W ( l + 2 ) D ( l + 1 ) ≤ α ϕ L − l − 1 t = l + 2 ∏ L s t W ( t ) . Let us consider M \mathrm{M} M ∥ M ∥ ≤ 1 \left\|\mathrm{M}\right\|\le1 ∥ M ∥ ≤ 1 M \mathrm{M} M
∥ φ ′ ′ ( α ~ ( l ) ( x ) ) ⊙ diag ( M ) ∥ ≤ ∥ φ ′ ′ ( α ~ ( l ) ( x ) ) ∥ ≤ β ϕ m l . \left\|\varphi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}(\mathrm{x})\right)\odot\operatorname{diag}(\mathrm{M})\right\|
\le\left\|\varphi^{\prime\prime}\left(\widetilde{\alpha}^{(l)}(\mathrm{x})\right)\right\|
\le\beta_{\phi}\sqrt{m_l}. φ ′′ ( α ( l ) ( x ) ) ⊙ diag ( M ) ≤ φ ′′ ( α ( l ) ( x ) ) ≤ β ϕ m l . Additionally, denoting
A = W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T , \mathrm{A}=\mathrm{W}^{(l+1)}\mathrm{D}^{(l)}\mathrm{M}\mathrm{D}^{(l)}\left(\mathrm{W}^{(l+1)}\right)^{\mathsf{T}}, A = W ( l + 1 ) D ( l ) M D ( l ) ( W ( l + 1 ) ) T , we have that ∥ T ( l + 1 ) ( A ) ∥ \left\|\mathrm{T}^{(l+1)}(\mathrm{A})\right\| T ( l + 1 ) ( A ) A = 0 \mathrm{A}=0 A = 0 A ≠ 0 \mathrm{A}\ne0 A = 0
∥ T ( l + 1 ) ( A ) ∥ = ∥ A ∥ ∥ A ∥ − 1 ∥ T ( l + 1 ) ( A ) ∥ = ∥ A ∥ ∥ ∥ A ∥ − 1 T ( l + 1 ) ( A ) ∥ = ∥ A ∥ ∥ T ( l + 1 ) ( ∥ A ∥ − 1 A ) ∥ ≤ ∥ A ∥ sup ∥ M ′ ∥ ≤ 1 ∥ T ( l + 1 ) ( M ′ ) ∥ ≤ α ϕ 2 ∥ W ( l + 1 ) ∥ 2 sup ∥ M ′ ∥ ≤ 1 ∥ T ( l + 1 ) ( M ′ ) ∥ . \begin{aligned}
\left\|\mathrm{T}^{(l+1)}(\mathrm{A})\right\|
=\left\|\mathrm{A}\right\|\left\|\mathrm{A}\right\|^{-1}\left\|\mathrm{T}^{(l+1)}(\mathrm{A})\right\| \\
&=\left\|\mathrm{A}\right\|\left\|\left\|\mathrm{A}\right\|^{-1}\mathrm{T}^{(l+1)}(\mathrm{A})\right\| \\
&=\left\|\mathrm{A}\right\|\left\|\mathrm{T}^{(l+1)}\left(\left\|\mathrm{A}\right\|^{-1}\mathrm{A}\right)\right\| \\
&\le\left\|\mathrm{A}\right\|\sup_{\left\|\mathrm{M}^{\prime}\right\|\le1}\left\|\mathrm{T}^{(l+1)}(\mathrm{M}^{\prime})\right\|
\le\alpha_{\phi}^{2}\left\|\mathrm{W}^{(l+1)}\right\|^{2}\sup_{\left\|\mathrm{M}^{\prime}\right\|\le1}\left\|\mathrm{T}^{(l+1)}(\mathrm{M}^{\prime})\right\|.
\end{aligned} T ( l + 1 ) ( A ) = ∥ A ∥ ∥ A ∥ − 1 T ( l + 1 ) ( A ) = ∥ A ∥ ∥ A ∥ − 1 T ( l + 1 ) ( A ) = ∥ A ∥ T ( l + 1 ) ( ∥ A ∥ − 1 A ) ≤ ∥ A ∥ ∥ M ′ ∥ ≤ 1 sup T ( l + 1 ) ( M ′ ) ≤ α ϕ 2 W ( l + 1 ) 2 ∥ M ′ ∥ ≤ 1 sup T ( l + 1 ) ( M ′ ) . Since the above holds for every M \mathrm{M} M ∥ M ∥ ≤ 1 \left\|\mathrm{M}\right\|\le1 ∥ M ∥ ≤ 1
sup ∥ M ∥ ≤ 1 ∥ T ( l ) ( M ) ∥ ≤ s l + 1 ( α ϕ L − l − 1 ∏ t = l + 2 L s t ∥ W ( t ) ∥ ) ∥ W ( l + 1 ) ∥ β ϕ m l + s l + 1 2 α ϕ 2 ∥ W ( l + 1 ) ∥ 2 sup ∥ M ′ ∥ ≤ 1 ∥ T ( l + 1 ) ( M ′ ) ∥ = β ϕ α ϕ L − l − 1 m l ∏ t = l + 1 L s t ∥ W ( t ) ∥ + s l + 1 2 α ϕ 2 ∥ W ( l + 1 ) ∥ 2 sup ∥ M ′ ∥ ≤ 1 ∥ T ( l + 1 ) ( M ′ ) ∥ . \begin{aligned}
\sup_{\left\|\mathrm{M}\right\|\le1}\left\|\mathrm{T}^{(l)}(\mathrm{M})\right\|
\le s_{l+1}\left(\alpha_{\phi}^{L-l-1}\prod_{t=l+2}^{L}s_t\left\|\mathrm{W}^{(t)}\right\|\right)\left\|\mathrm{W}^{(l+1)}\right\|\beta_{\phi}\sqrt{m_l}
+s_{l+1}^{2}\alpha_{\phi}^{2}\left\|\mathrm{W}^{(l+1)}\right\|^{2}\sup_{\left\|\mathrm{M}^{\prime}\right\|\le1}\left\|\mathrm{T}^{(l+1)}(\mathrm{M}^{\prime})\right\| \\
&=\beta_{\phi}\alpha_{\phi}^{L-l-1}\sqrt{m_l}\prod_{t=l+1}^{L}s_t\left\|\mathrm{W}^{(t)}\right\|
+s_{l+1}^{2}\alpha_{\phi}^{2}\left\|\mathrm{W}^{(l+1)}\right\|^{2}\sup_{\left\|\mathrm{M}^{\prime}\right\|\le1}\left\|\mathrm{T}^{(l+1)}(\mathrm{M}^{\prime})\right\|.
\end{aligned} ∥ M ∥ ≤ 1 sup T ( l ) ( M ) ≤ s l + 1 ( α ϕ L − l − 1 t = l + 2 ∏ L s t W ( t ) ) W ( l + 1 ) β ϕ m l + s l + 1 2 α ϕ 2 W ( l + 1 ) 2 ∥ M ′ ∥ ≤ 1 sup T ( l + 1 ) ( M ′ ) = β ϕ α ϕ L − l − 1 m l t = l + 1 ∏ L s t W ( t ) + s l + 1 2 α ϕ 2 W ( l + 1 ) 2 ∥ M ′ ∥ ≤ 1 sup T ( l + 1 ) ( M ′ ) . Unrolling this recurrence and absorbing powers of α ϕ \alpha_{\phi} α ϕ β ϕ \beta_{\phi} β ϕ C > 0 C>0 C > 0
sup ∥ M ∥ ≤ 1 ∥ T ( l ) ( M ) ∥ ≤ C L ∑ q = l L ( ∏ t = l + 1 q s t 2 ∥ W ( t ) ∥ 2 ) m q ( ∏ t = q + 1 L s t ∥ W ( t ) ∥ ) . \sup_{\left\|\mathrm{M}\right\|\le1}\left\|\mathrm{T}^{(l)}(\mathrm{M})\right\|
\le C^{L}\sum_{q=l}^{L}\left(\prod_{t=l+1}^{q}s_t^{2}\left\|\mathrm{W}^{(t)}\right\|^{2}\right)\sqrt{m_q}\left(\prod_{t=q+1}^{L}s_t\left\|\mathrm{W}^{(t)}\right\|\right). ∥ M ∥ ≤ 1 sup T ( l ) ( M ) ≤ C L q = l ∑ L ( t = l + 1 ∏ q s t 2 W ( t ) 2 ) m q ( t = q + 1 ∏ L s t W ( t ) ) . In particular, the above holds for M = I m l M = I_{m_l} M = I m l
∥ ∥ T ( l ) ( I m l ) ∥ ∥ L p ≤ ∥ ( C L ∑ q = l L ( ∏ t = l + 1 q s t 2 ∥ W ( t ) ∥ 2 ) m q ( ∏ t = q + 1 L s t ∥ W ( t ) ∥ ) ) ∥ L p ≤ C L ∑ q = l L ( ∏ t = l + 1 q s t 2 ∥ ∥ W ( t ) ∥ 2 ∥ L p ) m q ( ∏ t = q + 1 L s t ∥ ∥ W ( t ) ∥ ∥ L p ) \begin{aligned}
\left\|\left\|T^{(l)}\left(I_{m_l}\right)\right\|\right\|_{L_p}
&\le\left\|\left(C^L\sum_{q=l}^L\left(\prod_{t=l+1}^q s_t^2\left\|W^{(t)}\right\|^2\right)\sqrt{m_q}\left(\prod_{t=q+1}^L s_t\left\|W^{(t)}\right\|\right)\right)\right\|_{L_p}\\
&\le C^L\sum_{q=l}^L\left(\prod_{t=l+1}^q s_t^2\left\|\left\|W^{(t)}\right\|^2\right\|_{L_p}\right)\sqrt{m_q}\left(\prod_{t=q+1}^L s_t\left\|\left\|W^{(t)}\right\|\right\|_{L_p}\right)
\end{aligned} T ( l ) ( I m l ) L p ≤ C L q = l ∑ L ( t = l + 1 ∏ q s t 2 W ( t ) 2 ) m q ( t = q + 1 ∏ L s t W ( t ) ) L p ≤ C L q = l ∑ L ( t = l + 1 ∏ q s t 2 W ( t ) 2 L p ) m q ( t = q + 1 ∏ L s t W ( t ) L p ) ≤ C L ∑ q = l L ( ∏ t = l + 1 q c ′ v 0 2 s t 2 a t ( m t − 1 + m t + p ) ) m q ( ∏ t = q + 1 L c ′ v 0 2 s t 2 a t ( m t − 1 + m t + p ) ) , \begin{aligned}
&\le C^L\sum_{q=l}^L\left(\prod_{t=l+1}^q c'v_0^2\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)\sqrt{m_q}\left(\prod_{t=q+1}^L\sqrt{c'v_0^2\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)}\right),
\end{aligned} ≤ C L q = l ∑ L ( t = l + 1 ∏ q c ′ v 0 2 a t s t 2 ( m t − 1 + m t + p ) ) m q t = q + 1 ∏ L c ′ v 0 2 a t s t 2 ( m t − 1 + m t + p ) , where in ( i ) (i) ( i ) ( i i ) (ii) ( ii ) c ′ > 0 c' > 0 c ′ > 0 C C C
∥ ∥ T ( l ) ( I m l ) ∥ ∥ L p ≤ C 2 L ∑ q = l L ( ∏ t = l + 1 q s t 2 a t ( m t − 1 + m t + p ) ) m q ( ∏ t = q + 1 L s t 2 a t ( m t − 1 + m t + p ) ) . \left\|\left\|T^{(l)}\left(I_{m_l}\right)\right\|\right\|_{L_p}
\le C^{2L}\sum_{q=l}^L\left(\prod_{t=l+1}^q\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)\sqrt{m_q}\left(\prod_{t=q+1}^L\sqrt{\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)}\right). T ( l ) ( I m l ) L p ≤ C 2 L q = l ∑ L ( t = l + 1 ∏ q a t s t 2 ( m t − 1 + m t + p ) ) m q t = q + 1 ∏ L a t s t 2 ( m t − 1 + m t + p ) . Therefore,
∥ ∥ Tr ( ∂ 2 f ( θ ; x ) ∂ ( α ~ ( l ) ( x ) ) 2 ) ∥ L p ≤ c C 2 L ∑ q = l L ( ∏ t = l + 1 q s t 2 a t ( m t − 1 + m t + p ) ) m q ( ∏ t = q + 1 L s t 2 a t ( m t − 1 + m t + p ) ) s L + 1 2 a L + 1 p . \|\|\operatorname{Tr}(\frac{\partial^2 f(\theta;x)}{\partial(\widetilde{\alpha}^{(l)}(x))^2})\|_{L_p}
\le cC^{2L}\sum_{q=l}^L(\prod_{t=l+1}^q\frac{s_t^2}{a_t}(m_{t-1}+m_t+p))\sqrt{m_q}(\prod_{t=q+1}^L\sqrt{\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)})\sqrt{\frac{s_{L+1}^2}{a_{L+1}}}p. ∥∥ Tr ( ∂ ( α ( l ) ( x ) ) 2 ∂ 2 f ( θ ; x ) ) ∥ L p ≤ c C 2 L q = l ∑ L ( t = l + 1 ∏ q a t s t 2 ( m t − 1 + m t + p )) m q ( t = q + 1 ∏ L a t s t 2 ( m t − 1 + m t + p ) ) a L + 1 s L + 1 2 p . Recall from Proposition H.3 that there exists a constant c ′ ′ > 0 c'' > 0 c ′′ > 0
∥ ∥ α ( l − 1 ) ( x ) ∥ 2 ∥ L p ≤ ( c ′ ′ ) l − 1 ∥ x ∥ 2 ∏ t = 1 l − 1 s t 2 a t ( m t + p ) . \left\|\left\|\alpha^{(l-1)}(x)\right\|^2\right\|_{L_p}\le(c'')^{l-1}\|x\|^2\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+p). α ( l − 1 ) ( x ) 2 L p ≤ ( c ′′ ) l − 1 ∥ x ∥ 2 t = 1 ∏ l − 1 a t s t 2 ( m t + p ) . W.l.o.g. c ′ ′ = c = C c''=c=C c ′′ = c = C
∑ l = 1 L s l 2 ∥ ∥ α ( l − 1 ) ( x i ) ∥ 2 ∥ L 4 c β κ ˉ 2 β ∥ ∥ Tr ( ∂ 2 f i ( θ ) ∂ ( α ~ ( l ) ( x i ) ) 2 ) ∥ L 4 c β κ ˉ 2 β ∥ ≤ ∑ l = 1 L s l 2 c l − 1 ∥ x ∥ 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + 4 c β κ ˉ 2 β ) ) ⋅ C 2 L + 1 ∑ q = l L ( ∏ t = l + 1 q s t 2 a t ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) m q ⋅ ( ∏ t = q + 1 L s t 2 a t ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) s L + 1 2 a L + 1 4 c β κ ˉ 2 β ≲ C 3 L d ∑ l = 1 L ∑ q = l L ( ∏ t = 1 l − 1 s t 2 a t ( m t + m t − 1 + 4 c β κ ˉ 2 β ) ) s l 2 ( ∏ t = l + 1 q s t 2 a t ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) ⋅ m q ( ∏ t = q + 1 L s t 2 a t ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) s L + 1 2 a L + 1 4 c β κ ˉ 2 β , \begin{aligned}
&\sum_{l=1}^L s_l^2\left\|\left\|\alpha^{(l-1)}(x_i)\right\|^2\right\|_{L^{4c_\beta\bar{\kappa}^2\beta}}\left\|\left\|\operatorname{Tr}\left(\frac{\partial^2 f_i(\theta)}{\partial\left(\widetilde{\alpha}^{(l)}(x_i)\right)^2}\right)\right\|_{L^{4c_\beta\bar{\kappa}^2\beta}}\right\|\\
&\le\sum_{l=1}^L s_l^2c^{l-1}\|x\|^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+4c_\beta\bar{\kappa}^2\beta)\right)\\
&\qquad\cdot C^{2L+1}\sum_{q=l}^L\left(\prod_{t=l+1}^q\frac{s_t^2}{a_t}(m_{t-1}+m_t+4c_\beta\bar{\kappa}^2\beta)\right)\sqrt{m_q}\\
&\qquad\cdot\left(\prod_{t=q+1}^L\sqrt{\frac{s_t^2}{a_t}(m_{t-1}+m_t+4c_\beta\bar{\kappa}^2\beta)}\right)\sqrt{\frac{s_{L+1}^2}{a_{L+1}}4c_\beta\bar{\kappa}^2\beta}\\
&\lesssim C^{3L}d\sum_{l=1}^L\sum_{q=l}^L\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+m_{t-1}+4c_\beta\bar{\kappa}^2\beta)\right)s_l^2\left(\prod_{t=l+1}^q\frac{s_t^2}{a_t}(m_{t-1}+m_t+4c_\beta\bar{\kappa}^2\beta)\right)\\
&\qquad\cdot\sqrt{m_q}\left(\prod_{t=q+1}^L\sqrt{\frac{s_t^2}{a_t}(m_{t-1}+m_t+4c_\beta\bar{\kappa}^2\beta)}\right)\sqrt{\frac{s_{L+1}^2}{a_{L+1}}4c_\beta\bar{\kappa}^2\beta},
\end{aligned} l = 1 ∑ L s l 2 α ( l − 1 ) ( x i ) 2 L 4 c β κ ˉ 2 β Tr ( ∂ ( α ( l ) ( x i ) ) 2 ∂ 2 f i ( θ ) ) L 4 c β κ ˉ 2 β ≤ l = 1 ∑ L s l 2 c l − 1 ∥ x ∥ 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + 4 c β κ ˉ 2 β ) ) ⋅ C 2 L + 1 q = l ∑ L ( t = l + 1 ∏ q a t s t 2 ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) m q ⋅ t = q + 1 ∏ L a t s t 2 ( m t − 1 + m t + 4 c β κ ˉ 2 β ) a L + 1 s L + 1 2 4 c β κ ˉ 2 β ≲ C 3 L d l = 1 ∑ L q = l ∑ L ( t = 1 ∏ l − 1 a t s t 2 ( m t + m t − 1 + 4 c β κ ˉ 2 β ) ) s l 2 ( t = l + 1 ∏ q a t s t 2 ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) ⋅ m q t = q + 1 ∏ L a t s t 2 ( m t − 1 + m t + 4 c β κ ˉ 2 β ) a L + 1 s L + 1 2 4 c β κ ˉ 2 β , where in the last step we use Assumption 4.4. This is the deterministic upper bound we set out to find, which means it is a valid upper bound on Q Q Q
β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] ≲ C 3 L η β − 1 n d ⋅ ( ∑ l = 1 L ∑ q = l L ( ∏ t = l l − 1 s t 2 a t ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) s l 2 ( ∏ t = l + 1 q s t 2 a t ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) ) ⋅ m q ( ∏ t = q + 1 L s t 2 a t ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) s L + 1 2 a L + 1 c β κ ˉ 2 β E [ L ( θ t ) ] . \beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]
\lesssim C^{3L}\sqrt{\eta}\beta^{-1}nd
\cdot\left(\sum_{l=1}^{L}\sum_{q=l}^{L}\left(\prod_{t=l}^{l-1}\frac{s_t^2}{a_t}\left(m_{t-1}+m_t+4c_{\beta}\bar{\kappa}^{2}\beta\right)\right)s_l^2\left(\prod_{t=l+1}^{q}\frac{s_t^2}{a_t}\left(m_{t-1}+m_t+4c_{\beta}\bar{\kappa}^{2}\beta\right)\right)\right)
\cdot\sqrt{m_q}\left(\prod_{t=q+1}^{L}\sqrt{\frac{s_t^2}{a_t}\left(m_{t-1}+m_t+4c_{\beta}\bar{\kappa}^{2}\beta\right)}\right)\sqrt{\frac{s_{L+1}^2}{a_{L+1}}c_{\beta}\bar{\kappa}^{2}\beta}\sqrt{\mathbb{E}\left[\mathcal{L}(\theta_t)\right]}. β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] ≲ C 3 L η β − 1 n d ⋅ l = 1 ∑ L q = l ∑ L ( t = l ∏ l − 1 a t s t 2 ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) s l 2 ( t = l + 1 ∏ q a t s t 2 ( m t − 1 + m t + 4 c β κ ˉ 2 β ) ) ⋅ m q t = q + 1 ∏ L a t s t 2 ( m t − 1 + m t + 4 c β κ ˉ 2 β ) a L + 1 s L + 1 2 c β κ ˉ 2 β E [ L ( θ t ) ] . Fourth term
Let i ∈ [ n ] i \in[n] i ∈ [ n ]
∇ f i ( θ ) ⊤ Λ θ = ∑ l = 1 L ∂ f i ( θ ) ∂ vec ( W ( l ) ) ( a l v 0 2 I m l − 1 m l ) vec ( W ( l ) ) + ∂ f i ( θ ) ∂ v ( a L + 1 I m L ) v = 1 v 0 2 ∑ l = 1 L a l ∂ f i ( θ ) ∂ vec ( W ( l ) ) vec ( W ( l ) ) + a L + 1 ∂ f i ∂ v v . \begin{aligned}
\nabla f_i(\theta)^{\top}\Lambda\theta=\sum_{l=1}^{L}\frac{\partial f_i(\theta)}{\partial\operatorname{vec}\left(W^{(l)}\right)}\left(\frac{a_l}{v_0^2}I_{m_{l-1}m_l}\right)\operatorname{vec}\left(W^{(l)}\right)+\frac{\partial f_i(\theta)}{\partial v}\left(a_{L+1}I_{m_L}\right)v \\
&=\frac{1}{v_0^2}\sum_{l=1}^{L}a_l\frac{\partial f_i(\theta)}{\partial\operatorname{vec}\left(W^{(l)}\right)}\operatorname{vec}\left(W^{(l)}\right)+a_{L+1}\frac{\partial f_i}{\partial v}v.
\end{aligned} ∇ f i ( θ ) ⊤ Λ θ = l = 1 ∑ L ∂ vec ( W ( l ) ) ∂ f i ( θ ) ( v 0 2 a l I m l − 1 m l ) vec ( W ( l ) ) + ∂ v ∂ f i ( θ ) ( a L + 1 I m L ) v = v 0 2 1 l = 1 ∑ L a l ∂ vec ( W ( l ) ) ∂ f i ( θ ) vec ( W ( l ) ) + a L + 1 ∂ v ∂ f i v . For the first sum, observe that
∂ α ~ ( l ) ( x i ) ∂ vec ( W ( l ) ) vec ( W ( l ) ) = ∂ α ~ ( l ) ( x i ) ∂ W ( l ) ∂ W ( l ) ∂ vec ( W ( l ) ) vec ( W ( l ) ) = ∂ α ~ ( l ) ( x i ) ∂ W ( l ) W ( l ) = α ~ ( l ) ( x i ) . \begin{aligned}
\frac{\partial\widetilde{\alpha}^{(l)}(x_i)}{\partial\operatorname{vec}\left(W^{(l)}\right)}\operatorname{vec}\left(W^{(l)}\right)=\frac{\partial\widetilde{\alpha}^{(l)}(x_i)}{\partial W^{(l)}}\frac{\partial W^{(l)}}{\partial\operatorname{vec}\left(W^{(l)}\right)}\operatorname{vec}\left(W^{(l)}\right) \\
&=\frac{\partial\widetilde{\alpha}^{(l)}(x_i)}{\partial W^{(l)}}W^{(l)} \\
&=\widetilde{\alpha}^{(l)}(x_i).
\end{aligned} ∂ vec ( W ( l ) ) ∂ α ( l ) ( x i ) vec ( W ( l ) ) = ∂ W ( l ) ∂ α ( l ) ( x i ) ∂ vec ( W ( l ) ) ∂ W ( l ) vec ( W ( l ) ) = ∂ W ( l ) ∂ α ( l ) ( x i ) W ( l ) = α ( l ) ( x i ) . Using the above and denoting
h ( l ) ( x ) ≜ ∂ α ( L ) ( x ) ∂ α ~ ( l ) ( x ) α ~ ( l ) ( x ) , \mathbf{h}^{(l)}(x)\triangleq\frac{\partial\alpha^{(L)}(x)}{\partial\widetilde{\alpha}^{(l)}(x)}\widetilde{\alpha}^{(l)}(x), h ( l ) ( x ) ≜ ∂ α ( l ) ( x ) ∂ α ( L ) ( x ) α ( l ) ( x ) , we get
∂ f i ( θ ) ∂ vec ( W ( l ) ) vec ( W ( l ) ) = ∂ f i ( θ ) ∂ α ( L ) ( x i ) ∂ α ( L ) ( x i ) ∂ α ~ ( l ) ( x i ) ∂ α ~ ( l ) ( x i ) ∂ vec ( W ( l ) ) vec ( W ( l ) ) = s L + 1 v ⊤ h ( l ) ( x i ) . \begin{aligned}
\frac{\partial f_i(\theta)}{\partial\operatorname{vec}\left(W^{(l)}\right)}\operatorname{vec}\left(W^{(l)}\right)=\frac{\partial f_i(\theta)}{\partial\alpha^{(L)}(x_i)}\frac{\partial\alpha^{(L)}(x_i)}{\partial\widetilde{\alpha}^{(l)}(x_i)}\frac{\partial\widetilde{\alpha}^{(l)}(x_i)}{\partial\operatorname{vec}\left(W^{(l)}\right)}\operatorname{vec}\left(W^{(l)}\right) \\
&=s_{L+1}v^{\top}\mathbf{h}^{(l)}(x_i).
\end{aligned} ∂ vec ( W ( l ) ) ∂ f i ( θ ) vec ( W ( l ) ) = ∂ α ( L ) ( x i ) ∂ f i ( θ ) ∂ α ( l ) ( x i ) ∂ α ( L ) ( x i ) ∂ vec ( W ( l ) ) ∂ α ( l ) ( x i ) vec ( W ( l ) ) = s L + 1 v ⊤ h ( l ) ( x i ) . For the final term in ∇ f i ( θ ) ⊤ Λ θ \nabla f_i(\theta)^{\top}\Lambda\theta ∇ f i ( θ ) ⊤ Λ θ
a L + 1 ∂ f i ( θ ) ∂ v v = a L + 1 s L + 1 v ⊤ α ( L ) ( x i ) = a L + 1 y i + a L + 1 ( s L + 1 v ⊤ α ( L ) ( x i ) − y i ) = a L + 1 y i + a L + 1 r i ( θ ) . \begin{aligned}
a_{L+1}\frac{\partial f_i(\theta)}{\partial v}v=a_{L+1}s_{L+1}v^{\top}\alpha^{(L)}(x_i) \\
&=a_{L+1}y_i+a_{L+1}\left(s_{L+1}v^{\top}\alpha^{(L)}(x_i)-y_i\right)=a_{L+1}y_i+a_{L+1}r_i(\theta).
\end{aligned} a L + 1 ∂ v ∂ f i ( θ ) v = a L + 1 s L + 1 v ⊤ α ( L ) ( x i ) = a L + 1 y i + a L + 1 ( s L + 1 v ⊤ α ( L ) ( x i ) − y i ) = a L + 1 y i + a L + 1 r i ( θ ) . Thus, we get
∇ f i ( θ ) ⊤ Λ θ = 1 v 0 2 s L + 1 v ⊤ ( ∑ l = 1 L a l h ( l ) ( x i ) ) + a L + 1 y i + a L + 1 r i ( θ ) . \nabla f_i(\theta)^{\top}\Lambda\theta=\frac{1}{v_0^2}s_{L+1}v^{\top}\left(\sum_{l=1}^{L}a_l\mathbf{h}^{(l)}(x_i)\right)+a_{L+1}y_i+a_{L+1}r_i(\theta). ∇ f i ( θ ) ⊤ Λ θ = v 0 2 1 s L + 1 v ⊤ ( l = 1 ∑ L a l h ( l ) ( x i ) ) + a L + 1 y i + a L + 1 r i ( θ ) . Therefore,
r i ( θ ) ∇ f i ( θ ) ⊤ Λ θ = r i ( θ ) ( 1 v 0 2 s L + 1 v ⊤ ( ∑ l = 1 L a l h ( l ) ( x i ) ) + a L + 1 y i ) + 2 a L + 1 1 2 ∣ r i ( θ ) ∣ 2 . r_i(\theta)\nabla f_i(\theta)^{\top}\Lambda\theta=r_i(\theta)\left(\frac{1}{v_0^2}s_{L+1}v^{\top}\left(\sum_{l=1}^{L}a_l\mathbf{h}^{(l)}(x_i)\right)+a_{L+1}y_i\right)+2a_{L+1}\frac{1}{2}\left|r_i(\theta)\right|^2. r i ( θ ) ∇ f i ( θ ) ⊤ Λ θ = r i ( θ ) ( v 0 2 1 s L + 1 v ⊤ ( l = 1 ∑ L a l h ( l ) ( x i ) ) + a L + 1 y i ) + 2 a L + 1 2 1 ∣ r i ( θ ) ∣ 2 . Using this, we get
− β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Λ θ t ] = − β − 1 ∑ i = 1 n E [ r i ( θ t ) ∇ f i ( θ t ) T Λ θ t ] = − β − 1 ∑ i = 1 n E [ r i ( θ t ) ( 1 ν 0 2 s L + 1 v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) + a L + 1 y i ) ] − β − 1 ∑ i = 1 n E [ 2 a L + 1 + 1 2 ∣ r i ( θ t ) ∣ 2 ] ≤ β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) ( 1 ν 0 2 s L + 1 v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) + a L + 1 y i ) ∣ ] − β − 1 ∑ i = 1 n E [ 2 a L + 1 + 1 2 ∣ r i ( θ t ) ∣ 2 ] = β − 1 ∑ i = 1 n E [ ∣ r i ( θ t ) ( 1 ν 0 2 s L + 1 v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) + a L + 1 y i ) ∣ ] − 2 β − 1 n a L + 1 E [ L ( θ t ) ] . \begin{aligned}
-\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\mathsf{T}}\nabla f(\theta_t)\Lambda\theta_t\right]
&=-\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)\nabla f_i(\theta_t)^{\mathsf{T}}\Lambda\theta_t\right]\\
&=-\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)\left(\frac{1}{\nu_0^2}s_{L+1}\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)+a_{L+1}y_i\right)\right]\\
&\quad-\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[2a_{L+1}+\frac{1}{2}|r_i(\theta_t)|^2\right]\\
&\leq\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\left|r_i(\theta_t)\left(\frac{1}{\nu_0^2}s_{L+1}\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)+a_{L+1}y_i\right)\right|\right]\\
&\quad-\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[2a_{L+1}+\frac{1}{2}|r_i(\theta_t)|^2\right]\\
&=\beta^{-1}\sum_{i=1}^{n}\mathbb{E}\left[\left|r_i(\theta_t)\left(\frac{1}{\nu_0^2}s_{L+1}\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)+a_{L+1}y_i\right)\right|\right]\\
&\quad-2\beta^{-1}na_{L+1}\mathbb{E}\left[\mathcal{L}(\theta_t)\right].
\end{aligned} − β − 1 E [ r ( θ t ) T ∇ f ( θ t ) Λ θ t ] = − β − 1 i = 1 ∑ n E [ r i ( θ t ) ∇ f i ( θ t ) T Λ θ t ] = − β − 1 i = 1 ∑ n E [ r i ( θ t ) ( ν 0 2 1 s L + 1 v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) + a L + 1 y i ) ] − β − 1 i = 1 ∑ n E [ 2 a L + 1 + 2 1 ∣ r i ( θ t ) ∣ 2 ] ≤ β − 1 i = 1 ∑ n E [ r i ( θ t ) ( ν 0 2 1 s L + 1 v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) + a L + 1 y i ) ] − β − 1 i = 1 ∑ n E [ 2 a L + 1 + 2 1 ∣ r i ( θ t ) ∣ 2 ] = β − 1 i = 1 ∑ n E [ r i ( θ t ) ( ν 0 2 1 s L + 1 v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) + a L + 1 y i ) ] − 2 β − 1 n a L + 1 E [ L ( θ t ) ] . Focusing on the first term, we have that
∑ i = 1 n E [ ∣ r i ( θ t ) ( 1 ν 0 2 s L + 1 v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) + a L + 1 y i ) ∣ ] ≤ ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] E [ [ 1 ν 0 2 s L + 1 v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) + a L + 1 y i ] 2 ] ≤ ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] 2 ( a L + 1 2 y i 2 + 1 ν 0 4 s L + 1 2 E [ [ v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) ] 2 ] ) ≤ ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] ( 2 a L + 1 ∣ y i ∣ + 2 1 ν 0 2 s L + 1 E [ [ v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) ] 2 ] ) , \begin{aligned}
\sum_{i=1}^{n}\mathbb{E}\left[\left|r_i(\theta_t)\left(\frac{1}{\nu_0^2}s_{L+1}\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)+a_{L+1}y_i\right)\right|\right]
&\leq\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[|r_i(\theta_t)|^2\right]}\sqrt{\mathbb{E}\left[\left[\frac{1}{\nu_0^2}s_{L+1}\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)+a_{L+1}y_i\right]^2\right]}\\
&\leq\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[|r_i(\theta_t)|^2\right]}\sqrt{2\left(a_{L+1}^2y_i^2+\frac{1}{\nu_0^4}s_{L+1}^2\mathbb{E}\left[\left[\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)\right]^2\right]\right)}\\
&\leq\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[|r_i(\theta_t)|^2\right]}\left(2a_{L+1}|y_i|+2\frac{1}{\nu_0^2}s_{L+1}\sqrt{\mathbb{E}\left[\left[\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)\right]^2\right]}\right),
\end{aligned} i = 1 ∑ n E [ r i ( θ t ) ( ν 0 2 1 s L + 1 v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) + a L + 1 y i ) ] ≤ i = 1 ∑ n E [ ∣ r i ( θ t ) ∣ 2 ] E [ ν 0 2 1 s L + 1 v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) + a L + 1 y i ] 2 ≤ i = 1 ∑ n E [ ∣ r i ( θ t ) ∣ 2 ] 2 a L + 1 2 y i 2 + ν 0 4 1 s L + 1 2 E [ v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) ] 2 ≤ i = 1 ∑ n E [ ∣ r i ( θ t ) ∣ 2 ] 2 a L + 1 ∣ y i ∣ + 2 ν 0 2 1 s L + 1 E [ v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) ] 2 , where (i) uses Cauchy–Schwarz, (ii) uses ( a + b ) 2 ≤ 2 a 2 + 2 b 2 (a+b)^2\leq2a^2+2b^2 ( a + b ) 2 ≤ 2 a 2 + 2 b 2 a + b ≤ a + b \sqrt{a+b}\leq\sqrt{a}+\sqrt{b} a + b ≤ a + b Q Q Q i i i
E [ ( v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) ) 2 ] ≤ Q 2 , \mathbb{E}\left[\left(\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_lh^{(l)}(x_i)\right)\right)^2\right]\leq Q^2, E ( v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) ) 2 ≤ Q 2 , then we get
∑ i = 1 n E [ r i ( θ t ) ⊤ ( 1 v 0 2 s L + 1 v ( ∑ ℓ = 1 L a ℓ h ( ℓ ) ( x i ) ) + a L + 1 y i ) ] ≤ ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] ( 2 a L + 1 ∣ y i ∣ + 2 1 v 0 2 s L + 1 Q ) ≤ ( 2 a L + 1 R y + 2 1 v 0 2 s L + 1 Q ) n 1 n ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] \begin{aligned}
\sum_{i=1}^{n}\mathbb{E}\left[r_i(\theta_t)^\top\left(\frac{1}{v_0^2}s_{L+1}\mathbf{v}\left(\sum_{\ell=1}^{L}a_\ell\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right)+a_{L+1}y_i\right)\right] \\
&\le\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[|r_i(\theta_t)|^2\right]}\left(2a_{L+1}|y_i|+2\frac{1}{v_0^2}s_{L+1}Q\right) \\
&\le\left(2a_{L+1}R_y+2\frac{1}{v_0^2}s_{L+1}Q\right)n\frac{1}{n}\sum_{i=1}^{n}\sqrt{\mathbb{E}\left[|r_i(\theta_t)|^2\right]}
\end{aligned} i = 1 ∑ n E [ r i ( θ t ) ⊤ ( v 0 2 1 s L + 1 v ( ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) ) + a L + 1 y i ) ] ≤ i = 1 ∑ n E [ ∣ r i ( θ t ) ∣ 2 ] ( 2 a L + 1 ∣ y i ∣ + 2 v 0 2 1 s L + 1 Q ) ≤ ( 2 a L + 1 R y + 2 v 0 2 1 s L + 1 Q ) n n 1 i = 1 ∑ n E [ ∣ r i ( θ t ) ∣ 2 ] [Jensen] ≤ ( 2 a L + 1 R y + 2 1 v 0 2 s L + 1 Q ) n 1 n ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ] \le\left(2a_{L+1}R_y+2\frac{1}{v_0^2}s_{L+1}Q\right)n\sqrt{\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}\left[|r_i(\theta_t)|^2\right]} ≤ ( 2 a L + 1 R y + 2 v 0 2 1 s L + 1 Q ) n n 1 ∑ i = 1 n E [ ∣ r i ( θ t ) ∣ 2 ]
≤ 4 n E [ L ( θ t ) ] ( a L + 1 R y + 1 v 0 2 s L + 1 Q ) , \le4n\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}\left(a_{L+1}R_y+\frac{1}{v_0^2}s_{L+1}Q\right), ≤ 4 n E [ L ( θ t )] ( a L + 1 R y + v 0 2 1 s L + 1 Q ) , which yields
− β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ≤ 4 β − 1 n E [ L ( θ t ) ] ( a L + 1 R y + 1 v 0 2 s L + 1 Q ) -\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^\top\nabla f(\theta_t)\Lambda\theta_t\right]\le4\beta^{-1}n\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}\left(a_{L+1}R_y+\frac{1}{v_0^2}s_{L+1}Q\right) − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ≤ 4 β − 1 n E [ L ( θ t )] ( a L + 1 R y + v 0 2 1 s L + 1 Q ) − 2 β − 1 n a L + 1 E [ L ( θ t ) ] . -2\beta^{-1}na_{L+1}\mathbb{E}[\mathcal{L}(\theta_t)]. − 2 β − 1 n a L + 1 E [ L ( θ t )] . Let us find such Q Q Q i ∈ [ n ] i\in[n] i ∈ [ n ] F ≡ σ ( W ( 1 ) , … , W ( L ) ) \mathcal{F}\equiv\sigma\left(\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(L)}\right) F ≡ σ ( W ( 1 ) , … , W ( L ) ) σ \sigma σ σ ( W ( 1 ) , … , W ( L ) ) \sigma\left(\mathbf{W}^{(1)},\ldots,\mathbf{W}^{(L)}\right) σ ( W ( 1 ) , … , W ( L ) )
v ⊤ ( ∑ ℓ = 1 L a ℓ h ( ℓ ) ( x i ) ) ∣ F ∼ N ( 0 , 1 a L + 1 ∥ ∑ ℓ = 1 L a ℓ h ( ℓ ) ( x i ) ∥ 2 ) . \mathbf{v}^{\top}\left(\sum_{\ell=1}^{L}a_\ell\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right)\mid\mathcal{F}\sim\mathcal{N}\left(0,\frac{1}{a_{L+1}}\left\|\sum_{\ell=1}^{L}a_\ell\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right\|^2\right). v ⊤ ( ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) ) ∣ F ∼ N 0 , a L + 1 1 ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) 2 . Therefore,
∥ v ⊤ ( ∑ ℓ = 1 L a ℓ h ( ℓ ) ( x i ) ) ∥ L p ∣ F ≤ 1 a L + 1 ∥ ∑ ℓ = 1 L a ℓ h ( ℓ ) ( x i ) ∥ p . \left\|\mathbf{v}^{\top}\left(\sum_{\ell=1}^{L}a_\ell\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right)\right\|_{L_p|\mathcal{F}}\le\frac{1}{\sqrt{a_{L+1}}}\left\|\sum_{\ell=1}^{L}a_\ell\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right\|\sqrt{p}. v ⊤ ( ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) ) L p ∣ F ≤ a L + 1 1 ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) p . Using Claim E.4, applying the triangle inequality and applying Minkowski’s inequality, we get
∥ v ⊤ ( ∑ ℓ = 1 L a ℓ h ( ℓ ) ( x i ) ) ∥ L p ≤ 1 a L + 1 p ∥ ∥ ∑ ℓ = 1 L a ℓ h ( ℓ ) ( x i ) ∥ ∥ L p ≤ 1 a L + 1 p ∑ ℓ = 1 L a ℓ ∥ ∥ h ( ℓ ) ( x i ) ∥ ∥ L p ≤ 1 a L + 1 p ∑ ℓ = 1 L a ℓ ∥ ∥ h ( ℓ ) ( x i ) ∥ ∥ L p . \begin{aligned}
\left\|\mathbf{v}^{\top}\left(\sum_{\ell=1}^{L}a_\ell\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right)\right\|_{L_p}
\le\frac{1}{\sqrt{a_{L+1}}}\sqrt{p}\left\|\left\|\sum_{\ell=1}^{L}a_\ell\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right\|\right\|_{L_p} \\
&\le\frac{1}{\sqrt{a_{L+1}}}\sqrt{p}\sum_{\ell=1}^{L}a_\ell\left\|\left\|\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right\|\right\|_{L_p} \\
&\le\frac{1}{\sqrt{a_{L+1}}}\sqrt{p}\sum_{\ell=1}^{L}a_\ell\left\|\left\|\mathbf{h}^{(\ell)}(\mathbf{x}_i)\right\|\right\|_{L_p}.
\end{aligned} v ⊤ ( ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) ) L p ≤ a L + 1 1 p ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) L p ≤ a L + 1 1 p ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) L p ≤ a L + 1 1 p ℓ = 1 ∑ L a ℓ h ( ℓ ) ( x i ) L p . Observe that for every ℓ ∈ [ L ] \ell\in[L] ℓ ∈ [ L ]
h ( ℓ ) ( x i ) = ∂ α ( L ) ( x i ) ∂ α ~ ( ℓ ) ( x i ) α ~ ( ℓ ) ( x i ) = ( ∏ t = ℓ + 1 L ( ∂ α ( t ) ( x i ) ∂ α ( t − 1 ) ( x i ) ) ⊤ ) ∂ α ( L ) ( x i ) ∂ α ~ ( ℓ ) ( x i ) α ~ ( ℓ ) ( x i ) = ( ∏ t = ℓ + 1 L s t ( W ( t ) ) ⊤ diag ( ϕ ′ ( α ~ ( t ) ( x i ) ) ) ) ⊤ diag ( ϕ ′ ( α ~ ( ℓ ) ( x i ) ) ) α ~ ( ℓ ) ( x i ) . \begin{aligned}
\mathbf{h}^{(\ell)}(\mathbf{x}_i)=\frac{\partial\boldsymbol{\alpha}^{(L)}(\mathbf{x}_i)}{\partial\widetilde{\boldsymbol{\alpha}}^{(\ell)}(\mathbf{x}_i)}\widetilde{\boldsymbol{\alpha}}^{(\ell)}(\mathbf{x}_i) \\
&=\left(\prod_{t=\ell+1}^{L}\left(\frac{\partial\boldsymbol{\alpha}^{(t)}(\mathbf{x}_i)}{\partial\boldsymbol{\alpha}^{(t-1)}(\mathbf{x}_i)}\right)^{\top}\right)\frac{\partial\boldsymbol{\alpha}^{(L)}(\mathbf{x}_i)}{\partial\widetilde{\boldsymbol{\alpha}}^{(\ell)}(\mathbf{x}_i)}\widetilde{\boldsymbol{\alpha}}^{(\ell)}(\mathbf{x}_i) \\
&=\left(\prod_{t=\ell+1}^{L}s_t\left(\mathbf{W}^{(t)}\right)^{\top}\operatorname{diag}\left(\phi'\left(\widetilde{\boldsymbol{\alpha}}^{(t)}(\mathbf{x}_i)\right)\right)\right)^{\top}\operatorname{diag}\left(\phi'\left(\widetilde{\boldsymbol{\alpha}}^{(\ell)}(\mathbf{x}_i)\right)\right)\widetilde{\boldsymbol{\alpha}}^{(\ell)}(\mathbf{x}_i).
\end{aligned} h ( ℓ ) ( x i ) = ∂ α ( ℓ ) ( x i ) ∂ α ( L ) ( x i ) α ( ℓ ) ( x i ) = ( t = ℓ + 1 ∏ L ( ∂ α ( t − 1 ) ( x i ) ∂ α ( t ) ( x i ) ) ⊤ ) ∂ α ( ℓ ) ( x i ) ∂ α ( L ) ( x i ) α ( ℓ ) ( x i ) = ( t = ℓ + 1 ∏ L s t ( W ( t ) ) ⊤ diag ( ϕ ′ ( α ( t ) ( x i ) ) ) ) ⊤ diag ( ϕ ′ ( α ( ℓ ) ( x i ) ) ) α ( ℓ ) ( x i ) . Therefore,
∥ h ( l ) ( x i ) ∥ ≤ α ϕ L − l + 1 ( ∏ t = l + 1 L s t ∥ W ( t ) ∥ ) ∥ α ( l ) ( x i ) ∥ ≤ α ϕ L − l + 1 ( ∏ t = l L s t ∥ W ( t ) ∥ ) ∥ α ( l − 1 ) ( x i ) ∥ . \begin{aligned}
\left\|\mathbf{h}^{(l)}(x_i)\right\| \le\alpha_{\phi}^{L-l+1}\left(\prod_{t=l+1}^{L} s_t\left\|\mathbf{W}^{(t)}\right\|\right)\left\|\boldsymbol{\alpha}^{(l)}(x_i)\right\| \\
&\le\alpha_{\phi}^{L-l+1}\left(\prod_{t=l}^{L} s_t\left\|\mathbf{W}^{(t)}\right\|\right)\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|.
\end{aligned} h ( l ) ( x i ) ≤ α ϕ L − l + 1 ( t = l + 1 ∏ L s t W ( t ) ) α ( l ) ( x i ) ≤ α ϕ L − l + 1 ( t = l ∏ L s t W ( t ) ) α ( l − 1 ) ( x i ) . Since ∥ α ( l − 1 ) ( x i ) ∥ \left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\| α ( l − 1 ) ( x i )
∥ ∥ h ( l ) ( x i ) ∥ ∥ L p ≤ α ϕ L − l + 1 ( ∏ t = l L s t ∥ ∥ W ( t ) ∥ ∥ L p ) ∥ ∥ α ( l − 1 ) ( x i ) ∥ ∥ L p . \left\|\left\|\mathbf{h}^{(l)}(x_i)\right\|\right\|_{L^p} \le\alpha_{\phi}^{L-l+1}\left(\prod_{t=l}^{L} s_t\left\|\left\|\mathbf{W}^{(t)}\right\|\right\|_{L^p}\right)\left\|\left\|\boldsymbol{\alpha}^{(l-1)}(x_i)\right\|\right\|_{L^p}. h ( l ) ( x i ) L p ≤ α ϕ L − l + 1 ( t = l ∏ L s t W ( t ) L p ) α ( l − 1 ) ( x i ) L p . Using Propositions H.1 and H.3, we get that there exists a constant c > 0 c>0 c > 0
∥ ∥ h ( l ) ( x i ) ∥ ∥ L p ≤ ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + p ) ) . \left\|\left\|\mathbf{h}^{(l)}(x_i)\right\|\right\|_{L^p} \le\sqrt{\left(c\alpha_{\phi}^{2}v_0^2\right)^L\left\|x_i\right\|^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)}. h ( l ) ( x i ) L p ≤ ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + p ) ) . Thus, we get that
∥ ∥ v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) ∥ ∥ L p ≤ 1 a L + 1 p ∑ l = 1 L a l ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + p ) ) = ∑ l = 1 L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a l 2 1 a L + 1 p ( ∏ t = 1 l − 1 s t 2 a t ( m t + p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + p ) ) = L 1 L ∑ l = 1 L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a l 2 1 a L + 1 p ( ∏ t = 1 l − 1 s t 2 a t ( m t + p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + p ) ) ≤ L 1 L ∑ l = 1 L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a l 2 1 a L + 1 p ( ∏ t = 1 l − 1 s t 2 a t ( m t + p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + p ) ) = L 1 / 2 ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 1 a L + 1 p ∑ l = 1 L a l 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + p ) ) , \begin{aligned}
\left\|\left\|\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_l\mathbf{h}^{(l)}(x_i)\right)\right\|\right\|_{L^p} \\
&\le\frac{1}{\sqrt{a_{L+1}}}\sqrt{p}\sum_{l=1}^{L}a_l\sqrt{\left(c\alpha_{\phi}^{2}v_0^2\right)^L\left\|x_i\right\|^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)} \\
&= \sum_{l=1}^{L}\sqrt{\left(c\alpha_{\phi}^{2}v_0^2\right)^L\left\|x_i\right\|^2a_l^2\frac{1}{a_{L+1}}p\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)} \\
&= L\frac{1}{L}\sum_{l=1}^{L}\sqrt{\left(c\alpha_{\phi}^{2}v_0^2\right)^L\left\|x_i\right\|^2a_l^2\frac{1}{a_{L+1}}p\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)} \\
&\le L\sqrt{\frac{1}{L}\sum_{l=1}^{L}\left(c\alpha_{\phi}^{2}v_0^2\right)^L\left\|x_i\right\|^2a_l^2\frac{1}{a_{L+1}}p\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)} \\
&= L^{1/2}\sqrt{\left(c\alpha_{\phi}^{2}v_0^2\right)^L\left\|x_i\right\|^2\frac{1}{a_{L+1}}p\sum_{l=1}^{L}a_l^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+p)\right)},
\end{aligned} v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) L p ≤ a L + 1 1 p l = 1 ∑ L a l ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + p ) ) = l = 1 ∑ L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a l 2 a L + 1 1 p ( t = 1 ∏ l − 1 a t s t 2 ( m t + p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + p ) ) = L L 1 l = 1 ∑ L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a l 2 a L + 1 1 p ( t = 1 ∏ l − 1 a t s t 2 ( m t + p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + p ) ) ≤ L L 1 l = 1 ∑ L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a l 2 a L + 1 1 p ( t = 1 ∏ l − 1 a t s t 2 ( m t + p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + p ) ) = L 1/2 ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a L + 1 1 p l = 1 ∑ L a l 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + p ) ) , where in ( ∗ ) (*) ( ∗ )
∥ ( v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) ) 2 ∥ L p ≤ ∥ v T ( ∑ l = 1 L a l h ( l ) ( x i ) ) ∥ L 2 p 2 ≤ L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 1 a L + 1 2 p ∑ l = 1 L a l 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + 2 p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + 2 p ) ) ≤ c ϕ , σ 0 L ( c α ϕ 2 v 0 2 ) L d 1 a L + 1 2 p ∑ l = 1 L a l 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + 2 p ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + 2 p ) ) . \begin{aligned}
\left\|\left(\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_l\mathbf{h}^{(l)}(x_i)\right)\right)^2\right\|_{L^p} \\
&\le\left\|\mathbf{v}^{\mathsf{T}}\left(\sum_{l=1}^{L}a_l\mathbf{h}^{(l)}(x_i)\right)\right\|_{L^{2p}}^2 \\
&\le L\left(c\alpha_{\phi}^{2}v_0^2\right)^L\left\|x_i\right\|^2\frac{1}{a_{L+1}}2p\sum_{l=1}^{L}a_l^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+2p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+2p)\right) \\
&\le c_{\phi,\sigma_0}L\left(c\alpha_{\phi}^{2}v_0^2\right)^L d\frac{1}{a_{L+1}}2p\sum_{l=1}^{L}a_l^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}(m_t+2p)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}(m_{t-1}+m_t+2p)\right).
\end{aligned} ( v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) ) 2 L p ≤ v T ( l = 1 ∑ L a l h ( l ) ( x i ) ) L 2 p 2 ≤ L ( c α ϕ 2 v 0 2 ) L ∥ x i ∥ 2 a L + 1 1 2 p l = 1 ∑ L a l 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + 2 p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + 2 p ) ) ≤ c ϕ , σ 0 L ( c α ϕ 2 v 0 2 ) L d a L + 1 1 2 p l = 1 ∑ L a l 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + 2 p ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + 2 p ) ) . Together with Eq. (39) , we get
E p t [ ( v ⊤ ( ∑ l = 1 L a l h ( l ) ( x i ) ) ) 2 ] ≤ η c ϕ , σ 0 L ( c α ϕ 2 ν 0 2 ) L d a L + 1 2 c β κ ˉ 2 β ⋅ ∑ l = 1 L a l 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + 2 c β κ ˉ 2 β ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + 2 c β κ ˉ 2 β ) ) . \mathbb{E}_{p_t}\left[\left(\mathbf{v}^{\top}\left(\sum_{l=1}^{L}a_l\mathbf{h}^{(l)}(\mathbf{x}_i)\right)\right)^2\right]
\leq\eta c_{\phi,\sigma_0}L\left(c\alpha_{\phi}^{2}\nu_0^{2}\right)^L\frac{d}{a_{L+1}}2c_{\beta}\bar{\kappa}^{2}\beta
\cdot\sum_{l=1}^{L}a_l^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}\left(m_t+2c_{\beta}\bar{\kappa}^{2}\beta\right)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}\left(m_{t-1}+m_t+2c_{\beta}\bar{\kappa}^{2}\beta\right)\right). E p t ( v ⊤ ( l = 1 ∑ L a l h ( l ) ( x i ) ) ) 2 ≤ η c ϕ , σ 0 L ( c α ϕ 2 ν 0 2 ) L a L + 1 d 2 c β κ ˉ 2 β ⋅ l = 1 ∑ L a l 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + 2 c β κ ˉ 2 β ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + 2 c β κ ˉ 2 β ) ) . Thus, the right-hand side is a valid choice for Q 2 Q^2 Q 2
-\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}\nabla f(\theta_t)\Lambda\theta_t\right]
\lesssim\beta^{-1}n\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}
\cdot\left(a_{L+1}R_y+\frac{1}{\nu_0^2}s_{L+1}\sqrt{\eta L\left(c\alpha_{\phi}^{2}\nu_0^{2}\right)^L\frac{d}{a_{L+1}}2c_{\beta}\bar{\kappa}^{2}\beta
\cdot\sqrt{\sum_{l=1}^{L}a_l^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}\left(m_t+2c_{\beta}\bar{\kappa}^{2}\beta\right)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}\left(m_{t-1}+m_t+2c_{\beta}\bar{\kappa}^{2}\beta\right)\right)}\right)
-2\beta^{-1}na_{L+1}\mathbb{E}[\mathcal{L}(\theta_t)].}
Bounding the entire right-hand side
Recall that
d d t E [ L ( θ t ) ] = 1 n ( − n − 1 E [ r ( θ t ) ⊤ K N T K ( θ t ) r ( θ t ) ] ⏟ T 1 + β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] ⏟ T 2 + β − 1 E [ ∑ i = 1 n r i ( θ t ) Δ f i ( θ t ) ] ⏟ T 3 − β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] ⏟ T 4 ) . \frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}[\mathcal{L}(\theta_t)]
=
\frac{1}{n}\left(
\underbrace{-n^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}\mathbf{K}_{\mathrm{NTK}}(\theta_t)\mathbf{r}(\theta_t)\right]}_{T_1}
+
\underbrace{\beta^{-1}\mathbb{E}\left[\operatorname{Tr}\left(\nabla f(\theta_t)^{\top}\nabla f(\theta_t)\right)\right]}_{T_2}
+
\underbrace{\beta^{-1}\mathbb{E}\left[\sum_{i=1}^{n}r_i(\theta_t)\Delta f_i(\theta_t)\right]}_{T_3}
-
\underbrace{\beta^{-1}\mathbb{E}\left[\mathbf{r}(\theta_t)^{\top}\nabla f(\theta_t)\Lambda\theta_t\right]}_{T_4}
\right). d t d E [ L ( θ t )] = n 1 T 1 − n − 1 E [ r ( θ t ) ⊤ K NTK ( θ t ) r ( θ t ) ] + T 2 β − 1 E [ Tr ( ∇ f ( θ t ) ⊤ ∇ f ( θ t ) ) ] + T 3 β − 1 E [ i = 1 ∑ n r i ( θ t ) Δ f i ( θ t ) ] − T 4 β − 1 E [ r ( θ t ) ⊤ ∇ f ( θ t ) Λ θ t ] . We have shown that there exists a constant C > 0 C>0 C > 0 C C C
T 1 ≲ − λ 1 s L + 1 2 m E [ L ( θ t ) ] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 2 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) ; T_1 \lesssim-\lambda_1s_{L+1}^2m\mathbb{E}\left[\mathcal{L}(\theta_t)\right]+\eta\lambda_1s_{L+1}^2m\left(1+c_\beta\bar{\kappa}^2\beta\right)\exp\left(-\widetilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right); T 1 ≲ − λ 1 s L + 1 2 m E [ L ( θ t ) ] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 2 β ) exp ( − c 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) ; T 2 ≲ C 2 L η β − 1 n d s L + 1 2 ( m L + c β κ ˉ 2 β ) ⋅ ( s L 2 a L ∏ k = 1 L − 1 s k 2 a k ( m k + c β κ ˉ 2 β ) + 1 a L + 1 ∑ l = 1 L ( ∏ k = 1 l − 1 s k 2 a k ( m k + c β κ ˉ 2 β ) ) s l 2 ( ∏ k = l + 1 L s k 2 a k ( m k − 1 + m k + c β κ ˉ 2 β ) ) ) ; \begin{aligned}
T_2\lesssim{}&C^{2L}\eta\beta^{-1}nds_{L+1}^2\left(m_L+c_\beta\bar{\kappa}^2\beta\right)\\
&\cdot\left(\frac{s_L^2}{a_L}\prod_{k=1}^{L-1}\frac{s_k^2}{a_k}\left(m_k+c_\beta\bar{\kappa}^2\beta\right)\right.\\
&\qquad\left.+\frac{1}{a_{L+1}}\sum_{l=1}^{L}\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}\left(m_k+c_\beta\bar{\kappa}^2\beta\right)\right)s_l^2\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}\left(m_{k-1}+m_k+c_\beta\bar{\kappa}^2\beta\right)\right)\right);
\end{aligned} T 2 ≲ C 2 L η β − 1 n d s L + 1 2 ( m L + c β κ ˉ 2 β ) ⋅ ( a L s L 2 k = 1 ∏ L − 1 a k s k 2 ( m k + c β κ ˉ 2 β ) + a L + 1 1 l = 1 ∑ L ( k = 1 ∏ l − 1 a k s k 2 ( m k + c β κ ˉ 2 β ) ) s l 2 ( k = l + 1 ∏ L a k s k 2 ( m k − 1 + m k + c β κ ˉ 2 β ) ) ) ; T 3 ≲ C 3 L η β − 1 n d ⋅ ∑ l = 1 L ∑ q = l L ( ∏ t = 1 l − 1 s t + 1 2 a t + 1 ( m t + m t + 1 + 4 c β κ ˉ 2 β ) ) s l 2 ⋅ ( ∏ t = l q − 1 s t + 1 2 a t + 1 ( m t + m t + 1 + 4 c β κ ˉ 2 β ) ) ⋅ m q ( ∏ t = q L − 1 s t + 1 2 a t + 1 ( m t + m t + 1 + 4 c β κ ˉ 2 β ) ) 1 / 2 ⋅ s L + 1 2 a L + 1 c β κ ˉ 2 β E [ L ( θ t ) ] ; \begin{aligned}
T_3\lesssim{}&C^{3L}\sqrt{\eta\beta^{-1}nd}\\
&\cdot\sum_{l=1}^{L}\sum_{q=l}^{L}\left(\prod_{t=1}^{l-1}\frac{s_{t+1}^2}{a_{t+1}}\left(m_t+m_{t+1}+4c_\beta\bar{\kappa}^2\beta\right)\right)s_l^2\\
&\qquad\cdot\left(\prod_{t=l}^{q-1}\frac{s_{t+1}^2}{a_{t+1}}\left(m_t+m_{t+1}+4c_\beta\bar{\kappa}^2\beta\right)\right)\\
&\qquad\cdot\sqrt{m_q}\left(\prod_{t=q}^{L-1}\frac{s_{t+1}^2}{a_{t+1}}\left(m_t+m_{t+1}+4c_\beta\bar{\kappa}^2\beta\right)\right)^{1/2}\\
&\qquad\cdot\sqrt{\frac{s_{L+1}^2}{a_{L+1}}c_\beta\bar{\kappa}^2\beta\,\mathbb{E}\left[\mathcal{L}(\theta_t)\right]};
\end{aligned} T 3 ≲ C 3 L η β − 1 n d ⋅ l = 1 ∑ L q = l ∑ L ( t = 1 ∏ l − 1 a t + 1 s t + 1 2 ( m t + m t + 1 + 4 c β κ ˉ 2 β ) ) s l 2 ⋅ ( t = l ∏ q − 1 a t + 1 s t + 1 2 ( m t + m t + 1 + 4 c β κ ˉ 2 β ) ) ⋅ m q ( t = q ∏ L − 1 a t + 1 s t + 1 2 ( m t + m t + 1 + 4 c β κ ˉ 2 β ) ) 1/2 ⋅ a L + 1 s L + 1 2 c β κ ˉ 2 β E [ L ( θ t ) ] ; T 4 ≲ β − 1 n E [ L ( θ t ) ] ⋅ ( a L + 1 + s L + 1 η L C L d 1 a L + 1 2 c β κ ˉ 2 β ⋅ ( ∑ l = 1 L a l 2 ( ∏ t = 1 l − 1 s t 2 a t ( m t + 2 c β κ ˉ 2 β ) ) ) ( ∏ t = l L s t 2 a t ( m t − 1 + m t + 2 c β κ ˉ 2 β ) ) ) − β − 1 n a L + 1 E [ L ( θ t ) ] . \begin{aligned}
T_4\lesssim{}&\beta^{-1}n\sqrt{\mathbb{E}\left[\mathcal{L}(\theta_t)\right]}\\
&\cdot\left(a_{L+1}+s_{L+1}\sqrt{\eta L C^L d\frac{1}{a_{L+1}}2c_\beta\bar{\kappa}^2\beta}\right.\\
&\qquad\left.\cdot\sqrt{\left(\sum_{l=1}^{L}a_l^2\left(\prod_{t=1}^{l-1}\frac{s_t^2}{a_t}\left(m_t+2c_\beta\bar{\kappa}^2\beta\right)\right)\right)\left(\prod_{t=l}^{L}\frac{s_t^2}{a_t}\left(m_{t-1}+m_t+2c_\beta\bar{\kappa}^2\beta\right)\right)}\right)\\
&\qquad-\beta^{-1}na_{L+1}\mathbb{E}\left[\mathcal{L}(\theta_t)\right].
\end{aligned} T 4 ≲ β − 1 n E [ L ( θ t ) ] ⋅ ( a L + 1 + s L + 1 η L C L d a L + 1 1 2 c β κ ˉ 2 β ⋅ ( l = 1 ∑ L a l 2 ( t = 1 ∏ l − 1 a t s t 2 ( m t + 2 c β κ ˉ 2 β ) ) ) ( t = l ∏ L a t s t 2 ( m t − 1 + m t + 2 c β κ ˉ 2 β ) ) − β − 1 n a L + 1 E [ L ( θ t ) ] . From now on we shall assume m 1 = ⋯ = m L = m m_1=\cdots=m_L=m m 1 = ⋯ = m L = m d ≤ m d\leq m d ≤ m m > 4 c β κ ˉ 2 β m>4c_\beta\bar{\kappa}^2\beta m > 4 c β κ ˉ 2 β L L L C L C^L C L κ = s L + 1 2 m a L + 1 \kappa=\sqrt{\frac{s_{L+1}^2m}{a_{L+1}}} κ = a L + 1 s L + 1 2 m
T 1 ≲ − λ 1 s L + 1 2 m E [ L ( θ t ) ] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) ; T 2 ≲ η β − 1 n d s L + 1 2 m ( s L 2 a L ∏ k = 1 L − 1 s k 2 a k m + 1 a L + 1 ∑ l = 1 L ( ∏ k = 1 l − 1 s k 2 a k m ) s l 2 ( ∏ k = l + 1 L s k 2 a k m ) ) ; T 3 ≲ η c β κ ˉ 2 κ ˉ 2 β − 1 / 2 n d ⋅ ∑ l = 1 L ( ∏ t = 0 l − 2 s t + 1 2 a t + 1 m ) s l 2 ∑ q = l L ( ∏ t = l q − 1 s t + 1 2 a t + 1 m ) ( ∏ t = q L − 1 s t + 1 2 a t + 1 m ) E [ L ( θ t ) ] ; T 4 ≲ β − 1 n ( a L + 1 + η c β d κ 2 κ ˉ 2 β m ( ∏ t = 1 L s t 2 a t m ) ∑ l = 1 L a l 2 E [ L ( θ t ) ] − β − 1 n a L + 1 E [ L ( θ t ) ] ) . \begin{aligned}
T_1 &\lesssim-\lambda_1s_{L+1}^2m\mathbb{E}[\mathcal{L}(\theta_t)] + \eta\lambda_1s_{L+1}^2m\left(1+c_\beta\bar{\kappa}^4\beta\right)\exp\left(-\tilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right);\\
T_2 &\lesssim\eta\beta^{-1}nds_{L+1}^2m\left(\frac{s_L^2}{a_L}\prod_{k=1}^{L-1}\frac{s_k^2}{a_k}m+\frac{1}{a_{L+1}}\sum_{l=1}^{L}\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}m\right)s_l^2\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}m\right)\right);\\
T_3 &\lesssim\sqrt{\eta c_\beta\bar{\kappa}^2\bar{\kappa}^2\beta^{-1/2}nd}\\
&\qquad\cdot\sum_{l=1}^{L}\left(\prod_{t=0}^{l-2}\frac{s_{t+1}^2}{a_{t+1}}m\right)s_l^2\sum_{q=l}^{L}\left(\prod_{t=l}^{q-1}\frac{s_{t+1}^2}{a_{t+1}}m\right)\left(\prod_{t=q}^{L-1}\sqrt{\frac{s_{t+1}^2}{a_{t+1}}m}\right)\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]};\\
T_4 &\lesssim\beta^{-1}n\left(a_{L+1}+\sqrt{\eta c_\beta d\kappa^2\bar{\kappa}^2\frac{\beta}{m}}\sqrt{\left(\prod_{t=1}^{L}\frac{s_t^2}{a_t}m\right)\sum_{l=1}^{L}a_l^2}\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}\right.\\
&\qquad\left.-\beta^{-1}na_{L+1}\mathbb{E}[\mathcal{L}(\theta_t)]\right).
\end{aligned} T 1 T 2 T 3 T 4 ≲ − λ 1 s L + 1 2 m E [ L ( θ t )] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) ; ≲ η β − 1 n d s L + 1 2 m ( a L s L 2 k = 1 ∏ L − 1 a k s k 2 m + a L + 1 1 l = 1 ∑ L ( k = 1 ∏ l − 1 a k s k 2 m ) s l 2 ( k = l + 1 ∏ L a k s k 2 m ) ) ; ≲ η c β κ ˉ 2 κ ˉ 2 β − 1/2 n d ⋅ l = 1 ∑ L ( t = 0 ∏ l − 2 a t + 1 s t + 1 2 m ) s l 2 q = l ∑ L ( t = l ∏ q − 1 a t + 1 s t + 1 2 m ) t = q ∏ L − 1 a t + 1 s t + 1 2 m E [ L ( θ t )] ; ≲ β − 1 n a L + 1 + η c β d κ 2 κ ˉ 2 m β ( t = 1 ∏ L a t s t 2 m ) l = 1 ∑ L a l 2 E [ L ( θ t )] − β − 1 n a L + 1 E [ L ( θ t )] ) . Substituting these bounds, we get
d d t E [ L ( θ t ) ] ≲ 1 n ( − λ 1 s L + 1 2 m E [ L ( θ t ) ] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) + η β − 1 n d s L + 1 2 m ( s L 2 a L ∏ k = 1 L − 1 s k 2 a k m + 1 a L + 1 ∑ l = 1 L ( ∏ k = 1 l − 1 s k 2 a k m ) s l 2 ( ∏ k = l + 1 L s k 2 a k m ) ) + η c β κ ˉ 2 κ ˉ 2 β − 1 / 2 n d ∑ l = 1 L ( ∏ t = 0 l − 2 s t + 1 2 a t + 1 m ) s l 2 ∑ q = l L ( ∏ t = l q − 1 s t + 1 2 a t + 1 m ) ( ∏ t = q L − 1 s t + 1 2 a t + 1 m ) E [ L ( θ t ) ] + β − 1 n ( a L + 1 + η c β d κ 2 κ ˉ 2 β m ( ∏ t = 1 L s t 2 a t m ) ∑ l = 1 L a l 2 E [ L ( θ t ) ] − β − 1 n a L + 1 E [ L ( θ t ) ] ) ) . (40) \begin{aligned}
\frac{\mathrm{d}}{\mathrm{d}t}\mathbb{E}[\mathcal{L}(\theta_t)]
\lesssim\frac{1}{n}\bigg(&-\lambda_1s_{L+1}^2m\mathbb{E}[\mathcal{L}(\theta_t)]\\
&+\eta\lambda_1s_{L+1}^2m\left(1+c_\beta\bar{\kappa}^4\beta\right)\exp\left(-\tilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right)\\
&+\eta\beta^{-1}nds_{L+1}^2m\left(\frac{s_L^2}{a_L}\prod_{k=1}^{L-1}\frac{s_k^2}{a_k}m+\frac{1}{a_{L+1}}\sum_{l=1}^{L}\left(\prod_{k=1}^{l-1}\frac{s_k^2}{a_k}m\right)s_l^2\left(\prod_{k=l+1}^{L}\frac{s_k^2}{a_k}m\right)\right)\\
&+\sqrt{\eta c_\beta\bar{\kappa}^2\bar{\kappa}^2\beta^{-1/2}nd}\sum_{l=1}^{L}\left(\prod_{t=0}^{l-2}\frac{s_{t+1}^2}{a_{t+1}}m\right)s_l^2\sum_{q=l}^{L}\left(\prod_{t=l}^{q-1}\frac{s_{t+1}^2}{a_{t+1}}m\right)\left(\prod_{t=q}^{L-1}\sqrt{\frac{s_{t+1}^2}{a_{t+1}}m}\right)\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}\\
&+\beta^{-1}n\left(a_{L+1}+\sqrt{\eta c_\beta d\kappa^2\bar{\kappa}^2\frac{\beta}{m}}\sqrt{\left(\prod_{t=1}^{L}\frac{s_t^2}{a_t}m\right)\sum_{l=1}^{L}a_l^2}\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}\right.\\
&\qquad\left.-\beta^{-1}na_{L+1}\mathbb{E}[\mathcal{L}(\theta_t)]\right)\bigg). \tag*{(40)}
\end{aligned} d t d E [ L ( θ t )] ≲ n 1 ( − λ 1 s L + 1 2 m E [ L ( θ t )] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) + η β − 1 n d s L + 1 2 m ( a L s L 2 k = 1 ∏ L − 1 a k s k 2 m + a L + 1 1 l = 1 ∑ L ( k = 1 ∏ l − 1 a k s k 2 m ) s l 2 ( k = l + 1 ∏ L a k s k 2 m ) ) + η c β κ ˉ 2 κ ˉ 2 β − 1/2 n d l = 1 ∑ L ( t = 0 ∏ l − 2 a t + 1 s t + 1 2 m ) s l 2 q = l ∑ L ( t = l ∏ q − 1 a t + 1 s t + 1 2 m ) t = q ∏ L − 1 a t + 1 s t + 1 2 m E [ L ( θ t )] + β − 1 n a L + 1 + η c β d κ 2 κ ˉ 2 m β ( t = 1 ∏ L a t s t 2 m ) l = 1 ∑ L a l 2 E [ L ( θ t )] − β − 1 n a L + 1 E [ L ( θ t )] ) ) . (40) From now on we shall also assume a l s l 2 = m \frac{a_l}{s_l^2}=m s l 2 a l = m 2 ≤ l ≤ L 2\leq l\leq L 2 ≤ l ≤ L s l 2 a l \frac{s_l^2}{a_l} a l s l 2 (40) satisfies
η c β κ ˉ 2 κ ˉ 2 β − 1 / 2 n d ∑ l = 1 L ( ∏ t = 0 l − 2 s t + 1 2 a t + 1 m ) s l 2 ∑ q = l L ( ∏ t = l q − 1 s t + 1 2 a t + 1 m ) ( ∏ t = q L − 1 s t + 1 2 a t + 1 m ) E [ L ( θ t ) ] = η c β κ ˉ 2 κ ˉ 2 β − 1 / 2 n d ( s 1 2 + s 1 2 a 1 m ∑ l = 2 L s l 2 ) E [ L ( θ t ) ] . \begin{aligned}
\sqrt{\eta c_\beta\bar{\kappa}^2\bar{\kappa}^2\beta^{-1/2}nd}\sum_{l=1}^{L}\left(\prod_{t=0}^{l-2}\frac{s_{t+1}^2}{a_{t+1}}m\right)s_l^2\sum_{q=l}^{L}\left(\prod_{t=l}^{q-1}\frac{s_{t+1}^2}{a_{t+1}}m\right)\left(\prod_{t=q}^{L-1}\sqrt{\frac{s_{t+1}^2}{a_{t+1}}m}\right)\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]} \\
&=\sqrt{\eta c_\beta\bar{\kappa}^2\bar{\kappa}^2\beta^{-1/2}nd}\left(s_1^2+\frac{s_1^2}{a_1}m\sum_{l=2}^{L}s_l^2\right)\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}.
\end{aligned} η c β κ ˉ 2 κ ˉ 2 β − 1/2 n d l = 1 ∑ L ( t = 0 ∏ l − 2 a t + 1 s t + 1 2 m ) s l 2 q = l ∑ L ( t = l ∏ q − 1 a t + 1 s t + 1 2 m ) t = q ∏ L − 1 a t + 1 s t + 1 2 m E [ L ( θ t )] = η c β κ ˉ 2 κ ˉ 2 β − 1/2 n d ( s 1 2 + a 1 s 1 2 m l = 2 ∑ L s l 2 ) E [ L ( θ t )] . This allows us to simplify to
≲ 1 n ( − λ 1 s L + 1 2 m E [ L ( θ t ) ] ⏟ N 1 + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) + η β − 1 n d s L + 1 2 m ( s 1 2 a 1 + 1 a L + 1 s 1 2 + s 1 2 a 1 m 1 a L + 1 ∑ l = 2 L s l 2 ) + η c β κ 2 κ ˉ 2 β − 1 n d ( s 1 2 + s 1 2 a 1 m ∑ l = 2 L s l 2 ) E [ L ( θ t ) ] + β − 1 n ( a L + 1 + η c β d κ 2 κ ˉ 2 β m s 1 2 a 1 m ∑ l = 1 L a l 2 ) E [ L ( θ t ) ] − β − 1 n a L + 1 E [ L ( θ t ) ] ⏟ N 2 ) . \lesssim\frac{1}{n}\left(
\underbrace{-\lambda_1 s_{L+1}^2 m \mathbb{E}[\mathcal{L}(\theta_t)]}_{N_1}
+ \eta\lambda_1 s_{L+1}^2 m\left(1+c_\beta\bar{\kappa}^4\beta\right)\exp\left(-\tilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right)
+ \eta\beta^{-1}nd s_{L+1}^2m\left(\frac{s_1^2}{a_1}+\frac{1}{a_{L+1}}s_1^2+\frac{s_1^2}{a_1}m\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_l^2\right)
+ \sqrt{\eta c_\beta\kappa^2\bar{\kappa}^2\beta^{-1}}nd\left(s_1^2+\frac{s_1^2}{a_1}m\sum_{l=2}^{L}s_l^2\right)\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}
+ \beta^{-1}n\left(a_{L+1}+\sqrt{\eta c_\beta d\kappa^2\bar{\kappa}^2\frac{\beta}{m}}\sqrt{\frac{s_1^2}{a_1}m\sum_{l=1}^{L}a_l^2}\right)\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}
\underbrace{-\beta^{-1}na_{L+1}\mathbb{E}[\mathcal{L}(\theta_t)]}_{N_2}
\right). ≲ n 1 N 1 − λ 1 s L + 1 2 m E [ L ( θ t )] + η λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) + η β − 1 n d s L + 1 2 m ( a 1 s 1 2 + a L + 1 1 s 1 2 + a 1 s 1 2 m a L + 1 1 l = 2 ∑ L s l 2 ) + η c β κ 2 κ ˉ 2 β − 1 n d ( s 1 2 + a 1 s 1 2 m l = 2 ∑ L s l 2 ) E [ L ( θ t )] + β − 1 n a L + 1 + η c β d κ 2 κ ˉ 2 m β a 1 s 1 2 m l = 1 ∑ L a l 2 E [ L ( θ t )] N 2 − β − 1 n a L + 1 E [ L ( θ t )] . The negative regularization contribution N 2 N_2 N 2 N 2 N_2 N 2 β − 1 n a L + 1 E [ L ( θ t ) ] \beta^{-1}na_{L+1}\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]} β − 1 n a L + 1 E [ L ( θ t )] u = E [ L ( θ t ) ] ≤ 1 u=\mathbb{E}[\mathcal{L}(\theta_t)]\le1 u = E [ L ( θ t )] ≤ 1
− β − 1 n a L + 1 u + β − 1 n a L + 1 u -\beta^{-1}na_{L+1}u+\beta^{-1}na_{L+1}\sqrt{u} − β − 1 n a L + 1 u + β − 1 n a L + 1 u is nonnegative. We therefore close the inequality using only the NTK descent term N 1 N_1 N 1 u \sqrt{u} u N 2 N_2 N 2
⋯ ≲ 1 n [ − λ 1 s L + 1 2 m E [ L ( θ t ) ] + β − 1 n ( η c β κ ˉ 2 κ ˉ 2 β 1 / 2 d ( s 1 2 + s 1 2 a 1 m ∑ l = 2 L s l 2 ) + a L + 1 + η c β κ 2 κ ˉ 2 β 3 / 2 d s 1 2 a 1 d ∑ l = 1 L a l 2 ) ⋅ E [ L ( θ t ) ] + η β − 1 n d s L + 1 2 m ( s 1 2 a 1 + 1 a L + 1 s 1 2 + s 1 2 a 1 m 1 a L + 1 ∑ l = 2 L s l 2 ) + λ 1 η s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) ] . \cdots\lesssim\frac{1}{n}\left[
-\lambda_1s_{L+1}^2m\mathbb{E}[\mathcal{L}(\theta_t)]
+\beta^{-1}n\left(
\sqrt{\eta c_\beta\bar{\kappa}^2\bar{\kappa}^2\beta^{1/2}}d\left(s_1^2+\frac{s_1^2}{a_1}m\sum_{l=2}^{L}s_l^2\right)
+a_{L+1}
+\sqrt{\eta c_\beta\kappa^2\bar{\kappa}^2\beta^{3/2}d\frac{s_1^2}{a_1}d\sum_{l=1}^{L}a_l^2}
\right)\cdot\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]}
+\eta\beta^{-1}nd s_{L+1}^2m\left(\frac{s_1^2}{a_1}+\frac{1}{a_{L+1}}s_1^2+\frac{s_1^2}{a_1}m\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_l^2\right)
+\lambda_1\eta s_{L+1}^2m\left(1+c_\beta\bar{\kappa}^4\beta\right)\exp\left(-\tilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right)
\right]. ⋯ ≲ n 1 − λ 1 s L + 1 2 m E [ L ( θ t )] + β − 1 n η c β κ ˉ 2 κ ˉ 2 β 1/2 d ( s 1 2 + a 1 s 1 2 m l = 2 ∑ L s l 2 ) + a L + 1 + η c β κ 2 κ ˉ 2 β 3/2 d a 1 s 1 2 d l = 1 ∑ L a l 2 ⋅ E [ L ( θ t )] + η β − 1 n d s L + 1 2 m ( a 1 s 1 2 + a L + 1 1 s 1 2 + a 1 s 1 2 m a L + 1 1 l = 2 ∑ L s l 2 ) + λ 1 η s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) . Using Young’s inequality a b ≤ a 2 2 ϵ + ϵ 2 b 2 ab \le\frac{a^2}{2\epsilon}+\frac{\epsilon}{2}b^2 ab ≤ 2 ϵ a 2 + 2 ϵ b 2 ϵ = λ 1 s L + 1 2 m \epsilon=\lambda_1s_{L+1}^2m ϵ = λ 1 s L + 1 2 m b = E [ L ( θ t ) ] b=\sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]} b = E [ L ( θ t )] a a a E [ L ( θ t ) ] \sqrt{\mathbb{E}[\mathcal{L}(\theta_t)]} E [ L ( θ t )]
⋯ ≲ 1 n [ − λ 1 s L + 1 2 m 2 E [ L ( θ t ) ] + β − 2 n 2 2 λ 1 s L + 1 2 m ( η c β κ 2 κ ˉ 2 β 1 / 2 d ( s 1 2 + s 1 2 a 1 m ∑ l = 2 L s l 2 ) + a L + 1 + η c β κ 2 κ ˉ 2 β s 1 2 a 1 d ∑ l = 1 L a l 2 ) 2 + η β − 1 n d s L + 1 2 m ( s 1 2 a 1 + 1 a L + 1 s 1 2 + s 1 2 a 1 m 1 a L + 1 ∑ l = 2 L s l 2 ) + λ 1 η s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) ] \begin{aligned}
\cdots&\lesssim\frac{1}{n}\Bigg[-\frac{\lambda_1s_{L+1}^2m}{2}\mathbb{E}[\mathcal{L}(\theta_t)]\\
&\quad+\frac{\beta^{-2}n^2}{2\lambda_1s_{L+1}^2m}\Bigg(\sqrt{\eta c_\beta\kappa^2\bar{\kappa}^2\beta^{1/2}d\left(s_1^2+\frac{s_1^2}{a_1}m\sum_{l=2}^{L}s_l^2\right)}\\
&\qquad\qquad+a_{L+1}+\sqrt{\eta c_\beta\kappa^2\bar{\kappa}^2\beta\frac{s_1^2}{a_1}d\sum_{l=1}^{L}a_l^2}\Bigg)^2\\
&\quad+\eta\beta^{-1}nds_{L+1}^2m\left(\frac{s_1^2}{a_1}+\frac{1}{a_{L+1}}s_1^2+\frac{s_1^2}{a_1}m\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_l^2\right)\\
&\quad+\lambda_1\eta s_{L+1}^2m(1+c_\beta\bar{\kappa}^4\beta)\exp\left(-\tilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right)\Bigg]
\end{aligned} ⋯ ≲ n 1 [ − 2 λ 1 s L + 1 2 m E [ L ( θ t )] + 2 λ 1 s L + 1 2 m β − 2 n 2 ( η c β κ 2 κ ˉ 2 β 1/2 d ( s 1 2 + a 1 s 1 2 m l = 2 ∑ L s l 2 ) + a L + 1 + η c β κ 2 κ ˉ 2 β a 1 s 1 2 d l = 1 ∑ L a l 2 ) 2 + η β − 1 n d s L + 1 2 m ( a 1 s 1 2 + a L + 1 1 s 1 2 + a 1 s 1 2 m a L + 1 1 l = 2 ∑ L s l 2 ) + λ 1 η s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) ] ≤ 1 n [ − λ 1 s L + 1 2 m E [ L ( θ t ) ] + β − 2 n 2 2 λ 1 s L + 1 2 m ( c β κ 2 κ ˉ 2 β d 2 ( s 1 4 + ( s 1 2 a 1 m ) 2 ( ∑ l = 2 L s l 2 ) 2 ) + a L + 1 2 + c β κ 2 κ ˉ 2 β s 1 2 a 1 d ∑ l = 1 L a l 2 ) + β − 1 n s L + 1 2 m ( s 1 2 a 1 d + 1 a L + 1 s 1 2 d + s 1 2 a 1 d m 1 a L + 1 ∑ l = 2 L s l 2 ) + λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) ] . \begin{aligned}
\le\frac{1}{n}\Bigg[&-\lambda_1s_{L+1}^2m\mathbb{E}[\mathcal{L}(\theta_t)]\\
&+\frac{\beta^{-2}n^2}{2\lambda_1s_{L+1}^2m}\Bigg(c_\beta\kappa^2\bar{\kappa}^2\beta d^2\left(s_1^4+\left(\frac{s_1^2}{a_1}m\right)^2\left(\sum_{l=2}^{L}s_l^2\right)^2\right)\\
&\qquad\qquad+a_{L+1}^2+c_\beta\kappa^2\bar{\kappa}^2\beta\frac{s_1^2}{a_1}d\sum_{l=1}^{L}a_l^2\Bigg)\\
&+\beta^{-1}ns_{L+1}^2m\left(\frac{s_1^2}{a_1}d+\frac{1}{a_{L+1}}s_1^2d+\frac{s_1^2}{a_1}dm\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_l^2\right)\\
&+\lambda_1s_{L+1}^2m(1+c_\beta\bar{\kappa}^4\beta)\exp\left(-\tilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right)\Bigg].
\end{aligned} ≤ n 1 [ − λ 1 s L + 1 2 m E [ L ( θ t )] + 2 λ 1 s L + 1 2 m β − 2 n 2 ( c β κ 2 κ ˉ 2 β d 2 s 1 4 + ( a 1 s 1 2 m ) 2 ( l = 2 ∑ L s l 2 ) 2 + a L + 1 2 + c β κ 2 κ ˉ 2 β a 1 s 1 2 d l = 1 ∑ L a l 2 ) + β − 1 n s L + 1 2 m ( a 1 s 1 2 d + a L + 1 1 s 1 2 d + a 1 s 1 2 d m a L + 1 1 l = 2 ∑ L s l 2 ) + λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) ] . Then, from Claim E.9 we get
E [ L ( θ t ) ] ≤ ( E [ L ( θ 0 ) ] − L ∞ ) exp ( − 1 2 λ 1 s L + 1 2 m n t ) + L ∞ , \mathbb{E}[\mathcal{L}(\theta_t)]\le\left(\mathbb{E}[\mathcal{L}(\theta_0)]-\mathcal{L}_\infty\right)\exp\left(-\frac{1}{2}\lambda_1s_{L+1}^2\frac{m}{n}t\right)+\mathcal{L}_\infty, E [ L ( θ t )] ≤ ( E [ L ( θ 0 )] − L ∞ ) exp ( − 2 1 λ 1 s L + 1 2 n m t ) + L ∞ , where
L ∞ ≲ 1 2 λ 1 s L + 1 2 m [ β − 2 n 2 2 λ 1 s L + 1 2 m ( c β κ 2 κ ˉ 2 β d 2 ( s 1 4 + ( s 1 2 a 1 m ) 2 ( ∑ l = 2 L s l 2 ) 2 ) + a L + 1 2 + c β κ 2 κ ˉ 2 β s 1 2 a 1 d ∑ l = 1 L a l 2 ) + β − 1 n s L + 1 2 m ( s 1 2 a 1 d + 1 a L + 1 s 1 2 d + s 1 2 a 1 d m 1 a L + 1 ∑ l = 2 L s l 2 ) + λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β n λ 1 ln n λ 1 ) − 1 ) ] . \begin{aligned}
\mathcal{L}_\infty\lesssim\frac{1}{2\lambda_1s_{L+1}^2m}\Bigg[&\frac{\beta^{-2}n^2}{2\lambda_1s_{L+1}^2m}\Bigg(c_\beta\kappa^2\bar{\kappa}^2\beta d^2\left(s_1^4+\left(\frac{s_1^2}{a_1}m\right)^2\left(\sum_{l=2}^{L}s_l^2\right)^2\right)\\
&\qquad+a_{L+1}^2+c_\beta\kappa^2\bar{\kappa}^2\beta\frac{s_1^2}{a_1}d\sum_{l=1}^{L}a_l^2\Bigg)\\
&+\beta^{-1}ns_{L+1}^2m\left(\frac{s_1^2}{a_1}d+\frac{1}{a_{L+1}}s_1^2d+\frac{s_1^2}{a_1}dm\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_l^2\right)\\
&+\lambda_1s_{L+1}^2m(1+c_\beta\bar{\kappa}^4\beta)\exp\left(-\tilde{c}_1m\left(c_\beta\bar{\kappa}^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right)\Bigg].
\end{aligned} L ∞ ≲ 2 λ 1 s L + 1 2 m 1 [ 2 λ 1 s L + 1 2 m β − 2 n 2 ( c β κ 2 κ ˉ 2 β d 2 s 1 4 + ( a 1 s 1 2 m ) 2 ( l = 2 ∑ L s l 2 ) 2 + a L + 1 2 + c β κ 2 κ ˉ 2 β a 1 s 1 2 d l = 1 ∑ L a l 2 ) + β − 1 n s L + 1 2 m ( a 1 s 1 2 d + a L + 1 1 s 1 2 d + a 1 s 1 2 d m a L + 1 1 l = 2 ∑ L s l 2 ) + λ 1 s L + 1 2 m ( 1 + c β κ ˉ 4 β ) exp ( − c ~ 1 m ( c β κ ˉ 2 β λ 1 n ln λ 1 n ) − 1 ) ] . Opening the parenthesis, we get
L ∞ ≍ ( β − 1 n λ 1 s L + 1 2 m ) 2 ( c β κ 2 κ ˉ 2 β d 2 ( s 1 4 + ( s 1 2 a 1 m ) 2 ( ∑ l = 2 L s l 2 ) 2 ) + a L + 1 2 + c β κ 2 κ ˉ 2 β s 1 2 a 1 d ∑ l = 1 L a l 2 ) + β − 1 n λ 1 ( s 1 2 a 1 d + 1 a L + 1 s 1 2 d + s 1 2 a 1 d m 1 a L + 1 ∑ l = 2 L s l 2 ) + ( 1 + c β κ 4 β ) exp ( − c ~ 1 m ( c β κ 2 β n λ 1 ln n λ 1 ) − 1 ) . \begin{aligned}
\mathcal{L}_{\infty} \asymp\left(\frac{\beta^{-1}n}{\lambda_{1}s_{L+1}^{2}m}\right)^{2}\left(c_{\beta}\kappa^{2}\bar{\kappa}^{2}\beta d^{2}\left(s_{1}^{4}+\left(\frac{s_{1}^{2}}{a_{1}}m\right)^{2}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}\right)+a_{L+1}^{2}+c_{\beta}\kappa^{2}\bar{\kappa}^{2}\beta\frac{s_{1}^{2}}{a_{1}}d\sum_{l=1}^{L}a_{l}^{2}\right) \\
&+\frac{\beta^{-1}n}{\lambda_{1}}\left(\frac{s_{1}^{2}}{a_{1}}d+\frac{1}{a_{L+1}}s_{1}^{2}d+\frac{s_{1}^{2}}{a_{1}}dm\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_{l}^{2}\right) \\
&+\left(1+c_{\beta}\kappa^{4}\beta\right)\exp\left(-\tilde{c}_{1}m\left(c_{\beta}\kappa^{2}\beta\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)^{-1}\right).
\end{aligned} L ∞ ≍ ( λ 1 s L + 1 2 m β − 1 n ) 2 c β κ 2 κ ˉ 2 β d 2 s 1 4 + ( a 1 s 1 2 m ) 2 ( l = 2 ∑ L s l 2 ) 2 + a L + 1 2 + c β κ 2 κ ˉ 2 β a 1 s 1 2 d l = 1 ∑ L a l 2 + λ 1 β − 1 n ( a 1 s 1 2 d + a L + 1 1 s 1 2 d + a 1 s 1 2 d m a L + 1 1 l = 2 ∑ L s l 2 ) + ( 1 + c β κ 4 β ) exp ( − c ~ 1 m ( c β κ 2 β λ 1 n ln λ 1 n ) − 1 ) . Splitting the first term into two, we get
L ∞ ≍ ( β − 1 n a L + 1 λ 1 s L + 1 2 m ) 2 + ( β − 1 / 2 n λ 1 s L + 1 2 m ) 2 c β κ 2 κ ˉ 2 ( d 2 ( s 1 4 + ( s 1 2 a 1 m ) 2 ( ∑ l = 2 L s l 2 ) 2 ) + s 1 2 a 1 d ∑ l = 1 L a l 2 ) + β − 1 n λ 1 ( s 1 2 a 1 d + 1 a L + 1 s 1 2 d + s 1 2 a 1 d m 1 a L + 1 ∑ l = 2 L s l 2 ) + ( 1 + c β κ 4 β ) exp ( − c ~ 1 m ( c β κ 2 β n λ 1 ln n λ 1 ) − 1 ) . \begin{aligned}
\mathcal{L}_{\infty}\asymp\left(\frac{\beta^{-1}na_{L+1}}{\lambda_{1}s_{L+1}^{2}m}\right)^{2} \\
&+\left(\frac{\beta^{-1/2}n}{\lambda_{1}s_{L+1}^{2}m}\right)^{2}c_{\beta}\kappa^{2}\bar{\kappa}^{2}\left(d^{2}\left(s_{1}^{4}+\left(\frac{s_{1}^{2}}{a_{1}}m\right)^{2}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}\right)+\frac{s_{1}^{2}}{a_{1}}d\sum_{l=1}^{L}a_{l}^{2}\right) \\
&+\frac{\beta^{-1}n}{\lambda_{1}}\left(\frac{s_{1}^{2}}{a_{1}}d+\frac{1}{a_{L+1}}s_{1}^{2}d+\frac{s_{1}^{2}}{a_{1}}dm\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_{l}^{2}\right) \\
&+\left(1+c_{\beta}\kappa^{4}\beta\right)\exp\left(-\tilde{c}_{1}m\left(c_{\beta}\kappa^{2}\beta\frac{n}{\lambda_{1}}\ln\frac{n}{\lambda_{1}}\right)^{-1}\right).
\end{aligned} L ∞ ≍ ( λ 1 s L + 1 2 m β − 1 n a L + 1 ) 2 + ( λ 1 s L + 1 2 m β − 1/2 n ) 2 c β κ 2 κ ˉ 2 d 2 s 1 4 + ( a 1 s 1 2 m ) 2 ( l = 2 ∑ L s l 2 ) 2 + a 1 s 1 2 d l = 1 ∑ L a l 2 + λ 1 β − 1 n ( a 1 s 1 2 d + a L + 1 1 s 1 2 d + a 1 s 1 2 d m a L + 1 1 l = 2 ∑ L s l 2 ) + ( 1 + c β κ 4 β ) exp ( − c ~ 1 m ( c β κ 2 β λ 1 n ln λ 1 n ) − 1 ) . An equivalent formulation of the loss floor
Let us rewrite the loss floor by using more of our assumptions, substituting a 1 , … , a L + 1 a_{1},\ldots,a_{L+1} a 1 , … , a L + 1 κ \kappa κ ( β − 1 n λ 1 κ 2 ) 2 \left(\frac{\beta^{-1}n}{\lambda_{1}\kappa^{2}}\right)^{2} ( λ 1 κ 2 β − 1 n ) 2
( β − 1 / 2 n λ 1 s L + 1 2 m ) 2 c β κ 2 κ ˉ 2 ( d 2 ( s 1 4 + ( s 1 2 a 1 m ) 2 ( ∑ l = 2 L s l 2 ) 2 ) + s 1 2 a 1 d ∑ l = 1 L a l 2 ) = ( β − 1 / 2 n λ 1 ) 2 c β κ 2 κ ˉ 2 ( d 2 s L + 1 4 m 2 ( s 1 4 + ( s 1 2 a 1 m ) 2 ( ∑ l = 2 L s l 2 ) 2 ) + s 1 2 a 1 d s L + 1 4 m 2 ∑ l = 1 L a l 2 ) \begin{aligned}
\left(\frac{\beta^{-1/2}n}{\lambda_{1}s_{L+1}^{2}m}\right)^{2}c_{\beta}\kappa^{2}\bar{\kappa}^{2}\left(d^{2}\left(s_{1}^{4}+\left(\frac{s_{1}^{2}}{a_{1}}m\right)^{2}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}\right)+\frac{s_{1}^{2}}{a_{1}}d\sum_{l=1}^{L}a_{l}^{2}\right) \\
&=\left(\frac{\beta^{-1/2}n}{\lambda_{1}}\right)^{2}c_{\beta}\kappa^{2}\bar{\kappa}^{2}\left(\frac{d^{2}}{s_{L+1}^{4}m^{2}}\left(s_{1}^{4}+\left(\frac{s_{1}^{2}}{a_{1}}m\right)^{2}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}\right)+\frac{s_{1}^{2}}{a_{1}}\frac{d}{s_{L+1}^{4}m^{2}}\sum_{l=1}^{L}a_{l}^{2}\right)
\end{aligned} ( λ 1 s L + 1 2 m β − 1/2 n ) 2 c β κ 2 κ ˉ 2 d 2 s 1 4 + ( a 1 s 1 2 m ) 2 ( l = 2 ∑ L s l 2 ) 2 + a 1 s 1 2 d l = 1 ∑ L a l 2 = ( λ 1 β − 1/2 n ) 2 c β κ 2 κ ˉ 2 s L + 1 4 m 2 d 2 s 1 4 + ( a 1 s 1 2 m ) 2 ( l = 2 ∑ L s l 2 ) 2 + a 1 s 1 2 s L + 1 4 m 2 d l = 1 ∑ L a l 2 ≡ ( i ) ( β − 1 / 2 n λ 1 ) 2 c β κ 2 κ ˉ 2 ( d 2 s L + 1 4 m 2 ( s 1 4 + ( m d ) 2 ( ∑ l = 2 L s l 2 ) 2 ) + 1 s L + 1 4 m 2 ∑ l = 1 L a l 2 ) = ( β − 1 / 2 n λ 1 ) 2 c β κ 2 κ ˉ 2 ( s 1 4 s L + 1 4 ( d m ) 2 + 1 s L + 1 4 ( ∑ l = 2 L s l 2 ) 2 + 1 s L + 1 4 m 2 ∑ l = 1 L a l 2 ) \begin{aligned}
\overset{(i)}{\equiv}\left(\frac{\beta^{-1/2}n}{\lambda_{1}}\right)^{2}c_{\beta}\kappa^{2}\bar{\kappa}^{2}\left(\frac{d^{2}}{s_{L+1}^{4}m^{2}}\left(s_{1}^{4}+\left(\frac{m}{d}\right)^{2}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}\right)+\frac{1}{s_{L+1}^{4}m^{2}}\sum_{l=1}^{L}a_{l}^{2}\right) \\
&=\left(\frac{\beta^{-1/2}n}{\lambda_{1}}\right)^{2}c_{\beta}\kappa^{2}\bar{\kappa}^{2}\left(\frac{s_{1}^{4}}{s_{L+1}^{4}}\left(\frac{d}{m}\right)^{2}+\frac{1}{s_{L+1}^{4}}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}+\frac{1}{s_{L+1}^{4}m^{2}}\sum_{l=1}^{L}a_{l}^{2}\right)
\end{aligned} ≡ ( i ) ( λ 1 β − 1/2 n ) 2 c β κ 2 κ ˉ 2 s L + 1 4 m 2 d 2 s 1 4 + ( d m ) 2 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 m 2 1 l = 1 ∑ L a l 2 = ( λ 1 β − 1/2 n ) 2 c β κ 2 κ ˉ 2 s L + 1 4 s 1 4 ( m d ) 2 + s L + 1 4 1 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 m 2 1 l = 1 ∑ L a l 2 ≡ ( i i ) ( κ κ ˉ β − 1 / 2 n λ 1 ) 2 c β ( s 1 4 s L + 1 4 ( d m ) 2 + 1 s L + 1 4 ( ∑ l = 2 L s l 2 ) 2 + s 1 4 s L + 1 4 ( d m ) 2 + 1 s L + 1 4 ∑ l = 2 L s l 4 ) \overset{(ii)}{\equiv}\left(\kappa\bar{\kappa}\frac{\beta^{-1/2}n}{\lambda_{1}}\right)^{2}c_{\beta}\left(\frac{s_{1}^{4}}{s_{L+1}^{4}}\left(\frac{d}{m}\right)^{2}+\frac{1}{s_{L+1}^{4}}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}+\frac{s_{1}^{4}}{s_{L+1}^{4}}\left(\frac{d}{m}\right)^{2}+\frac{1}{s_{L+1}^{4}}\sum_{l=2}^{L}s_{l}^{4}\right) ≡ ( ii ) ( κ κ ˉ λ 1 β − 1/2 n ) 2 c β s L + 1 4 s 1 4 ( m d ) 2 + s L + 1 4 1 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 s 1 4 ( m d ) 2 + s L + 1 4 1 l = 2 ∑ L s l 4 ≲ ( κ κ ˉ β − 1 / 2 n λ 1 ) 2 c β ( s 1 4 s L + 1 4 ( d m ) 2 + 1 s L + 1 4 ( ∑ l = 2 L s l 2 ) 2 + 1 s L + 1 4 ∑ l = 2 L s l 4 ) , \lesssim\left(\kappa\bar{\kappa}\frac{\beta^{-1/2}n}{\lambda_{1}}\right)^{2}c_{\beta}\left(\frac{s_{1}^{4}}{s_{L+1}^{4}}\left(\frac{d}{m}\right)^{2}+\frac{1}{s_{L+1}^{4}}\left(\sum_{l=2}^{L}s_{l}^{2}\right)^{2}+\frac{1}{s_{L+1}^{4}}\sum_{l=2}^{L}s_{l}^{4}\right), ≲ ( κ κ ˉ λ 1 β − 1/2 n ) 2 c β s L + 1 4 s 1 4 ( m d ) 2 + s L + 1 4 1 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 1 l = 2 ∑ L s l 4 , where in ( i ) (i) ( i ) a 1 s 1 2 = m 0 = d \frac{a_1}{s_1^2}=m_0=d s 1 2 a 1 = m 0 = d ( i i ) (ii) ( ii ) a l = s l 2 m l − 1 a_l=s_l^2m_{l-1} a l = s l 2 m l − 1 l ∈ [ L ] l\in[L] l ∈ [ L ]
β − 1 n λ 1 ( s 1 2 a 1 d + 1 a L + 1 s 1 2 d + s 1 2 a 1 d m 1 a L + 1 ∑ l = 2 L s l 2 ) = ( i ) β − 1 n λ 1 ( 1 + 1 a L + 1 s 1 2 d + m 1 a L + 1 ∑ l = 2 L s l 2 ) = ( i i ) β − 1 n λ 1 ( 1 + κ 2 s 1 2 s L + 1 2 d m + κ 2 s L + 1 2 ∑ l = 2 L s l 2 ) , \frac{\beta^{-1}n}{\lambda_1}\left(\frac{s_1^2}{a_1}d+\frac{1}{a_{L+1}}s_1^2d+\frac{s_1^2}{a_1}dm\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_l^2\right)
\underset{(i)}{=}\frac{\beta^{-1}n}{\lambda_1}\left(1+\frac{1}{a_{L+1}}s_1^2d+m\frac{1}{a_{L+1}}\sum_{l=2}^{L}s_l^2\right)
\underset{(ii)}{=}\frac{\beta^{-1}n}{\lambda_1}\left(1+\kappa^2\frac{s_1^2}{s_{L+1}^2}\frac{d}{m}+\frac{\kappa^2}{s_{L+1}^2}\sum_{l=2}^{L}s_l^2\right), λ 1 β − 1 n ( a 1 s 1 2 d + a L + 1 1 s 1 2 d + a 1 s 1 2 d m a L + 1 1 l = 2 ∑ L s l 2 ) ( i ) = λ 1 β − 1 n ( 1 + a L + 1 1 s 1 2 d + m a L + 1 1 l = 2 ∑ L s l 2 ) ( ii ) = λ 1 β − 1 n ( 1 + κ 2 s L + 1 2 s 1 2 m d + s L + 1 2 κ 2 l = 2 ∑ L s l 2 ) , where ( i ) (i) ( i ) ( i i ) (ii) ( ii ) a L + 1 = s L + 1 2 m κ 2 a_{L+1}=\frac{s_{L+1}^2m}{\kappa^2} a L + 1 = κ 2 s L + 1 2 m
L ∞ ≲ ( β − 1 n λ 1 κ 2 ) 2 + ( κ κ ~ β − 1 / 2 n λ 1 ) 2 c β ( s 1 4 s L + 1 4 ( d m ) 2 + 1 s L + 1 4 ( ∑ l = 2 L s l 2 ) 2 + 1 s L + 1 4 ∑ l = 2 L s l 4 ) + β − 1 n λ 1 ( 1 + κ 2 s 1 2 s L + 1 2 d m + κ 2 s L + 1 2 ∑ l = 2 L s l 2 ) + ( 1 + c β κ 4 β ) exp ( − c ~ 1 m ( c β κ 2 β n λ 1 ln n λ 1 ) − 1 ) . \begin{aligned}
\mathcal{L}_\infty\lesssim{}&\left(\frac{\beta^{-1}n}{\lambda_1\kappa^2}\right)^2\\
&+\left(\kappa\tilde{\kappa}\frac{\beta^{-1/2}n}{\lambda_1}\right)^2c_\beta\left(\frac{s_1^4}{s_{L+1}^4}\left(\frac{d}{m}\right)^2+\frac{1}{s_{L+1}^4}\left(\sum_{l=2}^{L}s_l^2\right)^2+\frac{1}{s_{L+1}^4}\sum_{l=2}^{L}s_l^4\right)\\
&+\frac{\beta^{-1}n}{\lambda_1}\left(1+\kappa^2\frac{s_1^2}{s_{L+1}^2}\frac{d}{m}+\frac{\kappa^2}{s_{L+1}^2}\sum_{l=2}^{L}s_l^2\right)\\
&+(1+c_\beta\kappa^4\beta)\exp\left(-\tilde{c}_1m\left(c_\beta\kappa^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right).
\end{aligned} L ∞ ≲ ( λ 1 κ 2 β − 1 n ) 2 + ( κ κ ~ λ 1 β − 1/2 n ) 2 c β s L + 1 4 s 1 4 ( m d ) 2 + s L + 1 4 1 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 1 l = 2 ∑ L s l 4 + λ 1 β − 1 n ( 1 + κ 2 s L + 1 2 s 1 2 m d + s L + 1 2 κ 2 l = 2 ∑ L s l 2 ) + ( 1 + c β κ 4 β ) exp ( − c ~ 1 m ( c β κ 2 β λ 1 n ln λ 1 n ) − 1 ) . Finally, we observe that in order to achieve L ∞ < 1 \mathcal{L}_\infty<1 L ∞ < 1 β > n λ 1 \beta>\frac{n}{\lambda_1} β > λ 1 n κ 2 β ≥ β ≫ 1 \kappa^2\beta\geq\beta\gg1 κ 2 β ≥ β ≫ 1 c β = 2 c M ϕ 2 + 1 κ 2 β = Θ ( 1 ) c_\beta=2cM_\phi^2+\frac{1}{\kappa^2\beta}=\Theta(1) c β = 2 c M ϕ 2 + κ 2 β 1 = Θ ( 1 )
L ∞ ≲ ( β − 1 n λ 1 κ 2 ) 2 + ( κ κ ~ β − 1 / 2 n λ 1 ) 2 ( s 1 4 s L + 1 4 ( d m ) 2 + 1 s L + 1 4 ( ∑ l = 2 L s l 2 ) 2 + 1 s L + 1 4 ∑ l = 2 L s l 4 ) + β − 1 n λ 1 ( 1 + κ 2 s 1 2 s L + 1 2 d m + κ 2 s L + 1 2 ∑ l = 2 L s l 2 ) + ( 1 + κ 4 β ) exp ( − c ~ 1 m ( κ 2 β n λ 1 ln n λ 1 ) − 1 ) . \begin{aligned}
\mathcal{L}_\infty\lesssim{}&\left(\frac{\beta^{-1}n}{\lambda_1\kappa^2}\right)^2\\
&+\left(\kappa\tilde{\kappa}\frac{\beta^{-1/2}n}{\lambda_1}\right)^2\left(\frac{s_1^4}{s_{L+1}^4}\left(\frac{d}{m}\right)^2+\frac{1}{s_{L+1}^4}\left(\sum_{l=2}^{L}s_l^2\right)^2+\frac{1}{s_{L+1}^4}\sum_{l=2}^{L}s_l^4\right)\\
&+\frac{\beta^{-1}n}{\lambda_1}\left(1+\kappa^2\frac{s_1^2}{s_{L+1}^2}\frac{d}{m}+\frac{\kappa^2}{s_{L+1}^2}\sum_{l=2}^{L}s_l^2\right)\\
&+(1+\kappa^4\beta)\exp\left(-\tilde{c}_1m\left(\kappa^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right).
\end{aligned} L ∞ ≲ ( λ 1 κ 2 β − 1 n ) 2 + ( κ κ ~ λ 1 β − 1/2 n ) 2 s L + 1 4 s 1 4 ( m d ) 2 + s L + 1 4 1 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 1 l = 2 ∑ L s l 4 + λ 1 β − 1 n ( 1 + κ 2 s L + 1 2 s 1 2 m d + s L + 1 2 κ 2 l = 2 ∑ L s l 2 ) + ( 1 + κ 4 β ) exp ( − c ~ 1 m ( κ 2 β λ 1 n ln λ 1 n ) − 1 ) . Rearranging this, we get
L ∞ ≲ ( 1 + κ 4 β ) exp ( − c ~ 1 m ( κ 2 β n λ 1 ln n λ 1 ) − 1 ) + β − 1 n λ 1 ( 1 + κ 2 s 1 2 s L + 1 2 d m + κ 2 s L + 1 2 ∑ l = 2 L s l 2 ) + ( β − 1 / 2 n λ 1 ) 2 κ 2 κ ~ 2 ( s 1 4 s L + 1 4 ( d m ) 2 + 1 s L + 1 4 ( ∑ l = 2 L s l 2 ) 2 + 1 s L + 1 4 ∑ l = 2 L s l 4 ) + ( β − 1 n λ 1 ) 2 ( 1 κ 4 ) . \begin{aligned}
\mathcal{L}_\infty\lesssim{}&(1+\kappa^4\beta)\exp\left(-\tilde{c}_1m\left(\kappa^2\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right)\\
&+\frac{\beta^{-1}n}{\lambda_1}\left(1+\kappa^2\frac{s_1^2}{s_{L+1}^2}\frac{d}{m}+\frac{\kappa^2}{s_{L+1}^2}\sum_{l=2}^{L}s_l^2\right)\\
&+\left(\frac{\beta^{-1/2}n}{\lambda_1}\right)^2\kappa^2\tilde{\kappa}^2\left(\frac{s_1^4}{s_{L+1}^4}\left(\frac{d}{m}\right)^2+\frac{1}{s_{L+1}^4}\left(\sum_{l=2}^{L}s_l^2\right)^2+\frac{1}{s_{L+1}^4}\sum_{l=2}^{L}s_l^4\right)\\
&+\left(\frac{\beta^{-1}n}{\lambda_1}\right)^2\left(\frac{1}{\kappa^4}\right).
\end{aligned} L ∞ ≲ ( 1 + κ 4 β ) exp ( − c ~ 1 m ( κ 2 β λ 1 n ln λ 1 n ) − 1 ) + λ 1 β − 1 n ( 1 + κ 2 s L + 1 2 s 1 2 m d + s L + 1 2 κ 2 l = 2 ∑ L s l 2 ) + ( λ 1 β − 1/2 n ) 2 κ 2 κ ~ 2 s L + 1 4 s 1 4 ( m d ) 2 + s L + 1 4 1 ( l = 2 ∑ L s l 2 ) 2 + s L + 1 4 1 l = 2 ∑ L s l 4 + ( λ 1 β − 1 n ) 2 ( κ 4 1 ) . Standard parameterization
Under standard parameterization (∀ l ∈ [ L + 1 ] : s l = 1 , a 1 = d , ∀ l > 1 : a l = m , κ = 1 \forall l \in[L+1] : s_l = 1, a_1 = d, \forall l > 1 : a_l = m, \kappa= 1 ∀ l ∈ [ L + 1 ] : s l = 1 , a 1 = d , ∀ l > 1 : a l = m , κ = 1
L ∞ ≃ ( β − 1 n λ 1 ) 2 + ( β − 1 / 2 n λ 1 ) 2 ( ( d m ) 2 + L 2 + L ) + β − 1 n λ 1 ( 1 + d m + L ) + ( 1 + β ) exp ( − c ~ 1 m ( β n λ 1 ln n λ 1 ) − 1 ) ≃ ( β − 1 n λ 1 ) 2 + ( β − 1 / 2 n λ 1 ) 2 + β − 1 n λ 1 + ( 1 + β ) exp ( − c ~ 1 m ( β n λ 1 ln n λ 1 ) − 1 ) ≃ ( β − 1 n ) 2 + ( β − 1 / 2 n ) 2 + β − 1 n + ( 1 + β ) exp ( − c ~ 1 m β n ln n ) , \begin{aligned}
\mathcal{L}_{\infty} \simeq\left(\frac{\beta^{-1}n}{\lambda_1}\right)^2 + \left(\frac{\beta^{-1/2}n}{\lambda_1}\right)^2 \left(\left(\frac{d}{m}\right)^2 + L^2 + L\right)
+ \frac{\beta^{-1}n}{\lambda_1}\left(1+\frac{d}{m}+L\right) + (1+\beta)\exp\left(-\widetilde{c}_1m\left(\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right) \\
&\simeq\left(\frac{\beta^{-1}n}{\lambda_1}\right)^2 + \left(\frac{\beta^{-1/2}n}{\lambda_1}\right)^2 + \frac{\beta^{-1}n}{\lambda_1} + (1+\beta)\exp\left(-\widetilde{c}_1m\left(\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right) \\
&\simeq(\beta^{-1}n)^2 + (\beta^{-1/2}n)^2 + \beta^{-1}n + (1+\beta)\exp\left(-\widetilde{c}_1\frac{m}{\beta n\ln n}\right),
\end{aligned} L ∞ ≃ ( λ 1 β − 1 n ) 2 + ( λ 1 β − 1/2 n ) 2 ( ( m d ) 2 + L 2 + L ) + λ 1 β − 1 n ( 1 + m d + L ) + ( 1 + β ) exp ( − c 1 m ( β λ 1 n ln λ 1 n ) − 1 ) ≃ ( λ 1 β − 1 n ) 2 + ( λ 1 β − 1/2 n ) 2 + λ 1 β − 1 n + ( 1 + β ) exp ( − c 1 m ( β λ 1 n ln λ 1 n ) − 1 ) ≃ ( β − 1 n ) 2 + ( β − 1/2 n ) 2 + β − 1 n + ( 1 + β ) exp ( − c 1 β n ln n m ) , where in (i) we use d ≤ m d \le m d ≤ m λ 1 \lambda_1 λ 1 Θ ( 1 ) \Theta(1) Θ ( 1 ) L ∞ ≤ L g \mathcal{L}_{\infty} \le\mathcal{L}_g L ∞ ≤ L g L g ≤ 1 \mathcal{L}_g \le1 L g ≤ 1
m = Ω ( β n ln n ln ( m / L g ) ) m = \Omega\left(\beta n\ln n\ln\left(m/\mathcal{L}_g\right)\right) m = Ω ( β n ln n ln ( m / L g ) )
β = Ω ( n / L g ) \beta= \Omega\left(n/\mathcal{L}_g\right) β = Ω ( n / L g )
β = Ω ( n 2 / L g ) \beta= \Omega\left(n^2/\mathcal{L}_g\right) β = Ω ( n 2 / L g )
This can be achieved with m = Ω ~ ( n 3 / L g ) m = \widetilde{\Omega}\left(n^3/\mathcal{L}_g\right) m = Ω ( n 3 / L g ) β = Ω ~ ( n 2 / L g ) \beta= \widetilde{\Omega}\left(n^2/\mathcal{L}_g\right) β = Ω ( n 2 / L g )
Comparing the results to the case of shallow networks. For shallow networks, the analysis establishes convergence to a small loss floor under milder conditions. The stronger requirements in the deep-network case come from sums over layers l = 2 , … , L l = 2,\ldots,L l = 2 , … , L 0 0 0
NTK parameterization
Under NTK parameterization (s 1 = d − 1 / 2 s_1 = d^{-1/2} s 1 = d − 1/2 ∀ l > 1 : s l = m − 1 / 2 \forall l > 1 : s_l = m^{-1/2} ∀ l > 1 : s l = m − 1/2 ∀ l ∈ [ L + 1 ] : a l = 1 , κ = 1 \forall l \in[L+1] : a_l = 1, \kappa= 1 ∀ l ∈ [ L + 1 ] : a l = 1 , κ = 1
L ∞ ≃ ( β − 1 n λ 1 ) 2 + ( β − 1 / 2 n λ 1 ) 2 ( m 2 d 2 ( d m ) 2 + m 2 ( L m ) 2 + m 2 L m 2 ) + β − 1 n λ 1 ( 1 + m d d m + m L m ) + ( 1 + β ) exp ( − c ~ 1 m ( β n λ 1 ln n λ 1 ) − 1 ) ≃ ( β − 1 n λ 1 ) 2 + ( β − 1 / 2 n λ 1 ) 2 ( 1 + L 2 + L ) + β − 1 n λ 1 ( 1 + 1 + L ) + ( 1 + β ) exp ( − c ~ 1 m ( β n λ 1 ln n λ 1 ) − 1 ) ≃ ( β − 1 n ) 2 + ( β − 1 / 2 n ) 2 + β − 1 n + ( 1 + β ) exp ( − c ~ 1 m β n ln n ) \begin{aligned}
\mathcal{L}_{\infty} \simeq\left(\frac{\beta^{-1}n}{\lambda_1}\right)^2
+ \left(\frac{\beta^{-1/2}n}{\lambda_1}\right)^2 \left(\frac{m^2}{d^2}\left(\frac{d}{m}\right)^2 + m^2\left(\frac{L}{m}\right)^2 + m^2\frac{L}{m^2}\right)
+ \frac{\beta^{-1}n}{\lambda_1}\left(1+\frac{m}{d}\frac{d}{m}+m\frac{L}{m}\right)
+ (1+\beta)\exp\left(-\widetilde{c}_1m\left(\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right) \\
&\simeq\left(\frac{\beta^{-1}n}{\lambda_1}\right)^2 + \left(\frac{\beta^{-1/2}n}{\lambda_1}\right)^2(1+L^2+L) + \frac{\beta^{-1}n}{\lambda_1}(1+1+L)
+ (1+\beta)\exp\left(-\widetilde{c}_1m\left(\beta\frac{n}{\lambda_1}\ln\frac{n}{\lambda_1}\right)^{-1}\right) \\
&\simeq(\beta^{-1}n)^2 + (\beta^{-1/2}n)^2 + \beta^{-1}n + (1+\beta)\exp\left(-\widetilde{c}_1\frac{m}{\beta n\ln n}\right)
\end{aligned} L ∞ ≃ ( λ 1 β − 1 n ) 2 + ( λ 1 β − 1/2 n ) 2 ( d 2 m 2 ( m d ) 2 + m 2 ( m L ) 2 + m 2 m 2 L ) + λ 1 β − 1 n ( 1 + d m m d + m m L ) + ( 1 + β ) exp ( − c 1 m ( β λ 1 n ln λ 1 n ) − 1 ) ≃ ( λ 1 β − 1 n ) 2 + ( λ 1 β − 1/2 n ) 2 ( 1 + L 2 + L ) + λ 1 β − 1 n ( 1 + 1 + L ) + ( 1 + β ) exp ( − c 1 m ( β λ 1 n ln λ 1 n ) − 1 ) ≃ ( β − 1 n ) 2 + ( β − 1/2 n ) 2 + β − 1 n + ( 1 + β ) exp ( − c 1 β n ln n m ) where in the last step we treat λ 1 , L \lambda_1, L λ 1 , L Θ ( 1 ) \Theta(1) Θ ( 1 ) L ∞ ≤ L g \mathcal{L}_{\infty} \le\mathcal{L}_g L ∞ ≤ L g L g ≤ 1 \mathcal{L}_g \le1 L g ≤ 1
Convergence with gradient flow from stopped CLD
This appendix proves Proposition 5.2, the two-stage scheduling guarantee: regularized CLD is run until a deterministic time τ \tau τ
Recall Proposition 5.2 (Convergence to zero loss for shallow networks). For shallow networks, i.e. , L = 1 L=1 L = 1 τ = Ω ~ ( n s 2 2 λ 1 m ) \tau= \widetilde{\Omega}\left(\frac{n}{s_{2}^{2}\lambda_{1}m}\right) τ = Ω ( s 2 2 λ 1 m n ) (9) . Let K > 2 K > 2 K > 2 1 − η exp ( − c 4 m c β κ ˉ 2 β 2 λ 1 2 ln ( n λ 1 ) ) − 2 K 1-\eta\exp\left(-c_{4}\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta^{2}\lambda_{1}^{2}}\ln\left(\frac{n}{\lambda_{1}}\right)\right)-\frac{2}{K} 1 − η exp ( − c 4 c β κ ˉ 2 β 2 λ 1 2 m ln ( λ 1 n ) ) − K 2
L ∞ ( L = 1 ) ≲ s 2 2 m 2 K n 2 d 2 s 1 4 ( 1 d s 1 2 + κ 2 m s 2 2 ) ( κ 2 s 1 2 d s 2 2 + n m ) (41) \mathcal{L}_{\infty}^{(L=1)} \lesssim\frac{s_{2}^{2}m^{2}}{Kn^{2}d^{2}s_{1}^{4}\left(\frac{1}{ds_{1}^{2}}+\frac{\kappa^{2}}{ms_{2}^{2}}\right)\left(\frac{\kappa^{2}s_{1}^{2}d}{s_{2}^{2}}+\sqrt{nm}\right)} \tag*{(41)} L ∞ ( L = 1 ) ≲ K n 2 d 2 s 1 4 ( d s 1 2 1 + m s 2 2 κ 2 ) ( s 2 2 κ 2 s 1 2 d + nm ) s 2 2 m 2 (41) as in Eq. (10).
The recall above is not verbatim. The original statement of Proposition 5.2 closes with a worked example, which instantiates this demand under the standard parameterization and concludes that, with d = o ( n ) d=o\left(\sqrt{n}\right) d = o ( n ) K = Θ ( n / d ) K=\Theta\left(\sqrt{n}/d\right) K = Θ ( n / d ) 1 − o ( 1 ) 1-o(1) 1 − o ( 1 ) β = Ω ~ ( n ) \beta=\widetilde{\Omega}(n) β = Ω ( n ) m = Ω ~ ( n 2 ) m=\widetilde{\Omega}\left(n^{2}\right) m = Ω ( n 2 ) § 17 .§ 4 , immediately after the proof.
The proof combines a deterministic ingredient with a probabilistic one. The deterministic ingredient, developed in § 17 .§ 1 , is a criterion under which gradient flow started at a point θ 0 \theta_{0} θ 0 every time along the CLD trajectory, together with the expected-loss bound of Theorem 5.1. § 17 .§ 2 collects the initialization estimates that this transfer acts upon, and § 17 .§ 3 carries out the transfer, verifies the criterion at θ C L D ( τ ) \theta_{\mathrm{CLD}}(\tau) θ CLD ( τ ) § 17 .§ 4 then works out the resulting demand on the loss floor under the standard and NTK parameterizations.
Gradient flow convergence by proximity to good initialization
Our analysis here closely follows that of Song et al. [105 ] , with adaptation to our notations and settings. Importantly, we directly analyze without the bounded second layer assumption (which is part of their Assumption 2), similar to the (somewhat informal) treatment in their Appendix E.
We begin by citing some useful definitions and Lemmas from Song et al. [105 ] , adapted to our notations. Importantly, we restrict the smoothness definition to a specific convex set.
Definition I.1 (Smoothness). A function f : R N p → R n f:\mathbb{R}^{N_{p}}\to\mathbb{R}^{n} f : R N p → R n β f \beta_{f} β f Θ ⊆ R N p \Theta\subseteq\mathbb{R}^{N_{p}} Θ ⊆ R N p β f ≥ 0 \beta_{f}\geq0 β f ≥ 0 θ 1 , θ 2 ∈ Θ \theta_{1},\theta_{2}\in\Theta θ 1 , θ 2 ∈ Θ
∥ ∇ f ( θ 1 ) − ∇ f ( θ 2 ) ∥ 2 ≤ β f ∥ θ 1 − θ 2 ∥ \lVert\nabla f(\theta_{1})-\nabla f(\theta_{2})\rVert_{2}\leq\beta_{f}\lVert\theta_{1}-\theta_{2}\rVert ∥ ∇ f ( θ 1 ) − ∇ f ( θ 2 ) ∥ 2 ≤ β f ∥ θ 1 − θ 2 ∥ where ∇ f ( θ ) ∈ R n × N p \nabla f(\theta)\in\mathbb{R}^{n\times N_{p}} ∇ f ( θ ) ∈ R n × N p f f f θ \theta θ
Definition I.2 (Near-isometry). A matrix A A A ( μ , ν ) (\mu,\nu) ( μ , ν ) 0 < μ ≤ ν 0<\mu\leq\nu 0 < μ ≤ ν
μ ≤ σ min ( A ) ≤ σ max ( A ) ≤ ν , \mu\leq\sigma_{\min}(A)\leq\sigma_{\max}(A)\leq\nu, μ ≤ σ m i n ( A ) ≤ σ m a x ( A ) ≤ ν , where σ min ( A ) \sigma_{\min}(A) σ m i n ( A ) σ max ( A ) \sigma_{\max}(A) σ m a x ( A ) A A A
We now proceed to cite a slightly altered version of their Lemma 1, assuming smoothness only in the relevant ball. The proof is identical to theirs, and we do not repeat it here.
Lemma 1.3. Given some θ 0 ∈ R N p \theta_{0}\in\mathbb{R}^{N_{p}} θ 0 ∈ R N p f : R N p → R n f:\mathbb{R}^{N_{p}}\to\mathbb{R}^{n} f : R N p → R n β f \beta_{f} β f Θ = { θ : ∥ θ − θ 0 ∥ ≤ μ 2 β f } \Theta=\left\{\theta:\lVert\theta-\theta_{0}\rVert\leq\frac{\mu}{2\beta_{f}}\right\} Θ = { θ : ∥ θ − θ 0 ∥ ≤ 2 β f μ } ∇ f ( θ 0 ) \nabla f(\theta_{0}) ∇ f ( θ 0 ) ( μ , ν ) (\mu,\nu) ( μ , ν ) θ ∈ Θ \theta\in\Theta θ ∈ Θ ∇ f ( θ ) \nabla f(\theta) ∇ f ( θ ) ( 1 2 μ , 3 2 ν ) \left(\frac{1}{2}\mu,\frac{3}{2}\nu\right) ( 2 1 μ , 2 3 ν )
Definition I.4 (Polyak Łojasiewicz (PL) condition). A non-negative function F : R n → R \mathcal{F} : \mathbb{R}^{n} \to\mathbb{R} F : R n → R α F > 0 \alpha_{\mathcal{F}} > 0 α F > 0 z ∈ R n z \in\mathbb{R}^{n} z ∈ R n
F ( z ) ≤ ∥ ∇ F ( z ) ∥ 2 2 α F . \mathcal{F}(\mathbf{z}) \le\frac{\lVert\nabla\mathcal{F}(\mathbf{z}) \rVert^{2}}{2\alpha_{\mathcal{F}}}. F ( z ) ≤ 2 α F ∥ ∇ F ( z ) ∥ 2 . We now proceed by analyzing gradient flow of θ ∈ R N p \theta\in\mathbb{R}^{N_p} θ ∈ R N p L : R N p → R \mathcal{L} : \mathbb{R}^{N_p} \to\mathbb{R} L : R N p → R L ( θ ) ≜ F ( f ( θ ) ) \mathcal{L}(\theta) \triangleq\mathcal{F}(f(\theta)) L ( θ ) ≜ F ( f ( θ )) f : R N p → R n f : \mathbb{R}^{N_p} \to\mathbb{R}^{n} f : R N p → R n F : R n → R \mathcal{F} : \mathbb{R}^{n} \to\mathbb{R} F : R n → R
d d t θ ( t ) = − ∇ L ( θ ( t ) ) . (42) \frac{\mathrm{d}}{\mathrm{d}t}\theta(t) = -\nabla\mathcal{L}(\theta(t)). \tag*{(42)} d t d θ ( t ) = − ∇ L ( θ ( t )) . (42) We cite Song et al. [105, Lemma 2], adapted to our notations. The proof remains identical, hence we do not repeat it.
Lemma I.5. Let L ( θ ) ≜ F ( f ( θ ) ) \mathcal{L}(\theta) \triangleq\mathcal{F}(f(\theta)) L ( θ ) ≜ F ( f ( θ )) F ≥ 0 \mathcal{F} \ge0 F ≥ 0 α F > 0 \alpha_{\mathcal{F}} > 0 α F > 0 f f f β f \beta_f β f Θ = { θ : ∥ θ − θ 0 ∥ ≤ μ 2 β f } \Theta= \left\{ \theta: \lVert\theta- \theta_0 \rVert\le\frac{\mu}{2\beta_f} \right\} Θ = { θ : ∥ θ − θ 0 ∥ ≤ 2 β f μ } ∇ f ( θ 0 ) \nabla f(\theta_0) ∇ f ( θ 0 ) ( μ , ν ) (\mu,\nu) ( μ , ν ) θ ( t ) \theta(t) θ ( t ) (42) ) initialized by θ ( 0 ) = θ 0 \theta(0) = \theta_0 θ ( 0 ) = θ 0 τ b = inf { t > 0 : ∥ θ ( t ) − θ 0 ∥ > μ 2 β f } \tau_b = \inf\left\{t > 0 : \lVert\theta(t) - \theta_0 \rVert> \frac{\mu}{2\beta_f} \right\} τ b = inf { t > 0 : ∥ θ ( t ) − θ 0 ∥ > 2 β f μ } t < τ b t < \tau_b t < τ b
c ν L ( θ 0 ) μ 2 α F , c\frac{\nu\sqrt{\mathcal{L}(\theta_0)}}{\mu^2\sqrt{\alpha_{\mathcal{F}}}}, c μ 2 α F ν L ( θ 0 ) , where c > 0 c > 0 c > 0
We now extend Song et al. [105, Theorem 1], which proves convergence to a global minimum, by stating the convergence rate.
Theorem I.6. Let L ( θ ) ≜ F ( f ( θ ) ) \mathcal{L}(\theta) \triangleq\mathcal{F}(f(\theta)) L ( θ ) ≜ F ( f ( θ )) F ≥ 0 \mathcal{F} \ge0 F ≥ 0 α F > 0 \alpha_{\mathcal{F}} > 0 α F > 0 f f f β f \beta_f β f Θ = { θ : ∥ θ − θ 0 ∥ ≤ μ 2 β f } \Theta= \left\{ \theta: \lVert\theta- \theta_0 \rVert\le\frac{\mu}{2\beta_f} \right\} Θ = { θ : ∥ θ − θ 0 ∥ ≤ 2 β f μ } ∇ f ( θ 0 ) \nabla f(\theta_0) ∇ f ( θ 0 ) ( μ , ν ) (\mu,\nu) ( μ , ν ) L ( θ 0 ) ≤ 1 16 c 2 α F μ 6 β f 2 ν 2 \mathcal{L}(\theta_0) \le\frac{1}{16c^2}\frac{\alpha_{\mathcal{F}}\mu^6}{\beta_f^2\nu^2} L ( θ 0 ) ≤ 16 c 2 1 β f 2 ν 2 α F μ 6 c c c (42) ) initialized by θ ( 0 ) = θ 0 \theta(0) = \theta_0 θ ( 0 ) = θ 0
L ( θ ( t ) ) ≤ L ( θ 0 ) e − 1 2 α F μ 2 t . \mathcal{L}(\theta(t)) \le\mathcal{L}(\theta_0)e^{-\frac{1}{2}\alpha_{\mathcal{F}}\mu^2t}. L ( θ ( t )) ≤ L ( θ 0 ) e − 2 1 α F μ 2 t . Proof. We show that under the stated conditions, θ ( t ) \theta(t) θ ( t ) Θ \Theta Θ τ b = ∞ \tau_b = \infty τ b = ∞ τ b < ∞ \tau_b < \infty τ b < ∞ τ b \tau_b τ b
c ν L ( θ 0 ) μ 2 α F ≤ c ν 1 16 c 2 α F μ 6 β f 2 ν 2 μ 2 α F = μ 4 β f < μ 2 β f , c\frac{\nu\sqrt{\mathcal{L}(\theta_0)}}{\mu^2\sqrt{\alpha_{\mathcal{F}}}}
\le c\frac{\nu\sqrt{\frac{1}{16c^2}\frac{\alpha_{\mathcal{F}}\mu^6}{\beta_f^2\nu^2}}}{\mu^2\sqrt{\alpha_{\mathcal{F}}}}
= \frac{\mu}{4\beta_f} < \frac{\mu}{2\beta_f}, c μ 2 α F ν L ( θ 0 ) ≤ c μ 2 α F ν 16 c 2 1 β f 2 ν 2 α F μ 6 = 4 β f μ < 2 β f μ , where we employed the loss bound at θ 0 \theta_0 θ 0 τ b = ∞ \tau_b = \infty τ b = ∞
Then, by Lemma I.3, we have that for all t ≥ 0 t \ge0 t ≥ 0 ∇ f ( θ ( t ) ) \nabla f(\theta(t)) ∇ f ( θ ( t )) ( 1 2 μ , 3 2 ν ) \left(\frac{1}{2}\mu,\frac{3}{2}\nu\right) ( 2 1 μ , 2 3 ν )
d d t L ( θ ( t ) ) = ∇ L ( θ ( t ) ) ⊤ d d t θ ( t ) = − ∥ ∇ L ( θ ( t ) ) ∥ 2 = − ∥ ∇ F ( f ( θ ( t ) ) ) ⊤ ∇ f ( θ ( t ) ) ∥ 2 ≤ − ∥ ∇ F ( f ( θ ( t ) ) ) ∥ 2 σ min ( ∇ f ( θ ( t ) ) ) 2 ≤ − 2 α F F ( f ( θ ( t ) ) ) ( 1 2 μ ) 2 = − 1 2 α F μ 2 L ( θ ( t ) ) , \begin{aligned}
\frac{\mathrm{d}}{\mathrm{d}t}\mathcal{L}(\theta(t)) &= \nabla\mathcal{L}(\theta(t))^{\top}\frac{\mathrm{d}}{\mathrm{d}t}\theta(t) = -\lVert\nabla\mathcal{L}(\theta(t)) \rVert^2 \\
&= -\left\lVert\nabla\mathcal{F}(f(\theta(t)))^{\top}\nabla f(\theta(t)) \right\rVert^2 \\
&\le-\lVert\nabla\mathcal{F}(f(\theta(t))) \rVert^2 \sigma_{\min}(\nabla f(\theta(t)))^2 \\
&\le-2\alpha_{\mathcal{F}}\mathcal{F}(f(\theta(t)))\left(\frac{1}{2}\mu\right)^2 = -\frac{1}{2}\alpha_{\mathcal{F}}\mu^2\mathcal{L}(\theta(t)),
\end{aligned} d t d L ( θ ( t )) = ∇ L ( θ ( t ) ) ⊤ d t d θ ( t ) = − ∥ ∇ L ( θ ( t )) ∥ 2 = − ∇ F ( f ( θ ( t )) ) ⊤ ∇ f ( θ ( t )) 2 ≤ − ∥ ∇ F ( f ( θ ( t ))) ∥ 2 σ m i n ( ∇ f ( θ ( t )) ) 2 ≤ − 2 α F F ( f ( θ ( t ))) ( 2 1 μ ) 2 = − 2 1 α F μ 2 L ( θ ( t )) , where in the last inequality transition we used the PL condition satisfied by F \mathcal{F} F ∇ f ( θ ( t ) ) \nabla f(\theta(t)) ∇ f ( θ ( t )) ( 1 2 μ , 3 2 ν ) \left(\frac{1}{2}\mu,\frac{3}{2}\nu\right) ( 2 1 μ , 2 3 ν )
L ( θ ( t ) ) ≤ L ( θ 0 ) e − 1 2 α F μ 2 t . \mathcal{L}(\theta(t)) \le\mathcal{L}(\theta_0)e^{-\frac{1}{2}\alpha_{\mathcal{F}}\mu^2t}. L ( θ ( t )) ≤ L ( θ 0 ) e − 2 1 α F μ 2 t . Theorem I.6 is entirely deterministic, and asks for three things at the point from which gradient flow is started: a near-isometric Jacobian, smoothness of f f f μ 2 β f \frac{\mu}{2\beta_f} 2 β f μ 1 16 c 2 α F μ 6 β f 2 \frac{1}{16c^2}\frac{\alpha_F\mu^6}{\beta_f^2} 16 c 2 1 β f 2 α F μ 6 θ C L D ( τ ) \theta_{\mathrm{CLD}}(\tau) θ CLD ( τ )
The shallow network at initialization
We follow the shallow neural network setting, employing the assumptions and initialization presented in § 4 with L = 1 L=1 L = 1
f ( θ ) = ( s 2 v T ϕ ( s 1 W X T ) ) T , f(\theta)=\left(s_2\mathbf{v}^{\mathsf{T}}\phi\left(s_1\mathbf{W}\mathbf{X}^{\mathsf{T}}\right)\right)^{\mathsf{T}}, f ( θ ) = ( s 2 v T ϕ ( s 1 W X T ) ) T , where θ = { W , v } \theta=\{\mathbf{W},\mathbf{v}\} θ = { W , v }
L ~ ( z ) ≜ 1 2 n ∥ z − y ∥ 2 , \widetilde{\mathcal{L}}(z)\triangleq\frac{1}{2n}\lVert z-y\rVert^2, L ( z ) ≜ 2 n 1 ∥ z − y ∥ 2 , such that
L ( θ ) = L ~ ( f ( θ ) ) . \mathcal{L}(\theta)=\widetilde{\mathcal{L}}(f(\theta)). L ( θ ) = L ( f ( θ )) . Theorem I.6 is thus applied with F = L ~ \mathcal{F}=\widetilde{\mathcal{L}} F = L § 17 .§ 3 acts.
Claim I.7. The MSE loss L ~ : R n → R \widetilde{\mathcal{L}}:\mathbb{R}^n\to\mathbb{R} L : R n → R L ~ ( z ) ≜ 1 2 n ∥ z − y ∥ 2 \widetilde{\mathcal{L}}(z)\triangleq\frac{1}{2n}\lVert z-y\rVert^2 L ( z ) ≜ 2 n 1 ∥ z − y ∥ 2 y ∈ R n y\in\mathbb{R}^n y ∈ R n α L = n − 1 \alpha_{\mathcal{L}}=n^{-1} α L = n − 1
Proof. This is shown by simple algebra:
∇ L ~ ( z ) = 1 n ( z − y ) , \nabla\widetilde{\mathcal{L}}(z)=\frac{1}{n}(z-y), ∇ L ( z ) = n 1 ( z − y ) , ⟹ ∥ ∇ L ~ ( z ) ∥ 2 = 1 n 2 ∥ z − y ∥ 2 = 2 n 1 2 n ∥ z − y ∥ 2 = 2 n L ~ ( z ) \Longrightarrow\left\lVert\nabla\widetilde{\mathcal{L}}(z)\right\rVert^2=\frac{1}{n^2}\lVert z-y\rVert^2=\frac{2}{n}\frac{1}{2n}\lVert z-y\rVert^2=\frac{2}{n}\widetilde{\mathcal{L}}(z) ⟹ ∇ L ( z ) 2 = n 2 1 ∥ z − y ∥ 2 = n 2 2 n 1 ∥ z − y ∥ 2 = n 2 L ( z ) ⟹ L ~ ( z ) = ∥ ∇ L ~ ( z ) ∥ 2 2 n − 1 . \Longrightarrow\widetilde{\mathcal{L}}(z)=\frac{\left\lVert\nabla\widetilde{\mathcal{L}}(z)\right\rVert^2}{2n^{-1}}. ⟹ L ( z ) = 2 n − 1 ∇ L ( z ) 2 . The two near-isometry constants are read off from the spectrum of the Jacobian at initialization: the NTK minimum eigenvalue bound of Lemma 5.3 gives μ \mu μ ν \nu ν
Proposition I.8. At initialization, for m = Ω ( n λ 1 ln ( n λ 1 ) ln n ) m=\Omega\left(\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\ln n\right) m = Ω ( λ 1 n ln ( λ 1 n ) ln n )
1 − exp ( − c 1 m n λ 1 ln ( n λ 1 ) ) , 1-\exp\left(-c_1\frac{m}{\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)}\right), 1 − exp − c 1 λ 1 n ln ( λ 1 n ) m , μ ≜ σ min ( ∇ f ( θ 0 ) ) ≥ c ~ 1 λ 1 s 2 2 m \mu\triangleq\sigma_{\min}(\nabla f(\theta_0))\ge\widetilde{c}_1\sqrt{\lambda_1s_2^2m} μ ≜ σ m i n ( ∇ f ( θ 0 )) ≥ c 1 λ 1 s 2 2 m
Proof. From Lemma 5.3 , we have for m = Ω ( n λ 1 ln ( n λ 1 ) ln n ) m=\Omega\left(\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\ln n\right) m = Ω ( λ 1 n ln ( λ 1 n ) ln n ) 1 − exp ( − c 1 m n λ 1 ln ( n λ 1 ) ) 1-\exp\left(-c_1\frac{m}{\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)}\right) 1 − exp ( − c 1 λ 1 n l n ( λ 1 n ) m )
λ min ( K N T K ( θ 0 ) ) ≥ c 0 λ 1 s 2 2 m . \lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_0)\right)\ge c_0\lambda_1s_2^2m. λ m i n ( K NTK ( θ 0 ) ) ≥ c 0 λ 1 s 2 2 m . Since K N T K \mathbf{K}_{\mathrm{NTK}} K NTK
K N T K ( θ ) = ∇ f ( θ ) ∇ f ( θ ) T \mathbf{K}_{\mathrm{NTK}}(\theta)=\nabla f(\theta)\nabla f(\theta)^{\mathsf{T}} K NTK ( θ ) = ∇ f ( θ ) ∇ f ( θ ) T for all θ ∈ R N p \theta\in\mathbb{R}^{N_p} θ ∈ R N p
μ = σ min ( ∇ f ( θ 0 ) ) = λ min ( K N T K ( θ 0 ) ) ≥ c 0 λ 1 s 2 2 m = c ~ 1 λ 1 s 2 2 m . \mu=\sigma_{\min}(\nabla f(\theta_0))=\sqrt{\lambda_{\min}\left(\mathbf{K}_{\mathrm{NTK}}(\theta_0)\right)}\ge\sqrt{c_0\lambda_1s_2^2m}=\widetilde{c}_1\sqrt{\lambda_1s_2^2m}. μ = σ m i n ( ∇ f ( θ 0 )) = λ m i n ( K NTK ( θ 0 ) ) ≥ c 0 λ 1 s 2 2 m = c 1 λ 1 s 2 2 m . Proposition 1.9. At initialization, we have that with probability at least 1 − 4 n exp ( − c ~ 2 m ) 1-4n\exp(-\tilde{c}_2m) 1 − 4 n exp ( − c ~ 2 m )
ν ≜ σ max ( ∇ f ( θ 0 ) ) ≤ 4 α ϕ 2 c ϕ , σ 0 2 s 2 2 s 1 2 m d ( ν 0 2 a 1 + 1 a 2 ) n . \nu\triangleq\sigma_{\max}\left(\nabla f(\theta_0)\right) \leq\sqrt{4\alpha_{\phi}^{2}c_{\phi,\sigma_0}^{2}s_2^{2}s_1^{2}md\left(\frac{\nu_0^{2}}{a_1}+\frac{1}{a_2}\right)n}. ν ≜ σ m a x ( ∇ f ( θ 0 ) ) ≤ 4 α ϕ 2 c ϕ , σ 0 2 s 2 2 s 1 2 m d ( a 1 ν 0 2 + a 2 1 ) n . With m = Ω ( n λ 1 ln ( n λ 1 ) ln n ) m=\Omega\left(\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\ln n\right) m = Ω ( λ 1 n ln ( λ 1 n ) ln n ) c 2 c_2 c 2 1 − exp ( − c 2 m ) 1-\exp(-c_2m) 1 − exp ( − c 2 m )
Proof. Fix i ∈ [ n ] i\in[n] i ∈ [ n ]
∥ ∇ f i ( θ 0 ) ∥ 2 = ∥ s 2 ϕ ( s 1 W x i ) ∥ 2 + ∑ j = 1 m ∥ s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i ∥ 2 \left\|\nabla f_i(\theta_0)\right\|^2=\left\|s_2\phi\left(s_1\mathbf{W}x_i\right)\right\|^2+\sum_{j=1}^{m}\left\|s_2v_j\phi'\left(s_1\mathbf{w}_j^\mathsf{T}x_i\right)s_1x_i\right\|^2 ∥ ∇ f i ( θ 0 ) ∥ 2 = ∥ s 2 ϕ ( s 1 W x i ) ∥ 2 + j = 1 ∑ m s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i 2 [ ϕ ( ⋅ ) ≤ α ϕ ∣ ⋅ ∣ , ∣ ϕ ′ ∣ ≤ α ϕ ] ≤ α ϕ 2 s 2 2 s 1 2 ( ∥ W x i ∥ 2 + ∑ j = 1 m v j 2 ∥ x i ∥ 2 ) , \left[\phi(\cdot)\leq\alpha_\phi|\cdot|,\ |\phi'|\leq\alpha_\phi\right]\leq\alpha_\phi^2s_2^2s_1^2\left(\|\mathbf{W}x_i\|^2+\sum_{j=1}^{m}v_j^2\|x_i\|^2\right), [ ϕ ( ⋅ ) ≤ α ϕ ∣ ⋅ ∣ , ∣ ϕ ′ ∣ ≤ α ϕ ] ≤ α ϕ 2 s 2 2 s 1 2 ( ∥ W x i ∥ 2 + j = 1 ∑ m v j 2 ∥ x i ∥ 2 ) , and the RHS is distributed as
α ϕ 2 s 2 2 s 1 2 c ϕ , σ 0 d ( ν 0 2 a 1 ∥ z ( 1 ) ∥ 2 + 1 a 2 ∥ z ( 2 ) ∥ 2 ) , \alpha_\phi^2s_2^2s_1^2c_{\phi,\sigma_0}d\left(\frac{\nu_0^2}{a_1}\left\|\mathbf{z}^{(1)}\right\|^2+\frac{1}{a_2}\left\|\mathbf{z}^{(2)}\right\|^2\right), α ϕ 2 s 2 2 s 1 2 c ϕ , σ 0 d ( a 1 ν 0 2 z ( 1 ) 2 + a 2 1 z ( 2 ) 2 ) , where z ( 1 ) , z ( 2 ) ∼ N ( 0 m , I m ) \mathbf{z}^{(1)},\mathbf{z}^{(2)}\sim\mathcal{N}(\mathbf{0}_m,\mathbf{I}_m) z ( 1 ) , z ( 2 ) ∼ N ( 0 m , I m ) [111 ] , Theorem 3.1.1, we have
∣ ∥ z ∥ 2 − m ∣ ψ 2 ≤ C , \left|\|\mathbf{z}\|_2-\sqrt{m}\right|_{\psi_2}\leq C, ∥ z ∥ 2 − m ψ 2 ≤ C , thus, for any ϵ > 0 \epsilon>0 ϵ > 0
P ( ∣ ∥ z ∥ 2 − m ∣ ≥ ϵ ) ≤ 2 exp ( − c C 2 ϵ 2 ) = 2 exp ( − c ~ ϵ 2 ) . \mathbb{P}\left(\left|\|\mathbf{z}\|_2-\sqrt{m}\right|\geq\epsilon\right)\leq2\exp\left(-\frac{c}{C^2}\epsilon^2\right)=2\exp\left(-\tilde{c}\epsilon^2\right). P ( ∥ z ∥ 2 − m ≥ ϵ ) ≤ 2 exp ( − C 2 c ϵ 2 ) = 2 exp ( − c ~ ϵ 2 ) . Observe that
P ( ∣ ∥ z ∥ 2 − m ∣ ≥ ϵ ) ≥ P ( ∥ z ∥ 2 ≥ m + ϵ ) = P ( ∥ z ∥ 2 ≥ ( m + ϵ ) 2 ) . \mathbb{P}\left(\left|\|\mathbf{z}\|_2-\sqrt{m}\right|\geq\epsilon\right)\geq\mathbb{P}\left(\|\mathbf{z}\|_2\geq\sqrt{m}+\epsilon\right)=\mathbb{P}\left(\|\mathbf{z}\|^2\geq\left(\sqrt{m}+\epsilon\right)^2\right). P ( ∥ z ∥ 2 − m ≥ ϵ ) ≥ P ( ∥ z ∥ 2 ≥ m + ϵ ) = P ( ∥ z ∥ 2 ≥ ( m + ϵ ) 2 ) . Then, via union bound, with probability of at least 1 − 4 exp ( − c ~ ϵ 2 ) 1-4\exp\left(-\tilde{c}\epsilon^2\right) 1 − 4 exp ( − c ~ ϵ 2 )
∥ ∇ f i ( θ 0 ) ∥ 2 ≤ α ϕ 2 s 2 2 s 1 2 c ϕ , σ 0 d ( ν 0 2 a 1 + 1 a 2 ) ( m + ϵ ) 2 . \|\nabla f_i(\theta_0)\|^2\leq\alpha_\phi^2s_2^2s_1^2c_{\phi,\sigma_0}d\left(\frac{\nu_0^2}{a_1}+\frac{1}{a_2}\right)\left(\sqrt{m}+\epsilon\right)^2. ∥∇ f i ( θ 0 ) ∥ 2 ≤ α ϕ 2 s 2 2 s 1 2 c ϕ , σ 0 d ( a 1 ν 0 2 + a 2 1 ) ( m + ϵ ) 2 . Choosing ϵ = m \epsilon=\sqrt{m} ϵ = m 1 − 4 exp ( − c ~ m ) 1-4\exp(-\tilde{c}m) 1 − 4 exp ( − c ~ m ) i ∈ [ n ] i\in[n] i ∈ [ n ]
∥ ∇ f i ( θ 0 ) ∥ 2 ≤ 4 α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 m d ( ν 0 2 a 1 + 1 a 2 ) . \|\nabla f_i(\theta_0)\|^2\leq4\alpha_\phi^2c_{\phi,\sigma_0}s_2^2s_1^2md\left(\frac{\nu_0^2}{a_1}+\frac{1}{a_2}\right). ∥∇ f i ( θ 0 ) ∥ 2 ≤ 4 α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 m d ( a 1 ν 0 2 + a 2 1 ) . Via union bound, we get that with probability at least 1 − 4 n exp ( − c ~ m ) 1-4n\exp(-\tilde{c}m) 1 − 4 n exp ( − c ~ m )
ν = σ max ( ∇ f ( θ 0 ) ) = λ max ( ∇ f ( θ 0 ) T ∇ f ( θ 0 ) ) ≤ Tr ( ∇ f ( θ 0 ) T ∇ f ( θ 0 ) ) = ∑ j = 1 N p ∑ i = 1 n ( ( ∇ f i ( θ 0 ) j ) 2 ) = ∑ i = 1 n ∥ ∇ f i ( θ 0 ) ∥ 2 ≤ 4 α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 m d ( ν 0 2 a 1 + 1 a 2 ) n . \begin{aligned}
\nu=\sigma_{\max}\left(\nabla f(\theta_0)\right)=\sqrt{\lambda_{\max}\left(\nabla f(\theta_0)^\mathsf{T}\nabla f(\theta_0)\right)} \\
&\leq\sqrt{\operatorname{Tr}\left(\nabla f(\theta_0)^\mathsf{T}\nabla f(\theta_0)\right)}=\sqrt{\sum_{j=1}^{N_p}\sum_{i=1}^{n}\left(\left(\nabla f_i(\theta_0)_j\right)^2\right)} \\
&=\sqrt{\sum_{i=1}^{n}\|\nabla f_i(\theta_0)\|^2}\leq\sqrt{4\alpha_\phi^2c_{\phi,\sigma_0}s_2^2s_1^2md\left(\frac{\nu_0^2}{a_1}+\frac{1}{a_2}\right)n}.
\end{aligned} ν = σ m a x ( ∇ f ( θ 0 ) ) = λ m a x ( ∇ f ( θ 0 ) T ∇ f ( θ 0 ) ) ≤ Tr ( ∇ f ( θ 0 ) T ∇ f ( θ 0 ) ) = j = 1 ∑ N p i = 1 ∑ n ( ( ∇ f i ( θ 0 ) j ) 2 ) = i = 1 ∑ n ∥∇ f i ( θ 0 ) ∥ 2 ≤ 4 α ϕ 2 c ϕ , σ 0 s 2 2 s 1 2 m d ( a 1 ν 0 2 + a 2 1 ) n . For m = Ω ( n λ 1 ln ( n λ 1 ) ln n ) ≫ ln n m=\Omega\left(\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\ln n\right)\gg\ln n m = Ω ( λ 1 n ln ( λ 1 n ) ln n ) ≫ ln n 1 − exp ( − c 2 m ) 1-\exp(-c_2m) 1 − exp ( − c 2 m ) c 2 c_2 c 2 □ \square □
The smoothness constant is established in three steps. We first bound the Hessian of each network output at every parameter vector, in terms of the largest second-layer weight (Proposition I.10). We then convert that bound into smoothness on a ball whose radius is itself set by the smoothness constant; since the constant grows linearly with the radius, the two are pinned down together by a quadratic inequality in the radius (Proposition I.11). Finally, we bound the largest second-layer weight at initialization (Proposition I.12), the only one of the three steps that is probabilistic.
Proposition I.10. For all θ ∈ R N p \theta\in\mathbb{R}^{N_p} θ ∈ R N p i ∈ [ n ] i \in[n] i ∈ [ n ]
∥ ∇ 2 f i ( θ ) ∥ 2 ≤ s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d max j ∈ [ m ] ∣ v j ∣ ) . \left\lVert\nabla^{2} f_i(\theta) \right\rVert_2 \le s_2s_1\sqrt{c_{\phi,\sigma_0}d}\left(\alpha_{\phi}+\beta_{\phi}\sqrt{c_{\phi,\sigma_0}}s_1\sqrt{d}\max_{j\in[m]}|v_j|\right). ∇ 2 f i ( θ ) 2 ≤ s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d j ∈ [ m ] max ∣ v j ∣ ) . Proof. Observe the shallow network definition:
f i ( θ ) = s 2 v T ϕ ( s 1 W x i ) = ∑ j = 1 m s 2 v j ϕ ( s 1 w j T x i ) . f_i(\theta)=s_2\mathbf{v}^{\mathsf{T}}\phi\left(s_1\mathbf{W}\mathbf{x}_i\right)=\sum_{j=1}^{m}s_2v_j\phi\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right). f i ( θ ) = s 2 v T ϕ ( s 1 W x i ) = j = 1 ∑ m s 2 v j ϕ ( s 1 w j T x i ) . The network is additively separable over neurons, and each summand depends only on { v j , w j } \{v_j,\mathbf{w}_j\} { v j , w j } j ∈ [ m ] j\in[m] j ∈ [ m ]
f i ( j ) ( v j , w j ) = s 2 v j ϕ ( s 1 w j T x i ) , f_i^{(j)}(v_j,\mathbf{w}_j)=s_2v_j\phi\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right), f i ( j ) ( v j , w j ) = s 2 v j ϕ ( s 1 w j T x i ) , ∂ f i ( j ) ( v j , w j ) ∂ v j = s 2 ϕ ( s 1 w j T x i ) , \frac{\partial f_i^{(j)}(v_j,\mathbf{w}_j)}{\partial v_j}=s_2\phi\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right), ∂ v j ∂ f i ( j ) ( v j , w j ) = s 2 ϕ ( s 1 w j T x i ) , ∂ 2 f i ( j ) ( v j , w j ) ∂ v j 2 = 0 , \frac{\partial^2 f_i^{(j)}(v_j,\mathbf{w}_j)}{\partial v_j^2}=0, ∂ v j 2 ∂ 2 f i ( j ) ( v j , w j ) = 0 , ∂ 2 f i ( j ) ( v j , w j ) ∂ v j ∂ w j = s 2 ϕ ′ ( s 1 w j T x i ) s 1 x i , \frac{\partial^2 f_i^{(j)}(v_j,\mathbf{w}_j)}{\partial v_j\partial\mathbf{w}_j}=s_2\phi'\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)s_1\mathbf{x}_i, ∂ v j ∂ w j ∂ 2 f i ( j ) ( v j , w j ) = s 2 ϕ ′ ( s 1 w j T x i ) s 1 x i , ∂ f i ( j ) ( v j , w j ) ∂ w j = s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i , \frac{\partial f_i^{(j)}(v_j,\mathbf{w}_j)}{\partial\mathbf{w}_j}=s_2v_j\phi'\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)s_1\mathbf{x}_i, ∂ w j ∂ f i ( j ) ( v j , w j ) = s 2 v j ϕ ′ ( s 1 w j T x i ) s 1 x i , ∂ 2 f i ( j ) ( v j , w j ) ∂ w j 2 = s 2 s 1 2 v j ϕ ′ ′ ( s 1 w j T x i ) x i x i T . \frac{\partial^2 f_i^{(j)}(v_j,\mathbf{w}_j)}{\partial\mathbf{w}_j^2}=s_2s_1^2v_j\phi''\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\mathbf{x}_i\mathbf{x}_i^{\mathsf{T}}. ∂ w j 2 ∂ 2 f i ( j ) ( v j , w j ) = s 2 s 1 2 v j ϕ ′′ ( s 1 w j T x i ) x i x i T . The hessian block per neuron is then
∇ 2 f i ( j ) = [ A b b T 0 ] , \nabla^2 f_i^{(j)}=
\begin{bmatrix}
\mathbf{A} & \mathbf{b}\\
\mathbf{b}^{\mathsf{T}} & 0
\end{bmatrix}, ∇ 2 f i ( j ) = [ A b T b 0 ] , where
A = s 2 s 1 2 v j ϕ ′ ′ ( s 1 w j T x i ) x i x i T , \mathbf{A}=s_2s_1^2v_j\phi''\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\mathbf{x}_i\mathbf{x}_i^{\mathsf{T}}, A = s 2 s 1 2 v j ϕ ′′ ( s 1 w j T x i ) x i x i T , b = s 2 s 1 ϕ ′ ( s 1 w j T x i ) x i . \mathbf{b}=s_2s_1\phi'\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\mathbf{x}_i. b = s 2 s 1 ϕ ′ ( s 1 w j T x i ) x i . Hence,
∥ ∇ 2 f i ( j ) ∥ 2 ≤ ∥ [ A 0 0 T 0 ] ∥ 2 + ∥ [ 0 b b T 0 ] ∥ 2 = ∥ A ∥ 2 + ∥ b ∥ 2 = s 2 s 1 2 ∣ v j ∣ ∣ ϕ ′ ′ ( s 1 w j T x i ) ∣ ∥ x i x i T ∥ 2 + s 2 s 1 ∣ ϕ ′ ( s 1 w j T x i ) ∣ ∥ x i ∥ 2 = s 2 s 1 ∥ x i ∥ 2 ( s 1 ∣ v j ∣ ∣ ϕ ′ ′ ( s 1 w j T x i ) ∣ ∥ x i ∥ 2 + ∣ ϕ ′ ( s 1 w j T x i ) ∣ ) ≤ s 2 s 1 c ϕ , σ 0 d ( β ϕ s 1 c ϕ , σ 0 d ∣ v j ∣ + α ϕ ) . \begin{aligned}
\left\lVert\nabla^2 f_i^{(j)}\right\rVert_2
&\le\left\lVert
\begin{bmatrix}
\mathbf{A} & 0\\
0^{\mathsf{T}} & 0
\end{bmatrix}
\right\rVert_2+
\left\lVert
\begin{bmatrix}
0 & \mathbf{b}\\
\mathbf{b}^{\mathsf{T}} & 0
\end{bmatrix}
\right\rVert_2
=\lVert\mathbf{A}\rVert_2+\lVert\mathbf{b}\rVert_2\\
&=s_2s_1^2|v_j|\left|\phi''\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\right|\left\lVert\mathbf{x}_i\mathbf{x}_i^{\mathsf{T}}\right\rVert_2+s_2s_1\left|\phi'\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\right|\lVert\mathbf{x}_i\rVert_2\\
&=s_2s_1\lVert\mathbf{x}_i\rVert_2\left(s_1|v_j|\left|\phi''\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\right|\lVert\mathbf{x}_i\rVert_2+\left|\phi'\left(s_1\mathbf{w}_j^{\mathsf{T}}\mathbf{x}_i\right)\right|\right)\\
&\le s_2s_1\sqrt{c_{\phi,\sigma_0}d}\left(\beta_{\phi}s_1\sqrt{c_{\phi,\sigma_0}d}|v_j|+\alpha_{\phi}\right).
\end{aligned} ∇ 2 f i ( j ) 2 ≤ [ A 0 T 0 0 ] 2 + [ 0 b T b 0 ] 2 = ∥ A ∥ 2 + ∥ b ∥ 2 = s 2 s 1 2 ∣ v j ∣ ϕ ′′ ( s 1 w j T x i ) x i x i T 2 + s 2 s 1 ϕ ′ ( s 1 w j T x i ) ∥ x i ∥ 2 = s 2 s 1 ∥ x i ∥ 2 ( s 1 ∣ v j ∣ ϕ ′′ ( s 1 w j T x i ) ∥ x i ∥ 2 + ϕ ′ ( s 1 w j T x i ) ) ≤ s 2 s 1 c ϕ , σ 0 d ( β ϕ s 1 c ϕ , σ 0 d ∣ v j ∣ + α ϕ ) . Now, from the block structure of the entire Hessian,
∥ ∇ 2 f i ( θ ) ∥ 2 = max j ∈ [ m ] ∥ ∇ 2 f i ( j ) ∥ 2 ≤ s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d max j ∈ [ m ] ∣ v j ∣ ) . \left\lVert\nabla^2 f_i(\theta)\right\rVert_2=\max_{j\in[m]}\left\lVert\nabla^2 f_i^{(j)}\right\rVert_2\le s_2s_1\sqrt{c_{\phi,\sigma_0}d}\left(\alpha_{\phi}+\beta_{\phi}\sqrt{c_{\phi,\sigma_0}}s_1\sqrt{d}\max_{j\in[m]}|v_j|\right). ∇ 2 f i ( θ ) 2 = j ∈ [ m ] max ∇ 2 f i ( j ) 2 ≤ s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d j ∈ [ m ] max ∣ v j ∣ ) . Proposition I.11. Given some θ b \theta_b θ b which satisfies max j ∈ [ m ] ∣ v j ∣ ≤ ξ \max_{j\in[m]} |v_j| \le\xi max j ∈ [ m ] ∣ v j ∣ ≤ ξ for some ξ > 0 \xi> 0 ξ > 0 then for any μ > 0 \mu> 0 μ > 0 we have that f f f is β f \beta_f β f smooth in Θ = { θ : ∥ θ − θ b ∥ ≤ μ 2 β f } \Theta= \left\{\theta: \lVert\theta-\theta_b\rVert\le\frac{\mu}{2\beta_f}\right\} Θ = { θ : ∥ θ − θ b ∥ ≤ 2 β f μ } with
β f = n s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ξ ) + 1 2 n s 2 s 1 2 c ϕ , σ 0 d β ϕ μ . \beta_f = \sqrt{n}s_2s_1\sqrt{c_{\phi,\sigma_0}d}\left(\alpha_\phi+\beta_\phi\sqrt{c_{\phi,\sigma_0}}s_1\sqrt{d\xi}\right)+\sqrt{\frac{1}{2}\sqrt{n}s_2s_1^2c_{\phi,\sigma_0}d\beta_\phi\mu}. β f = n s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ξ ) + 2 1 n s 2 s 1 2 c ϕ , σ 0 d β ϕ μ . Proof. In a ball with radius ρ > 0 \rho>0 ρ > 0 θ b \theta_b θ b θ \theta θ ∥ θ − θ b ∥ ≤ ρ \lVert\theta-\theta_b\rVert\le\rho ∥ θ − θ b ∥ ≤ ρ max j ∈ [ m ] ∣ v j ∣ ≤ ξ + ρ \max_{j\in[m]}|v_j|\le\xi+\rho max j ∈ [ m ] ∣ v j ∣ ≤ ξ + ρ i ∈ [ n ] i\in[n] i ∈ [ n ]
∥ ∇ 2 f i ( θ ) ∥ 2 ≤ s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ( ξ + ρ ) ) . \lVert\nabla^2 f_i(\theta)\rVert_2\le s_2s_1\sqrt{c_{\phi,\sigma_0}d}\left(\alpha_\phi+\beta_\phi\sqrt{c_{\phi,\sigma_0}}s_1\sqrt{d}(\xi+\rho)\right). ∥ ∇ 2 f i ( θ ) ∥ 2 ≤ s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ( ξ + ρ ) ) . Then, for every θ 1 , θ 2 \theta_1,\theta_2 θ 1 , θ 2
∥ ∇ f ( θ 1 ) − ∇ f ( θ 2 ) ∥ 2 ≤ ∥ ∇ f ( θ 1 ) − ∇ f ( θ 2 ) ∥ F \lVert\nabla f(\theta_1)-\nabla f(\theta_2)\rVert_2\le\lVert\nabla f(\theta_1)-\nabla f(\theta_2)\rVert_F ∥ ∇ f ( θ 1 ) − ∇ f ( θ 2 ) ∥ 2 ≤ ∥ ∇ f ( θ 1 ) − ∇ f ( θ 2 ) ∥ F = ∑ i ∈ [ n ] ∥ ∇ f i ( θ 1 ) − ∇ f i ( θ 2 ) ∥ 2 =\sqrt{\sum_{i\in[n]}\lVert\nabla f_i(\theta_1)-\nabla f_i(\theta_2)\rVert^2} = i ∈ [ n ] ∑ ∥ ∇ f i ( θ 1 ) − ∇ f i ( θ 2 ) ∥ 2 ≤ n s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ( ξ + ρ ) ) ∥ θ 1 − θ 2 ∥ , \le\sqrt{n}s_2s_1\sqrt{c_{\phi,\sigma_0}d}\left(\alpha_\phi+\beta_\phi\sqrt{c_{\phi,\sigma_0}}s_1\sqrt{d}(\xi+\rho)\right)\lVert\theta_1-\theta_2\rVert, ≤ n s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ( ξ + ρ ) ) ∥ θ 1 − θ 2 ∥ , where we used that, on the ball { θ : ∥ θ − θ b ∥ ≤ ρ } \{\theta:\lVert\theta-\theta_b\rVert\le\rho\} { θ : ∥ θ − θ b ∥ ≤ ρ } ∥ ∇ 2 f i ( θ ) ∥ 2 \lVert\nabla^2f_i(\theta)\rVert_2 ∥ ∇ 2 f i ( θ ) ∥ 2 ∇ f i \nabla f_i ∇ f i
β 0 = n s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ξ ) , \beta_0=\sqrt{n}s_2s_1\sqrt{c_{\phi,\sigma_0}d}\left(\alpha_\phi+\beta_\phi\sqrt{c_{\phi,\sigma_0}}s_1\sqrt{d\xi}\right), β 0 = n s 2 s 1 c ϕ , σ 0 d ( α ϕ + β ϕ c ϕ , σ 0 s 1 d ξ ) , γ = n s 2 s 1 c ϕ , σ 0 d β ϕ c ϕ , σ 0 s 1 d , \gamma=\sqrt{n}s_2s_1\sqrt{c_{\phi,\sigma_0}d}\beta_\phi\sqrt{c_{\phi,\sigma_0}}s_1\sqrt{d}, γ = n s 2 s 1 c ϕ , σ 0 d β ϕ c ϕ , σ 0 s 1 d , we observe that f f f β ( ρ ) \beta(\rho) β ( ρ )
β ( ρ ) = β 0 + γ ρ . \beta(\rho)=\beta_0+\gamma\rho. β ( ρ ) = β 0 + γ ρ . We now demand that the ball’s radius is no larger than μ 2 β f \frac{\mu}{2\beta_f} 2 β f μ
ρ ≤ μ 2 ( β 0 + γ ρ ) \rho\le\frac{\mu}{2(\beta_0+\gamma\rho)} ρ ≤ 2 ( β 0 + γ ρ ) μ ⟹ 2 γ ρ 2 + 2 β 0 ρ − μ ≤ 0 \Longrightarrow2\gamma\rho^2+2\beta_0\rho-\mu\le0 ⟹ 2 γ ρ 2 + 2 β 0 ρ − μ ≤ 0 ⟹ ρ ≤ β 0 2 + 2 γ μ − β 0 2 γ . \Longrightarrow\rho\le\frac{\sqrt{\beta_0^2+2\gamma\mu}-\beta_0}{2\gamma}. ⟹ ρ ≤ 2 γ β 0 2 + 2 γ μ − β 0 . In the ball with radius ρ = μ 2 β ( ρ ) = β 0 2 + 2 γ μ − β 0 2 γ \rho=\frac{\mu}{2\beta(\rho)}=\frac{\sqrt{\beta_0^2+2\gamma\mu}-\beta_0}{2\gamma} ρ = 2 β ( ρ ) μ = 2 γ β 0 2 + 2 γ μ − β 0 f f f β ( ρ ) \beta(\rho) β ( ρ )
β ( ρ ) = β 0 + γ ρ = β 0 + γ β 0 2 + 2 γ μ − β 0 2 γ \beta(\rho)=\beta_0+\gamma\rho=\beta_0+\gamma\frac{\sqrt{\beta_0^2+2\gamma\mu}-\beta_0}{2\gamma} β ( ρ ) = β 0 + γ ρ = β 0 + γ 2 γ β 0 2 + 2 γ μ − β 0 ≤ β 0 + 1 2 ( β 0 + 2 γ μ − β 0 ) = β 0 + 1 2 γ μ . \le\beta_0+\frac{1}{2}\left(\beta_0+\sqrt{2\gamma\mu}-\beta_0\right)=\beta_0+\frac{1}{\sqrt{2}}\sqrt{\gamma\mu}. ≤ β 0 + 2 1 ( β 0 + 2 γ μ − β 0 ) = β 0 + 2 1 γ μ . Denoting β f = β 0 + 1 2 γ μ ≥ β ( ρ ) \beta_f=\beta_0+\frac{1}{\sqrt{2}}\sqrt{\gamma\mu}\ge\beta(\rho) β f = β 0 + 2 1 γ μ ≥ β ( ρ ) f f f β f \beta_f β f ρ = μ 2 β ( ρ ) \rho=\frac{\mu}{2\beta(\rho)} ρ = 2 β ( ρ ) μ f f f β f \beta_f β f μ 2 β f \frac{\mu}{2\beta_f} 2 β f μ
Proposition I.12. At initialization, we have that with probability at least 1 − exp ( − m n ) 1-\exp\left(-\frac{m}{n}\right) 1 − exp ( − n m )
max j ∈ [ m ] ∣ v j ∣ < 2 a 2 ( ln ( 2 m ) + m n ) . \max_{j\in[m]}|v_j|<\sqrt{\frac{2}{a_2}\left(\ln(2m)+\frac{m}{n}\right)}. j ∈ [ m ] max ∣ v j ∣ < a 2 2 ( ln ( 2 m ) + n m ) . With m = Ω ( n λ 1 ln ( n λ 1 ) ln n ) m=\Omega\left(\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\ln n\right) m = Ω ( λ 1 n ln ( λ 1 n ) ln n ) and for large enough n n n (satisfied by n ≫ 1 n\gg1 n ≫ 1 ), we have that with probability at least 1 − exp ( − m n ) 1-\exp\left(-\frac{m}{n}\right) 1 − exp ( − n m )
max j ∈ [ m ] ∣ v j ∣ ≤ 4 m a 2 n . \max_{j\in[m]}|v_j|\le\sqrt{\frac{4m}{a_2n}}. j ∈ [ m ] max ∣ v j ∣ ≤ a 2 n 4 m . Proof. At initialization, v j ∼ N ( 0 , 1 a 2 2 ) v_j \sim\mathcal{N}\left(0,\frac{1}{a_2^2}\right) v j ∼ N ( 0 , a 2 2 1 ) j ∈ [ m ] j \in[m] j ∈ [ m ] ϵ > 0 \epsilon> 0 ϵ > 0 [111 Exercise 2.6] $,
P ( max j ∈ [ m ] ∣ v j ∣ ≥ ϵ ) ≤ ∑ j ∈ [ m ] P ( ∣ v j ∣ ≥ ϵ ) ≤ 2 m exp ( − a 2 ϵ 2 2 ) . \mathbb{P}\left(\max_{j\in[m]} |v_j| \ge\epsilon\right) \le\sum_{j\in[m]} \mathbb{P}(|v_j| \ge\epsilon) \le2m\exp\left(-a_2\frac{\epsilon^2}{2}\right). P ( j ∈ [ m ] max ∣ v j ∣ ≥ ϵ ) ≤ j ∈ [ m ] ∑ P ( ∣ v j ∣ ≥ ϵ ) ≤ 2 m exp ( − a 2 2 ϵ 2 ) . Demanding this probability to be smaller than some δ > 0 \delta> 0 δ > 0
2 m exp ( − a 2 ϵ 2 2 ) ≤ δ 2m\exp\left(-a_2\frac{\epsilon^2}{2}\right) \le\delta 2 m exp ( − a 2 2 ϵ 2 ) ≤ δ ϵ 2 ≥ 2 a 2 ln ( 2 m δ ) . \epsilon^2 \ge\frac{2}{a_2}\ln\left(\frac{2m}{\delta}\right). ϵ 2 ≥ a 2 2 ln ( δ 2 m ) . Choosing δ = e − m / n \delta=e^{-m/n} δ = e − m / n 1 − exp ( − m n ) 1-\exp\left(-\frac{m}{n}\right) 1 − exp ( − n m )
max j ∈ [ m ] ∣ v j ∣ < 2 a 2 ( ln ( 2 m ) + m n ) . \max_{j\in[m]} |v_j| < \sqrt{\frac{2}{a_2}\left(\ln(2m)+\frac{m}{n}\right)}. j ∈ [ m ] max ∣ v j ∣ < a 2 2 ( ln ( 2 m ) + n m ) . We now assume m = Ω ( n λ 1 ln ( n λ 1 ) ln n ) = Ω ( n ( ln n ) 2 ) m=\Omega\left(\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\ln n\right)=\Omega\left(n(\ln n)^2\right) m = Ω ( λ 1 n ln ( λ 1 n ) ln n ) = Ω ( n ( ln n ) 2 ) λ 1 = O ( 1 ) \lambda_1=O(1) λ 1 = O ( 1 ) x > e x>e x > e ln x x \frac{\ln x}{x} x l n x n n n
ln ( 2 m ) = m ln ( 2 m ) m ≤ m ln ( 2 n ( ln n ) 2 ) n ( ln n ) 2 < m ln ( n 2 ) n ( ln n ) 2 = m 2 ln ( n ) n ( ln n ) 2 = m n 2 ln n < m n . \begin{aligned}
\ln(2m)=m\frac{\ln(2m)}{m}\le m\frac{\ln\left(2n(\ln n)^2\right)}{n(\ln n)^2} \\
&<m\frac{\ln(n^2)}{n(\ln n)^2}=m\frac{2\ln(n)}{n(\ln n)^2}=\frac{m}{n}\frac{2}{\ln n}<\frac{m}{n}.
\end{aligned} ln ( 2 m ) = m m ln ( 2 m ) ≤ m n ( ln n ) 2 ln ( 2 n ( ln n ) 2 ) < m n ( ln n ) 2 ln ( n 2 ) = m n ( ln n ) 2 2 ln ( n ) = n m ln n 2 < n m . Hence,
max j ∈ [ m ] ∣ v j ∣ < 2 a 2 ( ln ( 2 m ) + m n ) < 4 m a 2 n . \max_{j\in[m]} |v_j| < \sqrt{\frac{2}{a_2}\left(\ln(2m)+\frac{m}{n}\right)}<\sqrt{\frac{4m}{a_2n}}. j ∈ [ m ] max ∣ v j ∣ < a 2 2 ( ln ( 2 m ) + n m ) < a 2 n 4 m . Proof of the zero-loss guarantee by applying the second law change-of-measure inequality
Our change-of-measure method allows us to transfer initialization properties to arbitrary time properties with high probability, as well as upper bounding the expected loss at an arbitrary time by a finite value, by our main theorem Theorem 5.1. We will employ this method to show that with finite, lower bounded probability, necessary conditions for the application of Theorem I.6 hold when stopping CLD at a large enough deterministic time, allowing convergence to zero loss via gradient flow from this point.
Proof of Proposition 5.2. Step 1: transferring the initialization properties to time t t t For an event B B B δ 0 \delta_0 δ 0 (32) on the indicator function 1 B \mathbf{1}_B 1 B
p t ( B ) ≤ η δ 0 1 / ( c β κ ˉ 2 β ) p_t(B)\le\eta\delta_0^{1/(c_\beta\bar{\kappa}^2\beta)} p t ( B ) ≤ η δ 0 1/ ( c β κ ˉ 2 β ) for any t ≥ 0 t\ge0 t ≥ 0 (1) . Applying this to the complementary events of our high probability bounds, we get that at time t t t
By applying the second law with Proposition I.8, we get that with probability at least
1 − η exp ( − c 1 c β κ ˉ 2 β λ 1 ln ( n λ 1 ) ) , 1-\eta\exp\left(-c_1\frac{c_\beta\bar{\kappa}^2\beta}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\right), 1 − η exp ( − c 1 λ 1 c β κ ˉ 2 β ln ( λ 1 n ) ) , μ ≥ C 1 λ 1 s 2 2 m \mu\ge C_1\sqrt{\lambda_1s_2^2m} μ ≥ C 1 λ 1 s 2 2 m C 1 C_1 C 1
By applying the second law with Proposition I.9, we get that with probability at least
1 − η exp ( − c 2 c β κ ˉ 2 β m ) , 1-\eta\exp\left(-c_2\frac{c_\beta\bar{\kappa}^2\beta}{m}\right), 1 − η exp ( − c 2 m c β κ ˉ 2 β ) , ν ≤ C 2 s 2 2 s 1 2 m d a 2 a 1 ( a 1 + a 2 ) n \nu\le C_2\sqrt{\frac{s_2^2s_1^2md}{a_2a_1}}(a_1+a_2)n ν ≤ C 2 a 2 a 1 s 2 2 s 1 2 m d ( a 1 + a 2 ) n C 2 C_2 C 2
By applying the second law with Proposition I.12 , we get that with probability at least 1 − η exp ( − m c β κ ˉ 2 β ) 1-\eta\exp\left(-\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta}\right) 1 − η exp ( − c β κ ˉ 2 β m ) max j ∈ [ m ] ∣ v j ∣ ≤ 4 m a 2 n \max_{j\in[m]}|v_j|\le\sqrt{\frac{4m}{a_2n}} max j ∈ [ m ] ∣ v j ∣ ≤ a 2 n 4 m Proposition I.11 , we get that f f f β f \beta_f β f Θ = { θ : ∥ θ − θ C L D ( t ) ∥ ≤ C 1 λ 1 s 2 2 m 2 β f } \Theta=\left\{\theta:\|\theta-\theta_{\mathrm{CLD}}(t)\|\le\frac{C_1\sqrt{\lambda_1s_2^2m}}{2\beta_f}\right\} Θ = { θ : ∥ θ − θ CLD ( t ) ∥ ≤ 2 β f C 1 λ 1 s 2 2 m }
β f ≤ C 3 n s 2 s 1 d ( 1 + s 1 d m a 2 n ) + n s 2 s 1 d λ 1 s 2 2 m = C 3 s 2 s 1 d ( n + s 1 d m a 2 + ( λ 1 n m ) 1 / 4 ) , \begin{aligned}
\beta_f\le C_3\sqrt{n}s_2s_1\sqrt{d}\left(1+s_1\sqrt{d}\sqrt{\frac{m}{a_2n}}\right)+\sqrt{\sqrt{n}s_2s_1d\sqrt{\lambda_1s_2^2m}} \\
&= C_3s_2s_1\sqrt{d}\left(\sqrt{n}+s_1\sqrt{\frac{dm}{a_2}}+(\lambda_1nm)^{1/4}\right),
\end{aligned} β f ≤ C 3 n s 2 s 1 d ( 1 + s 1 d a 2 n m ) + n s 2 s 1 d λ 1 s 2 2 m = C 3 s 2 s 1 d ( n + s 1 a 2 d m + ( λ 1 nm ) 1/4 ) , where C 3 C_3 C 3
Via union bound, all of these will hold together with probability at least
1 − η ( exp ( − c 1 m c β κ ˉ 2 β n λ 1 ln ( n λ 1 ) ) + exp ( − c 2 m c β κ ˉ 2 β ) + exp ( − m c β κ ˉ 2 β n ) ) 1-\eta\left(\exp\left(-c_1\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)}\right)+\exp\left(-c_2\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta}\right)+\exp\left(-\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta n}\right)\right) 1 − η exp − c 1 c β κ ˉ 2 β λ 1 n ln ( λ 1 n ) m + exp ( − c 2 c β κ ˉ 2 β m ) + exp ( − c β κ ˉ 2 β n m ) ≥ 1 − η exp ( − c 4 m c β κ ˉ 2 β n λ 1 ln ( n λ 1 ) ) , \ge1-\eta\exp\left(-c_4\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)}\right), ≥ 1 − η exp − c 4 c β κ ˉ 2 β λ 1 n ln ( λ 1 n ) m , for a suitable constant c 4 c_4 c 4 m = Ω ( κ ˉ 2 β n λ 1 ln ( n λ 1 ) ) m=\Omega\left(\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)\right) m = Ω ( κ ˉ 2 β λ 1 n ln ( λ 1 n ) ) Theorem 5.1 .
Step 2: the loss at the stopping time. Our convergence theorem (Theorem 5.1 ) bounds the loss’s expected value (over the randomness of the initialization and the noise) at time t t t
E L ( θ C L D ( t ) ) ≤ ( E L ( θ 0 ) − L ∞ ) e − c s 2 2 λ 1 m n t + L ∞ . \mathbb{E}\mathcal{L}(\theta_{\mathrm{CLD}}(t))\le\left(\mathbb{E}\mathcal{L}(\theta_0)-\mathcal{L}_{\infty}\right)e^{-c s_2^2\frac{\lambda_1m}{n}t}+\mathcal{L}_{\infty}. E L ( θ CLD ( t )) ≤ ( E L ( θ 0 ) − L ∞ ) e − c s 2 2 n λ 1 m t + L ∞ . Throughout this appendix, L ∞ \mathcal{L}_{\infty} L ∞ L ∞ ( L = 1 ) \mathcal{L}_{\infty}^{(L=1)} L ∞ ( L = 1 ) (8) . Then, for any t ≥ 0 t\ge0 t ≥ 0 L g > 0 \mathcal{L}_g>0 L g > 0
P ( L ( θ C L D ( t ) ) ≥ L g ) ≤ E L ( θ C L D ( t ) ) L g ≤ ( E L ( θ 0 ) − L ∞ ) e − c s 2 2 λ 1 m n t + L ∞ L g . \begin{aligned}
\mathbb{P}\left(\mathcal{L}(\theta_{\mathrm{CLD}}(t))\ge\mathcal{L}_g\right)\le\frac{\mathbb{E}\mathcal{L}(\theta_{\mathrm{CLD}}(t))}{\mathcal{L}_g} \\
&\le\frac{\left(\mathbb{E}\mathcal{L}(\theta_0)-\mathcal{L}_{\infty}\right)e^{-c s_2^2\frac{\lambda_1m}{n}t}+\mathcal{L}_{\infty}}{\mathcal{L}_g}.
\end{aligned} P ( L ( θ CLD ( t )) ≥ L g ) ≤ L g E L ( θ CLD ( t )) ≤ L g ( E L ( θ 0 ) − L ∞ ) e − c s 2 2 n λ 1 m t + L ∞ . For t = τ t=\tau t = τ τ = Ω ~ ( n s 2 2 λ 1 m ) \tau=\widetilde{\Omega}\left(\frac{n}{s_2^2\lambda_1m}\right) τ = Ω ( s 2 2 λ 1 m n )
( E L ( θ 0 ) − L ∞ ) e − c s 2 2 λ 1 m n τ ≤ L ∞ , \left(\mathbb{E}\mathcal{L}(\theta_0)-\mathcal{L}_{\infty}\right)e^{-c s_2^2\frac{\lambda_1m}{n}\tau}\le\mathcal{L}_{\infty}, ( E L ( θ 0 ) − L ∞ ) e − c s 2 2 n λ 1 m τ ≤ L ∞ , then
E L ( θ C L D ( τ ) ) ≤ 2 L ∞ . \mathbb{E}\mathcal{L}(\theta_{\mathrm{CLD}}(\tau))\le2\mathcal{L}_{\infty}. E L ( θ CLD ( τ )) ≤ 2 L ∞ . By Markov’s inequality, we have that for any K > 2 K>2 K > 2
P ( L ( θ C L D ( τ ) ) > K L ∞ ) ≤ E L ( θ C L D ( τ ) ) K L ∞ ≤ 2 K . \mathbb{P}\left(\mathcal{L}(\theta_{\mathrm{CLD}}(\tau))>K\mathcal{L}_{\infty}\right)\le\frac{\mathbb{E}\mathcal{L}(\theta_{\mathrm{CLD}}(\tau))}{K\mathcal{L}_{\infty}}\le\frac{2}{K}. P ( L ( θ CLD ( τ )) > K L ∞ ) ≤ K L ∞ E L ( θ CLD ( τ )) ≤ K 2 . Step 3: verifying the criterion at θ C L D ( τ ) \theta_{\mathrm{CLD}}(\tau) θ CLD ( τ ) Via union bound, we will have at time τ \tau τ
1 − η exp ( − c 4 m c β κ ˉ 2 β n λ 1 ln ( n λ 1 ) ) − 2 K 1-\eta\exp\left(-c_4\frac{m}{c_{\beta}\bar{\kappa}^{2}\beta\frac{n}{\lambda_1}\ln\left(\frac{n}{\lambda_1}\right)}\right)-\frac{2}{K} 1 − η exp − c 4 c β κ ˉ 2 β λ 1 n ln ( λ 1 n ) m − K 2 over the randomness of the initialization and the noise, that f f f β f \beta_f β f Θ = { θ : ∥ θ − θ C L D ( τ ) ∥ ≤ μ 2 β f } \Theta= \left\{ \theta: \lVert\theta- \theta_{\mathrm{CLD}}(\tau) \rVert\le\frac{\mu}{2\beta_f} \right\} Θ = { θ : ∥ θ − θ CLD ( τ )∥ ≤ 2 β f μ } ∇ f ( θ C L D ( τ ) ) \nabla f(\theta_{\mathrm{CLD}}(\tau)) ∇ f ( θ CLD ( τ )) ( μ , ν ) (\mu,\nu) ( μ , ν )
L ( θ C L D ( τ ) ) < K L ∞ , \mathcal{L}\left(\theta_{\mathrm{CLD}}(\tau)\right) < K L_{\infty}, L ( θ CLD ( τ ) ) < K L ∞ , μ = C 1 λ 1 s 2 2 m , \mu= C_1\sqrt{\lambda_1 s_2^2m}, μ = C 1 λ 1 s 2 2 m , ν = C 2 s 2 2 s 1 2 m d a 2 a 1 ( a 1 + a 2 ) n , \nu= C_2\sqrt{\frac{s_2^2s_1^2md}{a_2a_1}}(a_1+a_2)n, ν = C 2 a 2 a 1 s 2 2 s 1 2 m d ( a 1 + a 2 ) n , β f = C 3 s 2 s 1 d ( n + s 1 d m a 2 + ( λ 1 n m ) 1 / 4 ) . \beta_f = C_3s_2s_1\sqrt{d}\left(\sqrt{n}+s_1\sqrt{\frac{dm}{a_2}}+(\lambda_1nm)^{1/4}\right). β f = C 3 s 2 s 1 d ( n + s 1 a 2 d m + ( λ 1 nm ) 1/4 ) . Together with L ~ \widetilde{\mathcal{L}} L α L = n − 1 \alpha_{\mathcal{L}}=n^{-1} α L = n − 1 θ C L D ( τ ) \theta_{\mathrm{CLD}}(\tau) θ CLD ( τ )
K L ∞ ≤ L G F = c ^ α L μ 6 β f 2 ν 2 , K L_{\infty} \leq L_{\mathrm{GF}} = \widehat{c}\frac{\alpha_{\mathcal{L}}\mu^6}{\beta_f^2\nu^2}, K L ∞ ≤ L GF = c β f 2 ν 2 α L μ 6 , where
L G F = c ^ α L μ 6 β f 2 ν 2 = c ^ n − 1 C 1 6 λ 1 3 s 2 6 m 3 C 3 2 s 2 2 s 1 2 d ( n + s 1 d m a 2 + ( λ 1 n m ) 1 / 4 ) 2 C 2 2 s 2 2 s 1 2 m d a 2 a 1 ( a 1 + a 2 ) n \begin{aligned}
L_{\mathrm{GF}} = \widehat{c}\frac{\alpha_{\mathcal{L}}\mu^6}{\beta_f^2\nu^2} \\
&= \widehat{c}\frac{n^{-1}C_1^6\lambda_1^3s_2^6m^3}{C_3^2s_2^2s_1^2d\left(\sqrt{n}+s_1\sqrt{\frac{dm}{a_2}}+(\lambda_1nm)^{1/4}\right)^2 C_2^2\frac{s_2^2s_1^2md}{a_2a_1}(a_1+a_2)n}
\end{aligned} L GF = c β f 2 ν 2 α L μ 6 = c C 3 2 s 2 2 s 1 2 d ( n + s 1 a 2 d m + ( λ 1 nm ) 1/4 ) 2 C 2 2 a 2 a 1 s 2 2 s 1 2 m d ( a 1 + a 2 ) n n − 1 C 1 6 λ 1 3 s 2 6 m 3 ≲ [ λ 1 = Θ ( 1 ) ] s 2 2 m 2 n 2 d 2 s 1 4 ( 1 a 1 + 1 a 2 ) ( n + s 1 2 d m a 2 + n m ) \underset{[\lambda_1=\Theta(1)]}{\lesssim}\frac{s_2^2m^2}{n^2d^2s_1^4\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\left(n+\frac{s_1^2dm}{a_2}+\sqrt{nm}\right)} [ λ 1 = Θ ( 1 )] ≲ n 2 d 2 s 1 4 ( a 1 1 + a 2 1 ) ( n + a 2 s 1 2 d m + nm ) s 2 2 m 2 ≲ s 2 2 m 2 n 2 d 2 s 1 4 ( 1 a 1 + 1 a 2 ) ( s 1 2 d m a 2 + n m ) , \lesssim\frac{s_2^2m^2}{n^2d^2s_1^4\left(\frac{1}{a_1}+\frac{1}{a_2}\right)\left(\frac{s_1^2dm}{a_2}+\sqrt{nm}\right)}, ≲ n 2 d 2 s 1 4 ( a 1 1 + a 2 1 ) ( a 2 s 1 2 d m + nm ) s 2 2 m 2 , where in the last transition we assumed m ≫ n m \gg n m ≫ n a 1 = d s 1 2 a_1=ds_1^2 a 1 = d s 1 2 a 2 = 1 κ 2 m s 2 2 a_2=\frac{1}{\kappa^2}ms_2^2 a 2 = κ 2 1 m s 2 2
L G F ≲ s 2 2 m 2 n 2 d 2 s 1 4 ( 1 d s 1 2 + κ 2 m s 2 2 ) ( κ 2 s 1 2 d s 2 2 + n m ) . L_{\mathrm{GF}} \lesssim\frac{s_2^2m^2}{n^2d^2s_1^4\left(\frac{1}{ds_1^2}+\frac{\kappa^2}{ms_2^2}\right)\left(\frac{\kappa^2s_1^2d}{s_2^2}+\sqrt{nm}\right)}. L GF ≲ n 2 d 2 s 1 4 ( d s 1 2 1 + m s 2 2 κ 2 ) ( s 2 2 κ 2 s 1 2 d + nm ) s 2 2 m 2 . Hence, it is enough to demand that
L ∞ ≲ s 2 2 m 2 K n 2 d 2 s 1 4 ( 1 d s 1 2 + κ 2 m s 2 2 ) ( κ 2 s 1 2 d s 2 2 + n m ) (43) L_{\infty}\lesssim\frac{s_2^2m^2}{Kn^2d^2s_1^4\left(\frac{1}{ds_1^2}+\frac{\kappa^2}{ms_2^2}\right)\left(\frac{\kappa^2s_1^2d}{s_2^2}+\sqrt{nm}\right)} \tag*{(43)} L ∞ ≲ K n 2 d 2 s 1 4 ( d s 1 2 1 + m s 2 2 κ 2 ) ( s 2 2 κ 2 s 1 2 d + nm ) s 2 2 m 2 (43) in order to achieve convergence via gradient flow from time t = τ t=\tau t = τ 1 − η exp ( − c 4 m c β κ 2 β λ 1 2 ln ( n λ 1 ) − 2 K ) 1-\eta\exp\left(-c_4\frac{m}{c_{\beta}\kappa^2\beta\lambda_1^2}\ln\left(\frac{n}{\lambda_1}\right)-\frac{2}{K}\right) 1 − η exp ( − c 4 c β κ 2 β λ 1 2 m ln ( λ 1 n ) − K 2 )
Examples of known parameterizations
We conclude by making the demand Eq. (43) on the loss floor explicit under the two common parameterizations, and reading off from Table 1 the width and inverse temperature that meet it.
Standard parameterization
Under standard parameterization, s 1 = s 2 = 1 s_1=s_2=1 s 1 = s 2 = 1 κ = 1 \kappa=1 κ = 1
L ∞ ≲ m 2 K n 2 d 2 ( 1 d + 1 m ) ( d + n m ) . L_{\infty}\lesssim\frac{m^2}{Kn^2d^2\left(\frac{1}{d}+\frac{1}{m}\right)\left(d+\sqrt{nm}\right)}. L ∞ ≲ K n 2 d 2 ( d 1 + m 1 ) ( d + nm ) m 2 . Under d ≪ n 3 / 2 d \ll n^{3/2} d ≪ n 3/2 n 2 ≲ m n^2 \lesssim m n 2 ≲ m § 5.2 , the above simplifies to
L ∞ ≲ m 3 / 2 K n 5 / 2 d . \mathcal{L}_{\infty} \lesssim\frac{m^{3/2}}{K n^{5/2} d}. L ∞ ≲ K n 5/2 d m 3/2 . From Table 1 , we have that this is achieved when demanding
β = Ω ( K n 7 / 2 d m 3 / 2 ) , m = Ω ~ ( K n 9 / 2 d m 3 / 2 ) . \beta= \Omega\left(\frac{K n^{7/2} d}{m^{3/2}}\right), \qquad m = \widetilde{\Omega}\left(\frac{K n^{9/2} d}{m^{3/2}}\right). β = Ω ( m 3/2 K n 7/2 d ) , m = Ω ( m 3/2 K n 9/2 d ) . To achieve this, it is enough to demand:
m = Ω ~ ( K 2 / 5 n 9 / 5 d 2 / 5 ) , β = Ω ( K 2 / 5 n 4 / 5 d 2 / 5 ) . m = \widetilde{\Omega}\left(K^{2/5} n^{9/5} d^{2/5}\right), \qquad\beta= \Omega\left(K^{2/5} n^{4/5} d^{2/5}\right). m = Ω ( K 2/5 n 9/5 d 2/5 ) , β = Ω ( K 2/5 n 4/5 d 2/5 ) . Since β = Ω ( n ) \beta= \Omega(n) β = Ω ( n ) K = O ( n d ) K = O\left(\sqrt{\frac{n}{d}}\right) K = O ( d n ) d ≪ n d \ll\sqrt{n} d ≪ n K = Θ ( n d ) K = \Theta\left(\sqrt{\frac{n}{d}}\right) K = Θ ( d n ) m = Ω ~ ( n 2 ) m = \widetilde{\Omega}\left(n^2\right) m = Ω ( n 2 ) β = Ω ( n ) \beta= \Omega(n) β = Ω ( n )
We further note that by allowing K = Θ ( 1 ) K = \Theta(1) K = Θ ( 1 ) ≳ 1 \gtrsim1 ≳ 1
m ≳ n 5 / 3 d 2 / 3 . m \gtrsim n^{5/3}d^{2/3}. m ≳ n 5/3 d 2/3 . This result matches the result of Song et al. [105 ] , up to d d d [105 ] the data is scaled such that ∀ i ∈ [ n ] \forall i \in[n] ∀ i ∈ [ n ] ∥ x i ∥ = 1 \lVert x_i\rVert= 1 ∥ x i ∥ = 1 ∥ x i ∥ ≍ d \lVert x_i\rVert\asymp\sqrt{d} ∥ x i ∥ ≍ d
NTK parameterization
Under NTK parameterization, s 1 = 1 d s_1 = \frac{1}{\sqrt{d}} s 1 = d 1 s 2 = 1 m s_2 = \frac{1}{\sqrt{m}} s 2 = m 1 κ = 1 \kappa= 1 κ = 1
L ∞ ≲ m K n 2 ( 1 + 1 ) ( m + n m ) , \mathcal{L}_{\infty} \lesssim\frac{m}{K n^2 (1+1)(m+\sqrt{nm})}, L ∞ ≲ K n 2 ( 1 + 1 ) ( m + nm ) m , which under m ≫ n m \gg n m ≫ n
L ∞ ≲ 1 K n 2 . \mathcal{L}_{\infty} \lesssim\frac{1}{K n^2}. L ∞ ≲ K n 2 1 . As elaborated in Table 1 , this is achieved when β = Ω ( K n 4 ) \beta= \Omega(Kn^4) β = Ω ( K n 4 ) m = Ω ~ ( β n ) = Ω ~ ( K n 5 ) m = \widetilde{\Omega}(\beta n) = \widetilde{\Omega}(K n^5) m = Ω ( β n ) = Ω ( K n 5 )
Experiments
Experimental details
The experiments are performed according to the following implementation. We sample n = 10 n = 10 n = 10 [3 , 94 ] , center the inputs to have 0 mean and scale them to have Euclidean norm d \sqrt{d} d d = 64 d = 64 d = 64 (3) with scales s 1 = s 2 = 1 s_1 = s_2 = 1 s 1 = s 2 = 1 ϕ = tanh \phi= \tanh ϕ = tanh (3) is the output dimension). The parameter initialization matches Eq. (5) with a 1 = d a_1 = d a 1 = d a 2 = m a_2 = m a 2 = m (1) with step size η = 10 − 5 \eta= 10^{-5} η = 1 0 − 5 m = 160,000 m = 160{,}000 m = 160 , 000 m = max ( d 2 , n β max log n log β max ) m = \max(d^2, n\beta_{\max}\log n\log\beta_{\max}) m = max ( d 2 , n β m a x log n log β m a x ) β max \beta_{\max} β m a x β \beta β β ∈ { 10 n , 50 n , 100 n , ∞ } \beta\in\{10n, 50n, 100n, \infty\} β ∈ { 10 n , 50 n , 100 n , ∞ } β = ∞ \beta= \infty β = ∞
The experiments were run on a dedicated server using 5 Nvidia RTX A6000 GPUs, each with 48GB of RAM. The run utilized 30GB of each GPU’s RAM and took 20 hours. The amount of required resources is due to the tracking of the neural network Jacobian. Its computation relies on the function torch.autograd.functional.jacobian, and storing its output takes a substantial amount of memory, as the number of parameters is n × m × ( d + 1 ) ≈ 10 8 n \times m \times(d + 1) \approx10^8 n × m × ( d + 1 ) ≈ 1 0 8
(a) Deviation from linearized dynamics in cosine distance.
(b) Deviation from linearized dynamics in L 2 L_2 L 2
(c) Training loss (NN in solid, linearized model in dashed).
Figure 2. The deviation from the linearized dynamics as a function of number of updates for β ∈ { 10 n , 50 n , 100 n , ∞ } \beta\in\{10n, 50n, 100n, \infty\} β ∈ { 10 n , 50 n , 100 n , ∞ } L 2 L_2 L 2
Additional results
In Figures 2a and 2b we measure the neural network’s deviation from its linearized dynamics : we construct the neural network’s linearization around its initialization, defined by
h ( θ ) = f ( θ 0 ) + ∇ θ f ( θ 0 ) ( θ − θ 0 ) h(\theta) = f(\theta_0) + \nabla_\theta f(\theta_0)(\theta- \theta_0) h ( θ ) = f ( θ 0 ) + ∇ θ f ( θ 0 ) ( θ − θ 0 ) f f f h h h 10 − 6 10^{-6} 1 0 − 6 β \beta β Figure 2c we observe that despite the similarity in parameter space, the loss of the neural network improves over that of the linearized model for sufficiently large noise. This suggests that the dynamics differ in the space of predictors .
In Figure 3 we again consider the Jacobian drift throughout training. Figure 3 a is identical to Figure 1 a and presented again for completion; Figure 3 b plots the distance in Frobenius norm, divided by the norm of the Jacobian at initialization.
Figure 3. The distance of the Jacobian from its value at initialization as a function of number of updates for β ∈ { 10 n , 50 n , 100 n , ∞ } \beta\in\{10n, 50n, 100n, \infty\} β ∈ { 10 n , 50 n , 100 n , ∞ } Figure 3 a the distance is the cosine distance; In Figure 3 b the distance is the Frobenius norm divided by the norm of the Jacobian at initialization.
In Figure 4 we consider a deep neural network defined as in Eq. (3) with L = 3 L = 3 L = 3 s l = 1 s_l = 1 s l = 1 l ∈ [ L + 1 ] l \in[L + 1] l ∈ [ L + 1 ] (5) with a l = m l − 1 a_l = m_{l-1} a l = m l − 1 l ∈ [ L + 1 ] l \in[L + 1] l ∈ [ L + 1 ] m = 10 , 000 m = 10,000 m = 10 , 000 m = max ( d 2 , n β max ) m = \max(d^2, n\beta_{\max}) m = max ( d 2 , n β m a x ) β max \beta_{\max} β m a x β \beta β
Figure 4. Deep network training statistics as a function of the number of updates for β ∈ { 10 n , 50 n , 100 n , ∞ } \beta\in\{10n, 50n, 100n, \infty\} β ∈ { 10 n , 50 n , 100 n , ∞ } Figure 4 a we plot the (cosine) Jacobian drift throughout training. In Figure 4 b we plot the training loss. In Figure 4 c we show the minimal eigenvalue of the feature Gram matrix, plotting each ( β , seed ) (\beta, \text{seed}) ( β , seed ) β = ∞ \beta= \infty β = ∞
