GitHub

딥러닝 모델을 학습할 때 배치 크기를 키우면 GPU를 더 효율적으로 활용할 수 있고, 여러 장비에 연산을 분산하기도 쉬워진다.

그렇다면 가능한 한 큰 배치를 사용하는 것이 항상 좋은 선택일까?

이번에 살펴본 논문은 ICLR 2017에 발표된 On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima이다.

이 논문은 Large Batch로 학습한 모델이 Small Batch로 학습한 모델과 비슷한 학습 정확도에 도달하면서도, 테스트 데이터에서는 더 낮은 성능을 보이는 현상을 분석한다.

여기에서 중요한 점은 Large Batch 모델이 학습에 실패한 것이 아니라는 것이다.

Large Batch 모델도 학습 데이터에는 매우 잘 맞았다. 문제는 학습이 아니라 일반화였다.

논문은 이러한 Generalization Gap이 Large Batch가 Sharp Minimum에 수렴하는 현상과 관련되어 있다고 주장한다. 반면 Small Batch는 Gradient 추정 과정에서 발생하는 Noise로 인해 상대적으로 Flat Minimum에 도달하는 경향을 보였다.

내가 이 논문을 읽으면서 가장 헷갈렸던 질문은 다음과 같았다.

Sharp Minimum과 Flat Minimum은 왜 일반화 성능에서 차이가 날까?
Mini-batch에서 Gradient에 Noise가 생긴다는 것은 정확히 무슨 뜻일까?
Large Batch의 문제는 Local Minimum에 빠지는 것이라고 알고 있었는데, Sharp와 Flat은 Local Minimum과 어떻게 다른 개념일까?

이 글에서는 논문의 실험과 수식을 따라가며 위 질문을 하나씩 정리한다.

1. 이 논문에서 관찰한 문제

신경망의 학습은 전체 학습 데이터에 대한 평균 손실을 최소화하는 문제로 표현할 수 있다.

f(x)=\frac{1}{M}\sum_{i=1}^{M}f_i(x)

여기에서 $M$ 은 전체 학습 데이터 수, $f_i$ 는 i번째 데이터에 대한 손실, $x$ 는 모델의 전체 파라미터를 의미한다.

전체 데이터를 한 번에 사용하면 정확한 Full Gradient를 계산할 수 있다.

\nabla f(x)=\frac{1}{M}\sum_{i=1}^{M}\nabla f_i(x)

하지만 데이터가 많으면 매번 전체 Gradient를 계산하는 비용이 매우 크다. 그래서 실제 딥러닝에서는 일부 데이터만 추출한 Mini-batch Gradient를 사용한다.

g_B(x)=\frac{1}{|B|}\sum_{i\in B}\nabla f_i(x)

논문에서는 Small Batch로 256개의 데이터를 사용했고, Large Batch에서는 전체 학습 데이터의 10%를 하나의 배치로 사용했다.

두 조건 모두 ADAM으로 충분히 학습하여 손실이 더 이상 개선되지 않을 때까지 진행했다.

학습 정확도는 비슷했지만 테스트 정확도는 달랐다

논문의 핵심 실험 결과는 다음과 같다.

Network	SB 학습 정확도	LB 학습 정확도	SB 테스트 정확도	LB 테스트 정확도	테스트 차이
F1	99.66%	99.92%	98.03%	97.81%	0.22%p
F2	99.99%	98.35%	64.02%	59.45%	4.57%p
C1	99.89%	99.66%	80.04%	77.26%	2.78%p
C2	99.99%	99.99%	89.24%	87.26%	1.98%p
C3	99.56%	99.88%	49.58%	46.45%	3.13%p
C4	99.10%	99.57%	63.08%	57.81%	5.27%p

Large Batch는 대부분의 실험에서 Small Batch와 비슷하거나 더 높은 학습 정확도를 기록했다.

그럼에도 테스트 정확도는 모든 네트워크에서 Large Batch가 낮았다.

Large Batch의 문제는 최적화 자체가 실패한 것이 아니라, 비슷한 학습 손실을 가진 다른 형태의 해에 도달했다는 것이다.

논문은 이를 일반적인 Over-training과도 구분한다.

