Direction Matters: On the Implicit Regularization Effect of - - PowerPoint PPT Presentation

Direction Matters: On the Implicit Regularization Effect of Stochastic Gradient Descent with Moderate Learning Rate

Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu Johns Hopkins University & UCLA

November 2020

Overview

• Background
• SGD vs. GD: Different Convergence Directions
• Small Learning Rate
• Moderate Learning Rate
• Direction Matters: SGD + Moderate LR is Good
• Proof Sketches

Implicit Regularization: SGD >> GD

CIFAR-10, ResNet-18, w/o weight decay, w/o data augmentation

SGD GD

Wu, Jingfeng, et al. "On the Noisy Gradient Descent that Generalizes as SGD." ICML 2020.

LS(w) = 1 n

n

X

i=1

`i(w)

Loss

Two More Figures about SGD (Less Relevant)

2500 5000 7500 10000 12500 15000 17500

LterDtLon

10 20 30 40 50 60 70

test DccurDcy (%)

GD (66.96) GLD const (66.66) GLD GynDmLc (69.25) GLD GLDg (67.96) 6GD (75.21)

Wilson, Ashia C., et al. "The marginal value of adaptive gradient methods in machine learning." NIPS 2017. Zhu, Zhanxing, et al. "The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Sharp Minima and Regularization Effects." ICML 2019.

SGD vs. GD: Learning Rate Matters!

Small LR Moderate LR GD L L SGD L J Q1: Small LR, SGD ≈ GD? Q2: Moderate LR, SGD >> GD? Q3: GD is bad anyhow?

In Theory, SGD ≈ GD or SGD ≠ GD ??

Theory disagrees with practice L

• “Easy” to prove SGD ≈ GD by concentration

<= e.g., small LR

• “Hard” to prove an inverse result

<= no concentration!

GD SGD

Unbiased noise (scales with 𝜃)

Small Learning Rate: SGD ≈ GD

dw = r LS(w)dt + pηΣ(w)

1 2 dBt

GD

Gradient Flow (GF)

SGD

Stochastic Modified Equation (SME)

Higher order term 𝜃 = 𝑒𝑢 → 0 𝜃 = 𝑒𝑢 → 0

Var[✏] = 2 Var[✏1 + · · · + ✏η] = ⌘2 ∼ O

• ⌘2

⇒ = O (√⌘)

Effects of Non-Small Learning Rate

SGD GD

𝜃 ≥ 2 𝐼 ! 𝜃 < 2 𝐼 ! 𝑀(𝑥) = 0.5 𝑥"𝐼𝑥 𝑥#$% = 𝐽 − 𝜃𝐼 ⋅ 𝑥#

LS(w) = 1 n

n

X

i=1

`i(w)

Always smooth! One of them can be non-smooth SGD + moderate and annealing LR

• Phase 1: moderate LR

=> fits smooth losses

• Phase 2: small LR

=> fits non-smooth losses

A 2-D Example

ℓ! 𝑥 = 0.5 𝑥"𝐼!𝑥, 𝐼! = 𝑒𝑗𝑏𝑕(2𝜆, 0) ℓ# 𝑥 = 0.5 𝑥"𝐼#𝑥, 𝐼# = 𝑒𝑗𝑏𝑕(0, 2) 𝑀 𝑥 = 0.5 𝑥"𝐼𝑥, 𝐼 = 𝑒𝑗𝑏𝑕(𝜆, 1) 𝜆 > 2 𝜃$ = 5 6 1.1 𝜆 , 𝑢 = 1, … , 𝑈

!

6 0.1 𝜆 , 𝑢 = 𝑈

! + 1, … , 𝑈#

𝜃$ = 6 0.1 𝜆 , 𝑢 = 1, … , 𝑈 Small LR Moderate LR

Same limits Different convergence directions

A High Dimensional Linear Regression

Setups

• Test data 𝑦 = 𝜂 ⋅ 𝜊 ∈ ℝ%, where A𝜂 ∈ 0, 1

𝜊 ∼ 𝒱(𝑇%&!)

• ℓ 𝑦; 𝑥 = 𝑥 − 𝑥∗ "𝑦𝑦" 𝑥 − 𝑥∗
• Training data X = 𝑦!, … , 𝑦( , 𝑗. 𝑗. 𝑒., 𝑒 ≫ 𝑜

WOLG

• Let 𝜇) =

𝑦) #

# ∈ (0,1]

• Assume 𝜇! ≥ 𝜇# ≥ ⋯ ≥ 𝜇(
• Let 𝑄 be the projection onto the column

space of 𝑌

• 𝑄* = 𝐽 − 𝑄

