Resumen Técnico: JEPA-DNA

Planteamiento del Problema

Los Modelos Fundacionales Genómicos (GFMs) han avanzado en el modelado de secuencias de ADN mediante la adaptación de arquitecturas de Modelos de Lenguaje de Gran Escala (LLM) a los datos genómicos. Sin embargo, los GFMs actuales dependen predominantemente de objetivos de auto-supervisión generativos, como el Modelado de Lenguaje Enmascarado (MLM) o la Predicción del Siguiente Token (NTP). Aunque son efectivos para capturar la sintaxis de la secuencia local e identificar motivos, estos paradigmas de reconstrucción a nivel de token suelen priorizar el "ruido" de secuencias de alta frecuencia (por ejemplo, elementos repetitivos) sobre el contexto funcional global. En consecuencia, los modelos pueden lograr una alta fidelidad sintáctica sin interiorizar la lógica regulatoria de alto nivel y las interacciones de largo alcance que gobiernan la función genómica. Existe una brecha crítica entre la sintaxis genómica y la semántica biológica que los objetivos generativos existentes no logan cerrar.

Metodología: Marco de Trabajo JEPA-DNA

Para abordar esta limitación, los autores presentan JEPA-DNA, un marco de entrenamiento continuo agnóstico al modelo que integra una Arquitectura Predictiva de Embebimiento Conjunto (JEPA) con objetivos generativos tradicionales. En lugar de predecir tokens brutos, JEPA-DNA obliga a los modelos a predecir las representaciones funcionales latentes de segmentos genómicos enmascarados.

Componentes de la Arquitectura

El marco aumenta un backbone estándar de GFM ($E_\theta$) con una rama JEPA que consta de:

1. Codificador de Contexto ($E_\theta$): Procesa una secuencia de entrada enmascarada $x$ para producir representaciones latentes contextuales $h_{ctx}$.
2. Codificador de Objetivo ($E_{\bar{\theta}}$): Un duplicado estructural del codificador de contexto con pesos actualizados mediante un Promedio Móvil Exponencial (EMA). Procesa la secuencia objetivo no enmascarada $y$ para generar objetivos latentes $z$.
3. Cabeza de Predicción ($P_\phi$): Mapea las representaciones de contexto al espacio latente del objetivo, aprendiendo a alinear las características derivadas del contexto con los embebimientos del codificador de objetivo.
4. Agregador de Secuencias ($A$): Comprime las secuencias latentes en representaciones globales (por ejemplo, usando tokens [CLS], el último token o el pooling de la media) para facilitar la alineación semántica global.

Estrategia de Entrenamiento

JEPA-DNA emplea un procedimiento de entrenamiento continuo de objetivos múltiples:

• Doble Enmascaramiento: Utiliza una estrategia de enmascaramiento basada en tramos (spans) con un programador que aumenta gradualmente las proporciones de enmascaramiento y las longitudes de los tramos, desafiando al modelo a inferir dependencias de largo alcance.
• Re-enmascaramiento: Reemplaza las posiciones enmascaradas en la salida del codificador de contexto con embebimientos latentes [MASK] aprendibles para evitar mapeos triviales.
• Función de Pérdida Compuesta: El modelo se optimiza minimizando una pérdida total ($L_{total}$) que equilibra cuatro componentes:
  1. Pérdida de LLM ($L_{llm}$): Pérdida estándar de MLM o NTP para mantener la precisión a nivel de token.
  2. Pérdida de Predicción Latente ($L_{jepa}$): Una pérdida de similitud de coseno que obliga a la representación latente global predicha ($\hat{z}_{glob}$) a alinearse con la representación latente global objetivo ($z_{glob}$).
  3. Pérdida de Varianza ($L_{var}$): Asegura que las dimensiones del embebimiento mantengan suficiente varianza a través del lote para prevenir el colapso de información.
  4. Pérdida de Covarianza ($L_{cov}$): Fomenta la descorrelación entre las dimensiones del embebimiento para evitar representaciones redundantes (basado en VICReg).

