Technische Samenvatting: JEPA-DNA

Probleemstelling

Genomische Fundamentele Modellen (GFMs) hebben de modellering van DNA-sequenties geavanceerd door architecturen van Large Language Models (LLMs) aan te passen aan genomische data. De huidige GFMs vertrouwen echter voornamelijk op generatieve zelfgesuperviseerde doelstellingen, zoals Masked Language Modeling (MLM) of Next-Token Prediction (NTP). Hoewel deze effectief zijn in het vastleggen van lokale sequentie-syntaxis en het identificeren van motieven, geven deze op token-niveau gebaseerde reconstructieparadigma's vaak de voorkeur aan hoogfrequente sequentie-"ruis" (bijv. repetitieve elementen) boven de globale functionele context. Bijgevolg kunnen modellen een hoge syntactische getrouwheid bereiken zonder de hoogwaardige regulatoire logica of langetermijninteracties te internaliseren die de genomische functie beheersen. Er bestaat een kritieke kloof tussen genomische syntaxis en biologische semantiek die bestaande generatieve doelstellingen niet weten te overbruggen.

Methodologie: JEPA-DNA Framework

Om deze beperking aan te pakken, introduceren de auteurs JEPA-DNA, een model-agnostisch framework voor continue training dat een Joint-Embedding Predictive Architecture (JEPA) integreert met traditionele generatieve doelstellingen. In plaats van ruwe tokens te voorspellen, dwingt JEPA-DNA modellen om de latente functionele representaties van gemaskeerde genomische segmenten te voorspellen.

Architectuurcomponenten

Het framework breidt een standaard GFM-backbone ($E_\theta$) uit met een JEPA-tak bestaande uit:

1. Context Encoder ($E_\theta$): Verwerkt een gemaskeerde inputsequentie $x$ om contextuele latente representaties $h_{ctx}$ te produceren.
2. Target Encoder ($E_{\bar{\theta}}$): Een structurele kopie van de context encoder met gewichten die worden bijgewerkt via Exponential Moving Average (EMA). Deze verwerkt de ongemaskeerde target-sequentie $y$ om latente targets $z$ te genereren.
3. Predictor Head ($P_\phi$): Brengt context-representaties in kaart naar de latente target-ruimte, waarbij geleerd wordt om context-afgeleide kenmerken af te stemmen op de embeddings van de target encoder.
4. Sequence Aggregator ($A$): Comprimeert latente sequenties tot globale representaties (bijv. via [CLS]-tokens, de laatste tokens of mean pooling) om globale semantische afstemming te faciliteren.

Trainingsstrategie

JEPA-DNA maakt gebruik van een multi-objective continue trainingsprocedure:

• Dual-Masking: Maakt gebruik van een span-gebaseerde maskeringstrategie met een scheduler die de maskeringratio's en span-lengtes geleidelijk verhoogt, wat het model uitdaagt om langetermijnafhankelijkheden te afleiden.
• Re-masking: Vervangt gemaskeerde posities in de output van de context encoder door leerbare latente [MASK]-embeddings om triviale mapping te voorkomen.
• Composiete Verliesfunctie: Het model wordt geoptimaliseerd door een totale loss ($L_{total}$) te minimaliseren die vier componenten balanceert:
  1. LLM Loss ($L_{llm}$): Standaard MLM of NTP loss om precisie op token-niveau te behouden.
  2. Latent Predictive Loss ($L_{jepa}$): Een cosine similarity loss die de voorspelde globale latente representatie ($\hat{z}_{glob}$) dwingt af te stemmen op de target globale latente representatie ($z_{glob}$).
  3. Variance Loss ($L_{var}$): Zorgt ervoor dat embedding-dimensies voldoende variantie behouden over de batch om informatieverlies door collapse te voorkomen.
  4. Covariance Loss ($L_{cov}$): Moedigt decorrelatie aan tussen embedding-dimensies om redundante representaties te voorkomen (gebaseerd op VICReg).

