Nucleotide Transforemr:building and evaluating robust foundation models for human genomics

논문 바로가기

INTRO

Fundation model은 대규모 데이터와 많은 파라미터로 학습되는 범용 모델이다. NLP에서는 BERT, GPT 등이 대표적인 예이다. 이러한 모델의 핵심 학습 방식은 MLM(Masked Language Modeling)으로 문장의 빈칸을 맞추는 방식으로 언어 패턴을 학습한다.

이러한 foundation model 접근 방식이 생물학 분야에 처음 적용된 사례는 단백질 서열에 대한 language model 학습이었다. 마스킹된 아미노산을 예측하도록 학습되었다. 이 방법은 구조 및 기능 예측에서 기존 방법보다 좋은 성능을 보여주었다.

DNA 서열에는 regulatory region, varint 영향, gene expression 같은 중요한 정보가 포함되어 있다. DNA는 유전체 데이터가 매우 방대하며 대부분 라벨이 없어 대규모 foundation model이 필요하다. 그래서 이 논문은 DNA 서열을 위한 foundation model인 Nucleotide Transformer(NT)를 제안한다.

---

Method

Model

유전체 데이터에 MLM을 적용하여 뉴클레오타이드 서열을 문장으로, k-mer(k = 6)를 단어로 간주하고 transformer 기반 언어 모델을 학습하였다.

Transformer는 입력 시퀀스의 각 위치를 임베딩 벡터로 변환하는 임베딩 레이어와, 이러한 임베딩을 점진적으로 정교화하는 sefl-attention 레이어 스택으로 구성된다. MLM을 이용해 transformer 언어 모델을 학습하는 대표적인 방법에 BERT가 있으며, BERT는 시퀀스의 모든 위치가 서로 참조할 수 있도록 양방향 attention을 사용하여 정보가 양 방향으로 흐를 수 있게 한다. 이러한 특성은 DNA 서열처럼 앞뒤 문맥 정보가 중요한 데이터에서 특히 중요하다. 학습 과정에서는 네트워크 최종 임베딩을 언어 모델 헤드에 입력하여 각 위치에 존재할 토큰에 대한 확률 분포를 예측한다.

Model Architecture

encoder-only transformer 구조를 따른다. 먼저 임베딩 레이어가 토큰 시퀀스를 임베딩 벡터 시퀀스로 변환한 후 각 임베딩에 위치 정보를 제공하기 위해 positional enoding을 추가한다(최대 1000개의 토큰을 처리할 수 있는 학습 가능한 positional encoding을 사용함).

토큰 단위는 6-mer 토큰을 사용하였고 이는 최대 6kb 길이의 서열을 처리하면서도 임베딩 크기를 적절하게 유지하기 위한 절충안이며, 다른 토큰 길이와 비교했을 때 가장 높은 성능을 보였기 때문이다.

이후 토큰 임베딩은 여러 개의 transformer 레이어로 구성된 스택을 통과한다. 각 transformer 레이어는 다음과 같은 연산을 수행한다.

• Layer normalization
• Multi-head self-attention
• Skip connection을 통해 attention 출력과 입력을 더함
• 또 다른 layer normalization
• GELU 활성화 함수를 사용하는 2층 퍼셉트론

각 모델의 attention head 수, 임베딩 차원, 퍼셉트론 은닉층의 뉴런 수, 저네 레이어 수는 위와 같다.

self-supervised 학습 동안 transformer 스택의 마지막 레이어에서 생성된 임베딩은 언어 모델 헤드를 통해 각 위치의 토큰에 대한 확률 분포로 변환된다.

NT-v2 모델
기존보다 더 효율적인 구조를 적용하였다.

• 기존의 학습 가능한 positional embedding 대신 rotary embedding을 사용
• swish 활성화 함수가 적용된 GLU 사용
• bias를 제거하여 계산 효율성 향상

또한 최대 2048개의 토큰을 입력으로 처리할 수 있으며, 이는 약 12kb 길이의 서열을 포함하는 context window에 해당한다.

Training

모델 학습은 BERT 방식을 따랐다. 각 학습 단계에서 토큰화된 시퀀스의 batch가 샘플링된다. 배치 크기는 사용 가능한 하드웨어와 모델 크기에 따라 조정하였다. 모든 실험은 A100 GPU 클러스터에서 수행되었으며, 500M 모델과 2.5B 모델을 학습할 때 각각 14개와 2개의 시퀀스를 배치로 사용하였다.

각 시퀀스에서 전체 토큰의 15%를 선택하여 마스킹하였다. 이 중

• 80%는 MASK 토큰으로 대체하고, 추가로 선택된 토큰의 10%는 무작위 토큰으로 대체하였다.

