DNACHUNKER: Learnable Tokenization for DNA Language Models

Motivation

• Known: 기존 DNA language model들은 single nucleotide, fixed-size k-mer, 혹은 BPE 기반 tokenization을 사용해왔으며, Nucleotide Transformer, DNABERT-2, GROVER, Caduceus, HyenaDNA, Evo 등 다양한 아키텍처가 이러한 고정 tokenization 위에서 학습되어 왔다.
• Gap: DNA는 자연어와 달리 canonical한 word 경계가 없기 때문에, 고정 tokenization은 shift, indel, local repeat 등 미세한 변이에 취약하고, k-mer는 사소한 mutation에도 tokenized output이 완전히 달라지며, BPE 같은 frequency 기반 방법은 반복적이고 비기능적인 서열을 우선시해 TF-binding motif나 cis-regulatory motif 같은 실제 기능적 substring을 조각내는 문제가 있다.
• Why: DNA 서열의 기능적 요소(promoter, enhancer 등)를 의미 있는 token으로 묶는 tokenization 전략은 장거리 상호작용을 포착하는 DNA language model의 표현력과 downstream 성능에 직접적인 영향을 미치므로, mutation-resilient하고 생물학적으로 유의미한 tokenizer 설계는 genomic sequence modeling 전반의 신뢰성과 효율성을 높이는 데 중요하다.
• Approach: 두 단계 hierarchical encoder 구조로, 먼저 lightweight bidirectional Mamba layer로 base-pair 임베딩을 처리한 후 cosine-similarity 기반 boundary decision으로 인접 위치를 chunk로 병합하고, 압축된 시퀀스를 더 expressive한 main network로 모델링한 뒤 bidirectional dechunking module로 다시 base-pair 해상도로 upsampling하는 masked DNA language model인 DNAChunker를 제안한다.

Achievement

Figure 2. DNALongBench. Performance on DNALongBench across five long-range genomic prediction tasks. We compare DNACHUN-

1. 다섯 개 벤치마크에서 일관된 성능 향상: NT benchmark(및 revised version), Genomic Benchmarks, BEND, DNALongBench에 걸쳐 DNAChunker가 strong fixed-tokenization baseline 대비 일관되게 우수한 성능을 보였으며, multi-species pretraining을 사용한 더 큰 baseline(Nucleotide Transformer 등)보다도 뛰어난 결과를 달성했다.
2. 생물학적으로 유의미하고 mutation-resilient한 segmentation 학습: 학습된 tokenizer가 TF-binding motif, cis-regulatory motif 등 기능적 motif를 단일 chunk로 보존하고, SNV, InDel, structural variant 등 다양한 변이 하에서도 안정적인 boundary를 유지함을 확인했다.
3. 긴 context에서의 효율성 개선: adaptive segmentation이 유효 시퀀스 길이를 줄여 FLOPs 분석 결과 long-context 처리 효율을 향상시켰다.
4. 각 구성 요소의 기여 검증: controlled ablation을 통해 제안된 chunking, dechunking 등 각 모듈이 성능 향상에 기여함을 확인했다.

How

Figure 1. Architecture, tokenizer robustness, and motif fragmentation analysis. (a) The architecture of DNACHUNKER. (b)

• Human reference genome(GRCh38/hg38)만을 사용해 사전학습을 수행하며, multi-species pretraining 없이도 경쟁력 있는 성능을 목표로 함.
• 2단계 hierarchical encoder: Stage 1에서는 lightweight bidirectional Mamba layers(Caduceus의 BiMamba 기반)로 raw base-pair embedding을 처리.
• Chunking 모듈이 cosine-similarity 기반 boundary decision을 통해 인접 위치의 표현이 충분히 유사하면 하나의 token으로 병합, dynamic segmentation procedure를 학습.
• 압축된(chunked) 시퀀스는 더 expressive한 main network(Stage 2 encoder)에서 처리되어 장거리 의존성(long-range dependency)을 포착.
• Dechunker(Stage 1, Stage 2)가 bidirectional dechunking module을 통해 압축된 표현을 다시 base-pair 해상도로 upsampling하며, encoder feature로부터의 residual connection을 활용해 세밀한 표현을 복원.
• 이 구조는 masked language modeling 목적함수로 학습되며, 기존 autoregressive 기반 dynamic chunking(H-Net 등)을 bidirectional masked DNA pretraining 세팅으로 확장.
• 다섯 개 벤치마크(NT benchmark, revised NT benchmark, Genomic Benchmarks, BEND, DNALongBench)에서 평가하며, motif fragmentation analysis, token size distribution per-chromosome, computation overhead(FLOPs) 분석 등 다양한 ablation과 분석을 수행.

Originality

• Autoregressive language modeling에서 주로 연구되던 dynamic/learnable tokenization(H-Net 등)을 masked DNA pretraining 세팅에 최초로 적용하여 bidirectional한 token boundary 학습을 가능하게 함.
• 기존 non-autoregressive 모델들이 gradient-based pooling이나 word-level external chunking처럼 사전 정의된 boundary heuristic에 의존했던 것과 달리, DNAChunker는 representation 학습과 segmentation을 완전히 end-to-end로 jointly 학습.
• Cosine-similarity 기반 boundary decision과 2단계 encoder-dechunker 구조를 결합해, DNA에 특화된 mutation-resilient하고 biologically-informed한 tokenization을 구현.

Limitation & Further Study

• 사전학습이 human reference genome(GRCh38/hg38)에 한정되어 있어, multi-species나 다양한 population variant를 포함하는 대규모 pretraining 대비 일반화 범위가 제한적일 수 있음.
• Cosine-similarity 기반 boundary decision의 임계값 설정이나 chunk 크기 분포에 대한 이론적 정당화가 충분히 제시되지 않을 가능성이 있어, 다양한 genomic 맥락(예: 비인간 genome, 비coding 영역)에서의 강건성 추가 검증이 필요.
• 발췌된 내용만으로는 최종 chunk 개수/길이에 대한 hyperparameter 민감도, 그리고 Evo2와 같은 초대형 autoregressive 모델과의 직접 비교가 충분히 다뤄졌는지 확인하기 어려움.
• 후속 연구로 multi-species pretraining과의 결합, 혹은 다른 생물학적 서열(RNA, 단백질)로의 확장 가능성을 탐구할 필요가 있음.

Evaluation

총평: DNA tokenization이라는 상대적으로 덜 주목받은 문제를 masked language model 세팅에서 학습 가능한 방식으로 해결한 참신하고 실용적인 연구로, 다양한 벤치마크와 생물학적 분석을 통해 접근법의 타당성을 잘 뒷받침하고 있다.