Slide 1

Slide 1 text

Convergence of mean field Langevin dynamics and its application to neural network feature learning 1 Taiji Suzuki The University of Tokyo / AIP-RIKEN 15th/Mar/2024 Workshop on Optimal Transport, Berlin (Deep learning theory team)

Slide 2

Slide 2 text

Outline of this talk 2 Convex β€’ 𝐹 is convex: β€’ Entropy regularization: Application: Training 2-layer NN in mean field regime. [Convergence] β€’ We introduce mean field Langevin dynamics (MFLD) to minimize β„’. β€’ We show its linear convergence under a log-Sobolev inequality condition. [Generalization error analysis] β€’ A generalization error analysis of 2-layer NN trained by MFLD is given. β€’ Separation from kernel methods is shown.

Slide 3

Slide 3 text

Feature learning of NN 3 Benefit of feature learning with optimization guarantee. β€’ [Computation] Suzuki, Wu, Nitanda: β€œConvergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction.” NeurIPS2023. β€’ [Generalization] ➒ Suzuki, Wu, Oko, Nitanda: β€œFeature learning via mean-field Langevin dynamics: Classifying sparse parities and beyond.” NeurIPS2023. ➒ Nitanda, Oko, Suzuki, Wu: β€œAnisotropy helps: improved statistical and computational complexity of the mean-field Langevin dynamics under structured data.” ICLR2024. Especially, we compare the generalization error between neural networks and kernel methods. Trade-off: Statistical complexity vs computational complexity Feature learning Optimization Neural network βœ“ Non-convex Kernel method Γ— Convex

Slide 4

Slide 4 text

Collaborators 4 Denny Wu (NYU/Flatiron Institute) Atsushi Nitanda (A*STAR) Kazusato Oko (The University of Tokyo/ RIKEN-AIP)

Slide 5

Slide 5 text

Noisy gradient descent 5 Noise Gradient descent Optimization of neural network is basically non-convex. ➒ Noisy gradient descent (e.g., SGD) is effective for non- convex optimization. Noisy perturbation is helpful to escape a local minimum. ➒ Likely converges to a flat global minimum.

Slide 6

Slide 6 text

Gradient Langevin Dynamics (GLD)6 (Gradient Langevin dynamics) (Non-convex) (Euler-Maruyama scheme) Discretization [Gelfand and Mitter (1991); Borkar and Mitter (1999); Welling and Teh (2011)] Stationary distribution: Can stay around the global minimum of 𝐹(π‘₯). Regularized loss:

Slide 7

Slide 7 text

GLD as a Wasserstein gradient flow7 : Distribution of 𝑋𝑑 (we can assume it has a density) PDE that describes πœ‡π‘‘β€™s dynamics [Fokker-Planck equation]: [linear w.r.t. 𝝁] = Stationary distribution This is the Wasserstein gradient flow to minimize the following objective: c.f., Donsker-Varadan duality formula β„’

Slide 8

Slide 8 text

2-layer NN in mean-field scaling 8 β€’ 2-layer neural network: Non-linear with respect to parameters π‘Ÿ 𝑗 , 𝑀𝑗 𝑗=1 𝑀 . where 𝑋(𝑗) = π‘Ÿπ‘— , 𝑀𝑗 and Regularized empirical risk: Non-convex Loss L2 regularization

Slide 9

Slide 9 text

Noisy gradient descent 9 Noisy gradient descent update (GLD): Does it converge? NaΓ―ve application of existing theory in gradient Langevin dynamics yields iteration complexity to achieve πœ– error. β†’ Cannot be applied to wide neural network. [Raginsky, Rakhlin and Telgarsky, 2017; Xu, Chen, Zou, and Gu, 2018; Erdogdu, Mackey and Shamir, 2018; Vempala and Wibisono, 2019] ⇔

Slide 10

Slide 10 text

Mean field limit 10 Loss function (empirical risk + regularization): 𝑀 β†’ ∞ … β˜…Mean field limit: Non-linear with respect to the parameters π‘Ÿπ‘— , 𝑀𝑗 𝑗=1 𝑀 . Convex w.r.t. πœ‡ if the loss ℓ𝑖 is convex (e.g., squared / logistic loss). [Nitanda&Suzuki, 2017][Chizat&Bach, 2018][Mei, Montanari&Nguyen, 2018][Rotskoff&Vanden-Eijnden, 2018] Linear with respect to πœ‡.

Slide 11

Slide 11 text