일반적인 Over-training에서는 학습이 진행될수록 학습 정확도는 높아지지만 테스트 정확도는 어느 시점부터 하락한다. 이 경우 Early Stopping이 도움이 될 수 있다.

그러나 논문의 Figure 2에서는 테스트 정확도가 정점에 도달한 뒤 하락하는 패턴이 관찰되지 않았다. Large Batch는 학습 초기부터 Small Batch보다 낮은 테스트 성능에 수렴했다.

따라서 이 논문에서 말하는 Generalization Gap은 단순히 너무 오래 학습해서 발생한 현상은 아니다.

2. Mini-batch Gradient에 Noise가 생긴다는 의미

Mini-batch의 Noise를 이해하려면 Full Gradient와 Mini-batch Gradient의 차이를 보면 된다.

Mini-batch Gradient를 Full Gradient와 오차 항의 합으로 표현하면 다음과 같다.

\xi_B(x)=g_B(x)-\nabla f(x)

여기에서 $ξ_B$ 가 논문에서 말하는 Gradient Noise이다.

이 Noise는 데이터에 오류를 추가하거나 정답 라벨을 훼손한다는 뜻이 아니다.

Gradient Noise란 일부 데이터만 보고 계산한 방향이 전체 데이터를 보고 계산한 방향과 조금씩 달라지는 현상이다.

예를 들어 전체 학습 데이터에서는 오른쪽으로 이동하는 것이 평균적으로 가장 좋은 방향이라고 해보자.

그런데 우연히 선택된 첫 번째 Mini-batch에는 특정 클래스의 데이터가 많을 수 있다. 이 배치에서는 오른쪽 위 방향이 더 좋은 것으로 계산될 수 있다.

다음 Mini-batch에는 다른 종류의 데이터가 많아 오른쪽 아래 방향이 계산될 수도 있다.

여러 Mini-batch의 방향을 평균하면 Full Gradient에 가까워지지만, 한 번의 업데이트만 보면 방향이 계속 흔들린다.

배치를 균일하게 무작위 추출한다고 가정하면 Mini-batch Gradient는 평균적으로 Full Gradient와 일치한다.

\mathbb{E}[g_B(x)]=\nabla f(x)

하지만 매번 같은 값이 나오는 것은 아니다. Gradient 추정값의 분산은 배치 크기가 커질수록 대체로 감소한다.

\operatorname{Var}(g_B)\approx\frac{\Sigma}{|B|}

따라서 배치 크기가 작을수록 Gradient 방향의 흔들림이 크고, 배치 크기가 커질수록 Full Gradient에 가까운 안정적인 방향이 계산된다.

모델의 업데이트 식은 다음과 같다.

x_{k+1}=x_k-\alpha_k g_{B_k}(x_k)

Small Batch에서는 $g_B$ 가 매번 조금씩 달라지므로 업데이트 경로가 매끄러운 직선이 아니라 좌우로 흔들리는 경로가 된다.

논문은 이러한 흔들림이 Sharp Minimum의 좁은 수렴 영역에서 모델을 밀어내고, 더 넓은 Flat Minimum을 탐색하도록 돕는다고 해석한다.

다만 Noise가 항상 좋은 것은 아니다. Noise가 지나치게 크면 수렴이 느려지거나 최솟값 주변에서 계속 흔들릴 수 있다.

이 논문의 주장은 Noise 자체가 무조건 일반화를 향상시킨다는 것이 아니라, 적절한 크기의 Gradient Noise가 좁은 수렴 영역에서 빠져나오는 탐색 능력을 제공할 수 있다는 것이다.

3. Sharp Minimum과 Flat Minimum은 무엇인가

Sharp Minimum은 최솟값 주변의 손실 함수가 가파르게 증가하는 지점을 의미한다. Flat Minimum은 주변의 넓은 범위에서 손실이 천천히 변하는 지점을 의미한다.

Sharp Minimum

최솟값 주변의 골짜기가 좁다.
파라미터가 조금만 이동해도 손실이 빠르게 증가한다.
특정 파라미터 방향에서 곡률이 크다.
동일한 성능을 유지하려면 파라미터를 높은 정밀도로 지정해야 한다.