El marco está diseñado para ser compatible con diversas arquitecturas (Transformers, Modelos de Espacio de Estados como HyenaDNA) y métodos de tokenización, requiriendo únicamente un operador de agregación para comprimir las secuencias en representaciones latentes.

Contribuciones Clave

1. Nuevo Marco de Entrenamiento: Introducción de un marco de entrenamiento continuo que desplaza el objetivo de aprendizaje de la reconstrucción literal de tokens hacia la alineación de características latentes vía JEPA.
2. Agnosticismo de Arquitectura: Demostración de que JEPA-DNA actúa como un mecanismo universal de "mejora representacional", compatible con los principales GFMs (DNABERT-2, NTv3, HyenaDNA) y con objetivos tanto de MLM como de NTP.
3. Rendimiento de Vanguardia (SOTA): Establecimiento de nuevos techos de rendimiento en tareas de sondaje lineal supervisado al unir la precisión generativa con el anclaje semántico latente.
4. Análisis Sinérgico: Caracterización de la interacción entre los objetivos generativos y latentes, mostrando que un enfoque híbrido supera a los métodos puramente reconstructivos o puramente contrastivos.

Resultados Experimentales

El marco fue evaluado en 17 tareas de referencia genómica diversas utilizando protocolos de sondaje lineal y zero-shot.

• Sondaje Lineal (Linear Probing): JEPA-DNA demostró ganancias consistentes en las tres arquitecturas de backbone probadas.
  • DNABERT-2: Mejoró el rendimiento en 6 de 9 tareas, con ganancias de AUROC de hasta +8.0%.
  • NTv3: Mejoró el rendimiento en 8 de 9 tareas.
  • HyenaDNA: Mejoró el rendimiento en 7 de 9 tareas, con ganancias notables en el reconocimiento de promotores (+11.22%) y la clasificación de sitios de splicing (+10.28%).
  • Especificidad de la Tarea: Las ganancias fueron más pronunciadas en tareas que dependen de un contexto regulatorio más amplio (promotores, sitios de splicing, meQTL, eQTL causales), mientras que las tareas impulsadas por cambios locales finos (patogenicidad de codificación) mostraron resultados más variables.
• Evaluación Zero-Shot: Utilizando DNABERT-2, JE-PA-DNA mejoró el rendimiento zero-shot en 6 de 8 benchmarks clínicos y de puntuación de fitness, incluyendo una ganancia de AUROC de +12.06% en la predicción de variantes de enfermedad.
• Estudios de Ablación:
  • Regularización VICReg: Identificada como esencial para desbloquear los beneficios del paradigma híbrido.
  • Arquitectura del Predictor: Los predictores basados en atención (4 capas) superaron a los MLPs y a las arquitecturas más profundas.
  • Composición de la Pérdida: La alineación en el espacio latente proporcionó una señal más informativa (+3.86% de ganancia) que la reconstrucción de tokens por sí sola (+2.06%), pero la configuración óptima requirió un híbrido de ambas.

Significado

El artículo afirma que JEPA-DNA establece un nuevo estado del arte para los Modelos Fundacionales Genómicos al lograr anclar con éxito sus representaciones en la semántica biológica. Al desplazar la señal de aprendizaje de la recuperación de tokens hacia la alineación semántica, el marco fuerza a los modelos a aprender características funcionales abstractas que son invariantes al ruido de secuencia de bajo nivel. Este enfoque cierra la brecha entre la sintaxis genómica y la lógica biológica, ofreciendo un camino escalable para el anclaje de modelos fundacionales sin la carga computacional del alineamiento de múltiples secuencias o los requisitos de datos del minado de negativos contrastivos. Los autores posicionan a JEPA-DNA no como un reemplazo de los GFMs existentes, sino como una fase de entrenamiento continuo versátil que mejora la calidad representacional de cualquier modelo fundacional genómico existente.