Understanding Transfer Learning of RNA Foundation Models on Downstream Tasks

Motivation

• Known: NLP와 protein science에서 대규모 unlabeled sequence 데이터로 사전학습된 foundation model(BERT, GPT, ESM 등)이 다양한 다운스트림 태스크에 효과적으로 전이됨이 알려져 있으며, RNA 도메인에서도 RNABERT, RNA-FM 등 transformer 기반 RNA FM이 구조·기능 관련 태스크에서 SOTA 대비 우수한 성능을 보인다고 보고되어 왔다.
• Gap: 그러나 RNA FM의 pretraining이 실제로 왜, 그리고 어떤 메커니즘으로 다운스트림 성능을 향상시키는지에 대한 체계적 검증은 이루어진 바 없으며, computer vision에서 확립된 hierarchical feature reuse 패러다임이 RNA FM에도 그대로 적용되는지 불분명하다. 또한 protein FM과 genomic FM에 대한 최근 연구들이 서로 상반된 결과(feature reuse 확인 vs. pretraining의 이점 없음)를 보고하여 RNA FM에 대한 독립적 검증이 필요하다.
• Why: RNA FM이 구조·기능 예측 등 생물학적으로 중요한 응용에 광범위하게 사용되고 있음에도, feature reuse 가정에 기반해 모델을 맹목적으로 스케일업하는 것은 비생산적일 수 있으므로, transfer learning의 실제 작동 원리를 규명하는 것은 향후 RNA FM 설계 및 활용에 중요한 지침을 제공한다.
• Approach: 다양한 크기와 사전학습 데이터셋을 가진 여러 transformer 기반 RNA FM을 구조 예측(SSP, SSI)과 기능 예측 등 폭넓은 다운스트림 태스크에 대해 frozen representation 평가, layer-wise ℓ2 distance 추적, loss landscape(flatness) 분석을 결합하여 체계적으로 비교하는 실증적 접근을 취한다.

Achievement

Figure 1. For each task, the orange bars (PT-FT) represent the performance of RNA foundation models fine-tuned from pret

1. Frozen representation의 태스크 의존성 규명: 사전학습된 RNA FM의 frozen last-layer representation이 구조 예측 태스크에서는 baseline보다 우수하지만, 기능 예측 태스크에서는 단순한 one-hot encoding보다 못한 성능을 보임을 실험적으로 입증했다.
2. Hierarchical feature reuse 미적용 발견: layer별 ℓ2 distance 변화 추적을 통해, computer vision과 달리 RNA FM에서는 얕은 층의 안정적인 저수준 특징 재사용이 보편적이지 않으며 세 가지 상이한 지식 전이 패턴이 존재함을 확인했다.
3. Pretraining을 optimization prior로 재해석: 사전학습과 다운스트림 목적함수가 잘 정렬될 때, pretraining이 feature reuse보다는 flatter minima로의 수렴을 가속하는 유리한 초기화(optimization initialization) 역할을 한다는 것을 loss curve 및 generalization 분석으로 입증했다.
4. 가장 포괄적인 RNA FM transfer learning 벤치마크 구축: transformer 기반 RNA FM에 대해 구조·기능 예측을 아우르는 가장 광범위한 transfer learning 평가를 수행하여, 향후 RNA FM 설계에 원칙적 지침을 제공했다.

How

Figure 3. Block-wise l2 distance evolution during training on the

• OmniGenBench, BEACON, Archive2, bpRNA 등 다양한 데이터셋을 활용해 secondary structure prediction(SSP)과 structural score imputation(SSI)을 포함한 구조 예측 태스크와 다양한 기능 예측 태스크를 curation
• 서로 다른 크기와 pretraining corpus를 가진 여러 RNA FM(PT-FT, 즉 pretrained-then-finetuned)을 무작위 초기화(RI-FT) baseline 및 one-hot encoding baseline과 비교
• frozen representation을 이용한 downstream 성능(Fig 1)을 태스크별로 비교하여 구조 vs 기능 태스크 간 성능 패턴 차이 분석
• fine-tuning 과정에서 block-wise(layer-wise) hidden state의 초기 상태 대비 ℓ2 distance 변화를 추적(Fig 3)하여 지식 전이 패턴을 세 가지로 분류
• PT-FT와 RI-FT(random init-finetune) 설정 하의 training loss curve 비교(Fig 4)를 통해 수렴 속도와 flat minima 도달 여부를 분석
• loss landscape의 flatness와 generalization 간 관계를 정량적으로 평가하여 pretraining의 optimization prior로서의 역할을 검증

Originality

• computer vision 중심으로 논의되어 온 transfer learning의 feature reuse 패러다임을 RNA FM이라는 새로운 도메인에 처음으로 체계적으로 적용하여 검증한 점
• frozen representation 성능뿐 아니라 layer-wise ℓ2 distance, loss landscape flatness 분석을 결합한 다각적 방법론으로 pretraining의 메커니즘을 규명한 점
• pretraining의 역할을 "feature reuse"가 아닌 "optimization initialization/prior"로 재정의하여, task alignment와 model capacity가 transferability를 결정짓는 핵심 요인이라는 새로운 개념적 프레임을 제시한 점

Limitation & Further Study

• 본 연구는 상대적으로 small-to-medium 크기의 transformer 기반 encoder-only RNA FM에 국한되어 있어, 향후 대규모 RNA FM이나 다른 아키텍처(예: state-space model, decoder 기반)로의 일반화 가능성은 추가 검증이 필요하다
• MLM(masked language modeling) 목적함수에 기반한 사전학습만을 다루므로, contrastive learning이나 multi-modal pretraining 등 다른 사전학습 전략에 대한 결론은 별도로 확인되어야 한다
• flat minima와 generalization 간의 인과관계에 대한 이론적 근거가 실험적 상관관계 수준에 머물러 있어, 보다 엄밀한 이론적 분석이 후속 연구로 요구된다
• task alignment를 정량적으로 사전에 예측할 수 있는 지표나 방법론이 제시되지 않아, 실무에서 어떤 다운스트림 태스크가 pretraining으로부터 이득을 볼지 사전에 판단하기 어렵다는 한계가 있다

Evaluation

총평: RNA foundation model의 transfer learning 메커니즘에 대한 최초의 체계적이고 다각적인 실증 연구로, computer vision에서 통용되던 feature reuse 가정을 RNA 도메인에서 재검토하고 pretraining을 optimization prior로 재해석한 점에서 학술적 기여도가 높으며, 향후 RNA FM 설계에 실질적 지침을 제공한다.