그래디언트 클리핑

Gradient clipping; 그래디언트 클리핑
기울기의 크기가 정해 둔 한계를 넘으면 잘라 내거나 줄여서, 한 번의 갱신이 너무 크게 일어나지 않게 막는 기법

그래디언트 클리핑(gradient clipping, 기울기 클리핑)은 역전파로 구한 기울기가 너무 클 때 그 크기를 임계값으로 제한하는 기법이다. 기울기 폭주로 한 번에 가중치가 크게 튀어 학습이 발산하는 것을 막는다. 원래 순환 신경망 학습의 문제를 풀려고 쓰였고, 지금은 트랜스포머 기반 대형 언어 모델 학습에서도 기본 설정으로 들어간다. GPT-3는 기울기 전체 노름을 1.0으로 잘랐고,^[1] 라마 2도 클리핑 값 1.0을 썼다.^[2]

방식

기울기 벡터를 $g$ , 임계값을 $c$ 라 하자.

방식	식	방향 보존	비고
값 클리핑(clip by value)	$g_{i} \leftarrow \min (\max (g_{i}, - c), c)$	아니오	원소마다 독립적으로 자름. 미콜로프가 RNN 언어 모델에 쓴 방식
노름 클리핑(clip by norm)	$g \leftarrow g \cdot \min (1, \frac{c}{‖ g ‖})$	예	파스카누 등이 제안. 현재 가장 기본
전역 노름 클리핑	모든 매개변수의 기울기를 이어 붙인 벡터 하나의 노름 $‖ g ‖ = \sqrt{\sum_{p} ‖ g_{p} ‖^{2}}$ 로 계산	예(모델 전체)	PyTorch `clip_grad_norm_`의 동작
적응형 클리핑(AGC)	단위별로 $‖ G_{i} ‖ / ‖ W_{i} ‖ > λ$ 이면 $G_{i} \leftarrow λ \frac{‖ W_{i} ‖^{⋆}}{‖ G_{i} ‖} G_{i}$	단위마다	가중치 크기에 비례한 한계. NFNet에서 제안

값 클리핑은 큰 성분만 잘라 기울기 방향이 바뀐다. 노름 클리핑은 방향은 그대로 두고 길이만 $c$ 로 줄이므로, 학습률을 그 단계에서만 $c / ‖ g ‖$ 배 줄인 것과 같다. 임계값보다 작은 기울기는 건드리지 않는다.

배경과 이론

2013년 파스카누(Pascanu), 미콜로프(Mikolov), 벤지오(Bengio)는 RNN의 기울기 폭주와 소실을 분석하면서 오차 곡면에 가파른 "벽"(절벽)이 생긴다고 설명했다. 기울기가 폭주하는 방향은 곡률이 큰 방향과 겹치므로, 보통의 경사 하강 단계로는 벽에 닿는 순간 아주 멀리 튕겨 나가 그동안의 학습이 망가진다. 기울기 노름이 임계값을 넘을 때 길이를 줄이면 벽 근처에 머물며 다른 방향으로 내려갈 수 있다. 이들은 미콜로프가 박사 논문에서 쓴 원소별 클리핑을 바탕으로 노름 클리핑(알고리즘 1)을 제안했고, 임계값은 충분히 많은 갱신 동안의 평균 노름을 보고 정하면 되며 이 값에 학습이 그리 민감하지 않다고 했다.^[3]

2020년 장징자오(Jingzhao Zhang) 등은 실제 신경망 학습 경로에서 국소 매끄러움(헤세 행렬 크기)이 기울기 노름과 함께 커진다는 관찰에서 $(L_{0}, L_{1})$ -매끄러움이라는 완화된 조건을 세우고, 이 조건에서는 클리핑한 경사 하강이 고정 학습률 경사 하강보다 임의로 빠르게 수렴할 수 있음을 증명했다.^[4] 클리핑이 단순한 안전장치가 아니라 일종의 적응형 학습률이라는 해석이다.

적응형 그래디언트 클리핑(AGC)

딥마인드의 브록(Brock) 등은 배치 정규화 없는 ResNet(NFNet)을 큰 배치로 학습하기 위해 AGC를 제안했다. 기울기 노름과 가중치 노름의 비 $‖ G ‖_{F} / ‖ W ‖_{F}$ 는 한 번의 갱신이 가중치를 상대적으로 얼마나 바꾸는지를 나타내므로, 이 비가 $λ$ 를 넘으면 자른다. 층 단위보다 출력 단위(가중치 행렬의 행) 단위로 계산하는 편이 실험에서 더 좋았고, $‖ W_{i} ‖^{⋆} = \max (‖ W_{i} ‖, ϵ)$ (기본 $ϵ = 1 0^{- 3}$ )로 둬서 0으로 초기화된 가중치의 기울기가 항상 0으로 잘리는 것을 막는다.^[5]

PyTorch 구현

torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=2.0, error_if_nonfinite=False, foreach=None)은 모든 매개변수 기울기의 노름을 하나의 벡터처럼 계산해 제자리(in-place)에서 줄이고, 자르기 전의 전체 노름을 돌려준다.^[6] 값 클리핑은 clip_grad_value_(parameters, clip_value)다.

for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    loss = loss_fn(model(x), y)
    loss.backward()
    total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    log("grad_norm", total_norm.item())   # 자르기 전 노름을 기록해 두면 이상 징후를 잡기 좋다

같은 일을 직접 쓰면 다음과 같다. 위 함수와 결과가 같은 것을 확인했다.