Flat Minimum

최솟값 주변의 골짜기가 넓다.
파라미터가 조금 이동해도 손실이 크게 증가하지 않는다.
주변의 넓은 파라미터 영역이 비슷한 성능을 낸다.
상대적으로 낮은 정밀도로도 해를 표현할 수 있다.

최솟값 근처에서 Gradient가 거의 0이라고 하면 손실 함수의 변화는 헤시안 행렬을 이용해 다음과 같이 근사할 수 있다.

f(x+\Delta)\approx f(x)+\frac{1}{2}\Delta^{\mathsf T}\nabla^2 f(x)\Delta

헤시안의 큰 고윳값은 해당 방향으로 손실 함수가 빠르게 휘어진다는 의미이다.

따라서 Sharp Minimum에서는 헤시안에 큰 양의 고윳값이 존재하는 경향이 있다. Flat Minimum에서는 상대적으로 작은 고윳값이 많이 나타난다.

여기에서 한 가지 주의할 점이 있다.

Sharp Minimum이라고 해서 모든 방향으로 손실이 가파르게 증가하는 것은 아니다.

논문은 Large Batch 해의 주변을 조사한 결과, 전체 방향 중 일부 저차원 부분공간에서만 손실이 매우 가파르게 증가하고 나머지 방향에서는 비교적 평평할 수 있다고 설명한다.

즉, 고차원 공간의 Sharp Minimum은 단순한 뾰족한 원뿔 모양이라기보다, 몇몇 방향으로만 매우 좁은 골짜기에 가깝다.

4. Flat Minimum은 왜 일반화에 유리하다고 보는가

Sharp Minimum과 Flat Minimum이 동일한 학습 손실을 가진다고 가정해보자.

Sharp Minimum에서는 파라미터가 조금만 바뀌어도 손실이 크게 증가한다. 반면 Flat Minimum에서는 어느 정도 파라미터가 바뀌어도 손실이 비슷하게 유지된다.

학습 데이터와 테스트 데이터는 완전히 동일한 분포의 동일한 표본이 아니다. 새로운 데이터가 들어오면 모델이 보게 되는 입력과 손실 함수도 조금씩 달라진다.

이때 학습한 해가 매우 민감한 Sharp Minimum이라면 이러한 작은 변화가 테스트 손실의 큰 변화로 이어질 가능성이 있다.

반대로 Flat Minimum은 주변 변화에 상대적으로 둔감하기 때문에 학습 데이터에 없던 새로운 표본에서도 비슷한 예측을 유지할 가능성이 높다는 것이 논문의 직관이다.

파라미터 정밀도 관점

논문은 Minimum Description Length 관점도 제시한다.

Sharp Minimum에서 낮은 손실을 유지하려면 모델의 파라미터 값을 매우 정밀하게 지정해야 한다. 파라미터가 조금만 변해도 손실이 커지기 때문이다.

Flat Minimum에서는 넓은 범위의 파라미터가 비슷한 손실을 내므로 상대적으로 낮은 정밀도로 모델을 표현할 수 있다.

Minimum Description Length 관점에서는 더 적은 정보로 설명할 수 있는 단순한 모델이 일반화에 유리하다고 본다. 따라서 낮은 정밀도로 기술할 수 있는 Flat Minimum이 더 좋은 일반화와 연결될 수 있다.

다만 이 논문이 Sharp Minimum이면 반드시 일반화가 나쁘다는 수학적 증명을 제공한 것은 아니다.

논문은 Large Batch의 낮은 테스트 성능과 높은 Sharpness 사이의 강한 상관관계를 수치 실험으로 제시한다.

5. 논문은 Sharpness를 어떻게 측정했는가

딥러닝 모델의 전체 헤시안을 계산하고 모든 고윳값을 구하는 것은 계산 비용이 매우 크다.

그래서 논문은 최솟값 주변에 작은 파라미터 변화를 주었을 때 손실이 얼마나 증가하는지를 Sharpness의 휴리스틱 지표로 사용했다.

\phi_{x,f}(\epsilon,A)=\frac{\max_{y\in C_{\epsilon}}f(x+Ay)-f(x)}{1+f(x)}\times100

