[논문 리뷰] Training Compute-Optimal Protein Language Models

Xingyi Cheng, Bo Chen외 4명|arXiv (Cornell University)|2024. 11. 04.

Machine Learning in Bioinformatics인용 수 4

한 줄 요약

이 논문은 단백질 서열 데이터에 맞춘 스케일링 법칙을 수립하여 고성능 단백질 언어 모델을 위한 계산 최적화된 훈련 전략을 제안한다. 1940억 토큰 데이터셋을 사용하고 300개 이상의 모델을 훈련하여, 고정된 계산 예산 내에서 데이터 및 모델 크기의 비선형적 증가가 우수한 성능을 낳음을 입증한다. CLM 사전 훈련을 통해 추정된 효과적 전이 토큰 수를 통해 MLM 작업으로의 효과적인 전이가 가능하며, 동일한 계산 자원 조건에서 ESM-2 및 PROGEN2와 같은 기준 모델보다 뛰어난 성능을 보였다.

ABSTRACT

We explore optimally training protein language models, an area of significant interest in biological research where guidance on best practices is limited. Most models are trained with extensive compute resources until performance gains plateau, focusing primarily on increasing model sizes rather than optimizing the efficient compute frontier that balances performance and compute budgets. Our investigation is grounded in a massive dataset consisting of 939 million protein sequences. We trained over 300 models ranging from 3.5 million to 10.7 billion parameters on 5 to 200 billion unique tokens, to investigate the relations between model sizes, training token numbers, and objectives. First, we observed the effect of diminishing returns for the Causal Language Model (CLM) and that of overfitting for the Masked Language Model~(MLM) when repeating the commonly used Uniref database. To address this, we included metagenomic protein sequences in the training set to increase the diversity and avoid the plateau or overfitting effects. Second, we obtained the scaling laws of CLM and MLM on Transformer, tailored to the specific characteristics of protein sequence data. Third, we observe a transfer scaling phenomenon from CLM to MLM, further demonstrating the effectiveness of transfer through scaling behaviors based on estimated Effectively Transferred Tokens. Finally, to validate our scaling laws, we compare the large-scale versions of ESM-2 and PROGEN2 on downstream tasks, encompassing evaluations of protein generation as well as structure- and function-related tasks, all within less or equivalent pre-training compute budgets.

연구 동기 및 목표

• 고정된 예산 조건에서 단백질 언어 모델 훈련을 위한 최적의 계산 자원 할당 전략을 규명하기 위해.
• Uniref를 초과하는 메타게놈 서열을 활용한 훈련 데이터 확장을 통해 CLM 및 MLM에서의 과적합 및 성능 정체 문제를 해결하기 위해.
• Causal Language Model(클라우지언 언어 모델, CLM) 및 Masked Language Model(마스킹 언어 모델, MLM) 목적함수에 대해 단백질 서열 데이터에 특화된 스케일링 법칙을 도출하기 위해.
• 전이 학습을 통한 CLM 사전 훈련의 효과성을 검증하기 위해.
• 기존의 모델 설정(예: ESM-2, PROGEN2)을 유도된 스케일링 법칙을 기반으로 재최적화하여 동일한 계산 예산 내에서 성능을 향상시키기 위해.

제안 방법

• Uniref와 메타게놈 서열을 조합하여 총 1940억 토큰, 9억 3900만 개 서열으로 구성된 데이터셋(uniMeta200B)을 구축하여 다양성을 향상시키고 과적합을 감소시켰다.
• 50억에서 2000억 개의 고유 토큰까지 다양한 훈련 데이터량과 350만에서 107억 파라미터까지의 다양한 모델 크기를 갖춘 300개 이상의 모델을 고정된 계산 예산 조건에서 훈련하여 스케일링 행동을 분석했다.
• 모델 크기와 훈련 토큰 수에 따른 손실 및 성능에 대해 거듭제곱 법칙 스케일링 함수를 fitting하여 CLM(0.57)과 MLM(0.77)의 스케일링 지수를 추정했다.
• CLM에서 MLM으로의 전이 효율을 수량화하기 위해 효과적 전이 토큰 수(Dₜ) 개념을 도입하여 두 목적함수 간의 최적 계산 자원 할당을 가능하게 하였다.
• 유도된 스케일링 법칙을 기반으로 ESM-2(3B) 및 PROGEN2-xlarge의 모델 크기와 훈련 토큰 수를 재할당하여 각각 72억 및 107억 파라미터의 새로운 모델을 훈련시켰다.
• A100 GPU에서 데이터 병렬 처리를 활용하고 FlashAttention, bfloat16, 최적화된 학습률 스케줄링(코사인 감쇠, 2.5% 웜업)을 적용하여 훈련 속도를 향상시켰다.

실험 결과

연구 질문

• RQ1 단백질 언어 모델링에서 CLM과 MLM의 스케일링 법칙은 어떻게 다를 수 있으며, 각각의 거듭제곱 법칙 지수는 무엇인가?
• RQ2 Uniref를 초월하는 훈련 데이터 다양성 증가로 CLM 및 MLM에서의 과적합과 성능 정체 문제를 완화시킬 수 있는가?
• RQ3 고정된 계산 예산 조건에서 CLM 사전 훈련이 MLM 성능 향상에 얼마나 효과적으로 전이될 수 있는가?
• RQ4 CLM 및 MLM 사전 훈련이 모두 필요할 경우, 최적의 계산 자원 할당 방식은 무엇인가?
• RQ5 스케일링 법칙을 기반으로 모델 크기와 훈련 토큰 수를 재할당하면, ESM-2 및 PROGEN2와 같은 기존 모델 대비 후행 작업에서 성능 향상을 이룰 수 있는가?

주요 결과

• MLM의 스케일링 법칙은 CLM(0.57)보다 높은 지수(0.77)를 보이며, 모델 크기 증가에 따른 계산 자원 증가에 더 크게 기여함을 시사한다.
• 계산 자원이 10배 증가하면, MLM의 경우 모델 크기가 4배 증가하고 데이터 양이 3배 증가하는 데 비해, CLM의 경우 모델 크기가 6배 증가하고 데이터 양이 70% 증가한다.
• CLM 사전 훈련은 MLM으로의 효과적인 전이를 가능하게 하며, 이 전이 효율은 효과적 전이 토큰 수(Dₜ)로 정량화되어 두 목적함수 간의 최적 계산 자원 할당이 가능하다.
• ESM-2(3B) 및 PROGEN2-xlarge 예산 기반으로 자원 재할당을 통해 각각 72억 및 107억 파라미터의 모델을 도출하였으며, 이는 후행 작업에서 원래 모델보다 뛰어난 성능을 보였다.
• 새로운 전략에 따라 훈련된 107억 파라미터 모델은 동일한 계산 예산 내에서 단백질 생성 및 구조/기능 예측 작업에서 최고 성능을 기록했다.
• 본 연구는 메타게놈 서열을 통한 데이터 다양성이 CLM 및 MLM 목적함수에서 과적합 및 성능 포화를 방지하는 데 핵심적임을 확인했다.