Theorem 0:

• There are multiple minima for 𝑀𝒯 𝑥 = !

( 𝑥 − 𝑥∗ "𝑌𝑌"(𝑥 − 𝑥∗)

• The iterates of gradient methods belong to a hypothesis class ℋ𝒯 = {𝑥: 𝑄*𝑥 = 𝑄*𝑥,}
• If gradient methods find a global minima, then it is the one closest to initialization

Remark: this is also known as “minimal-norm solution” since the initialization is usually zero

A High Dimensional Linear Regression

Theorem 1 (informal): Consider SGD with moderate LR, 𝜃& = 4𝜃 ∈ 6 1 𝜇% + 𝑝 1 , 6 1 𝜇! − 𝑝(1) , 𝑢 = 1, … , 𝑈

%

𝑝 1 , 𝑢 = 𝑈

% + 1, … , 𝑈!

then 𝑄 𝑥 − 𝑥∗ ‖ ‖ 𝑄(𝑥 − 𝑥∗) ! → 𝑤% ± 𝑝(1) Theorem 2 (informal): Consider GD with moderate or small LR, 𝜃& ∈ 0, 6 𝑜 2𝜇! − 𝑝(1) , 𝑢 = 1, … , 𝑈! then 𝑄 𝑥 − 𝑥∗ ‖ ‖ 𝑄(𝑥 − 𝑥∗) ! → 𝑤( ± 𝑝(1) Rayleigh quotient: 𝑆 𝑌𝑌", 𝑣 =

)!**!) )!)

Remark: 𝑤% (𝑤() is the largest (smallest) eigen vector of 𝑌𝑌"

Convergence Direction Matters

LD(walg)−inf

w LD(w) = LD(walg) −

inf

w02HS LD(w0)

| {z }

∆(walg), estimation error

+ inf

w02HS LD(w0) − inf w LD(w)

| {z }

approximation error

intrinsic error, not improvable determined by the algorithms and hyperparameters

• 𝛽-level set: 𝒳
• = {𝑥 ∈ ℋ𝒯: 𝑀𝒯 𝑥 = 𝛽}
• Optimal estimation error within a level set: Δ-

∗ = min .∈𝒳! Δ(𝑥)

Theorem 3:

• For SGD with moderate LR, Δ 𝑥12% < 1 + 𝑝 1

⋅ Δ-

∗

• For GD with moderate or small LR, Δ 𝑥2% >

⁄

3" 3# − 𝑝 1

⋅ Δ-

∗

Remark: 𝛿% (𝛿() is the largest (smallest) eigenvalue of 𝑌𝑌"

Proof Sketch of Theorem 3

Test Loss Training Loss Within a level set of the training loss, larger eigenvalue direction ⇒ smaller test loss

Proof Sketch of Theorem 2

• 𝑀𝒯 𝑥 = !

( 𝑥"𝑌𝑌"𝑥

• GD: 𝑥$4! = 𝐽 − #5

( 𝑌𝑌" 𝑥$

⇒ 𝑥$ = 𝐽 − #5

( 𝑌𝑌" $

𝑥, ⇒ 𝑣$ = 𝐽 − #5

( Γ $

𝑣,

• Moderate/small LR: 0 ≤ 1 − #53"

( $

≪ ⋯ ≪ 1 − #53#

( $

≪ 1

• Hence GD converges slower along the small eigenvalue directions

Proof Sketch of Theorem 1

• ℓ) 𝑥 = 𝑥"𝑦)𝑦)

"𝑥

• SGD in one epoch: 𝑥6,84! = 𝐽 − 2𝜃𝑦9 8 𝑦9 8

"

𝑥6,8 ⇒ 𝑥64! = ∏8:!

(

𝐽 − 2𝜃𝑦9 8 𝑦9 8

"

𝑥6

• Bounding the spectrum of ∏8:!

(

𝐽 − 2𝜃𝑦9 8 𝑦9 8

"

as projected onto 𝑌&! = (𝑦#, … , 𝑦() and its complement 6

; ;$"

• Bounding the updates of 𝑥6 as projected onto 𝑌&! and 6

; ;$"

• Reverse engineering Phase 2 and Phase 1…

Remark: with moderate LR, concentration

• f matrix products turns out to be vacuous

Henriksen, Amelia, and Rachel Ward. "Concentration inequalities for random matrix products." Linear Algebra and its Applications 594 (2020): 81-94.

Take Home

• SGD + moderate LR: converges along large eigenvalue directions
• GD or SGD + small LR: converge along small eigenvalue directions
• The former directional bias benefits generalization
• The analysis is “anti-concentration”

Get the paper ->