이 지표는 다음과 같이 해석할 수 있다.

현재 해 $x$ 주변에 작은 파라미터 변화를 준다.
허용된 범위에서 손실을 가장 많이 증가시키는 방향을 찾는다.
기존 손실 대비 손실이 얼마나 증가했는지를 계산한다.
값이 클수록 주변 변화에 민감한 Sharp Minimum으로 본다.

전체 파라미터 공간에서 측정한 결과, Large Batch의 Sharpness는 Small Batch보다 대부분 1~2자릿수 차수가 높았다.

Network	SB Sharpness	LB Sharpness
F1	1.23	205.14
F2	1.39	310.64
C1	28.58	707.23
C2	8.68	925.32
C3	29.85	258.75
C4	12.83	421.84

ε = 10⁻³

, 전체 파라미터 공간에서 측정한 Sharpness

예를 들어 C2 네트워크에서 Small Batch 해의 Sharpness는 약 8.68이었지만 Large Batch 해는 약 925.32였다.

무작위로 선택한 100차원 부분공간에서도 Large Batch의 Sharpness가 훨씬 높게 나타났다.

논문의 Figure 3에서는 Small Batch 해와 Large Batch 해를 연결하는 선을 따라 손실을 측정했다.

α = 0은 Small Batch 해이고, α = 1은 Large Batch 해이다. 대부분의 네트워크에서 α = 1 주변의 손실 곡선이 훨씬 좁고 가파른 모습을 보였다.

6. Large Batch는 왜 Sharp Minimum에 머무르는가

논문의 설명은 Gradient Noise의 크기와 수렴 영역의 폭을 연결한다.

Sharp Minimum은 주변의 손실이 빠르게 증가하는 좁은 골짜기이다. Small Batch의 Gradient Noise가 충분히 크면 업데이트 방향이 계속 흔들리면서 이 좁은 골짜기 밖으로 빠져나올 수 있다.

반면 Large Batch는 Full Gradient에 가까운 안정적인 방향을 사용한다. 따라서 한 번 Sharp Minimum의 수렴 영역에 들어가면 바깥으로 밀어낼 만한 확률적 변동이 부족할 수 있다.

넓은 Flat Minimum에서는 상황이 다르다.

어느 정도 Gradient Noise가 존재하더라도 골짜기의 바닥이 넓기 때문에 쉽게 밖으로 빠져나가지 않는다. Small Batch는 좁은 골짜기에서는 튕겨 나오지만, 넓은 골짜기에서는 머무를 수 있다는 설명이다.

Batch Size Threshold

논문은 Batch Size를 점차 증가시키며 테스트 정확도와 Sharpness를 측정했다.

그 결과 특정 Batch Size를 넘는 시점부터 테스트 정확도가 크게 하락하고 Sharpness가 빠르게 증가하는 임계점이 관찰되었다.

F2 네트워크에서는 Batch Size 약 15,000 부근
C1 네트워크에서는 Batch Size 약 500 부근

이 임계점 아래에서는 Gradient Noise가 Sharp한 수렴 영역에서 빠져나올 만큼 충분했지만, 임계점 위에서는 Noise가 부족해졌다는 것이 논문의 해석이다.

Small Batch로 먼저 탐색하고 Large Batch로 수렴시키기

논문에서는 Small Batch로 일정 기간 학습한 모델을 시작점으로 사용한 뒤, Large Batch 학습으로 전환하는 실험도 수행했다.

Small Batch 학습을 몇 Epoch만 수행하고 Large Batch로 전환하면 여전히 Sharpness가 높고 테스트 성능도 낮았다.

반면 Small Batch로 충분히 학습하여 Flat한 영역을 발견한 뒤 Large Batch로 전환하면, Large Batch도 높은 테스트 정확도를 유지하면서 해당 영역에 수렴할 수 있었다.

Small Batch는 탐색을 담당하고, Large Batch는 이미 발견된 좋은 영역에서 빠르게 수렴하는 역할을 할 수 있다.

이 결과는 학습 초반에는 작은 배치를 사용하고, 학습이 진행될수록 배치 크기를 키우는 Dynamic Batch Size 전략의 가능성을 보여준다.

