Resumo Técnico: JEPA-DNA

Declaração do Problema

Os Modelos de Fundação Genômica (GFMs) avançaram na modelagem de sequências de DNA ao adaptar arquiteturas de Grandes Modelos de Linguagem (LLMs) para dados genômicos. No entanto, os GFMs atuais dependem predominantemente de objetivos autossupervisionados generativos, como o Modelagem de Linguagem Mascarada (MLM) ou a Predição de Próximo Token (NTP). Embora eficazes na captura da sintaxe de sequência local e na identificação de motivos, esses paradigmas de reconstrução ao nível de token frequentemente priorizam o "ruído" de sequência de alta frequência (ex: elementos repetitivos) em detrimento do contexto funcional global. Consequentemente, os modelos podem alcançar alta fidelidade sintática sem internalizar a lógica regulatória de alto nível e as interações de longo alcance que governam a função genômica. Existe uma lacuna crítica entre a sintaxe genômica e a semântica biológica que os objetivos generativos existentes falham em preencher.

Metodologia: Estrutura JEPA-DNA

Para abordar essa limitação, os autores introduzem o JEPA-DNA, uma estrutura de treinamento contínuo agnóstica ao modelo que integra uma Arquitetura Preditiva de Embedding Conjunto (JEoma - Joint-Embedding Predictive Architecture) com objetivos generativos tradicionais. Em vez de prever tokens brutos, o JEPA-DNA força os modelos a prever as representações funcionais latentes de segmentos genômicos mascarados.

Componentes da Arquitetura

A estrutura aumenta o backbone padrão de um GFM ($E_\theta$) com um ramo JEPA consistindo em:

1. Codificador de Contexto ($E_\theta$): Processa uma sequência de entrada mascarada $x$ para produzir representações latentes contextuais $h_{ctx}$.
2. Codificador de Alvo ($E_{\bar{\theta}}$): Uma duplicata estrutural do codificador de contexto com pesos atualizados via Média Móvel Exponencial (EMA). Ele processa a sequência alvo não mascarada $y$ para gerar alvos latentes $z$.
3. Cabeça de Predição ($P_\phi$): Mapeia as representações de contexto para o espaço latente do alvo, aprendendo a alinhar as características derivadas do contexto com os embeddings do codificador de alvo.
4. Agregador de Sequência ($A$): Comprime sequências latentes em representações globais (ex: usando tokens [CLS], último token ou pooling de média) para facilitar o alinhamento semântico global.

Estratégia de Treinamento

O JEPA-DNA emprega um procedimento de treinamento contínuo de múltiplos objetivos:

• Mascaramento Duplo (Dual-Masking): Utiliza uma estratégia de mascaramento baseada em spans com um scheduler que aumenta gradualmente as proporções de mascaramento e os comprimentos dos spans, desafiando o modelo a inferir dependências de longo alcance.
• Remascaramento (Re-masking): Substitui as posições mascaradas na saída do codificador de contexto por embeddings [MASK] latentes aprendíveis para evitar mapeamentos triviais.
• Função de Perda Composta: O modelo é otimizado minimizando uma perda total ($L_{total}$) que equilibra quatro componentes:
  1. Perda LLM ($L_{llm}$): Perda padrão de MLM ou NTP para manter a precisão ao nível de token.
  2. Perda de Predição Latente ($L_{jepa}$): Uma perda de similaridade de cosseno que força a representação latente global predita ($\hat{z}_{glob}$) a se alinhar com a representação latente global alvo ($z_{glob}$).
  3. Perda de Variância ($L_{var}$): Garante que as dimensões do embedding mantenham variância suficiente através do batch para evitar o colapso informacional.
  4. Perda de Covariância ($L_{cov}$): Incentiva a descorrelação entre as dimensões do embedding para evitar representações redundantes (baseado em VICReg).

A estrutura é projetada para ser compatível com diversas arquiteturas (Transformers, State Space Models como HyenaDNA) e métodos de tokenização, exigindo apenas um operador de agregação para comprimir sequências em representações latentes.

Principais Contribuições

1. Nova Estrutura de Treinamento: Introdução de uma estrutura de treinamento contínuo que desloca o objetivo de aprendizado da reconstrução literal de tokens para o alinhamento de características latentes via JEPA.
2. Agnosticismo de Arquitetura: Demonstração de que o JEPA-DNA atua como um mecanismo universal de "elevação representacional", compatível com os principais GFMs (DNABERT-2, NTv3, HyenaDNA) e ambos os objetivos MLM e NTP.
3. Desempenho de Estado da Arte: Estabelecimento de novos tetos de desempenho em tarefas de sondagem linear supervisionada ao unir a precisão generativa com o embasamento semântico latente.
4. Análise Sinérgica: Caracterização da interação entre objetivos generativos e latentes, mostrando que uma abordagem híbrida supera métodos puramente de reconstrução ou puramente contrastivos.

Resultados Experimentais

A estrutura foi avaliada em 17 tarefas de benchmark genômico diversas usando protocolos de sondagem linear e zero-shot.

• Sondagem Linear (Linear Probing): O JEPA-DNA demonstrou ganhos consistentes em todas as três arquiteturas de backbone testadas.
  • DNABERT-2: Melhorou o desempenho em 6 a 8 de 9 tarefas, com ganhos de AUROC de até +8,0%.
  • NTv3: Melhorou o desempenho em 8 de 9 tarefas.
  • HyenaDNA: Melhorou o desempenho em 7 de 9 tarefas, com ganhos notáveis em reconhecimento de promotores (+11,22%) e classificação de sítios de splicing (+10,28%).
  • Especificidade de Tarefa: Os ganhos foram mais pronunciados em tarefas que dependem de um contexto regulatório mais amplo (promotores, sítios de splicing, meQTL, eQTL causal), enquanto tarefas impulsionadas por mudanças locais refinadas (patogenicidade de codificação) mostraram resultados mais variáveis.
• Avaliação Zero-Shot: Usando o DNABERT-2, o JEPA-DNA melhorou o desempenho zero-shot em 6 de 8 benchmarks clínicos e de pontuação de fitness, incluindo um ganho de +12,06% de AUROC na predição de variantes de doenças.
• Estudos de Ablação:
  • Regularização VICReg: Identificada como essencial para desbloquear os benefícios do paradigma híbrido.
  • Arquitetura do Preditor: Preditores baseados em atenção (4 camadas) superaram MLPs e arquiteturas mais profundas.
  • Composição de Perda: O alinhamento do espaço latente forneceu um sinal mais informativo (+3,86% de ganho) do que a reconstrução de token sozinha (+2,06%), mas a configuração ideal exigiu um híbrido de ambos.

Significância

O artigo afirma que o JEPA-DNA estabelece um novo estado da arte para os Modelos de Fundação Genômica ao conseguir fundamentá-los na semântica biológica. Ao deslocar o sinal de aprendizado da recuperação de tokens para o alinhamento semântico, a estrutura força os modelos a aprender características funcionais abstratas que são invariantes ao ruído de sequência de baixo nível. Esta abordagem preenche a lacuna entre a sintaxe genômica e a lógica biológica, oferecendo um caminho escalável para fundamentar modelos de fundação sem o custo computacional do alinhamento de múltiplas sequências ou os requisitos de dados de mineração de negativos contrastivos. Os autores posicionam o JEPA-DNA não como um substituto para os GFMs existentes, mas como uma fase de treinamento contínuo versátil que melhora a qualidade representacional de qualquer modelo de fundação genômica existente.