import torch

def clip_by_global_norm(params, max_norm, eps=1e-6):
    grads = [p.grad for p in params if p.grad is not None]
    total = torch.linalg.vector_norm(
        torch.stack([torch.linalg.vector_norm(g) for g in grads]))
    coef = (max_norm / (total + eps)).clamp(max=1.0)   # 임계값 이하면 1
    for g in grads:
        g.mul_(coef)
    return total

함께 쓸 때 주의할 점

FP16 혼합 정밀도: GradScaler는 손실에 스케일을 곱해 역전파하므로, 클리핑 전에 scaler.unscale_(optimizer)로 원래 크기로 되돌려야 한다. 그렇지 않으면 스케일된 노름을 기준으로 잘라 버린다.^[7]
기울기 누적: 여러 마이크로배치의 기울기를 다 더한 뒤, optimizer.step() 직전에 한 번만 자른다.
분산 학습: 전역 노름은 모든 GPU에 흩어진 기울기 조각의 노름을 합쳐야 한다. 데이터 병렬(DDP)은 올 리듀스 뒤 기울기가 같으므로 그대로 쓰면 된다. FSDP1은 모델 메서드 model.clip_grad_norm_(max_norm)을 쓰고, DTensor 기반 FSDP2와 텐서 병렬은 torch.nn.utils.clip_grad_norm_이 DTensor를 처리한다. DeepSpeed는 설정 파일의 "gradient_clipping" 값으로 엔진이 처리한다.^[8]
다른 프레임워크: 허깅 페이스 Transformers의 TrainingArguments.max_grad_norm 기본값은 1.0이다. 케라스 옵티마이저는 clipnorm(변수별 노름), global_clipnorm(전역 노름), clipvalue(값) 인자를 받는다.

임계값 정하기

클리핑 없이 짧게 학습하면서 기울기 노름을 기록하고, 평소 값 근처나 조금 위로 정한다(파스카누 등의 제안).^[3]
트랜스포머 언어 모델은 전역 노름 1.0이 GPT-3, 라마 2 등 공개된 학습 설정에서 쓰인 값이다.^[1]^[2]
학습 내내 거의 모든 단계가 잘린다면 임계값이 너무 낮아 사실상 학습률을 줄인 셈이다. 학습률을 다시 보는 것이 낫다.
클리핑은 갑작스러운 손실 급등(loss spike)을 완화하지만 원인을 없애지는 않는다. 노름이 자주 치솟으면 학습률, 워밍업, 가중치 초기화, 정밀도, 데이터 이상치를 함께 점검한다.

같이 보기

각주

↑ ^1.0 ^1.1 Tom B. Brown 외, 〈Language Models are Few-Shot Learners〉, 2020, 부록 B. arXiv:2005.14165
↑ ^2.0 ^2.1 Hugo Touvron 외, 〈Llama 2: Open Foundation and Fine-Tuned Chat Models〉, 2023. arXiv:2307.09288
↑ ^3.0 ^3.1 Razvan Pascanu, Tomas Mikolov, Yoshua Bengio, 〈On the difficulty of training Recurrent Neural Networks〉, ICML 2013. arXiv:1211.5063
↑ Jingzhao Zhang, Tianxing He, Suvrit Sra, Ali Jadbabaie, 〈Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity〉, ICLR 2020. arXiv:1905.11881
↑ Andrew Brock, Soham De, Samuel L. Smith, Karen Simonyan, 〈High-Performance Large-Scale Image Recognition Without Normalization〉, ICML 2021. arXiv:2102.06171
↑ PyTorch, 〈torch.nn.utils.clip_grad_norm_〉, 2.14 문서. docs.pytorch.org
↑ PyTorch, 〈Automatic Mixed Precision examples: Gradient clipping〉, 2.14 문서. docs.pytorch.org
↑ DeepSpeed, 〈DeepSpeed Configuration JSON〉. deepspeed.ai

참고 문헌

위키독스, 〈딥 러닝을 이용한 자연어 처리 입문〉: 그래디언트 클리핑

[gpt3-1] 1.0 ^1.1 Tom B. Brown 외, 〈Language Models are Few-Shot Learners〉, 2020, 부록 B. arXiv:2005.14165

[llama2-2] 2.0 ^2.1 Hugo Touvron 외, 〈Llama 2: Open Foundation and Fine-Tuned Chat Models〉, 2023. arXiv:2307.09288

[pascanu13-3] 3.0 ^3.1 Razvan Pascanu, Tomas Mikolov, Yoshua Bengio, 〈On the difficulty of training Recurrent Neural Networks〉, ICML 2013. arXiv:1211.5063

[zhang20-4] Jingzhao Zhang, Tianxing He, Suvrit Sra, Ali Jadbabaie, 〈Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity〉, ICLR 2020. arXiv:1905.11881

[nfnet-5] Andrew Brock, Soham De, Samuel L. Smith, Karen Simonyan, 〈High-Performance Large-Scale Image Recognition Without Normalization〉, ICML 2021. arXiv:2102.06171

[ptclip-6] PyTorch, 〈torch.nn.utils.clip_grad_norm_〉, 2.14 문서. docs.pytorch.org

[ampex-7] PyTorch, 〈Automatic Mixed Precision examples: Gradient clipping〉, 2.14 문서. docs.pytorch.org

[dscfg-8] DeepSpeed, 〈DeepSpeed Configuration JSON〉. deepspeed.ai

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]