[논문 리뷰] RITA: a Study on Scaling Up Generative Protein Sequence Models

Daniel Hesslow, Niccoló Zanichelli외 3명|arXiv (Cornell University)|2022. 05. 11.

Machine Learning in Bioinformatics인용 수 59

한 줄 요약

RITA는 1.2B 매개변수까지의 자기회귀 단백질 서열 모델을 도입하고 2억80M UniRef-100 시퀀스에서 학습하며, 확장 규모가 다음 아미노산 예측, 적합도 예측, 그리고 효소 기능 예측에서 성능에 어떻게 영향을 주는지 체계적으로 연구한다.

ABSTRACT

In this work we introduce RITA: a suite of autoregressive generative models for protein sequences, with up to 1.2 billion parameters, trained on over 280 million protein sequences belonging to the UniRef-100 database. Such generative models hold the promise of greatly accelerating protein design. We conduct the first systematic study of how capabilities evolve with model size for autoregressive transformers in the protein domain: we evaluate RITA models in next amino acid prediction, zero-shot fitness, and enzyme function prediction, showing benefits from increased scale. We release the RITA models openly, to the benefit of the research community.

연구 동기 및 목표

• 단백질 설계를 가속화하기 위해 대규모 자기회귀 단백질 서열 모델의 설계를 동기부여한다.
• 모델 크기가 증가함에 따라 성능을 분석하여 단백질 서열 모델링의 확장 법칙을 확립한다.
• 다운스트림 작업에서 RITA를 평가한다: 다음 아미노산 예측, zero-shot 적합도(ProteinGym), 그리고 효소 기능 예측.
• 커뮤니티에 모델을 공개하여 개방형 연구와 벤치마킹을 가능하게 한다.

제안 방법

• 사전 학습 데이터로 150B 아미노산을 사용하여 1.2B 매개변수까지의 4가지 크기(Small to XLarge)의 디코더-만 트랜스포머 모델을 학습한다.
• AliBi로 소거 분석 후 RoPE 로타리 위치 임베딩을 사용한다; 언어 모델링 손실이 더 낮은 RoPE를 선택한다.
• 주로 UniRef-100에서 사전학습하고 보완 데이터로 Metaclust와 MGnify를 사용하되 사전 클러스터링 없이 수행한다.
• Megatron-Deepspeed로 150B 아미노산에 걸쳐 수백 대의 GPU에서 Adam 옵티마이저, 배치 크기 512, 컨텍스트 크기 1024로 학습한다.
• UniRef-100, MGnify, Metaclust에서 보류된 단백질 가족들에 대한 perplexity를 평가하고 ProtGPT2 및 ProtXLNet와 비교한다.
• 다운스트림 작업을 평가한다: (i) ProteinGym 치환을 통한 변이 효과 예측; (ii) SwissProt 주석을 사용한 효소 기능 예측; (iii) 표적 단백질 계통군을 생성하기 위한 프롬프트 튜닝.

실험 결과

연구 질문

• RQ1 자기회귀 트랜스포머에서 모델 크기가 커질수록 단백질 서열 모델의 능력이 어떻게 확장되는가?
• RQ2 더 큰 RITA 모델이 기저선과 비교해 다음 아미노산 예측, 적합도 예측, 및 효소 기능 예측을 향상시키는가?
• RQ3 프롬프트 튜닝이 보류된 단백질 계통군에 대한 생성 품질을 더 개선할 수 있는가?
• RQ4 컴퓨팅 제약하에서 단백질 서열 모델링의 확장 법칙 특성(예: 지수)은 무엇인가?

주요 결과

• 모든 테스트 데이터셋에서 모델 크기가 커질수록 perplexity가 개선된다; RITA-XL이 UniRef-100, Metaclust, MGnify, Pfam 보류에서 가장 좋은 perplexity를 달성한다.
• 변이 효과(ProteinGym) 예측에서 RITA 모델의 성능은 크기가 커질수록 증가하며 Large 및 XLarge 변형에서 ESM-1v를 능가할 수 있다.
• 효소 기능 예측에서 더 큰 모델이 더 높은 top-k 정확도를 보이며, XLarge가 최상의 결과를 달성한다(상위 1위에서 10위까지의 값이 연구에 제시됨).
• 프롬프트 튜닝은 기본 모델에 비해 보류된 단백질 계통PF03272에 대해 perplexity를 크게 감소시키며 제어 가능한 생성을 보여준다.
• 다운스트림 작업 전반에 걸쳐 RITA 모델은 과제별 베이스라인에 근접하거나 이를 능가하며 NLP 확장 법칙과 유사한 compute 주도형 확장 행동을 보이지만 자원 vocab 크기와 훈련 역학의 차이로 주목할 만한 편차가 있다.