[97] Sharpness-Aware Minimization for Efficiently Improving Generalization (SAM)

[dhkim0225]

paper code

Probably Approximately Correct (PAC)

John Mount 의 영상 설명
전상혁님 블로그
전상혁님의 PAC learning 글을 5번은 보았던 것 같다. 4년 전에는 까막눈으로 보았었고, 매년 PAC learning 개념이 논문에서 사용될 때마다 보았었는데, 잘 이해가 안 갔었다. 기초가 쌓이기 시작한 이제서야 글이 대에에충 이해가 간다.

높은 확률로 고른 hypothesis (모델) 가 low error 를 갖게 하도록 만들자!

PAC 용 notation

Training data sample 수 $m$
Train error 와 population error 의 gap $\epsilon$.

$$ {\huge error_{true}(h) \le error_{train}(h) + \epsilon} $$

Finite hypothesis space 의 complexity $|H|$
confidence of the relation : at least $(1-\delta)$

PAC Bound

finite hypothesis space 에 대해 PAC bound 는 다음과 같이 주어진다.

$$ {\huge Pr[error_{true}(h) \le error_{train}(h) + \epsilon] \le ||H||exp(-2m\epsilon^2)} $$

즉, gap 의 upper bound 는 training data $m$ 과 complexity $|H|$, generalization error $\epsilon$, 그리고 confidence $\delta$ 에 의해 정의된다.
$\epsilon$ 은 generalization 과 overfitting 과 직접적인 연관이 있다.
$\epsilon$ 이 직접적인 generalization term이고,
$\epsilon$ 이 크면 클 수록 더 좋은 대체 모델이 존재할 확률이 커지므로, overfitting 과도 연관이 된다.

PAC Bayesian Generalization Bound (Appendix A.1)

gaussian perturbation 을 줘도 population loss L_D 가 test error 를 감소시키지 않을 것이라는 제한을 두고 theorem 증명일 진행한다.
증명은 못 알아먹겠다. 생략.

Notation

Model ${\huge w \in \mathcal{W} \in \mathbb{R}^d}$
Per-data-point loss function ${\huge l: \mathcal{W} \times \mathcal{X} \times \mathcal{Y} \rightarrow \mathbb{R}_{+}}$

training set loss

population loss

SAM

population loss는 알 수 없으니 training set loss 로 이를 근사하고, optimizing 하는 게 일반적이다. 근데, sub-optimal 로 빠지기 쉽기 때문에 새로운 방식으로 SAM을 제시한다.

시작은 간소화된 PAC Bayesian Generalization Bound 이다.
간소화된 수식에 ${\huge -L_S (w)+ L_S (w)}$ 를 중간 텀에 넣어주면, sharpness 와 train loss, regularization 을 뽑아낼 수 있다.

즉, train loss 에 regularaization 을 더해서 사용하는 일반적인 경우, sharpness 는 신경 안 쓰면서 PAC Learning 을 하고 있는 느낌.
저자들은, Sharpness 와 regularziation term 을 활용하는 방향으로 학습을 진행시킨다.
여기서 $\rho$ 값은 hyperparmeter 인데 (epsilon 의 크기), batch-size 에 민감해서 잘 조정해야 한다.
Norm 차원 p 값은 2 를 사용하였다. (실험적으로 결정했다고 한다. appendix c.5)

아무튼 제시한 방법을 컴퓨터로 풀어내야 하니, 수식을 좀 더 전개를 해보자.
${\huge L_S (w+\epsilon)}$ 을 최대화 하는 ${\huge \epsilon}$ 을 ${\huge \epsilon^*}$ 이라 하자.
${\huge \epsilon}$ 값이 작기 때문에, 테일러 급수를 적용할 수 있고, ${\huge L_S (w)}$ 는 epsilon 에 무관한 값이니까, 제외시킨다.
그러면, dual norm 문제가 하나 툭 튀어 나온다.
깔끔하게 풀어내면, 최종 Loss term 에 hessian term 이 더해진 형태가 나오게 된다.

Hessian 은 빼 버리는게 국룰.

알고리즘을 정리하면 다음과 같다.

그림의 1, 2, 3 순으로

1. epsilon 만큼 이동 시키고,
2. 이동된 곳에서 loss 를 구하고
3. 현재 weight에 적용.

Results

batch size m 과 rho 값을 적당히 가져가는게 성능에 주요했다.

second order 까지 실험을 했는데, first order 실험결과가 더 좋다. 이유는 저자들도 모르겠다고 한다.

성능. 나쁘지 않다.

augmentation 이 강해져도 효과가 있음을 보였다.