기술 요약: JEPA-DNA

문제 정의

유전체 파운데이션 모델(GFMs)은 LLM 아키텍처를 유전체 데이터에 적응시킴으로써 DNA 서열 모델링을 발전시켜 왔다. 그러나 현재의 GFM들은 주로 마스크드 언어 모델링(MLM)이나 다음 토큰 예측(NTP)과 같은 생성적 자기지도 학습 목적 함수에 의존한다. 이러한 방식은 국소적인 서열 구문(syntax)을 포착하고 모티프를 식별하는 데는 효과적이지만, 토큰 수준의 재구성 패러다임은 종종 전역적인 기능적 맥락보다 고빈도 서열 "노이즈"(예: 반복 요소)를 우선시하는 경향이 있다. 결과적으로, 모델은 높은 구문적 충실도를 달ek할 수는 있으나, 유전체 기능을 규정하는 고차원적 조절 로직이나 장거리 상호작용을 내부화하지 못할 수 있다. 즉, 기존의 생성적 목적 함수가 해결하지 못하는 유전체 구문(syntax)과 생물학적 의미론(semantics) 사이의 결정적인 간극이 존재한다.

방법론: JEPA-DNA 프레임워크

이러한 한계를 해결하기 위해, 저자들은 결합 임베딩 예측 아키텍처(JEPA)를 전통적인 생성적 목적 함수와 통합한 모델 불가지론적(model-agnostic) 지속 학습 프레임워크인 JEPA-DNA를 소개한다. JEPA-DNA는 단순히 원시 토큰을 예측하는 대신, 마스킹된 유전체 세그먼트의 잠재적 기능 표현(latent functional representations)을 예측하도록 모델을 강제한다.

아키텍처 구성 요소

본 프레임워크는 표준 GFM 백본($E_\theta$)을 다음과 같은 JEPA 브랜치로 확장한다:

1. 컨텍스트 인코더 ($E_\theta$): 마스킹된 입력 서열 $x$를 처리하여 맥락적 잠재 표현 $h_{ctx}$를 생성한다.
2. 타겟 인코더 ($E_{\bar{\theta}}$): 지수 이동 평균(EMA)을 통해 가중치가 업데이트되는 구조적 복제본으로, 마스킹되지 않은 타겟 서열 $y$를 처리하여 잠재 타겟 $z$를 생성한다.
3. 예측 헤드 ($P_\phi$): 컨텍스트 표현을 타겟 잠재 공간으로 매핑하며, 컨텍스트 유래 특징을 타겟 인코더의 임베딩과 정렬하도록 학습한다.
4. 서열 어그리게이터 ($A$): 글로벌 표현(예: [CLS] 토큰, 마지막 토큰, 또는 평균 풀링 사용)으로 잠재 서열을 압축하여 전역적 의미론적 정렬을 용이하게 한다.

학습 전략

JEPA-DNA는 다중 목적 지속 학습(multi-objective continual training) 절차를 채택한다:

• 이중 마스킹(Dual-Masking): 스팬(span) 기반 마스킹 전략을 사용하여 마스킹 비율과 스팬 길이를 점진적으로 증가시키는 스케줄러를 활용함으로써, 모델이 장거리 의존성을 추론하도록 도전 과제를 부여한다.
• 재마스킹(Re-masking): 컨텍스트 인코더의 출력에서 마스킹된 위치를 학습 가능한 잠재 [MASK] 임베딩으로 대체하여 사소한 매핑(trivial mapping)을 방지한다.
• 복합 손실 함수(Composite Loss Function): 모델은 다음 네 가지 구성 요소를 균형 있게 조절하는 총 손실($L_{total}$)을 최소화하도록 최적화된다:
  1. LLM 손실 ($L_{llm}$): 토큰 수준의 정밀도를 유지하기 위한 표준 MLM 또는 NTP 손실.
  2. 잠재 예측 손실 ($L_{jepa}$): 예측된 글로벌 잠재 표현($\hat{z}_{glob}$)이 타겟 글로벌 잠재 표현($z_{glob}$)과 일치하도록 강제하는 코사인 유사도 손실.
  3. 분산 손실 ($L_{var}$): 정보 붕괴(informational collapse)를 방지하기 위해 임베딩 차원이 배치 전반에 걸쳐 충분한 분산을 유지하도록 보장한다.
  4. 공분산 손실 ($L_{cov}$): 중복된 표현을 방지하기 위해 임베딩 차원 간의 비상관성을 촉진한다(VICReg 기반).

