그래디언트 축적은 신경망 및 기타 머신러닝 모델, 특히 딥러닝을 훈련할 때 메모리 제약으로 인해 모든 샘플을 한 번에 처리할 수 없는 상황에서 큰 배치 크기의 효과를 근사하는 데 사용되는 기법이다. 옵티마이저는 단일 순전파 및 역전파 과정에서 큰 배치에 대한 그래디언트를 계산하는 대신, 여러 개의 작은 미니 배치에 걸쳐 그래디언트를 축적하고 축적된 그래디언트가 원하는 유효 배치 크기에 도달한 후에만 가중치 업데이트를 수행한다. 이 접근 방식을 통해 실무자는 단일 장치(예: GPU 또는 TPU)의 메모리 용량을 초과하는 배치 크기로 모델을 훈련할 수 있다.
이 방법은 1950년대 Robbins–Monro 알고리즘을 통해 도입된 반복 최적화 알고리즘인 확률적 경사 하강법(SGD)과 밀접한 관련이 있다. SGD에서는 전체 데이터셋에 대해 계산된 목적 함수의 실제 그래디언트를 무작위로 선택된 데이터 하위 집합에 대한 그래디언트로 근사한다. 일반적인 절충안은 각 단계에서 소수의 샘플 집합에 대해 그래디언트를 계산하는 미니 배치 접근 방식이다. 그래디언트 축적은 모델 매개변수를 업데이트하기 전에 여러 미니 배치에 걸쳐 그래디언트를 평균화하여 이 아이디어를 확장하며, 배치 크기를 메모리 사용량과 효과적으로 분리한다.
동기 및 메모리 제약
대규모 언어 모델 및 트랜스포머와 같은 대규모 모델을 훈련하려면 활성화, 그래디언트 및 옵티마이저 상태를 저장하기 위한 상당한 메모리가 필요하다. 배치 크기는 메모리 소비에 직접적인 영향을 미친다. 배치가 클수록 중간 계산에 더 많은 메모리가 필요하다. 소비자용 GPU나 엣지 장치와 같이 메모리가 제한된 하드웨어에서는 최대 실현 가능한 배치 크기가 작을 수 있다. 그러나 작은 배치 크기는 그래디언트 추정치의 노이즈가 커지고 훈련 역학이 불안정해질 수 있다. 그래디언트 축적은 순전파 및 역전파에 작은 배치를 사용하면서 여러 단계에 걸쳐 그래디언트를 축적하여 더 큰 유효 배치 크기를 달성함으로써 이 문제에 대한 해결책을 제공하며, 최대 메모리 사용량을 늘리지 않고 그래디언트 안정성을 개선한다.
그래디언트 축적의 작동 방식
표준 미니 배치 훈련에서 모델은 크기 \(b\)의 미니 배치를 처리하고, 손실을 계산하고, 역전파를 수행하여 그래디언트를 얻은 다음, SGD 또는 Adam과 같은 옵티마이저를 사용하여 즉시 가중치를 업데이트한다. 그래디언트 축적을 사용하면 프로세스가 수정된다. 모델은 \(k\)개의 미니 배치를 순차적으로 처리하고 각 역전파 후에 그래디언트를 축적한다(일반적으로 합산 또는 평균화). \(k\)단계 후에 축적된 그래디언트를 사용하여 모델 가중치를 업데이트하고 그래디언트 축적기를 0으로 재설정한다. 유효 배치 크기는 \(k \times b\)이다. 이 접근 방식은 손실 함수가 샘플에 대한 합 또는 평균일 때 배치 정규화 레이어를 주의 깊게 처리한다는 전제 하에 \(k \times b\)의 배치 크기로 훈련하는 것과 수학적으로 동일하다.
장점 및 사용 사례
그래디언트 축적은 큰 배치 크기가 유리하지만 메모리 제한으로 인해 실현 불가능한 시나리오에서 널리 사용된다. 예를 들어 생성 모델을 훈련하거나 대규모 언어 모델을 미세 조정하려면 훈련을 안정화하고 수렴을 개선하기 위해 수백 또는 수천 개의 배치 크기가 필요한 경우가 많다. 그래디언트를 축적함으로써 연구자들은 단일 GPU 또는 CPU 기반 시스템과 같이 메모리가 제한된 하드웨어에서 이러한 큰 배치를 시뮬레이션할 수 있다. 또한 그래디언트 축적은 매 단계마다 그래디언트를 동기화하지 않고도 여러 장치에서 훈련할 수 있게 하여 분산 훈련에서 통신 오버헤드를 줄일 수 있다. 이는 다양한 메모리 용량의 GPU 인스턴스를 제공하는 Amazon Web Services, Microsoft Azure, Google Cloud와 같은 클라우드 기반 플랫폼에서 특히 관련이 있다.
배치 크기 및 학습률과의 관계
그래디언트 축적을 사용하면 유효 배치 크기가 증가하며, 이는 종종 학습률을 조정해야 한다. 실제로 대규모 배치 훈련에 대한 기존 연구에서 제안된 바와 같이 학습률을 배치 크기에 비례하여 선형적으로 확장하는 것이 일반적인 휴리스틱이다. 그러나 최적의 확장은 모델 아키텍처, 옵티마이저 및 데이터셋에 따라 달라진다. 그래디언트 축적은 가중치 업데이트당 표시되는 샘플 수를 변경하지 않으며 그래디언트가 집계되는 방식만 변경한다. 따라서 유효 배치 크기가 증가함에 따라 총 훈련 단계 수가 감소하여 수렴 속도와 최종 성능에 영향을 미칠 수 있다. 실무자는 실제 대규모 배치 훈련과 유사한 결과를 얻기 위해 학습률과 축적 단계를 조정해야 한다.
구현 고려 사항
그래디언트 축적을 구현하려면 훈련 루프를 수정해야 한다. 각 순전파 및 역전파 후에 그래디언트는 모델의 그래디언트 버퍼에 축적된다(예: PyTorch의 accumulate_grad 또는 TensorFlow의 GradientTape와 영구 변수 사용). 옵티마이저 단계는 원하는 축적 단계 수 후에만 수행된다. 손실을 적절히 스케일링하는 것이 중요하다. 손실이 미니 배치에 대해 평균화된 경우 축적된 그래디언트를 축적 단계 수로 나누어 올바른 그래디언트 크기를 유지해야 한다. 대안으로 손실이 합산된 경우 스케일링이 필요하지 않다. 또한 배치 정규화 레이어는 현재 미니 배치에 대한 통계를 계산하므로 특별한 주의가 필요하다. 그래디언트 축적을 사용하면 정규화를 위한 유효 배치 크기는 여전히 미니 배치 크기이므로 실제 대규모 배치 훈련과 다른 동작을 초래할 수 있다. 일부 프레임워크는 장치 간 동기화된 배치 정규화를 통해 이를 자동으로 처리하는 기능을 제공한다.
다른 기법과의 비교
그래디언트 축적은 역전파 중에 활성화를 재계산하여 메모리를 줄이는 그래디언트 체크포인팅과 낮은 정밀도 산술을 사용하여 메모리를 줄이는 혼합 정밀도 훈련과 종종 비교된다. 이러한 기법은 결합될 수 있다. 예를 들어 그래디언트 축적을 혼합 정밀도와 함께 사용하면 더 큰 유효 배치 크기로 훈련할 수 있다. 여러 GPU에 걸친 분산 훈련과 달리 그래디언트 축적은 장치 간 통신이 필요하지 않으므로 구현이 더 간단하고 동기화 오버헤드가 적다. 그러나 순차적 단계 수가 증가하므로 미니 배치 크기가 매우 작고 여러 역전파의 오버헤드가 클 경우 훈련 속도가 느려질 수 있다.
한계 및 대안
그래디언트 축적은 실제 대규모 배치 훈련의 완벽한 대체물은 아니다. 그래디언트 추정치는 서로 다른 미니 배치에서 계산되므로 데이터 셔플링 및 배치 정규화 통계로 인해 약간의 차이가 발생할 수 있다. 어떤 경우에는 이로 인해 수렴 동작이 달라질 수 있다. 대안으로는 높은 메모리 대역폭과 용량을 제공하는 Cerebras 시스템이나 Groq 가속기와 같은 더 큰 메모리 장치를 사용하거나 모델 병렬 처리와 파이프라인 병렬 처리를 사용하여 배치를 여러 장치에 분산하는 방법이 있다. 또한 LAMB 또는 LARS와 같은 일부 옵티마이저는 큰 배치 크기를 직접 처리하도록 설계되어 축적의 필요성을 줄일 수 있다.
역사적 배경 및 채택
여러 미니 배치에 걸쳐 그래디언트를 축적하는 개념은 현대 딥러닝 이전에 존재했으며 1980년대의 "번치 모드 역전파 알고리즘"에 뿌리를 두고 있다. 2010년대에 딥러닝 모델의 크기가 커지고 메모리 제약이 병목 현상이 되면서 이 개념은 두각을 나타냈다. 오늘날 그래디언트 축적은 PyTorch, TensorFlow 및 JAX를 포함한 주요 딥러닝 프레임워크의 표준 기능이며 OpenAI, Anthropic 및 Google DeepMind와 같은 조직에서 최첨단 모델을 훈련하는 데 널리 사용된다. 또한 제한된 하드웨어에서 대규모 모델을 미세 조정하는 연구 및 프로덕션 환경에서 일반적인 기법이다.
결론
그래디언트 축적은 메모리 제약 하에서 대규모 모델을 훈련하기 위한 실용적이고 널리 채택된 기법이다. 여러 미니 배치에 걸쳐 그래디언트를 축적함으로써 비례적인 메모리 증가 없이 큰 배치 크기를 시뮬레이션할 수 있다. 일부 오버헤드를 도입하고 학습률 및 정규화 레이어의 세심한 조정이 필요하지만, 특히 대규모 훈련 및 미세 조정 작업에서 딥러닝 실무자의 필수 도구로 남아 있다.