다만 1000 Genomes 데이터셋을 이용한 학습에서는 이러한 추가 데이터 증강을 적용하지 않았다. 무작위 토큰 치환이 인간 유전체에서 자연적으로 발생하는 돌연변이 빈도보다 더 큰 노이즈를 생성할 수 있기 때문이다.

각 배치의 손실 함수는 선택된 위치에서 예측된 토큰 확률 분포와 실제 토큰 사이의 cross-entropy loss의 합으로 계산하였다. 효과적인 학습을 위해 gradient를 누적하여 배치당 약 100만 개의 토큰이 사용되도록 하였다.

최적화에는 Adam optimizer를 사용하였고 다음과 같은 하이퍼파라미터를 사용하였다

• $\beta_1 = 0.9$
• $\beta_2 = 0.999$
• $\epsilon = 1 \times 10^{-8}$

학습 초기에 warmup 단계를 두어 learning rate를 $5 \times 10^{-5}$에서 $1 \times 10^{-4}$까지 16,000 step 동안 선형적으로 증가시킨 뒤, 이후에 square-root decay 방식으로 감소시켰다.

Probing

Probing은 사전학습된 모델이 생성한 임베딩이 다운스트림 작업을 얼마나 잘 수행할 수 있는지 평가하기 위한 방법이다. 이를 위해 다운스트림 작업에 사용되는 nucleotide 서열 데이터셋을 모델에 입력하고, 모델의 10개 레이어에서 생성된 임베딩을 각각 추출하여 저장한다. 이후 각 레이어의 임베딩을 입력으로 사용하여 다운스트림 모델을 학습시켜 성능을 비교함으로써, 어떤 레이어의 표현이 가장 유용한지 분석한다. 다운스트림 모델로는 scikit-learn의 기본 설정을 사용한 logistic regression과 MLP를 사용하였다. MLP의 경우 학습률, 활성화 함수, 은닉층 수와 같은 하이퍼파라미터가 성능에 영향을 미치기 때문에 하이퍼파라미터 탐색을 수행하였다.

모델 평가는 10-fold 검증을 통해 진행되었다. 전체 데이터셋을 90%의 학습 데이터와 10%의 검증 데이터로 나누는 과정을 서로 다른 셔플로 10번 반복하고, 동일한 하이퍼파라미터 설정에서 10개의 모델을 학습한 뒤 검증 성능의 평균을 계산하였다. 이러한 과정은 Tree-structured Parzen Estimator(TPE) 기반 탐색을 사용하여 총 100번 반복되었으며, 이후 가장 성능이 좋은 모델을 선택하여 테스트 데이터에서 최종 성능을 평가하였다. 이 과정에서 각 사전학습 모델의 10개 레이어에 대해 성능을 측정하였으며, 전체 probing 실험을 통해 약 76만 개의 다운스트림 모델이 학습되었다. 이를 통해 모델의 레이어별 표현 능력과 다운스트림 작업 성능 간의 관계를 상세히 분석하였다.

또한 비교를 위해 baseline 실험도 수행하였다. Transformer 레이어를 거치지 않은 상태에서 토큰화된 서열을 그대로 입력으로 사용하는 logistic regression 모델을 학습하여 성능을 평가하였다. 추가로 토큰 ID를 one-hot 인코딩한 뒤 시퀀스 길이 축에 대해 합이나 평균으로 pooling하는 방식도 실험하였으나, 이러한 방법보다 원시 토큰 시퀀스를 그대로 사용하는 방식이 더 좋은 성능을 보였다

Finetuning

probing처럼 임베딩을 추출해 사용하는 것뿐 아니라, IA3(PEFT)를 이용해 모델을 미세조정하였다. 기존의 언어 모델 헤드를 제거하고, 작업의 종류에 따라 classification 또는 regression 헤드로 교체한다. Transformer 레이어와 임베딩 레이어의 기존 가중치는 모두 고정하고, 대신 각 Transformer 레이어에 새로운 학습 가능한 벡터 $l_k, l_v, l_{ff}$를 추가하여 self-attention과 feed-forward 네트워크 내부에서 rescaling 역할을 하도록 한다. 이러한 벡터들은 attention의 key와 value, 그리고 feed-forward 연산에 elemet-wise 곱 형태로 적용되어, downstream 작업에 더 적합한 표현을 만들도록 transformer 레이어의 출력을 조정한다. (추가되는 파라미터 수는 $L(d_k+d_v+d_{ff})$이며, 여기서 $L$은 Transformer 레이어 수이다.)

이 방식으로 추가되는 파라미터는 전체 모델 파라미터의 약 0.1% 수준에 불과하며 전체 모델을 다시 학습하는 기존 fine-tuning보다 학습 속도와 공간 측면에서 훨씬 효율적이다.