본 프레임 lack은 다양한 아키텍처(Transformer, HyenaDNA와 같은 State Space Model) 및 토크나이제이션 방식과 호환되도록 설계되었으며, 서열을 잠재 표현으로 압축하기 위한 어그리게이션 연산자만을 필요로 한다.

주요 기여

1. 새로운 학습 프레임워크: 학습 목적을 단순한 토큰 복구에서 JEPA를 통한 잠재 특징 정렬로 전환하는 새로운 지속 학습 프레임워크를 도입하였다.
2. 아키텍처 불가지론성: JEPA-DNA가 다양한 GFM(DNABERT-2, NTv3, HyenaDNA) 및 MLM과 NTP 목적 함수 모두와 호환되는 보편적인 "표현력 향상(representational uplift)" 메커니즘임을 입증하였다.
3. 최첨단 성능(SOTA): 생성적 정밀도와 잠재적 의미론적 접지를 결합함으로써 지도 학습 선형 프로빙(linear probing) 작업에서 새로운 성능 한계를 구축하였다.
4. 시너지 분석: 생성적 목적 함수와 잠재적 목적 함수 사이의 상호작용을 규명하여, 하이브리드 접근 방식이 순수 재구성 방식이나 순수 대조 학습 방식보다 우수함을 보여주었다.

실험 결과

프레임워크는 선형 프로빙 및 제로샷 프로토콜을 사용하여 17개의 다양한 유전체 벤치마크 작업에 대해 평가되었다.

• 선형 프로빙 (Linear Probing): JEPA-DNA는 테스트된 세 가지 백본 아키텍처 모두에서 일관된 성능 향상을 보였다.
  • DNABERT-2: 9개 작업 중 6~8개 작업에서 성능이 개선되었으며, AUROC 상승폭은 최대 +8.0%에 달했다.
  • NTv3: 9개 작업 중 8개 작업에서 성능이 개선되었다.
  • HyenaDNA: 9개 작업 중 7개 작업에서 성능이 개선되었으며, 특히 프로모터 인식(+11.22%)과 스플라이스 사이트 분류(+10.28%)에서 주목할 만한 이득을 보였다.
  • 작업 특이성: 성능 향상은 광범위한 조절 맥락(프로모터, 스플라이스 사이트, meQTL, 인과적 eQTL)에 의존하는 작업에서 가장 두드러졌으며, 미세한 국소적 변화(코딩 병원성)에 의해 주도되는 작업에서는 결과가 더 가변적이었다.
• 제로샷 평가 (Zero-Shot Evaluation): DNABERT-2를 사용하여, JE way-DNA는 질병 변이 예측에서의 +12.06% AUROC 상승을 포함하여 8개의 임상 및 적합도 점수 벤치마크 중 6개에서 제로샷 성능을 향상시켰다.
• 절제 연구 (Ablation Studies):
  • VICReg 정규화: 하이브리드 패러다임의 이점을 실현하는 데 필수적인 요소로 확인되었다.
  • 예측기 아키텍처: 어텐션 기반 예측기(4개 레이어)가 MLP 및 더 깊은 구조보다 우수한 성능을 보였다.
  • 손실 구성: 잠재 공간 정렬이 토큰 재구성 단독 방식(+2.06%)보다 더 유익한 신호(+3.86% 이득)를 제공했으나, 최적의 구성은 두 방식의 하이브리드 형태를 필요로 했다.

의의

본 논문은 JEPA-DNA가 유전체 모델을 생물학적 의미론에 성공적으로 접지시킴으로써 유전체 파운데이션 모델의 새로운 최첨단(SOTA)을 구축했다고 주장한다. 학습 신호를 토큰 복구에서 의미론적 정렬로 전환함으로써, 프레임워크는 모델이 저수준의 서열 노이즈에 불변한 추상적이고 기능적인 특징을 학습하도록 강제한다. 이 접근 방식은 유전체 구문과 생물학적 로직 사이의 간극을 메우며, 다중 서열 정렬의 계산적 오버헤드나 대조적 네거티브 마이닝의 데이터 요구 사항 없이도 파운데이션 모델을 접지시킬 수 있는 확장 가능한 경로를 제공한다. 저자들은 JEPA-DNA를 기존 GFM을 대체하는 것이 아니라, 기존의 모든 유전체 파운데이션 모델의 표현 품질을 향상시킬 수 있는 다용도 지속 학습 단계로 정의한다.