간단하고 확장 가능한 전략을 통한 대규모 언어 모델의 지속적 사전 학습

Simple and Scalable Strategies to Continually Pre-train Large Language Models

March 13, 2024

저자: Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats L. Richter, Quentin Anthony, Timothée Lesort, Eugene Belilovsky, Irina Rish

cs.AI

초록

대규모 언어 모델(LLMs)은 일반적으로 수십억 개의 토큰으로 사전 학습되며, 새로운 데이터가 사용 가능해지면 이 과정을 다시 시작합니다. 이러한 모델을 지속적으로 사전 학습하는 것이 훨씬 더 효율적인 해결책이며, 재학습에 비해 상당한 계산 자원을 절약할 수 있습니다. 그러나 새로운 데이터로 인한 분포 변화는 일반적으로 이전 데이터에 대한 성능 저하나 새로운 데이터에 대한 적응 실패로 이어집니다. 본 연구에서는 학습률(LR) 재가열, 학습률 재감소, 그리고 이전 데이터의 재생이라는 간단하고 확장 가능한 조합이 모든 사용 가능한 데이터를 처음부터 재학습한 경우의 성능과 일치함을 보여줍니다. 이는 최종 손실 및 언어 모델(LM) 평가 벤치마크를 통해 측정되었습니다. 구체적으로, 우리는 두 가지 일반적으로 사용되는 LLM 사전 학습 데이터셋 간의 약하지만 현실적인 분포 변화(영어→영어)와 더 강한 분포 변화(영어→독일어)를 405M 파라미터 규모의 모델과 대규모 데이터셋(수천억 개의 토큰)에서 실험했습니다. 더 큰 규모의 실험을 위해 약하지만 현실적인 분포 변화를 선택한 결과, 우리의 지속 학습 전략이 10B 파라미터 LLM에 대한 재학습 기준과 일치함을 확인했습니다. 우리의 결과는 LLM이 간단하고 확장 가능한 지속 학습 전략을 통해 성공적으로 업데이트될 수 있으며, 재학습 기준과 일치하는 성능을 훨씬 적은 계산 자원으로 달성할 수 있음을 보여줍니다. 마지막으로, 이전 연구에서 영감을 받아, 우리는 학습률 재가열로 인한 망각을 방지하고 고정된 토큰 예산에 구속되지 않는 코사인 학습률 스케줄의 대안을 제안합니다.

English

Large language models (LLMs) are routinely pre-trained on billions of tokens, only to start the process over again once new data becomes available. A much more efficient solution is to continually pre-train these models, saving significant compute compared to re-training. However, the distribution shift induced by new data typically results in degraded performance on previous data or poor adaptation to the new data. In this work, we show that a simple and scalable combination of learning rate (LR) re-warming, LR re-decaying, and replay of previous data is sufficient to match the performance of fully re-training from scratch on all available data, as measured by final loss and language model (LM) evaluation benchmarks. Specifically, we show this for a weak but realistic distribution shift between two commonly used LLM pre-training datasets (EnglishrightarrowEnglish) and a stronger distribution shift (EnglishrightarrowGerman) at the 405M parameter model scale with large dataset sizes (hundreds of billions of tokens). Selecting the weak but realistic shift for larger-scale experiments, we also find that our continual learning strategies match the re-training baseline for a 10B parameter LLM. Our results demonstrate that LLMs can be successfully updated via simple and scalable continual learning strategies, matching the re-training baseline using only a fraction of the compute. Finally, inspired by previous work, we propose alternatives to the cosine learning rate schedule that help circumvent forgetting induced by LR re-warming and that are not bound to a fixed token budget.