LDARNet: DNA Adaptive Representation Network with Learnable Tokenization for Genomic Modeling

기술 요약: LDARNet

문제 정의

유전체 파운데이션 모델은 점차 대규모 언어 모델(LLM) 아키텍처를 채택하고 있으나, 주로 고정된 토큰화 방식(예: k-mer, Byte-Pair Encoding 또는 단일 뉴클레오타이드)에 의존하고 있습니다. 이러한 고정된 방식은 생물학적 근거가 부족하여 생물학적으로 유의미한 구조를 가릴 수 있는 임의적인 시퀀스 경계를 부과합니다. 최근 H-Net과 같은 연구를 통해 적응형(adaptive) 및 동적(dynamic) 토큰화가 자기회귀(autoregressive, AR) 생성에서는 가능하다는 것이 입증되었으나, 이러한 적응형 토큰화가 유전체 파운데이션 모델의 표준인 마스크드 언어 모델링(Masked Language Modeling, MLM) 패러다임 하에서도 경쟁력 있는 임베딩을 생성하는지는 여전히 불분명합니다. 또한, 기존의 적응형 방식들은 학습된 경계가 정형적인 유전체 랜드마크와 일치하는지 검증하거나, 학습된 라우팅과 성능 향상 사이의 인과 관계를 입증하는 데 실패하는 경우가 많습니다.

방법론

저자들은 AR 생성에서 MLM 사전 학습으로 H-Net 아키텍처를 적응시킨 120M 파라미터 규모의 계층적 유전체 파운데이션 모델인 LDARNet(Learnable DNA Adaptive Representation Network)을 제시합니다.

아키텍처

LDARNet은 H-Net의 인코더-백본-디코더 계층 구조를 유지하면서, 양방향 유전체 모델링에 적합하도록 네 가지 핵심적인 아키텍처 수정을 도입했습니다:

1. BiMamba-2 백본: 인과적(Causal) Mamba 레이어를 BiMamba-2 블록(Caduceus의 BiMamba 설계를 Mamba-2로 확장)으로 교체했습니다. 이 블록은 파라미터 수를 두 배로 늘리지 않고도 전체 컨텍스트 조건화를 가능하게 하기 위해 순방향 및 역방향 패스에 공유 가중치를 사용합니다.
2. 비인과적 어텐션(Non-Causal Attention): 모든 어텐션 메커니즘은 양방향 컨텍스트를 지원하기 위해 비인과적(non-causal)입니다.
3. 로컬 어텐션: 인코더는 미세한 패턴 인식을 돕기 위해 단일 로컬 어텐션 레이어를 포함하며, 백본과 디코더는 전적으로 BiMamba-2에 의존합니다.
4. 양방향 라우팅 및 디청킹(Dechunking):
   • 양방향 라우팅: H-Net의 단방향 라우터와 달리, LDARNet은 인접한 쌍의 순방향 및 역방향 방향성을 모두 사용하여 경계 증거를 추정합니다. 이는 대칭적으로 경계를 감지하기 위해 양방향 유사도 지표를 활용하며, 이를 통해 MLM에서 사용되는 15% 마스킹에 대해 견고한 프로세스를 구축합니다.
   • 양방향 디청킹: 디청커는 양방향 지수 이동 평균(Exponential Moving Average, EMA) 스무더를 통합합니다. 이는 양방향으로부터 정보를 전달하여 가변 길이 스팬(span)의 안정적인 재구성을 보장합니다.

학습 목표

모델은 MLM 손실과 비율 기반 정규화 항($L_{ratio}$)을 결과로 최적화됩니다. 이 정규화 항은 모든 단계에서 목표 압축 비율($N=4$)을 강제함으로써 퇴보적인 압축 솔루션(예: 사소한 압축)을 방지합니다. 학습 코퍼스는 인간 참조 게놈과 Nucleotide Transformer의 다종 데이터를 결합하며, 생물학적 대칭성을 강제하기 위해 상보 가닥 증강(reverse-complement augmentation)을 활용합니다.

주요 기여

본 논문은 네 가지 주요 기여를 명시합니다:

1. 아키텍처 적응: 적응형 토큰화 프레임워크를 자기회귀 생성에서 마스크드 언어 모델링으로 성공적으로 적응시켰습니다.
2. 성능 효율성: 동적 경계 예측을 갖춘 120M 파라미터 규모의 컴팩트한 모델이 4~20배 더 큰 베이스라인(최대 2.5B 파라미터)과 대등하거나 이를 능가할 수 있음을 입증하였으며, 5가지 히스톤 변형 작업에서 최첨단(SOTA) 결과를 달-성했습니다.
3. 교차 종 일반화: 학습 가능한 다중 스케일 압축을 통한 아키텍처의 일반성이 도메인 특화 최적화보다 교차 종 전이 가능성을 더 잘 보존함을 입상했습니다. LDARNet은 Nucleotide Transformer 벤치마크에서 11/18 승리를 거두며, 비슷한 규모의 인간 게놈 특화 모델들을 능가했습니다.
4. 학습된 경계에 대한 인과적 증거: 학습된 라우팅(단순히 계층적 스캐폴딩뿐만 아니라)이 성능 향상을 주도한다는 직접적인 증거를 제공합니다. FLOPs를 맞춘 통제 실험 결과, 학습된 경계가 고정 그리드 경계보다 히스톤 작업에서 최대 14.3%포인트 더 높은 성능을 보였습니다. 또한, 뉴클레오타이드 해상도 분석을 통해 학습된 경계가 감독 없이도 정형적인 프로모터 모티프 및 스플라이스 접합부(splice junctions)와 일치함을 밝혔습니다.

결과

Nucleotide Transformer (NT) 벤치마크

18개의 다양한 작업에서 LDARNet은 컴팩트 모델(<300M 파라미터) 중 11승을 거두었으며, 이는 다른 컴팩트 모델들(각각 1~2승 달성)을 크게 앞지른 결과입니다. 특히 5가지 히스톤 변형 작업(H3K4me1, H3K4me2, H3K4me3, H3K79me3, H4ac)에서 경쟁 모델들을 최대 2.5B 파라미터까지 압도하며 최고의 성적을 기록했습니다.

유전체 벤치마크 (GB)

9가지 조절 유전체 작업에서 LDARNet은 DNABERT-2와 함께 컴팩트 모델 중 최고 성능(각 3승)을 기록하며 동률을 이루었습니다. GB 작업은 높은 베이스라인 정확도로 인해 천장 효과(ceiling effect)가 나타나기도 하지만, LDARNet은 "Human non-TATA Promoters" 작업에서 96.3%의 최고 정확도를 달성하며 Generator-1.2B와 같은 더 큰 모델을 능가했습니다.

통제 실험

LDARNet(학습된 라우팅)과 4개 뉴클레오타이드마다 고정된 경계($N=4$)를 가진 변형 모델을 비교한 FLOPs 매칭 실험을 통해 성능 향상의 원인을 분리했습니다:

• 히스톤 작업: LDARNet은 고정 경계 변형 모델에 비해 H3K4me1에서 +14.3pp, H3K36me3에서 +6.7pp 높은 성능을 보였습니다. 고정 경계 모델은 이러한 작업에서 훨씬 높은 분산을 보였습니다.
• 스플라이스 사이트(Splice Sites): 스플라이스 사이트 분류에서는 경향이 역전되어, 고정 경계(0.961)가 학습된 라우팅(0.901)보다 우수한 성능을 보였습니다. 저자들은 이를 스플라이스 모티프(GT/AG)의 국소적 특성 때문으로 설명하며, 고정 그리드가 학습된 라우팅보다 균일한 위치 정보를 더 잘 보존한다고 분석했습니다.

의의

본 논문은 LDARNet이 유전체 파운데이션 모델에서 적응형 토큰화에 대한 생물학적 해석을 제공한다고 주장합니다. 감독 없이도 기능적 유전체 요소(프로모터, 스플라이스 접합부)와 일치하는 경계를 학습함으로써, 이 모델은 적응형 토큰화가 고정된 방식보다 기능적 신호를 더 충실하게 포착할 수 있음을 보여줍니다. 이 연구는 학습된 라우팅이 컴팩트한 모델이 히스톤 변형 예측과 같은 장거리 후성유전학적 모델링 작업에서 더 큰 아키텍처와 경쟁할 수 있게 만드는 핵심 요소임을 강조합니다.