7. Local Minimum에 대한 오해 정리

기존에는 Large Batch를 사용하면 Local Minimum에 빠져 학습이 잘되지 않는다고 이해하고 있었다.

하지만 이 논문의 결과를 기준으로 보면 이 설명은 정확하지 않다.

Local Minimum 자체가 문제는 아니다

딥러닝의 손실 공간에는 많은 Local Minimum이 존재할 수 있다.

중요한 것은 Local Minimum에 도달했느냐가 아니라, 어떤 형태의 Local Minimum에 도달했느냐이다.

Sharp Minimum도 Local Minimum이다.
Flat Minimum도 Local Minimum이다.
두 해는 비슷한 학습 손실을 가질 수 있다.
하지만 주변 변화에 대한 민감도와 테스트 성능은 다를 수 있다.

논문의 Large Batch 모델은 학습 정확도가 99%에 가까웠다. 따라서 Local Minimum에 빠져 학습 자체가 실패했다고 보기 어렵다.

오히려 Large Batch는 학습 데이터에는 잘 맞는 Sharp Local Minimum에 수렴했지만, 새로운 데이터에 대한 일반화 성능이 낮았다고 보는 것이 정확하다.

Saddle Point와도 구분해야 한다

Saddle Point는 어떤 방향에서는 손실이 증가하고 다른 방향에서는 감소하는 정지점이다.

논문은 Large Batch의 Generalization Gap이 Saddle Point에 갇혔기 때문이라는 설명보다, Small Batch와 Large Batch가 질적으로 다른 Local Minimum에 도달한다는 설명을 지지한다.

정확한 표현

따라서 다음 표현은 다소 부정확하다.

Large Batch는 Local Minimum에 빠져서 학습되지 않는다.

논문의 주장에 더 가까운 표현은 다음과 같다.

Large Batch는 Gradient Noise가 작아 Sharp Minimum의 수렴 영역에서 빠져나오지 못할 수 있으며, 그 결과 학습 정확도는 높지만 테스트 일반화 성능이 낮아질 수 있다.

8. 내가 헷갈렸던 세 가지 질문 정리

8.1. Sharp와 Flat은 왜 일반화에서 차이가 나는가?

Sharp Minimum은 파라미터와 데이터의 작은 변화에도 손실이 크게 변하는 민감한 해이다.

Flat Minimum은 주변의 넓은 파라미터 범위에서 비슷한 손실을 유지한다.

따라서 학습 데이터와 조금 다른 새로운 데이터가 들어왔을 때 Flat Minimum이 예측 성능을 더 안정적으로 유지할 수 있다는 것이 논문의 일반화 직관이다.

다만 논문은 이를 수학적으로 증명한 것이 아니라, Sharpness와 테스트 정확도 사이의 수치적 상관관계를 제시했다.

8.2. Mini-batch의 Noise는 정확히 무엇인가?

데이터에 인위적인 잡음을 추가한다는 뜻이 아니다.

일부 데이터만 사용해 계산한 Gradient가 전체 데이터의 Gradient와 다른 현상을 Noise라고 부른다.

배치가 작을수록 어떤 데이터가 선택되었는지에 따라 Gradient 방향의 변동이 커진다. 배치가 커질수록 여러 데이터의 차이가 평균화되어 Full Gradient에 가까워진다.

논문은 이 확률적 변동이 좁은 Sharp Minimum을 탈출하고 넓은 Flat Minimum을 탐색하는 데 도움을 준다고 설명한다.

8.3. Large Batch의 문제는 Local Minimum인가?

단순히 Local Minimum에 도달하는 것이 문제는 아니다.

Small Batch와 Large Batch 모두 Local Minimum에 도달할 수 있다. 차이는 Small Batch가 Flat한 Local Minimum에, Large Batch가 Sharp한 Local Minimum에 도달하는 경향을 보였다는 것이다.

Large Batch 모델의 학습 정확도가 높았다는 점에서도 학습 실패나 나쁜 최적화만으로 설명하기 어렵다.

따라서 이 논문의 핵심 문제는 Optimization Failure가 아니라 Solution Geometry와 Generalization이다.

9. Large Batch 문제를 개선하려는 시도