General form of mean field LD 11 ➒ SDE the Fokker-Planck equation of which corresponds to the Wasserstein GF: 𝐹 Gradient convex strictly convex = strictly convex + Mean field Langevin dynamics: The first variation 𝛿𝐹 π›Ώπœ‡ : 𝒫 Γ— ℝ𝑑 β†’ ℝ is defined as a continuous functional such as Definition (first variation) GLD: , ➒ ➒

Slide 12

Slide 12 text

MF-LD to optimize mean field NN 12 Loss function: (distribution of π‘‹π‘˜ ) Neuron β„Žπ‘₯ (β‹…) π‘₯ Discrete time MFLD:

Slide 13

Slide 13 text

Proximal Gibbs measure 13 𝐹 Gradient Minimizer Proximal Gibbs measure ➒The proximal Gibbs measure is a kind of β€œtentative” target. ➒It plays important role in the convergence analysis. Linearized objective at 𝝁:

Slide 14

Slide 14 text

Convergence rate 14 Proximal Gibbs measure: Theorem (Linear convergence) [Nitanda, Wu, Suzuki (AISTATS2022)][Chizat (2022)] Assumption (Log-Sobolev inequality) KL-div Fisher-div There exists 𝛼 > 0 such that for any probability measure 𝜈 (abs. cont. w.r.t. π‘πœ‡), If π‘πœ‡π‘‘ satisfies the LSI condition for any 𝑑 β‰₯ 0, then This is a non-linear extension of well known GLD convergence analysis. c.f., Polyak-Lojasiewicz condition 𝑓 π‘₯ βˆ’ 𝑓 π‘₯βˆ— ≀ 𝐢 𝛻𝑓 π‘₯ 2 The rate of convergence is characterized by LSI constant

Slide 15

Slide 15 text

Log-Sobolev inequality 15 L2-regularized loss function for mean field 2-layer NN: Proximal Gibbs: If sup 𝑧 ℓ𝑖 β€² π‘“πœ‡ (β‹…) β„Žπ‘₯ (β‹…) ≀ 𝐡, the proximal Gibbs measure π‘πœ‡ satisfies the LSI with a constant 𝛼 with ∡ Bakry-Emery criterion (1985) and Holley-Strook bounded perturbation lemma (1987) Bounded (≀ 𝐡) Strongly convex where Gaussian Bounded perturbation

Slide 16

Slide 16 text

Finite particles & discrete time algorithm 16

Slide 17

Slide 17 text

We have obtained a convergence of infinite width and continuous time dynamics. Question: Can we evaluate a finite particles & discrete time approximation errors? 17 (distribution of 𝑋𝑑) Neuron π‘₯ (vector field) (Finite particle approximation)

Slide 18

Slide 18 text

Difficulty β€’ SDE of interacting particles (McKean, Kac,…, 60’) 18 𝑑 = 1 𝑑 = 2 𝑑 = 3 𝑑 = 4 Finite particle approximation error can be amplified through time. β†’ It is difficult to bound the perturbation uniformly over time. The particles behave as if they are independent as the number of particles increases to infinity. Propagation of chaos [Sznitman, 1991; Lacker, 2021]: β€’ A naΓ―ve evaluation gives exponential growth on time: ➒ Weak interaction/Strong regularization in existing work exp 𝑑 /𝑀 [Mei et al. (2018, Theorem 3)]

Slide 19

Slide 19 text

