LAMB(레이어별 적응형 모멘트 배치 훈련)는 심층 신경망 훈련, 특히 대규모 배치 분산 컴퓨팅 환경에서 훈련을 위해 설계된 최적화 알고리즘입니다. 이 알고리즘은 2019년 Google과 토론토 대학교의 연구자들이 수천 개의 가속기로 훈련을 확장하면서 모델 정확도와 수렴 속도를 유지하는 문제를 해결하기 위해 도입했습니다.
이 알고리즘은 Adam 옵티마이저를 확장하여 매개변수별이 아닌 레이어별로 적응형 학습률을 계산합니다. 이러한 레이어별 적응은 LAMB가 서로 다른 네트워크 레이어 간의 다양한 기울기 규모를 더 효과적으로 처리할 수 있게 하며, 이는 트랜스포머 및 잔차 네트워크와 같은 심층 구조에서 특히 중요합니다. 레이어의 가중치와 기울기의 노름을 기반으로 업데이트를 정규화함으로써 LAMB는 매우 큰 배치 크기에서도 안정적이고 효율적인 훈련을 보장합니다.
배경 및 동기
대규모 신경망 훈련은 일반적으로 많은 GPU 또는 TPU에 분산된 막대한 계산 리소스를 필요로 합니다. 배치 크기를 늘리는 것은 이러한 리소스를 효율적으로 활용하는 일반적인 전략이지만, 종종 모델 성능 저하나 수렴 속도 저하로 이어집니다. 확률적 경사 하강법(SGD) 및 Adam과 같은 전통적인 옵티마이저는 레이어 간 기울기 규모의 이질성을 고려하지 않는 전역 학습률에 의존하기 때문에 큰 배치 크기에서 어려움을 겪습니다.
LAMB 옵티마이저는 이러한 한계를 극복하기 위해 개발되었습니다. 그 설계는 이전에 합성곱 네트워크의 대규모 배치 훈련에 사용된 레이어별 적응형 비율 스케일링(LARS) 알고리즘에서 영감을 얻었습니다. LAMB는 이 개념을 적응형 모멘트 추정과 결합하여 일반화하며, 두 접근 방식의 이점을 모두 활용합니다.
알고리즘 세부 사항
LAMB는 레이어의 가중치 노름과 기울기 노름의 비율을 기반으로 각 레이어에 대한 업데이트를 계산합니다. 단계 t에서 매개변수 텐서에 대한 핵심 업데이트 규칙은 다음과 같습니다:
- Adam에서와 같이 첫 번째 및 두 번째 모멘트 추정치(기울기의 평균과 분산)를 계산합니다.
- 업데이트 방향을 두 번째 모멘트의 제곱근에 작은 엡실론을 더한 값으로 나눈 모멘트 보정 기울기로 계산합니다.
- 이 방향을 레이어의 가중치 노름과 업데이트 방향 노름의 비율로 스케일링합니다.
- 전역 학습률을 곱하고 업데이트를 적용합니다.
이 레이어별 스케일링은 가중치 노름이 큰 레이어가 비례적으로 더 큰 업데이트를 받는 반면, 노름이 작은 레이어는 보수적으로 업데이트되도록 보장합니다. 또한 알고리즘은 기울기 클리핑 기술과 유사하게 극단적인 업데이트를 방지하기 위해 클리핑할 수 있는 신뢰 비율을 통합합니다.
저자들은 LAMB가 배치 크기 32,768로 ImageNet에서 ResNet-50을 훈련하면서 배치 크기 256의 기준선과 동일한 정확도를 달성하지만 훨씬 적은 단계로 달성할 수 있음을 입증했습니다.这使得它非常适合在数百或数千个加速器上进行分布式训练。
응용 및 영향
LAMB는 대규모 언어 모델 및 기타 딥러닝 모델 훈련에 널리 채택되었습니다. 예를 들어, BERT 및 기타 트랜스포머 기반 모델을 대규모로 훈련하는 데 사용되어 훈련 시간을 며칠에서 몇 시간으로 줄였습니다. 이 알고리즘은 AWS 및 Microsoft Azure 클라우드 플랫폼과 같은 하드웨어 리소스가 풍부한 환경과 AWS Trainium 및 Graphcore IPU와 같은 특수 AI 하드웨어에서 특히 가치가 있습니다.
LAMB의 변형 및 후속 모델과 같은 많은 후속 옵티마이저가 그 원리를 기반으로 구축되었습니다. 또한 학습률 스케줄 및 적응형 최적화 방법 연구에도 영향을 미쳤습니다. TensorFlow 및 PyTorch와 같은 프레임워크의 오픈 소스 구현은 더 넓은 머신러닝 커뮤니티에서 접근 가능하게 만들었습니다.
다른 옵티마이저와의 비교
Adam과 비교할 때, LAMB는 일반적으로 큰 배치 크기를 사용할 때 더 빠른 수렴과 더 나은 최종 성능을 달성합니다. Adam의 전역 학습률은 종종 세심한 튜닝이 필요하며 큰 배치에서 불안정성을 초래할 수 있습니다. LAMB의 레이어별 적응은 이러한 문제를 완화하여 더 공격적인 스케일링을 가능하게 합니다.
모멘텀이 있는 SGD용으로 설계된 LARS와 비교할 때, LAMB는 적응형 모멘트 추정을 통합하여 노이즈가 많은 기울기와 희소 특징에 더 강건합니다.这使得LAMB成为各种架构的更通用选择，包括序列到序列模型和编码器-解码器框架。
한계 및 고려 사항
장점에도 불구하고 LAMB는 한계가 없지 않습니다. 이 알고리즘은 신뢰 비율 클리핑 임계값 및 엡실론 항과 같은 추가 하이퍼파라미터를 도입하며, 특정 작업에 대해 튜닝이 필요할 수 있습니다. 또한 레이어별 스케일링이 유익하다고 가정하지만, 이는 레이어가 고도로 상관된 구조나 특정 가중치 초기화 방식을 사용할 때 항상 성립하지 않을 수 있습니다.
또한 LAMB는 대규모 배치 설정에서 뛰어나지만, 단순한 옵티마이저인 Adam으로 충분한 작은 배치 크기에서는 이점이 줄어듭니다. 연구자들은 알고리즘의 성능이 전역 학습률 선택에 민감할 수 있으며 최적 결과를 얻기 위해 학습률 워밍업이 필요할 수 있다고 지적했습니다.
결론
LAMB는 대규모 딥러닝 훈련을 위한 최적화에서 중요한 발전을 나타냅니다. 레이어별 적응과 적응형 모멘트를 결합함으로써 대규모 배치 크기에서 효율적이고 안정적인 훈련을 가능하게 하며, 생성형 AI 및 인공지능 연구 시대의 초석 기술이 되었습니다. 그 영향은 원래 응용을 넘어 후속 옵티마이저 및 훈련 방법론의 발전을 형성했습니다.