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

Motivation

• Known: DNABERT, Nucleotide Transformer, GENA-LM, Caduceus, HyenaDNA, JanusDNA 등 다양한 genomic foundation model들이 대규모 사전학습을 통해 promoter, enhancer, splice site 등 다운스트림 과제에 전이 가능한 임베딩을 생성함이 알려져 있으며, 이들 대부분은 k-mer나 BPE 같은 고정 tokenization 방식에 의존한다.
• Gap: H-Net이 DNA에 대해 동적 계층적 tokenization을 도입해 perplexity 감소를 보였지만 다운스트림 분류 벤치마크에서 평가되지 않았고, MxDNA·VQDNA 같은 적응형 tokenization 연구들도 학습된 경계 자체를 생물학적 랜드마크와 대조 검증하지 않아, adaptive tokenization이 MLM 패러다임에서 확립된 genomic foundation model과 경쟁력 있는 임베딩을 만드는지, 그리고 그 경계가 생물학적으로 의미 있는지가 미해결 문제로 남아 있었다.
• Why: 고정 tokenization은 DNA 서열에 임의적 경계를 부과해 생물학적으로 중요한 구조(프로모터 모티프, splice junction 등)를 가릴 수 있으므로, 학습 가능한 tokenization이 더 작은 모델로도 더 나은 성능과 해석 가능성을 제공할 수 있음을 보이는 것은 genomic foundation model 설계에 중요한 함의를 갖는다.
• Approach: H-Net 아키텍처를 MLM에 맞게 적응시켜, BiMamba-2 state-space layer, local attention, bidirectional routing, ratio 기반 정규화를 결합함으로써 비지도 방식으로 적응형 토큰 경계를 유도하는 계층적 인코더-백본-디코더 구조(LDARNet)를 설계했다.

Achievement

Figure 2. Boundary probabilities around canonical promoter motifs. Mean boundary probability with 95% confidence interva

1. 성능: Nucleotide Transformer 및 Genomic Benchmarks의 27개 과제에 대한 fine-tuning에서, 300M 파라미터 미만 compact model 중 18개 과제 중 11개에서 우승했으며, 5개 histone modification 과제에서 20배 큰 모델(2.5B 파라미터)을 능가하는 state-of-the-art 성능을 달성했다.
2. 인과적 검증: FLOPs를 맞춘 통제 실험을 통해 학습된 routing이 성능 향상의 원인임을 입증했다—동일 연산량에서 학습된 경계가 고정 격자 경계보다 histone 과제에서 최대 14.3 percentage point 우수했다.
3. 생물학적 해석: 뉴클레오타이드 해상도 분석에서 학습된 경계가 비지도 방식으로도 정규 프로모터 모티프 및 splice junction과 정렬됨을 보여, adaptive tokenization의 생물학적 타당성을 제시했다.
4. 교차종 전이성: NT 벤치마크에서 인간 게놈에 특화 최적화된 유사 규모 모델(1-2/18 승)에 비해 LDARNet이 아키텍처적 일반성을 통해 더 우수한 교차종 전이성(11/18 승)을 유지함을 보였다.

How

Figure 1. Model overview. Left: the LDARNet architecture with BiMamba-2 outer layers and a BiMamba-2 backbone operating

• H-Net의 인코더-백본-디코더 계층 구조를 유지하되 Mamba 레이어를 Caduceus의 BiMamba 설계를 Mamba-2로 확장한 BiMamba-2 블록(가중치 공유)으로 교체
• 모든 attention을 non-causal로 변경하여 양방향 DNA 구조에 맞춤
• 인코더에 세밀한 패턴 인식을 위한 단일 local attention layer 추가, 백본과 디코더는 전적으로 BiMamba-2 블록으로 구성
• router와 dechunking 모듈을 bidirectional 변형으로 확장하여 비지도 경계 예측 수행
• ratio 기반 정규화를 도입해 청크 비율을 제어하며 학습 안정화(Figure 4)
• 27개 다운스트림 과제(Nucleotide Transformer, Genomic Benchmarks 스위트)에서 fine-tuning 평가 수행
• FLOPs를 맞춘 고정 격자(fixed-grid) 경계 대비 controlled ablation으로 학습된 routing의 기여를 분리
• 뉴클레오타이드 해상도에서 경계 확률을 promoter motif(Figure 2) 및 splice junction(Figure 3) 위치와 대조 분석

Originality

• 기존 H-Net은 autoregressive 생성 전용으로 설계되었으나 이를 MLM 패러다임에 최초로 적응시켜 DNA의 양방향적 특성을 반영한 bidirectional routing/dechunking 및 BiMamba-2 기반 구조로 확장한 점
• 기존 adaptive tokenization 연구(MxDNA, VQDNA)와 달리 학습된 경계 자체를 명시적 boundary-placement 문제로 정식화하고, 이를 생물학적 랜드마크(프로모터 모티프, splice junction)와 직접 대조 검증한 최초의 시도
• FLOPs-matched controlled experiment를 통해 성능 향상의 원인이 계층적 구조 자체가 아니라 학습된 routing 메커니즘임을 인과적으로 분리해낸 실험 설계

Limitation & Further Study

• MxDNA, VQDNA 등 주요 경쟁 adaptive tokenization 방법들이 구현 상의 이유로 다운스트림 baseline 비교에서 제외되어, 직접적인 성능 비교가 제한적임
• 27개 과제 중 18개 compact model 비교에서 11승에 그쳐 나머지 7개 과제에서는 우위를 점하지 못했으며, 이에 대한 실패 원인 분석이 본문 발췌에서는 충분히 제시되지 않음
• 단일 120M 파라미터 규모에서의 실험으로, 더 큰 스케일에서 학습된 routing의 이점이 유지되는지에 대한 스케일링 검증이 부족할 수 있음
• 후속 연구로 다양한 종·과제에 대한 일반화 검증, 더 큰 모델 스케일에서의 재현, 다른 adaptive tokenization 방법과의 직접적 재구현 비교가 필요함

Evaluation

총평: H-Net의 동적 청킹을 MLM 패러다임으로 성공적으로 이식하고, FLOPs-matched 실험과 뉴클레오타이드 해상도 분석을 통해 학습된 tokenization의 성능적·생물학적 타당성을 인과적으로 입증한 견고한 연구로, compact genomic foundation model 설계에 실질적 기여를 한다.