[논문 리뷰] Are Protein Language Models Compute Optimal?

Yaiza Serrano, Álvaro Ciudad외 1명|arXiv (Cornell University)|2024. 06. 11.

Machine Learning in Bioinformatics인용 수 4

한 줄 요약

이 연구는 NLP 스케일링 법칙을 pLM에 적응시켜 계산 최적의 단백질 언어 모델(pLMs)을 조사하며, 모델 크기나 데이터 스케일이 어떻게 되든 성능이 고정된 손실에서 정점에 도달함을 밝혀냈다. 단일 패assing 학습을 통해 35M 파라미터 모델이 11.43의 퍼플렉서티를 달성했으며, 이는 ESM-2(15B)나 xTrimoPGLM(100B)와 유사한 성능을 보이며, 최적의 학습을 통해 더 작은 모델이 더 계산 효율적일 수 있음을 보여준다.

ABSTRACT

While protein language models (pLMs) have transformed biological research, the scaling laws governing their improvement remain underexplored. By adapting methodologies from NLP scaling laws, we investigated the optimal ratio between model parameters and training tokens within a fixed compute budget. Our study reveals that pLM sizes scale sublinearly with compute budget, showing diminishing returns in performance as model size increases, and we identify a performance plateau in training loss comparable to the one found in relevant works in the field. Our findings suggest that widely-used pLMs might not be compute-optimal, indicating that larger models could achieve convergence more efficiently. Training a 35M model on a reduced token set, we attained perplexity results comparable to larger models like ESM-2 (15B) and xTrimoPGLM (100B) with a single dataset pass. This work paves the way towards more compute-efficient pLMs, democratizing their training and practical application in computational biology.

연구 동기 및 목표

• 일반적으로 사용되는 단백질 언어 모델(pLMs)이 고정된 FLOP 예산 하에서 계산 최적인지 조사하기.
• 고정된 계산 예산 내에서 학습 손실을 최소화하기 위해 모델 파라미터와 훈련 토큰 수의 최적 비율을 규명하기.
• ESM-2 및 xTrimoPGLM와 같은 현재의 pLMs가 계산 최적의 스케일링 법칙을 따르는지 여부 또는 수익 감소 현상이 나타나는지 평가하기.
• 퍼플렉서티를 핵심 지표로 사용해 단일 패assing으로 학습된 소형 모델(35M 파라미터)의 성능을 더 큰 모델들과 비교하기.
• 다양한 계산 수준에서 손실을 최소화하는 데 최적의 모델 크기와 데이터셋 크기의 스케일링 법칙 유도하기.

제안 방법

• Hoffmann 등(2022)과 Kaplan 등(2020)의 NLP 스케일링 법칙을 적응시켜 고정된 FLOP 예산 하에서 모델 크기(N), 데이터셋 크기(D), 그리고 학습 손실 간의 관계를 모델링함.
• UniRef50의 무작위로 샘플된 부분집합을 대상으로 5M에서 650M 파라미터까지 다양한 pLM을 훈련시키며, 토큰 수를 6.5B에서 20.8B로 변화시켜 다양한 계산 할당을 시뮬레이션함.
• 학습 손실 곡선의 스퍼플라인 보간을 통해 각 계산 수준에서 FLOPs와 최적의 모델 크기 및 토큰 수를 매핑함으로써 정밀한 최적 구성 식별 가능.
• 결합된 스케일링 법칙을 피팅: $ L(N,D) = ig[ (N_c/N)^{\alpha_N/\alpha_D} + (D_c/D) \big]^{\alpha_D} $, 이를 통해 손실을 모델 및 데이터 스케일의 함수로 모델링함.
• 손실 정점의 다항식 피팅을 통해 고정된 계산 예산(예: $10^{17}$ FLOPs)에 대한 최적 모델 크기를 식별하며, $10^{17}$ 및 $10^{18}$ FLOPs에서 각각 약 50.6B 및 157B 파라미터의 최적 크기를 예측함.
• 검증 세트에서 퍼플렉서티를 사용해 모델 성능을 평가하며, 단일 패assing으로 학습된 35M 파라미터 모델을 ESM-2(15B) 및 xTrimoPGLM(100B) 등 더 큰 모델들과 비교함.

Figure 1: FLOPs vs loss with token-based learning rate decay.

실험 결과

연구 질문

• RQ1 고정된 계산 예산 하에서 단백질 언어 모델의 모델 파라미터와 훈련 토큰 수의 최적 비율은 무엇인가요?
• RQ2 ESM-2 및 xTrimoPGLM와 같은 널리 사용되는 pLMs는 계산 최적인지, 아니면 수익 감소 현상을 보이는가요?
• RQ3 최적의 계산 할당과 조기 정지 기법을 적용할 경우 소형 pLM이 큰 모델과 유사한 성능을 달성할 수 있는가요?
• RQ4 pLM 사전학습에서 계산 예산이 증가함에 따라 최적의 모델 크기는 어떻게 변화하는가요?
• RQ5 동일한 데이터를 다중 패assing으로 학습하면 단일 패assing에서 관찰된 손실 정점 수준을 초월해 성능 향상이 크게 이루어지는가요?

주요 결과

• UniRef50의 20M 시퀀스 서브셋에서 단일 패assing으로 학습된 35M 파라미터 모델이 퍼플렉서티 11.43을 달성했으며, 이는 ESM-2(15B) 및 xTrimoPGLM(100B)와 비교해 계산 효율성이 뛰어나다.
• 모델 크기에 관계없이 일관된 손실 정점이 관찰되어, 성능 향상이 모델 크기나 데이터 스케일과 무관하게 정점에 도달함을 시사하며, 현재 pLM 학습의 근본적인 한계가 있음을 시사한다.
• FLOP 예산이 $10^{17}$ FLOPs일 경우 최적의 모델 크기는 약 50.6억 파라미터이며, $10^{18}$ FLOPs일 경우 157억 파라미터로 증가함을 확인했으며, 이는 최적 모델 크기가 계산과 함께 비선형적으로 증가함을 보여준다.
• 더 큰 모델은 더 높은 샘플 효율성을 보이며 더 적은 FLOPs로 손실 정점에 도달함을 확인했으며, 이는 더 높은 용량이 비록 파라미터 수가 많아도 계산 효율성을 향상시킨다는 것을 의미한다.
• 최적의 모델 크기는 계산 예산과 함께 비선형적으로 증가하며, 최적의 토큰 수 역시 비선형적으로 증가함을 확인했으며, 이는 모델 크기가 커질수록 성능 향상의 수익 감소 현상이 나타남을 시사한다.
• 결과적으로 ESM-2 및 xTrimoPGLM와 같은 현재의 pLMs는 계산 최적이 아니며, 더 작은, 더 잘 최적화된 모델이 동일한 성능 정점에 도달하기 위해 훨씬 더 적은 FLOPs를 소비함을 시사한다.

Figure 2: FLOPs vs loss without learning rate decay.