Het framework is ontworpen om compatibel te zijn met diverse architecturen (Transformers, State Space Models zoals HyenaDNA) en tokenisatiemethoden, waarbij enkel een aggregatie-operator nodig is om sequenties in latente representaties te comprimeren.

Belangrijkste Bijdragen

1. Nieuw Trainingsframework: Introductie van een framework voor continue training dat de leerdoelstelling verschuift van letterlijke tokenreconstructie naar latente kenmerk-afstemming via JEPA.
2. Architectuur-agnosticisme: Demonstratie dat JEPA-DNA fungeert als een universeel mechanisme voor "representational uplift", compatibel met leidende GFMs (DNABERT-2, NTv3, HyenaDNA) en zowel MLM als NTP doelstellingen.
3. State-of-the-Art Prestaties: Vaststelling van nieuwe prestatieplafonds over gesuperviseerde linear probing taken door generatieve precisie te combineren met latente semantische gronding.
4. Synergetische Analyse: Karakterisering van de interactie tussen generatieve en latente doelstellingen, waarbij wordt aangetoond dat een hybride aanpak beter presteert dan pure reconstructie of pure contrastieve methoden.

Experimentele Resultaten

Het framework werd geëvalueerd op 17 diverse genomische benchmark-taken met behulp van linear probing en zero-shot protocollen.

• Linear Probing: JEPA-DNA vertoonde consistente winst over alle drie de geteste backbone-architecturen.
  • DNABERT-2: Verbeterde de prestaties in 6–8 van de 9 taken, met AUROC-winsten tot +8,0%.
  • NTv3: Verbeterde de prestaties in 8 van de 9 taken.
  • HyenaDNA: Verbeterde de prestaties in 7 van de 9 taken, met opvallende winsten in promotorherkenning (+11,22%) en splice site classificatie (+10,28%).
  • Taakspecifiteit: De winsten waren het meest uitgesproken in taken die steunen op bredere regulatoire context (promotoren, splice sites, meQTL, causale eQTL), terwijl taken die gedreven worden door fijnmazige lokale veranderingen (coding pathogenicity) meer variabele resultaten lieten zien.
• Zero-Shot Evaluatie: Gebruikmakend van DNABERT-2 verbeterde JEPA-DNA de zero-shot prestaties in 6 van de 8 klinische en fitness-scoring benchmarks, inclusclusief een +12,06% AUROC-winst op ziektevariantvoorspelling.
• Ablatieonderzoeken:
  • VICReg Regularisatie: Geïdentificeerd als essentieel om de voordelen van het hybride paradigma te ontsluiten.
  • Predictor Architectuur: Attention-gebaseerde predictors (4 lagen) presteerden beter dan MLPs en diepere architecturen.
  • Loss Samenstelling: Afstemming in de latente ruimte leverde een informatiever signaal op (+3,86% winst) dan tokenreconstructie alleen (+2,06%), maar de optimale configuratie vereiste een hybride van beide.

Betekenis

Het artikel stelt dat JEPA-DNA een nieuwe state-of-the-art vestigt voor Genomische Fundamentele Modellen door ze succesvol te funderen in biologische semantiek. Door het leersignaal te verschuiven van tokenherstel naar semantische afstemming, dwingt het framework modellen om abstracte, functionele kenmerken te leren die invariant zijn voor laag-niveau sequentieruis. Deze aanpak overbrugt de kloof tussen genomische syntaxis en biologische logica, en biedt een schaalbaar pad naar het funderen van foundation modellen zonder de computationele overhead van multi-sequence alignment of de datavereisten van contrastieve negative mining. De auteurs positioneren JEPA-DNA niet als een vervanging voor bestaande GFMs, maar als een veelzijdig fase van continue training die de kwaliteit van de representatie van elk bestaand genomisch fundamenteel model verbetert.