[논문 리뷰] RiNALMo: General-Purpose RNA Language Models Can Generalize Well on Structure Prediction Tasks

Rafael Josip Penić, Tin Vlašić외 3명|arXiv (Cornell University)|2024. 02. 29.

RNA and protein synthesis mechanisms인용 수 21

한 줄 요약

RiNALMo는 650M 매개변수를 가진 RNA 언어 모델로, 3600만 개의 ncRNA 시퀀스에서 사전 학습되어 RNA 구조 및 기능 작업에 대해 강한 일반화 성능을 보이며, unseen RNA 가족을 포함한 여러 데이터셋에서 기존 모델을 능가한다.

ABSTRACT

While RNA has recently been recognized as an interesting small-molecule drug target, many challenges remain to be addressed before we take full advantage of it. This emphasizes the necessity to improve our understanding of its structures and functions. Over the years, sequencing technologies have produced an enormous amount of unlabeled RNA data, which hides a huge potential. Motivated by the successes of protein language models, we introduce RiboNucleic Acid Language Model (RiNALMo) to unveil the hidden code of RNA. RiNALMo is the largest RNA language model to date, with 650M parameters pre-trained on 36M non-coding RNA sequences from several databases. It can extract hidden knowledge and capture the underlying structure information implicitly embedded within the RNA sequences. RiNALMo achieves state-of-the-art results on several downstream tasks. Notably, we show that its generalization capabilities overcome the inability of other deep learning methods for secondary structure prediction to generalize on unseen RNA families.

연구 동기 및 목표

• 대규모 비라벨링 RNA 데이터를 활용하여 일반-purpose RNA 표현 학습.
• 현대적 아키텍처 개선을 적용한 RNA 중심 트랜스포머 인코더 사전 학습.
• 구조와 기능 일반화를 평가하기 위해 다양한 다운스트림 태스크에서 RiNALMo 평가.
• 보지 못한 RNA 가족에서의 평가를 통한 가족 간 일반화 시연.
• 더 넓은 채용성과 재현성을 위해 코드의 공개 제공.

제안 방법

• RNAcentral, Rfam, nt, Ensembl의 3600만 ncRNA 시퀀스에 대해 마스킹 언어 모델링을 사용하여 BERT 스타일 인코더로 RiNALMo (650M 매개변수) 사전 학습.
• RoPE 위치 인코딩, SwiGLU 활성화, FlashAttention-2, 블록당 20개 어텐션 헤드를 갖춘 33개의 트랜스포머 블록 사용.
• 입력 길이를 1024 토큰으로 제한하고 15% 손상(80% 마스킹, 10% 임의 토큰, 10% 변경 없음)으로 교차 엔트로피 손실로 학습.
• 시퀀스를 1280-d 임베딩 차원으로 표현; CLS 토큰 선행, EOS 토큰 후행; 다운스트림 태스크에서 간단한 예측 헤드로 엔드-투-엔드 미세 조정.
• 학습된 임베딩을 구조 과제에 대해 단순한 ResNet 기반 예측기 부착 및 스플라이스-사이트와 리보솜 로딩 과제에 대해 MLP 헤드로 평가.

Figure 1: RiNALMo pre-training and applications. In the pre-training stage, RiNALMo is trained on unlabeled RNA sequences from several databases using masked language modeling (MLM). The final language model comprises $33$ Transformer blocks with the embedding dimension of $1280$ . Once pre-trained,

실험 결과

연구 질문

• RQ1 RiNALMo가 훈련 동안 보지 못한 RNA 가족에 대해 2차 구조 예측에서 일반화할 수 있는가?
• RQ2 다종 종의 스플라이스-사이트 예측에서 RiNALMo의 성능은 기존 RNA LM 및 전통적 방법과 비교하여 어떠한가?
• RQ3 평균 리보솜 로딩 예측이 정확하고, 무작위-origin 학습에서 인간 시퀀스로 일반화할 수 있는가?

주요 결과

• RiNALMo는 가족 간 2차 구조 예측에서 다른 방법보다 크게 우수하며, 대부분의 시험된 가족에서 더 높은 F1 점수를 달성.
• 가족 간 일반화에서 RiNALMo는 9개 가족 중 8개에서 열역학 기반 및 딥러닝 베이스라인을 능가(텔로메라스 RNA는 예외).
• 다종 종 스플라이스-사이트 예측에서 RiNALMo는 물고기, 파리, 식물, 벌레 데이터세트에서 최첨단 F1 점수 달성.
• 평균 리보솜 로딩 예측에서 RiNALMo는 Random7600과 Human7600 데이터셋에서 Uni-RNA, RNA-FM, Optimus 5-Prime보다 더 높은 R^2 점수 달성.
• RiNALMo의 가족 간 구조 임베딩은 t-SNE 시각화에서 RNA 가족으로 명확하게 군집화되어 임베딩에 풍부한 구조 정보를 시사.
• RiNALMo의 평균 가족 간 2차 구조 F1 점수는 0.72로 RNAstructure(0.59), CONTRAfold(0.61) 및 기타 베이스라인보다 높음(표 2).
• 도너/수용체 스플라이스-사이트 예측에서 RiNALMo의 평균 F1 점수는 97.70(어류), 96.11(곤충), 96.25(식물), 95.63(웜).
• RiNALMo의 스플라이스-사이트 결과는 보고된 종 분할에서 SpliceBERT, Uni-RNA, RNA-FM, Spliceator를 능가.
• RiNALMo는 Random-origin 시퀀스로 학습했을 때도 인간 UTR MRL 데이터에서 거의 최첨단 성능에 근접한 일반화를 보여줌(Random7600 vs Human7600 표 4).

Figure 2: t-SNE visualization of sequence embeddings produced by RiNALMo for RNAs from the inter-family generalization evaluation dataset.