Practical algorithm 19 β€’ Time discretization: 𝑑 β†’ π‘˜πœ‚ (πœ‚: step size, π‘˜: # of steps) β€’ Space discretization: πœ‡π‘‘ is approximatd by 𝑀 particles β€’ Stochastic gradient: 𝛻 𝛿𝐹 πœ‡ π›Ώπœ‡ β†’ π‘£π‘˜ 𝑖 πœ‡π‘‘ β†’ ො πœ‡π‘˜ = 1 𝑀 βˆ‘π›Ώ 𝑋 π‘˜ (𝑖) where and (stochastic gradient) (space discretization) (time discretization) ➒ Noisy gradient descent on 2-layer NN with finite width. 𝑀 particles 𝑋 π‘˜ 𝑖 𝑖=1 𝑀

Slide 20

Slide 20 text

Convergence analysis 20 Time discr. Space discr. Stochastic approx. Under smoothness and boundedness of the loss function, it holds that Suppose that π‘πœ‡ satisfies log-Sobolev inequality with a constant 𝛼. Theorem (One-step update) [Suzuki, Wu, Nitanda (2023)] : proximal Gibbs measure 1. 𝐹: 𝒫 β†’ ℝ is convex and has a form of 𝑭 𝝁 = 𝑳 𝝁 + π€πŸ 𝔼𝝁 𝒙 𝟐 . 2. (smoothness) 𝛻𝛿𝐿 πœ‡ π›Ώπœ‡ π‘₯ βˆ’π›»π›ΏπΏ 𝜈 π›Ώπœ‡ 𝑦 ≀ 𝐢(π‘Š2 πœ‡, 𝜈 + π‘₯ βˆ’ 𝑦 ) and (boundedness) 𝛻𝛿𝐿 πœ‡ π›Ώπœ‡ π‘₯ ≀ 𝑅. Assumption: (+ second order differentiability) NaΓ―ve bound: [Suzuki, Wu, Nitanda: Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction. arXiv:2306.07221] 𝐎(𝟏/𝑴)

Slide 21

Slide 21 text

Uniform log-Sobolev inequality 21 𝑋 π‘˜ (1) 𝑋 π‘˜ (2) 𝑋 π‘˜ (𝑁) π’³π‘˜ = 𝑋 π‘˜ 𝑖 𝑖=1 𝑀 ∼ πœ‡ π‘˜ 𝑀 : Joint distribution of 𝑀 particles. Potential of the joint distribution 𝝁 π’Œ (𝑴) on ℝ𝒅×𝑴 : where (Fisher divergence) where ➒ The finite particle dynamics is the Wasserstein gradient flow that minimizes . (Approximate) Uniform log-Sobolev inequality [Chen et al. 2022] Recall [Chen, Ren, Wang. Uniform-in-time propagation of chaos for mean field Langevin dynamics. arXiv:2212.03050, 2022.] For any 𝑴, Reference

Slide 22

Slide 22 text

Computational complexity 22 Time discr. Space discr. Stochastic approx. SG-MFLD Iteration complexity: to achieve πœ– + 𝑂(1/(πœ†2 𝛼𝑁)) accuracy. By setting , the iteration complexity becomes ➒ 𝐡 = 1/(πœ†2 π›Όπœ–) is the optimal mini-batch size. β†’ π‘˜ = 𝑂 Ξ€ log πœ–βˆ’1 πœ– . (finite sum), (stochastic gradient) (Mini-batch size = 𝐡) ➒Approximation errors are uniform in time. ➒No exponential dependency on 𝑴 (number of neurons).

Slide 23

Slide 23 text

Numerical experiment 23 Test error v.s. Number of steps (regularization term: π‘Ÿ π‘₯ = π‘₯ 2) 𝑀 β†’ ∞

Slide 24

Slide 24 text

Generalization error analysis 24

Slide 25

Slide 25 text

Generalization error analysis So far, we have obtained convergence of MFLD. β‡’ How effective is the feature learning of MFLD in terms of generalization error? 25 β€’ Benefit of feature learning? Neural network vs Kernel method (NTK vs mean field)

Slide 26

Slide 26 text

Classification task 26 Problem setting (classification): ➒Logistic loss: β„“ 𝑦𝑓 = log(1 + exp(βˆ’π‘¦π‘“)) ➒tanh activation: β„Žπ‘₯ 𝑧 = ΰ΄€ 𝑅 β‹… [tanh π‘₯1 , 𝑧 + π‘₯2 + 2 β‹… tanh π‘₯3 ]/3 Loss function and model: (+1) (-1)

Slide 27

Slide 27 text

Assumptions There exists πœ‡βˆ— such that 1. KL 𝜈, πœ‡βˆ— ≀ 𝑅, 2. π‘Œπ‘“πœ‡βˆ— 𝑍 β‰₯ 𝑐0 for some constants 𝑅, 𝑐0 > 0. 27 Assumption where 𝜈 = 𝑁(0, πœ†/(2πœ†1 )). Objective of MFLD: The Bayes classifier is attained by πβˆ— with a bounded KL-div from 𝝂. (a. s. ), KL-regularization 𝑐0 π‘Œ = 1 π‘Œ = βˆ’1 π‘“πœ‡βˆ—(𝑧) 𝑧 supp(𝑃𝑍 ) supp(𝑃𝑍 ) (+ classification calibration condition)

Slide 28

Slide 28 text

Main theorem 28 Suppose that πœ† = Θ( Ξ€ 1 𝑅), then it holds that with probability 1 βˆ’ exp βˆ’π‘‘ . Class. error Theorem 1 O ΰ΄€ 𝑅2𝑅 𝑛 β€’ β„Žπ‘₯ 𝑧 = ΰ΄€ 𝑅 β‹… [tanh π‘₯1 , 𝑧 + π‘₯2 + 2 β‹… tanh π‘₯3 ]/3 β€’ πœ‡βˆ—: KL 𝜈, πœ‡βˆ— ≀ 𝑅, π‘Œπ‘“πœ‡βˆ— 𝑍 β‰₯ 𝑐0 Existing bound: Chen et al. (2020); Nitanda, Wu, Suzuki (2021) Class. Error ≀ O 1 𝑛 . (Rademacher complexity bound) β€’ Our bound provides fast learning rate (faster than 1/ 𝑛). O Ξ€ 𝑅 𝑛 β‰ͺ O Ξ€ 1 𝑛

Slide 29

Slide 29 text

Main theorem 2 29 Theorem 2 𝔼[Class. Error] ≀ O 𝑅 𝑛 . then it holds that with probability Theorem 1: 𝔼[Class. Error] ≀ O exp(βˆ’O(𝑛/𝑅2)) if 𝑛 β‰₯ 𝑅2. Theorem 2: Suppose that πœ† = Θ(1/𝑅) and If we have sufficiently large training data, we have exponential convergence of test error. We only need to evaluate 𝑅 to obtain a test error bound.

Slide 30

Slide 30 text

Example: k-sparse parity problem 30 β€’ π‘˜-sparse parity problem on high dimensional data ➒ 𝑍 ∼ Unif( βˆ’1,1 𝑑) (up to freedom of rotation) ➒ π‘Œ = ς𝑗=1 π‘˜ 𝑍𝑗 Table 1 of [Telgarsky: Feature selection and low test error in shallow low-rotation ReLu networks, ICLR2023]. Q: Can we learn sparse π’Œ-parity with GD? Is there any benefit of neural network? β€» Suppose that we don’t know which coordinate 𝑍𝑗 is aligned to. π‘˜ = 2: XOR problem 𝑑 = 3, π‘˜ = 2 Complexity to learn XOR function (π‘˜ = 2) Only the first π‘˜-coordinates are informative.

Slide 31

Slide 31 text

Generalization bound 31 𝔼[Class. Error] ≀ O 𝑅 𝑛 . Theorem 1: 𝔼[Class. Error] ≀ O exp(βˆ’O(𝑛/𝑅2)) Theorem 2: if 𝑛 β‰₯ 𝑅2. πœ‡βˆ—: KL 𝜈, πœ‡βˆ— ≀ 𝑅, π‘Œπ‘“πœ‡βˆ— 𝑍 β‰₯ 𝑐0 (perfect classifier with margin 𝑐0) Suppose that there exists πœ‡βˆ— such that For the π‘˜-parity problem, we may take Then, Lemma We can evaluate 𝑅 required for the π‘˜-sparse parity problem: Reminder

Slide 32

Slide 32 text

Generalization error bound β€’ Setting 2: 𝑛 > 𝑑2 32 β€’ Setting 1: 𝑛 > 𝑑 ➒ Test error (classification error) = 𝐎(exp(βˆ’π’/π’…πŸ)) ➒ Test error (classification error) = 𝐎( Ξ€ 𝒅 𝒏) These are better than NTK (kernel method); Sample complexity of NTK 𝒏 = 𝛀 π’…π’Œ vs NN 𝒏 = 𝐎(𝒅) Trade-off between computational complexity and sample complexity. Our analysis provides β€’ better sample complexity β€’ discrete-time/finite-width analysis β€’ 𝑑 and π‘˜ are β€œdecoupled.” Corollary (Test accuracy of MFLD) (Computational complexity is exp O 𝑑 (But, can be relaxed to O(1) if X is anisotropic))

Slide 33

Slide 33 text

Anisotropic data structure β€’ Isotropic data: ➒Test error bound: O( Ξ€ 𝑑 𝑛). ➒Computational complexity is O(exp 𝑑 ). 33 If data has anisotropic covariance, sample / computational complexities can be much improved. True signal True signal Isotropic data distribution Anisotropic data distribution # of iterations: (𝑅 = 𝑑 yields exp(𝑑)) The data structure affects the complexities.

Slide 34

Slide 34 text

Anisotropic k-parity setting 34 Input Label (k-sparse parity) +1 +1 -1 -1 +1 +1 -1 -1 Example: coordinate wise scaling where 𝑠𝑗 is the scaling factor s.t. βˆ‘ 𝑗=1 𝑑 𝑠𝑗 2 = 1. β€’ Power low decay: 𝑠𝑗 2 ≍ π‘—βˆ’π›Ό/𝑑1βˆ’π›Ό β€’ Spiked covariance: 𝑠𝑗 2 ≍ π‘‘π›Όβˆ’1 𝑗 ∈ π‘˜ 𝑠𝑗 2 ≍ π‘‘βˆ’1 𝑗 ∈ π‘˜ + 1, 𝑛 (𝛼 ∈ [0,1]) (𝑍𝑗 = ±𝑠𝑗 )

Slide 35

Slide 35 text

Generalization error bound 35 Anisotropic π’Œ-sparse parity with coordinate wise scaling (we assume βˆ‘ 𝑗=1 𝑑 𝑠𝑗 2 = 1) Input: Label: β€’ Setting 2: 𝑛 > π‘†π‘˜ 2 β€’ Setting 1: 𝑛 > π‘†π‘˜ ➒ Test error (classification error) = ➒ Test error (classification error) = Corollary 1. 𝔼[Class. Error] ≀ O 𝑅 𝑛 . 2. 𝔼[Class. Error] ≀ O exp(βˆ’O(𝑛/𝑅2)) if 𝑛 β‰₯ 𝑅2. Lemma

Slide 36

Slide 36 text

Example 36 (with βˆ‘ 𝑗=1 𝑑 𝑠𝑗 2 = 1) Input: β€’ Setting 2: 𝑛 > π‘†π‘˜ 2 β€’ Setting 1: 𝑛 > π‘†π‘˜ ➒ Test error (classification error) = ➒ Test error (classification error) = Test error = β€’ Isotropic setting (𝒔𝒋 𝟐 = 𝟏/𝒅): Test error = 1. Power low decay (𝒔𝒋 𝟐 ≍ π’‹βˆ’πœΆ): 2. Spiked covariance Test error = Test error β€’ Anisotropic setting:

Slide 37

Slide 37 text

37 Anisotropic Isotropic k-signal components Noisy components k-signal components Noisy components Large signal Small noise Small signal Large noise π’‹βˆ’πœΆ

Slide 38

Slide 38 text

Computational complexity When πœ† = 𝑂( Ξ€ 1 𝑅), the number of iterations can be bounded by 38 By substituting, then the # of iterations can be summarized as β€’ Isotropic setting: β€’ Anisotropic setting: Anisotropic structure mitigate the computational complexity. Especially, there is no exponential dependency on 𝒅 when 𝜢 = 𝟏. (assuming π‘˜ = 𝑂(1))

Slide 39

Slide 39 text

Kernel lower bound 39 When π‘˜ = π‘˜βˆ—, Mean field: Kernel: Mean field NN can β€œdecouple” π’Œ and 𝒅, while kernel has exponential relation between them. Setting: Thm For arbitrary 𝛿 > 0, the sample complexity of kernel methods is lower bounded as (kernel)

Slide 40

Slide 40 text

Coordinate transform 40 β€’ When the input 𝑍 is isotropic, we may estimate the β€œinformative direction” by the gradients at the initialization. β€’ Then, 𝐺 estimates the informative direction. By the following coordinate transformation, we may take 𝑅 independent of 𝑑 (exp 𝑑 β†’ exp(π‘˜)): True signal True signal Isotropic input (avoiding curse of dim)

Slide 41

Slide 41 text

Sample complexity to compute G 41 By using training data with size Then, 𝑅 can be modified as we can compute 𝐺 as Without 𝐺 estimation: 𝑅 = Ξ©(𝑑). Isotropic setting:

Slide 42

Slide 42 text

Summary of the result 42 Upper bound for NN (our result) Lower bound for kernel method (our result) Improve sample complexity

Slide 43

Slide 43 text

Discussion β€’ The CSQ lower bound states that O π‘‘π‘˜βˆ’1 sample complexity is optimal for methods with polynomial order computational complexity. [Abbe et al. (2023); Refinetti et al. (2021); Ben Arous et al. (2022); Damian et al. (2022)] β€’ On the other hand, our analysis is about full-batch GD. 43 Minibatch size # of iterations Sample complexity Our analysis 𝒏 𝒆𝒅 𝒅 SGD (CSQ-lower bound) 1 π‘‘π‘˜βˆ’1 π‘‘π‘˜βˆ’1 We obtain a better sample complexity than O(π‘‘π‘˜βˆ’1) with higher computational complexity. β†’ We can obtain a polynomial order method with MFLD for anisotropic input.

Slide 44

Slide 44 text

Conclusion β€’ Mean field Langevin dynamics ➒ Mean field representation of 2-layer NNs ➒ Optimizing convex functional ➒ Convergence guarantee (Wasserstein gradient flow, Uniform-in-time propagation of chaos) β€’ Generalization error of mean field 2-layer NN ➒ Fast learning rate ➒ Sparse π‘˜-parity problem ➒Better sample complexity than kernel methods ➒Structure of data (anisotropic covariance) can improve the complexities. 44 Kernel Mean field Mean field Kernel lower bound Kernel lower bound