[논문 리뷰] Staged Training for Transformer Language Models

Sheng Shen, Pete Walsh외 4명|arXiv (Cornell University)|2022. 03. 11.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 변환기 언어 모델의 단계적 훈련을 제안하며, 먼저 작은 모델을 훈련한 후 손실과 훈련 동역학을 유지하는 학습된 연산자를 사용해 깊이 또는 너비를 단계적으로 증가시킨다. 이전에 사용한 계산 자원을 재사용하고 스케일링 법칙을 통해 단계 전환을 최적화함으로써, 처음부터 훈련하는 것에 비해 최대 22%의 계산 자원 절감을 달성하면서 성능은 그대로 유지하거나 향상시킨다.

ABSTRACT

The current standard approach to scaling transformer language models trains each model size from a different random initialization. As an alternative, we consider a staged training setup that begins with a small model and incrementally increases the amount of compute used for training by applying a "growth operator" to increase the model depth and width. By initializing each stage with the output of the previous one, the training process effectively re-uses the compute from prior stages and becomes more efficient. Our growth operators each take as input the entire training state (including model parameters, optimizer state, learning rate schedule, etc.) and output a new training state from which training continues. We identify two important properties of these growth operators, namely that they preserve both the loss and the "training dynamics" after applying the operator. While the loss-preserving property has been discussed previously, to the best of our knowledge this work is the first to identify the importance of preserving the training dynamics (the rate of decrease of the loss during training). To find the optimal schedule for stages, we use the scaling laws from (Kaplan et al., 2020) to find a precise schedule that gives the most compute saving by starting a new stage when training efficiency starts decreasing. We empirically validate our growth operators and staged training for autoregressive language models, showing up to 22% compute savings compared to a strong baseline trained from scratch. Our code is available at https://github.com/allenai/staged-training.

연구 동기 및 목표

• 대규모 변환기 언어 모델을 훈련하는 데 드는 높은 계산 비용을 줄이기 위해.
• 표준 훈련 방식에서 더 큰 모델이 초기에는 계산 효율성이 낮고 나중에야 효율성이 향상되는 비효율성을 해결하기 위해.
• 모델 크기 증가 시 손실과 훈련 동역학을 모두 유지하는 성장 연산자를 설계하기 위해.
• 스케일링 법칙을 활용해 계산 효율성을 극대화할 수 있는 체계적인 단계 전환 스케줄을 개발하기 위해.
• Wikitext-103 및 LAMBADA와 같은 표준 벤치마크를 사용해 GPT-2 스타일의 자동회귀 모델에서 이 방법을 실증적으로 검증하기 위해.

제안 방법

• 이 방법은 전체 훈련 상태(가중치, 옵티마이저 상태, 학습률 스케줄)를 입력으로 받아 새로운 더 큰 모델 상태를 출력하는 깊이 및 너비 성장 연산자를 사용한다.
• 성장 연산자는 손실 유지 및 훈련 동역학 유지 기능을 갖추고 있어, 성장 후 손실 감소 속도가 처음부터 훈련한 모델과 일치함을 보장한다.
• 단계 스케줄은 Kaplan 등(2020)의 스케일링 법칙을 사용해 결정되며, 이는 훈련 효율성이 떨어지기 시작하고 새로운 단계를 시작해야 할 시점을 예측한다.
• 정확한 스케일링 법칙이 알려져 있지 않은 실제 환경에서 최적의 스케줄에 대한 실용적 근사치를 도출한다.
• 표준 벤치마크인 Wikitext-103과 LAMBADA를 사용해 GPT-2 스타일의 자동회귀 언어 모델에서 이 방법을 검증한다.
• 이 방법은 이전 훈련 단계를 재사용하고 중복 계산을 피함으로써 대규모 모델을 훨씬 적은 계산 자원으로 훈련할 수 있도록 한다.

실험 결과

연구 질문

• RQ1 모델 성장을 포함한 단계적 훈련이 대규모 변환기 언어 모델의 훈련에 드는 총 계산 비용을 줄일 수 있는가?
• RQ2 안정적이고 효율적인 훈련을 위해 성장 연산자가 만족해야 할 조건은 무엇인가?
• RQ3 최적의 단계 전환 스케줄을 결정하기 위해 어떤 방법을 사용할 수 있는가?
• RQ4 옵티마이저와 학습률 스케줄을 포함한 전체 훈련 상태를 모델 성장 중에 얼마나 잘 유지할 수 있는가?
• RQ5 훈련 동역학을 유지하면 단순한 성장 방법에 비해 더 나은 수렴과 성능을 달성할 수 있는가?

주요 결과

• 제안된 성장 연산자는 손실과 훈련 동역학을 모두 유지하여, 성장 후 손실 감소 속도가 처음부터 훈련한 모델과 동일함을 보장한다.
• 스케일링 법칙을 사용해 단계 전환을 결정함으로써, 처음부터 훈련하는 것에 비해 최대 22%의 계산 자원 절감을 달성한다.
• GPT-2 스타일 모델에 대한 실증 결과에 따르면, 성장된 모델는 훨씬 적은 계산 자원으로도 유사하거나 더 낮은 검증 손실과 제로샷 퍼플렉서티를 달성한다.
• Wikitext-103에서 GPT-2 베이스 모델의 37.5 PPL을 유지하면서 전체 GPT-2 베이스 모델의 15%의 계산 자원만을 사용했고, LAMBADA 정확도가 동일한 경우 33%의 계산 자원만을 소비했다.
• 너무 이르거나 너무 늦게 성장을 시작하면 계산 자원 절감 효과가 떨어지므로, 정밀한 단계 스케줄링의 중요성을 확인한다.
• 이 방법은 다양한 모델 크기와 훈련 목표에서 성능를 유지함으로써 강건하고 일반화 가능함을 입증한다.