ICML 2022

PoF: Post-Training of Feature Extractor
for Improving Generalization
Ikuro Sato*12, Ryota Yamada*1,Masayuki Tanaka1, Nakamasa Inoue1, Rei Kawakami12*Equal contribution

1Tokyo Institute of Technology, Japan
2Denso IT Laboratory, Inc., Japan

Background

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

2 /19

Loss landscape and generalization ability
Previous studies showed that flatter minima tend to generalize better.

Δℒ

Curvature-based
loss increment

Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Dziugaite & Roy, 2017; Jiang et al., 2020; Dinh et al., 2017.

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

3 /19

Related work
SAM: Sharpness Aware Minimization (P. Foret et al., 2021)• Can find a flatter minimum within a “ball” of fixed radius.
• High performance gains“ball” of fixed radius

Table: Top-1 test error rates.

Flattened region

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

4 /19

MethodI. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

5 /19

Aim of this work
To develop an optimization method that optimizes feature extractor s.t.
loss landscape in the classifier parameter space becomes flatter.
Assumption: Sharp landscape tends to be more harmful in higher-layer parameter space.

Feature extractor𝑥𝑥

𝐹𝐹𝜙𝜙 𝑥𝑥

testClassifier𝐶𝐶𝜃𝜃 𝐹𝐹𝜙𝜙 𝑥𝑥

ℒℒΔℒ

𝜃𝜃0

training

Classifier parameter space

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

6 /19

PoF: Post-training of Feature-extractor
Post-trains the feature-extractor part of DNN with parameter-perturbed classifiers.
Training loss landscapeℒ 𝐶𝐶𝜃𝜃0 𝐹𝐹𝜙𝜙 𝑥𝑥

𝜃𝜃0I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

7 /19

PoF: Post-training of Feature-extractor
Post-trains the feature-extractor part of DNN with parameter-perturbed classifiers.
Training loss landscapeℒ 𝐶𝐶𝜃𝜃0 +Δ𝜃𝜃 𝐹𝐹𝜙𝜙 𝑥𝑥

𝜃𝜃0 + Δ𝜃𝜃
I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

8 /19

PoF: Post-training of Feature-extractor
Post-trains the feature-extractor part of DNN with parameter-perturbed classifiers.
Training loss landscapeℒ 𝐶𝐶𝜃𝜃0 +Δ𝜃𝜃 𝐹𝐹𝜙𝜙+Δ𝜙𝜙 𝑥𝑥
𝜃𝜃0 + Δ𝜃𝜃
I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

9 /19

PoF: Post-training of Feature-extractor
Post-trains the feature-extractor part of DNN with parameter-perturbed classifiers.
Training loss landscapeinitialfinalFlatter𝜃𝜃0I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

10 /19

Range of parameter perturbation
How can one estimate good perturbation range?

ℒ 𝑤𝑤

Training setTest set𝜃𝜃0?

𝜃𝜃I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

11 /19

Range of parameter perturbation
How can one estimate good perturbation range?
 Use mini-batch statistics!

ℒ 𝑤𝑤

Training set
Training mini-batch
Test set
Test mini-batch𝜃𝜃0𝜉𝜉𝐵𝐵⋆𝜃𝜃

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

12 /19

Range of parameter perturbation
Distribution of 𝜉𝜉𝐵𝐵⋆

2×from toy experiment The peak for test set is roughly 2x as that for training set.

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

13 /19

Algorithm

Classifier-parameter
perturbation
× 𝑛𝑛Feature-extractor
update𝜉𝜉ℬ⋆ = arg min ℒℬ 𝜙𝜙 𝑡𝑡 , 𝜃𝜃0 − 𝜉𝜉 ℒ̂ 𝐵𝐵′

Normalizedmini-batch
gradient𝜉𝜉≥0⋆ ̂′𝜕𝜕ℒ𝜙𝜙,
𝜃𝜃
−2𝜉𝜉
0𝐵𝐵�
ℬ ℒ 𝐵𝐵𝑡𝑡+1𝑡𝑡, 𝜂𝜂 > 0𝜙𝜙= 𝜙𝜙 − 𝜂𝜂𝜕𝜕𝜕𝜕I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

14 /19Theory
With certain assumption, we can derive effective loss,Hessianℒ𝒟𝒟 𝜙𝜙, 𝜃𝜃0 − 2𝜉𝜉𝐵𝐵⋆ ℒ̂ 𝐵𝐵′ℒ 𝑤𝑤⋆ ̂′ ⊤2𝜉𝜉𝐵𝐵 ℒ𝐵𝐵 ℋ𝒟𝒟 𝜙𝜙, 𝜃𝜃0 ℒ̂ 𝐵𝐵′≃ ℒ𝒟𝒟 𝜙𝜙, 𝜃𝜃0 +Rough estimate of ΔℒΔℒ𝜃𝜃0Training set
Test set𝜃𝜃

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

15 /19

Evaluation

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

16 /19

Generalization
PoF can improve generalization of network trained by SAM on 3/4 datasets.
Table: Test error rates
SAM

PoF

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

17 /19Loss curvatures
The largest Hessian component gets halved by PoF.

× 1�2I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

18 /19

Summary• Proposed PoF: Post-training of Feature extractor• The flattening range is controlled
in a data-driven manner so thatthe algorithm rough reduces Δℒ.

ℒ 𝑤𝑤Δℒ• Confirmed generalization improvement
over SAM on 3/4 datasets.

I. Sato, R. Yamada, et al., PoF: Post-Training of Feature Extractor, ICML 2022

𝜃𝜃0Training setTest set𝜃𝜃

19 /19