논문은 Large Batch의 Generalization Gap을 줄이기 위한 몇 가지 방법을 실험했다.

Data Augmentation

이미지를 반전하고 회전하거나 이동하여 학습 데이터를 증가시켰다.

일부 네트워크에서는 Large Batch의 테스트 정확도가 Small Batch와 비슷한 수준까지 개선되었다. 하지만 학습된 해의 Sharpness는 여전히 높았다.

Conservative Training

현재 파라미터에서 너무 멀리 이동하지 않도록 제약을 추가하면서 하나의 Large Batch를 더 충분히 활용하는 방법을 적용했다.

테스트 정확도는 일부 개선되었지만 Sharpness 문제를 해결하지는 못했다.

Adversarial·Robust Training

파라미터나 입력의 작은 변화에 대해 최악의 손실을 줄이는 Robust Optimization도 검토했다.

그러나 논문의 실험에서는 Adversarial Training과 Stability Training이 Large Batch의 일반화 성능을 뚜렷하게 개선하지 못했다.

Dynamic Batch Size

논문에서 가장 가능성이 높게 제시된 방향 중 하나는 학습 초반에는 Small Batch를 사용하고, 학습이 진행됨에 따라 Batch Size를 증가시키는 방법이다.

초기에는 Gradient Noise를 이용해 넓은 영역을 탐색하고, 이후에는 Large Batch를 이용해 계산 효율과 안정적인 수렴을 확보하는 방식이다.

10. 논문을 해석할 때 주의할 점

논문은 Large Batch가 반드시 Sharp Minimum에 수렴한다는 이론적 증명을 제시하지 않았다.
6개의 네트워크와 제한된 학습 설정에서 수치적 증거를 제시했다.
사용한 Sharpness Metric은 계산 가능한 휴리스틱이며, Sharpness의 완전한 정의는 아니다.
Sharpness와 일반화 사이의 관계는 인과관계가 아니라 강한 상관관계로 제시되었다.
Large Batch가 항상 나쁘다는 결론보다는, Batch Size가 Optimization Path와 도달하는 해의 성질을 바꿀 수 있다는 점이 핵심이다.

논문도 결론에서 Large Batch와 Sharp Minimum의 관계를 증명할 수 있는지, Sharp와 Flat Minimum의 밀도는 어떻게 다른지, Large Batch를 Flat Minimum으로 유도할 수 있는지를 열린 질문으로 남겼다.

11. 정리

Large Batch 모델은 학습 정확도가 낮아서 문제가 된 것이 아니다.
Small Batch와 비슷한 학습 손실에 도달했지만 테스트 성능이 낮은 Generalization Gap이 발생했다.
논문에서는 Large Batch가 Sharp Minimum에, Small Batch가 Flat Minimum에 수렴하는 경향을 관찰했다.
Sharp Minimum은 일부 파라미터 방향의 작은 변화에도 손실이 빠르게 증가하는 민감한 해이다.
Flat Minimum은 주변의 넓은 범위에서 손실이 비슷하게 유지되는 상대적으로 안정적인 해이다.
Mini-batch Noise는 데이터에 잡음을 넣는 것이 아니라, 일부 표본으로 계산한 Gradient와 Full Gradient 사이의 차이를 의미한다.
Small Batch의 Gradient Noise는 Sharp한 수렴 영역에서 빠져나와 Flat한 영역을 탐색하도록 도울 수 있다.
Large Batch의 문제를 단순히 Local Minimum에 빠진다고 표현하는 것은 부정확하다.
Sharp Minimum과 Flat Minimum은 모두 Local Minimum이며, 문제는 어떤 형태의 해에 도달했는가이다.
이 논문의 핵심은 Batch Size가 단순한 연산 단위가 아니라 Optimization Path와 일반화 성능을 바꿀 수 있다는 점이다.

Small Batch의 Noise는 학습을 방해하는 오차이기도 하지만, 동시에 좁은 해를 탈출하고 넓은 해를 탐색하게 만드는 확률적 탐색 장치가 될 수 있다.

참고 논문

Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., & Tang, P. T. P. (2017). On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima. International Conference on Learning Representations.

Jay's <Devlog />