---

Result

기존 유전체 모델은 대부분 task-specific supervised model이므로 새로운 task마다 모델을 새로 학습해야 한다는 문제가 있었다. 이를 해결하기 위해 DNA foundation model인 Nucleotide Transformer를 제안하게 되었다.

• 6kb 길이 DNA 서열 사용
• Transformer 기반 language model
• self-supervised pretraining

방법을 사용하였다.

모델 크기와 데이터 다양성이 성능에 미치는 영향을 분석하기 위해 4개의 NT 모델을 구축하였다.

모델명	특징
Human ref 500M	- human reference genome - 500M parameters
1000G 500M	- 3202 human genomes - 500M parameters
1000G 2.5B	-3202 human genomes - 2.5B parameters
Multispecies 2.5B	- 850 species genomes - 2.5B parameters

위 그래프 결과를 통해서 단순히 model size 증가보다 데이터 다양성 증가가 더 중요할 수 있음을 보여준다.

NT 모델을 평가하기 위해 18개의 genomic prediction task dataset을 구축하였다.
Task 종류는 splice site prediction, promoter prediction, histone modification, enhancer prediction이고 데이터 출처는 GENCODE, EPD, ENCODE이고 10-fold cross validation으로 평가하였다.

NT 모델 평가는 두 가지 방식으로 평가하였다.

1. Probing
   pretrained embedding을 사용하여 간단한 classifier 학습
   모델 ▶️ logistic regression, MLP
   

2. Fine-tuning
   LM head를 classification head or regression head로 교체후 PEFT로 재학습 하였다. 전체 파라미터의 0.1%만 업데이트 한 것이 특징이다. 이러한 특징은 하나의 pretrained 모델을 여러 genomics task에 재사용이 가능함을 시사한다.

NT 모델을 기존 genomics 모델과 비교하였다.
baseline 모델은 BPNet(CNN)이며 MCC = 0.665, Larget BPNet = 0.683으로 supervised CNN도 매우 강력한 baseline임을 보여준다.

결과에서 보이듯이 NT 모델은 많은 task에서 baseline을 능가하는 것을 보여준다

또한 기존 SOTA 모델과 비교했을 때, NT 모델은 specialized DL model과 거의 동일한 성능을 보여주었다. 이는 foundation model이 task-specific model과 경쟁이 가능함을 보여준다.

이러한 결과를 바탕으로 다음과 같은 해석이 가능하다. NT 모델은 unsupervised training만으로 단순한 sequence 패턴이 아니라 실제 genomic elements를 학습할 수 있었다. 모델은 intron, exon, intergenic region, UTR을 구분하였고 attention 분석 결과 TF binding site, promoter, enchancer 등에 attention이 집중됨을 확인할 수 있었다. 즉, 모델이 gene structure를 학습할 수 있음을 보여주었다.

또한 NT 모델은 mutation effect도 예측 가능했다.

모델 최적화 실험을 통해 더 작은 모델로도 성능 유지가 가능했음을 보여주었다.

마지막으로 Long context 효과가 있음을 보여주었다.

1. 본 연구는 DNA 서열을 이용한 transformer 기반 Nucleotide Transformer(NT) 모델을 구축하고, 사전학습 데이터의 종류와 모델 크기가 유전체 예측 성능에 어떤 영향을 미치는지를 분석하였다. 특히 동일한 모델 구조를 사용하더라도 어떤 데이터로 사전학습하느냐에 따라 성능이 크게 달라진다는 점이 확인되었다. 실험 결과 인간 유전체만을 사용하여 학습한 모델보다 여러 종의 유전체 데이터를 함께 사용한 multispecies 모델이 대부분의 인간 유전체 예측 작업에서 더 높은 성능을 보였다. 이는 서로 다른 종 사이에서도 보존된 기능적 유전체 패턴을 모델이 학습했기 때문으로 해석되며, 다양한 종의 데이터를 활용하는 것이 모델의 일반화 성능을 높이는 데 중요한 요소임을 시사한다.

2. 또한 본 연구는 모델 크기와 성능 사이의 관계를 분석하여 genomics 분야에서도 scaling law가 존재함을 확인하였다. 5천만에서 25억 개의 파라미터까지 다양한 크기의 모델을 비교한 결과, 모델 규모가 커질수록 전반적인 예측 성능이 향상되는 경향이 나타났다. 다만 이러한 대규모 모델을 학습하기 위해서는 상당한 계산 자원이 필요하며, 실제로 본 연구에서는 수백 개의 GPU를 사용하여 장기간 학습을 수행하였다. 그러나 일단 모델이 학습된 이후에는 다양한 downstream 분석에서 비교적 낮은 비용으로 활용할 수 있어 연구 활용성이 높다는 점을 강조한다.

