LAMB 옵티마이저(Layer-wise Adaptive Moments for Batch training)는 심층 신경망 훈련을 위한 확률적 최적화 알고리즘으로, Adam 옵티마이저에 레이어별 정규화 단계를 추가하여 확장한 것이다. 2019년 Google 연구진(특히 Yang You, Jing Li, Jonathan Hseu 등)이 소개한 LAMB는 모델 정확도를 저하시키거나 방대한 하이퍼파라미터 튜닝을 요구하지 않으면서 매우 큰 미니배치 크기(예: 32,768 이상)를 효율적으로 사용할 수 있도록 설계되었다. 이는 각 레이어의 가중치 노름과 업데이트 노름의 비율을 기반으로 레이어별 업데이트 크기를 조정하여 학습률을 각 레이어 파라미터의 규모로부터 사실상 분리함으로써 이를 달성한다.
LAMB는 특히 Transformer (architecture) 기반 모델 훈련에 큰 영향을 미쳤으며, 초기 대규모 언어 모델(LLM)과 ResNet과 같은 비전 아키텍처를 포함한다. TensorFlow(tf.keras.optimizers.LAMB를 통해) 및 PyTorch(NVIDIA, Hugging Face, FairScale과 같은 라이브러리의 LAMB 구현을 통해)와 같은 분산 훈련 프레임워크에서의 채택은 LAMB를 훈련 실행 규모를 확장하기 위한 표준 도구로 만들었다. 더 큰 배치를 허용함으로써 LAMB는 최첨단 모델 훈련에 필요한 벽시계 시간을 줄여주며, 이는 대규모 컴퓨팅 클러스터에 의존하는 OpenAI, Anthropic, Google DeepMind와 같은 조직에 중요하다.
동기 및 배경
심층 신경망을 큰 미니배치로 훈련하면 에포크당 계산 단계 수가 줄어들지만, 배치 크기를 단순히 확장하면 종종 일반화 성능 저하와 불안정한 수렴으로 이어진다. 이를 "대형 배치 훈련 문제"라고 한다. 모멘텀이 있는 확률적 경사 하강법(SGD) 또는 Adam과 같은 표준 옵티마이저는 배치 크기가 증가할 때 학습률을 신중하게 조정해야 하며, 그렇게 하더라도 정확도가 저하되는 경우가 많다. LAMB 옵티마이저는 배치 크기 변화에 더 강건한 옵티마이저를 만들어 이 문제를 해결하기 위해 개발되었다.
LAMB의 핵심 통찰은 심층 네트워크의 서로 다른 레이어가 매우 다른 규모의 그래디언트와 가중치 노름을 나타낸다는 것이다. 예를 들어, 초기 합성곱 레이어는 작은 가중치를 갖는 반면, 후기 완전 연결 레이어는 큰 가중치를 갖는다. Adam의 단일 전역 학습률은 일부 레이어에 대해 너무 큰 업데이트(발산 유발)를 유발하거나 다른 레이어에 대해 너무 작은 업데이트(수렴 지연)를 유발할 수 있다. LAMB는 레이어의 가중치와 그래디언트 노름의 비율을 기반으로 업데이트를 정규화하는 레이어별 적응형 비율을 도입하여 각 레이어가 그 크기에 비례하여 안정적인 속도로 움직이도록 보장한다.
알고리즘 설명
LAMB는 추가 정규화 단계가 있는 Adam의 변형으로 볼 수 있다. \(\theta_t\)를 반복 \(t\)에서의 파라미터, \(g_t\)를 \(\theta_t\)에 대한 손실의 그래디언트라고 하자. LAMB는 Adam과 유사하게 그래디언트의 1차 및 2차 모멘트(\(m_t\) 및 \(v_t\))를 지수 감쇠율 \(\beta_1\) 및 \(\beta_2\)(일반적으로 0.9 및 0.999)로 유지한다. 편향 보정 후 Adam 업데이트 \(\frac{m_t}{\sqrt{v_t} + \epsilon}\)를 계산한다.
결정적인 차이는 신뢰 비율 \(\phi\)이다. 각 레이어 \(i\)에 대해 \(\phi_i = \frac{||\theta_{t,i}||}{||r_{t,i}||}\)이며, 여기서 \(r_{t,i} = \frac{m_{t,i}}{\sqrt{v_{t,i}} + \epsilon}\)는 해당 레이어의 Adam 업데이트(학습률 제외)이고, \(||\cdot||\)는 L2 노름을 나타낸다. 레이어 \(i\)의 최종 업데이트는 \(\theta_{t+1,i} = \theta_{t,i} - \eta \cdot \phi_i \cdot r_{t,i}\)이며, 여기서 \(\eta\)는 전역 학습률이다. 이 신뢰 비율은 레이어의 가중치 노름에 비례하여 업데이트를 조정하므로 작은 레이어는 작지만 무시할 수 없는 업데이트를 받고, 큰 레이어는 더 크지만 안정적인 업데이트를 받는다.
실제로는 0으로 나누는 것을 방지하기 위해 분모에 작은 상수(예: 1e-6)가 추가된다. 이 알고리즘은 또한 AdamW에서 사용되는 분리된 가중치 감쇠 접근 방식을 따라 업데이트에 통합된 선택적 가중치 감쇠(L2 정규화)를 포함한다. 모든 레이어에 대해 신뢰 비율이 1로 설정되면 LAMB는 표준 Adam(편향 보정 포함)으로 축소된다.
하이퍼파라미터 및 튜닝
LAMB는 Adam의 대부분의 하이퍼파라미터를 상속한다: \(\beta_1\)(모멘텀), \(\beta_2\)(분산 감쇠), \(\epsilon\)(수치 안정성), 및 가중치 감쇠율. 주요 새로운 하이퍼파라미터는 전역 학습률 \(\eta\)이며, 대형 배치 훈련의 경우 일반적으로 Adam(예: 1e-3)에 비해 상당히 높은 0.01-0.1 범위로 설정된다. 저자들은 매우 큰 배치(예: BERT의 경우 32,768)의 경우 첫 10% 단계에 걸친 선형 워밍업과 함께 0.01의 학습률이 잘 작동한다는 것을 발견했으며, 나머지 단계에 걸쳐 코사인 감쇠 학습률 스케줄을 사용할 것을 권장했다.
또한 \(\beta_2\)의 선택은 안정성에 영향을 줄 수 있으며, 희소 그래디언트가 있는 모델의 경우 더 높은 \(\beta_2\)(예: 0.99)를 사용할 수 있다. 저자들은 또한 배치 크기를 학습률에 비례하여 확장할 수 있다고 제안했으며(선형 스케일링 규칙), 이는 LAMB와 잘 작동하는 지침이다. 예를 들어, 배치 크기가 두 배가 되면 정확도 손실 없이 학습률도 두 배로 늘릴 수 있다.
성능 및 벤치마크
원래 논문에서 LAMB는 ImageNet(이미지 분류)에서 ResNet-50 훈련과 마스크 언어 모델링을 위한 BERT(트랜스포머 기반 언어 모델) 훈련이라는 두 가지 주요 작업에서 평가되었다. LAMB를 사용하여 저자들은 배치 크기 32,768로 단 2,048회 반복 만에 ImageNet top-1 정확도 76.0%를 달성했으며, 이는 훨씬 적은 에포크로 더 작은 배치(예: 256)로 달성한 최첨단 정확도와 일치한다. BERT의 경우 1,024개의 TPU를 사용하여 약 3.5분 만에 동일한 정확도(SQUAD에서 F1 점수 1.0)로 모델을 훈련시켰으며, 이는 이전 방법보다 10배 빠른 속도 향상이다.
이후 LAMB는 Google 내부 워크플로우에서 BERT 기반 모델 훈련을 위한 기본 옵티마이저가 되었다. 이 논문은 배치 크기를 1,024에서 65,536으로 확장할 때 LAMB가 Adam과 모멘텀이 있는 SGD보다 성능이 우수하다고 보고했다. 저자들은 또한 LAMB가 Gradient Clipping(폭발적 그래디언트 방지에 사용)과 잘 작동하며 NVIDIA GPU 및 Google TPU와 같은 최신 하드웨어에서 사용되는 혼합 정밀도 훈련과 호환된다는 것을 보여주었다.
대규모 훈련에서의 응용
LAMB의 주요 응용 분야는 배치 크기가 단일 장치 메모리에 맞추기에는 너무 큰 분산 훈련이다. 이러한 설정에서 그래디언트는 데이터 병렬 처리를 사용하여 여러 GPU 또는 TPU에 걸쳐 평균화된다. 예를 들어, OpenAI와 Google DeepMind는 수천 개의 시퀀스 길이를 가진 대형 Transformer (architecture) 모델을 훈련할 때 LAMB와 유사한 옵티마이저를 사용한다. LAMB(및 그 후속작인 LAMB2)와 같은 최신 옵티마이저가 제안되었지만, LAMB는 비전 트랜스포머 및 LLM 훈련을 포함한 많은 오픈소스 노력에서 AI21 Labs 및 SambaNova와 같은 연구 그룹과 회사에 의해 여전히 신뢰할 수 있는 선택으로 남아 있다.
머신러닝이 Amazon Web Services(AWS Trainium 하드웨어 포함)에서 사용되는 맥락에서 LAMB는 효율성을 위해 맞춤형 커널에서 지원된다. 마찬가지로 Intel과 AMD는 자사 가속기에서 LAMB를 벤치마킹했다. 이 옵티마이저가 극단적인 배치 크기를 처리할 수 있는 능력은 단일 에포크 비용이 높고 에포크 수를 줄이는 것이 가장 중요한 대규모 데이터셋에서 모델을 사전 훈련하는 데 가치가 있다.
다른 옵티마이저와의 관계
LAMB는 SGD 변형, Adam, 및 AdamW(분리된 가중치 감쇠) 및 LARS(Layer-wise Adaptive Rate Scaling)와 같은 후속 옵티마이저를 포함하는 적응형 옵티마이저 계열의 일부이다. 2017년 You 등이 CNN의 대형 배치 훈련을 위해 도입한 LARS는 유사한 레이어별 신뢰 비율을 사용하지만 2차 모멘트를 유지하지 않으며 1차 모멘트(모멘텀)와 그래디언트 노름에 의존한다. LAMB는 LARS의 이점(레이어별 스케일링)과 Adam의 적응형 파라미터별 학습률을 결합하여 희소 그래디언트(트랜스포머와 같은)가 있는 모델에 더 강건하게 만든다.
또 다른 밀접하게 관련된 옵티마이저는 분산 감소를 통합하는 NVLAMB(Nvidia 제공)이다. 그러나 LAMB는 더 단순하고 널리 사용된다. Beam Search를 사용하는 Sequence-to-Sequence (Seq2Seq) 작업의 경우 LAMB는 추론에 직접적인 영향을 미치지 않지만 훈련 수렴을 도와 시퀀스 디코딩을 간접적으로 개선한다.
확장 및 변형
소개 이후 여러 변형이 제안되었다. LAMB2(Google 제공)는 그래디언트 분산을 기반으로 정규화 요소를 추가하여 특정 문제에 대한 안정성을 개선한다. 편향 보정이 있는 LARS 및 기타 수정도 일반적이다. 실제로 많은 프레임워크는 처음 몇 단계 동안 유용한 모멘트에 대한 선택적 편향 보정과 함께 LAMB를 구현한다. PyTorch의 torch.optim.Lamb(torch_optimizer 패키지)와 같은 일부 구현에서는 신뢰 비율 매개변수를 조정하거나 사용자 지정 레이어별 학습률을 사용할 수 있다.
1cycle 스케줄과 같은 다양한 스케일링 전략을 가진 AdamW와 같은 새로운 옵티마이저의 등장에도 불구하고 LAMB는 대형 배치 훈련을 위한 강력한 기준선으로 남아 있다. 연구 커뮤니티는 일반화를 더욱 개선하기 위해 LAMB를 데이터 증강 및 그래디언트 클리핑과 결합하는 것을 탐구했다.
실용적 고려 사항 및 한계
LAMB는 대형 배치 설정에서 탁월하지만 작은 배치 크기(예: 1,024 미만)에서는 항상 최선의 선택은 아니다. 이러한 영역에서는 표준 Adam 또는 모멘텀이 있는 SGD가 더 간단하고 동등하게 효과적일 수 있다. LAMB는 레이어별 노름 계산의 계산 오버헤드를 추가하며, 이는 최신 하드웨어에서는 무시할 수 있지만 많은 작은 레이어(예: U-Net 아키텍처)가 있는 모델에서는 사소하지 않을 수 있다.
또 다른 한계는 일부 레이어의 가중치 노름이 매우 작은(0에 가까운) 경우 LAMB의 신뢰 비율이 때때로 불안정한 훈련으로 이어질 수 있다는 것이다. 이는 일반적으로 분모에 엡실론 항을 추가하고 가중치가 0으로 흘러가는 것을 방지하는 가중치 감쇠를 사용하여 완화된다. 또한 LAMB는 학습률과 워밍업 단계의 신중한 튜닝이 필요하며, 부적절한 스케줄은 발산으로 이어질 수 있다.
메모리 사용량은 Adam(파라미터당 두 개의 모멘트 벡터)과 유사하므로 메모리 집약적이지 않다. 매우 큰 모델의 경우 모델 가지치기 또는 그래디언트 클리핑을 LAMB와 함께 사용할 수 있지만 이는 독립적인 기술이다.
결론
LAMB는 대규모 딥러닝을 위한 최적화 알고리즘 도구 상자에서 초석이 되었다. 거대한 미니배치로 효과적인 훈련을 가능하게 함으로써 많은 벤치마크 모델의 개발을 가속화하고 실험 비용을 줄였다. 레이어별 적응 원리는 이후 옵티마이저 설계에 영향을 미쳤으며 분산 훈련의 문제에 직면한 실무자에게 실용적이고 잘 이해된 솔루션으로 남아 있다. 인공지능이 계속 성장함에 따라 LAMB와 같은 옵티마이저는 진화할 가능성이 높지만 레이어별 신뢰와 적응형 모멘트라는 핵심 아이디어는 계속 유지될 것이다.