3. 기존의 생물학 언어 모델 연구에서는 주로 transformer의 마지막 layer embedding을 이용한 probing 방식으로 downstream 성능을 평가하는 경우가 많았다. 그러나 본 연구에서는 다양한 layer와 downstream 모델을 체계적으로 비교한 결과, 최종 layer보다 중간 layer에서 더 좋은 성능이 나타나는 경우가 많다는 사실을 확인하였다. 또한 probing 방식보다 parameter-efficient fine-tuning 방식이 더 높은 성능과 안정성을 보였으며, 전체 파라미터의 약 0.1%만 학습해도 효과적인 결과를 얻을 수 있었다. 이러한 결과는 실제 활용 측면에서 fine-tuning 방식이 더 실용적임을 보여준다.

4. 연구 결과, DNA language model 접근법은 기존의 task-specific supervised 모델과 경쟁 가능한 성능을 보였다. histone modification, splicing site, regulatory element 예측 등 다양한 유전체 분석 작업에서 NT 모델은 기존의 특화된 딥러닝 모델들과 유사하거나 때로는 더 높은 성능을 달성하였다. 특히 데이터셋 규모가 작은 경우에는 supervised 모델을 처음부터 학습하는 것보다 대규모 사전학습 모델을 활용하는 방식이 더 효과적일 수 있음을 확인하였다.

5. 또한 다양한 해석 분석을 통해 NT 모델이 단순한 서열 패턴을 넘어 실제 생물학적 기능 정보를 학습하고 있음을 확인하였다. attention map 분석과 embedding 분석을 통해 모델이 exon, intron, promoter, enhancer, transcription factor binding site와 같은 주요 유전체 요소를 자동으로 인식한다는 사실이 밝혀졌다. 이는 self-supervised 방식으로 학습된 모델이 유전자 구조와 조절 요소와 같은 생물학적 특징을 내재적으로 학습하고 있음을 보여준다.

6. 본 연구는 또한 NT 모델이 genetic variant의 기능적 영향 예측에도 활용될 수 있음을 보여주었다. zero-shot 기반 점수를 이용하여 변이의 심각도와 기능적 중요도를 평가한 결과, 모델이 variant severity와 일정 수준의 상관관계를 보였으며 eQTL, meQTL, ClinVar 변이와 같은 기능적 변이를 분류하는 데에도 활용 가능함을 확인하였다. 특히 이러한 접근은 기능적 주석 정보가 부족한 비인간 생물 종에서도 유용하게 활용될 가능성이 있다.

7. 연구 후반에서는 모델 구조를 개선한 NT-v2 모델을 제안하여 더 효율적인 모델 설계 가능성을 보여주었다. NT-v2는 더 긴 입력 서열(12kb)을 처리할 수 있도록 설계되었으며, 개선된 아키텍처를 통해 모델 크기를 줄이면서도 성능을 향상시켰다. 실제로 2.5B 파라미터 모델보다 약 10배 작은 250M 모델이 더 높은 성능을 보이는 결과가 나타났으며, 이는 모델 구조와 학습 전략의 개선이 성능 향상과 계산 효율성을 동시에 달성할 수 있음을 의미한다.

8. 마지막으로 저자들은 transformer 기반 모델이 여전히 긴 유전체 서열을 처리하는 데 한계가 있다는 점을 지적한다. 유전자 발현은 수십 kb 이상 떨어진 조절 요소에 의해 영향을 받을 수 있지만, 일반적인 transformer 구조는 서열 길이가 증가할수록 계산량이 급격히 증가하는 문제가 있다. 따라서 향후 연구에서는 긴 서열을 효율적으로 처리하면서도 높은 성능을 유지할 수 있는 새로운 모델 구조를 개발하는 것이 중요한 연구 방향이 될 것으로 제시된다.

종합적으로 이 연구는 대규모 DNA 서열을 이용한 genomics foundation model 접근법이 다양한 유전체 분석 문제에서 효과적으로 활용될 수 있음을 보여준다. 특히 다양한 종의 유전체 데이터를 이용한 사전학습이 모델의 일반화 성능을 향상시키며, 이러한 모델은 유전자 구조 인식, 기능적 요소 탐지, 변이 영향 예측 등 다양한 생물학적 문제에 적용될 수 있다. 저자들은 이러한 접근이 앞으로 genomics 분야에서도 자연어처리와 유사한 foundation model 중심의 연구 패러다임을 확산시키는 계기가 될 것이라고 결론짓는다.