Source-linked AI summary
JEPA-DNA: Grounding Genomic Foundation Models through Joint-Embedding Predictive Architectures
Ariel Larey, Elay Dahan, Amit Bleiweiss, Raizy Kellerman, Guy Leib, Omri Nayshool, Dan Ofer, Tal Zinger, Dan Dominissini, Gideon Rechavi, Nicole Bussola, Simon Lee, Shane O'Connell, Dung Hoang, Marissa Wirth, Alexander W. Charney, Nati Daniel, Yoli Shavit
TL;DR
Genomic generative objectives capture local syntax but emphasize token reconstruction over high-level functional context. JEPA-DNA adds latent prediction to generative training; across 17 benchmark tasks, it reports consistent linear-probing and zero-shot gains and a new state of the art for supervised linear probing.
Problem
Generative objectives can identify local motifs and sequence patterns but focus on token reconstruction rather than high-level functional context.
Source: §1 · p. 2Method
JEPA-DNA is a continual-training framework that combines token-level recovery with prediction of masked-span latent representations to encourage abstract functional features.
Source: §1 · p. 2Results
Across 17 genomic benchmark tasks, JEPA-DNA reports consistent linear-probing and zero-shot gains and a new state of the art for supervised linear probing.
Source: §1 · p. 2Takeaways & Limitations
JEPA-DNA is reported to capture high-level biological context more effectively than generative objectives alone while boosting backbone performance across supervised and zero-shot tasks.
Source: §5 · p. 9The end-to-end comparison covers only one backbone and a subset of tasks, and broader full-fine-tuning and parameter-efficient adaptation studies are deferred.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Genomic Foundation Models (GFMs) typically rely on Masked Language Modeling (MLM) or Next-Token Prediction (NTP) to learn the "Laws of Nature". While effective at capturing local syntax, these generative paradigms prioritize token-level reconstruction over high-level functional context. We introduce JEPA-DNA, a model-agnostic continual training framework that integrates a Joint-Embedding Predictive Architecture (JEPA) with traditional generative objectives. By supervising global sequence embeddings in a latent space, JEPA-DNA forces models to predict the functional representations of masked genomic segments, shifting the learning signal from token recovery to semantic alignment. We evaluate JEPA-DNA on 17 diverse genomic benchmark tasks, demonstrating consistent gains in linear probing and zero-shot performance regardless of the underlying GFM architecture or generative objective, and establishing a new state-of-the-art for GFMs under the linear-probing protocol. Through extensive ablation studies, we further characterize the synergistic interplay between generative and latent objectives. Our code is provided in the supplementary material.
1 Introduction
Genomic generative objectives capture local sequence patterns but emphasize token reconstruction over high-level functional context. JEPA-DNA adds latent prediction to generative training and is evaluated across 17 tasks, with reported gains in linear probing and zero-shot performance.
1 Introduction
- The framework targets the gap between local motif and sequence-pattern recognition and the high-level regulatory logic and long-range interactions governing genomic function.The motivation draws on the view that input-space reconstruction can prioritize low-level details over abstract semantic structure [7].Source: §1 · p. 2
- JEPA-DNA combines token recovery with latent-space prediction to encourage abstract, functionally meaningful features that are less sensitive to low-level sequence noise.Its predictor learns to recover the latent representation of a masked span from surrounding context.Source: §1 · p. 2
- Across 17 benchmark tasks, JEPA-DNA reports consistent linear-probing and zero-shot gains and a new state of the art for supervised linear probing.Ablations attribute gains to core framework components rather than simply additional training.Source: §1 · p. 2
- JEPA-DNA is presented as a continual-training framework that shifts learning from literal token reconstruction toward latent feature alignment.The evaluation spans DNABERT-2, NTv3, and HyenaDNA [4], tokenization methods, and MLM and NTP objectives.Source: §1 · p. 2
2 Related Work
Related work describes generative genomic models and alternatives to token reconstruction, then situates JEPA among joint-embedding approaches. The section contrasts these approaches’ scope and trade-offs with the multiscale sequence demands of genomic modeling.
2.1 Pre-training Paradigms of Genomic Foundation Models
- Genomic foundation models span Transformer, long-convolution, state-space, and U-Net-style architectures, yet nearly all rely on MLM or NTP token reconstruction.These models target longer contexts and DNA syntax, while prior work notes limits in capturing functional state.Source: §2.1 · p. 2
- MSA-based GPN-MSA [15] uses cross-species conservation, while DNASimCLR uses augmented sequence views; the former requires costly preprocessing and alignments, and the latter may require augmentations or negative mining.JEPA-DNA is described as capturing functional invariants directly from raw sequence.Source: §2.1 · p. 3
2.2 Joint-Embedding Predictive Architecture (JEPA)
- JEPA predicts masked-segment representations rather than literal input, and prior work extends the paradigm to language [17] and single-cell transcriptomics.Unlike gene-expression-vector models, genomic JEPA must handle sequential motifs and long-range dependencies.Source: §2.2 · p. 3
3 Method
JEPA-DNA augments a genomic foundation model with a joint-embedding branch for learning local syntax and global semantics. Its architecture pairs context-derived predictions with target latents and aggregates sequence embeddings into global representations.
3 Method
- JEPA-DNA augments a standard genomic foundation model with a joint-embedding branch, enabling multi-objective learning over local sequence syntax and global semantics.The framework treats the genome as a structured signal rather than only a token sequence.Source: §3 · p. 3
3.1 Architecture Components
- The context encoder processes masked input, an EMA-updated target encoder embeds the unmasked target, and a predictor maps context latents toward target latents.A backbone-native aggregator, such as a [CLS] token, last token, or mean pooling, produces global target and predicted representations; auxiliary causal and variance encoders address autoregressive backbones and mode collapse.Source: §3.1 · p. 3
3.2 Masking and Re-masking Strategy
The masking strategy hides contiguous target spans from the context encoder and predictor. A schedule increases masking difficulty during training, while re-masking prevents the predictor from directly using encoded content at masked positions.
3.2 Masking and Re-masking Strategy
- Training masks multiple contiguous spans, typically exceeding a 20% aggregate ratio, and gradually increases both the masking ratio and span count.The schedule progresses from mild corruption to more aggressive masking; the same configuration applies to JEPA and MLM losses.Source: §3.2 · p. 4
- At masked positions, context latents are replaced by a learnable latent mask embedding before prediction, preventing a trivial mapping to target representations.The predictor receives the full context sequence, allowing it to use localized information and a global sequence summary.Source: §3.2 · p. 4
3.3 Multi-Objective Continual training
JEPA-DNA continues pre-trained GFM training with a composite objective that retains token-level language modeling while adding latent prediction and embedding-diversity constraints.
3.3.1 LLM Loss (Lllm)
- The LLM loss retains standard MLM or NTP token prediction, using masked or causally selected token positions.For NTP, it is computed with an additional causal forward pass through an encoder sharing weights with the context encoder.Source: §3.3.1 · p. 4
3.3.2 Latent Predictive Loss (Ljepa)
- The JEPA loss aligns the predictor’s global output with the target encoder’s embedding using cosine similarity.Minimization encourages the predicted context representation to align with the target encoder’s functional embedding.Source: §3.3.2 · p. 5
3.3.3 Variance and Covariance Regularization
- Variance regularization maintains sufficient variation in each embedding dimension across the batch, with γ typically set to 1.It is applied to prevent informational collapse.Source: §3.3.3 · p. 5
- Covariance regularization penalizes squared off-diagonal covariance values to encourage embedding dimensions to encode independent information.The covariance is computed from batch-centered embeddings.Source: §3.3.3 · p. 5
3.3.4 Total Objective
- The final objective balances language-modeling, latent-prediction, variance, and covariance losses using hyperparameter weights.The framework begins with a frozen-backbone predictor-training stage, then jointly updates the encoder and predictor.Source: §3.3 · p. 4, §3.3.4 · p. 5
3.4 Compatibility across Model Architectures and Objectives
JEPA-DNA is compatible with MLM and NTP across model architectures, provided the sequence can be aggregated into a latent representation. It uses architecture-appropriate sequence embeddings to supervise that global representation.
3.4 Compatibility across Model Architectures and Objectives
- JEPA-DNA works with MLM and NTP objectives across Transformer, state-space, and long-convolution backbones by requiring an operator that aggregates sequences into latent representations.The framework uses the last [SEP] embedding for Transformer decoders, SSMs, and long-convolution models, or mean pooling for Nucleotide Transformer-style models.Source: §3.4 · pp. 5–6
4 Experimental Results
JEPA-DNA improves linear-probing performance across three backbone families and also improves zero-shot results for DNABERT-2. Ablations indicate that gains depend on the framework’s components, while additional experiments compare it with contrastive learning and examine design choices.
4.1 Experimental Setup
- The setup tests three backbones on data subsets from their original pre-training corpora and evaluates nine supervised and eight zero-shot tasks.Linear probing trains a single-layer classifier on a frozen backbone; zero-shot evaluation ranks variants by cosine similarity between reference and variant embeddings.Source: §4.1 · p. 6, §4.2 · p. 7, §4.3 · p. 8
4.2 Linear Evaluation
- JEPA-DNA reaches state-of-the-art linear-probing performance on 7 of 9 tasks and improves 6–8 of 9 tasks for each evaluated backbone.The reported AUROC gains over original checkpoints reach +7.96% for DNABERT-2, +6.77% for NTv3, and +11.22% for HyenaDNA.Source: Table 1 · p. 7, §4.2 · pp. 7–8
- JEPA-DNA beats step-matched and compute-matched continual-training baselines on 22 of 27 supervised task–backbone comparisons, with mean relative gains of +2.52% and +2.46%.Native-objective continuation improves only 15 of 27 pairs on average, and its gains are small and inconsistent.Source: §4.2 · p. 7
4.3 Zero-Shot Evaluation
- DNABERT-2 improves on 6 of 8 zero-shot tasks, with a mean AUROC gain of +4.1% over its original checkpoint.Continual training alone provides little to no benefit; JEPA-DNA beats matched baselines on 17 of 24 task–backbone pairs.Source: §4.3 · p. 8, Table 2 · p. 8
4.4 Ablations and Additional Experiments
- In DNABERT-2 ablations, VICReg regularization is essential, and peak performance combines regularized latent prediction with Scheduled Masking.The factorial grid compares step-matched and compute-matched budgets and includes single-component elimination controls.Source: §4.4 · p. 9, Table 3 · p. 9
- Latent-space alignment provides a stronger training signal than token reconstruction, but optimal results retain both MLM tokens and latent sequence alignment.The reported gains are +3.86% for latent-space alignment and +2.06% for token reconstruction.Source: §4.4 · p. 9
- An attention-based predictor outperforms the MLP variant, performance peaks at four layers, and SGDM supports training stability where Adam or AdamW often degrade performance.The target encoder is also used for inference.Source: §4.4 · p. 9
- JEPA-DNA exceeds DNASimCLR and three matched contrastive variants on all evaluated tasks, with +7.57% versus +4.06% for the strongest variant.Additional experiments cover embedding analysis, full fine-tuning, sequence-length stability, and loss-weight sensitivity.Source: §4.4 · p. 9
5 Conclusion
The paper concludes that adding latent prediction to token reconstruction improves biological-context representations and boosts performance across supervised and zero-shot tasks. It notes that full fine-tuning and training from scratch remain areas for future study.
5 Conclusion
- JEPA-DNA combines latent-representation prediction with token reconstruction and consistently boosts backbone performance across supervised and zero-shot tasks.The authors conclude it captures high-level biological context more effectively than generative objectives alone.Source: §5 · p. 9
- Full fine-tuning is only partly addressed: the end-to-end comparison covers one backbone and a subset of tasks, while broader fine-tuning and parameter-efficient adaptation are deferred.The study focuses on continual pre-training of existing state-of-the-art backbones because large-scale genomic training is computationally demanding.Source: §5 · p. 9
- Training JEPA from scratch on massive, multi-species datasets remains future work, including testing whether it accelerates the emergence of biological priors.The authors also identify LeJEPA as a possible route to better scaling for extremely long-range dependencies.Source: §5 · p. 10
A Appendix
The appendix supplements the paper’s findings and methods with cost estimates, extended experiments, and reproducibility details.
A Appendix
- The appendix provides supplementary information supporting the paper’s findings and methodology.Source: A Appendix · p. 13
- Section A.1 quantifies JEPA-DNA’s training and inference cost.Source: A Appendix · p. 13
- Section A.2 adds robustness results, performance breakdowns, ablations, baseline comparisons, a fine-tuning study, and sensitivity analyses.Source: A Appendix · p. 13
- Section A.3 documents the experimental setup, datasets, evaluation protocols, and implementation and training details for reproducibility.Source: A Appendix · p. 13
A.1 Computational Cost
JEPA-DNA adds computation during continual pre-training, but the reported inference runtime and memory match the original backbone.
A.1 Computational Cost
- Pre-training takes 1.09× to 7.88× the time per step of matched native continual training, depending on the backbone.The added cost comes from forward passes through the EMA target encoder and predictor and is confined to pre-training.Source: §A.1 · p. 13, Table 4 · p. 13
- At inference, runtime and memory are identical to the original backbone because only the EMA target encoder is retained.The predictor and VICReg heads are discarded.Source: §A.1 · p. 13
A.2 Extended Results
Extended results show JEPA-DNA establishes a linear-probing performance ceiling across most tasks, with gains also examined across seeds, zero-shot settings, ablations, and additional evaluations. These analyses identify complementary roles for latent and token-level objectives and probe robustness under varied training and evaluation conditions.
Robustness Analysis
- Across five seeds, JEPA-DNA improves 6 of 9 DNABERT-2 tasks, 8 of 9 NTv3 tasks, and 6 of 9 HyenaDNA tasks, with some task-specific regressions.The reported regressions include locally driven motif tasks for HyenaDNA; the NTv3 sQTL change is within one standard deviation.Source: §A.2 · p. 13, Table 6 · p. 15
Benchmarking Against Leading GFMs
- JEPA-DNA outperforms the best baseline GFM on a majority of linear-probing tasks, despite the baseline comparison taking the maximum across seven models.Table 8 compares the best original baseline with the best JEPA-DNA extension across tasks.Source: §A.2 · p. 14, Table 8 · p. 16
Zero-Shot Performance for Additional Backbones
- Zero-shot gains are consistent across DNABERT-2, NTv3, and HyenaDNA, spanning different architectures, pre-training objectives, and datasets.The supplied zero-shot table caption identifies variant-effect performance as AUROC and describes comparisons with original and continual-training baselines.Source: §A.2 · p. 14, Table 9 · p. 16, Table 10 · p. 17
Extended Ablation Studies
- Ablations show latent prediction and token-level MLM are complementary: latent prediction without token loss reaches +3.85%, while the combined configuration reaches +6.22%.VICReg is an enabling condition: without it, MLM+JEPA falls below MLM alone (+1.63% versus +2.16%); applying JEPA at token level instead of retaining MLM also reduces the gain (+3.86% versus +6.22%).Source: §A.2 · pp. 14–16, Table 11 · p. 17
Predictor and Optimizer Ablations
- An attention predictor outperforms the MLP variant and peaks at four layers, while SGDM achieves +6.22% versus −12.24% for Adam and −4.45% for AdamW.AdaBelief yields a marginal +1.22% gain in the optimizer comparison.Source: §A.2 · p. 17, Table 12 · p. 18, Table 13 · p. 18
DNASimCLR Comparison
- JEPA-DNA with DNABERT-2 exceeds DNASimCLR across evaluated tasks, with a 7.57% average AUROC improvement; controlled contrastive variants remain below it.Using the same pretrained DNABERT-2 backbone with contrastive loss recovers +4.06%, while adding MLM to contrastive learning yields −0.34%.Source: §A.2 · p. 18, Table 14 · p. 19
Embedding Analysis and Full Fine-Tuning
- On chromosome 22 embeddings, a k-nearest-neighbor classifier reaches weighted F1 0.903 for the target encoder and 0.877 for the context encoder; full fine-tuning retains an advantage on four DNABERT-2 tasks.On saturated GUE benchmarks, the fine-tuning margin is under 0.005 absolute AUROC; on LRB causal eQTL it is +8.5% relative.Source: §A.2 · pp. 18–19, Table 15 · p. 20
Sequence-Length and Loss-Weight Sensitivity
- Beyond DNABERT-2’s 512-token training context, JEPA-DNA leads on Enhancer–Target at every tested length and degrades 3.7× more slowly on BRCA1; ±50% loss-weight perturbations remain above both baselines.Across eight loss-weight perturbations on GUE Splice Site, the AUROC spread is 0.021.Source: §A.2 · pp. 19–20, Table 16 · p. 20
A.3 Extended Implementation Details
The implementation details describe the tested backbones, evaluation datasets and protocols, and a shared continual-training setup. JEPA-DNA uses masked-span prediction with latent and native generative losses, alongside predictor, variance, and covariance components.
GFM Backbones Details
- The evaluated backbones use distinct sequence representations: DNABERT-2 uses BPE and [CLS], NTv3 uses single-nucleotide tokens and mean pooling, and HyenaDNA uses the last [SEP] embedding.DNABERT-2 supports up to 512 tokens; NTv3 is a roughly 100M-parameter hybrid, while HyenaDNA is an approximately 600K-parameter autoregressive model.Source: §A.3 · p. 20
Continual Training Datasets
- Continual-training data follow a data-parity regime, reusing each backbone’s pre-training genomic sources rather than adding genomic information.The corpora comprise multi-species genome windows for DNABERT-2, Open Genome 2 records for NTv3, and hg38 windows for HyenaDNA.Source: §A.3 · p. 21
Evaluation Protocol
- Evaluation combines frozen linear probing with zero-shot variant-effect ranking based on reference–variant embedding similarity or angular distance.Linear probing trains only a classifier; zero-shot scores are evaluated with AUROC and AUPRC.Source: §A.3 · p. 22
Training Schedule and Masking
- Training progressively increases masked contiguous spans from 1–2 regions covering 10–30% of each sequence to 1–4 regions covering 30–50%.The context encoder processes the masked sequence, while the EMA-updated target encoder processes the full unmasked sequence.Source: §A.3 · p. 22
Training Schedule and Optimization
- Training warms up the predictor with the encoder frozen for 1,000 steps, then jointly updates them with a learning-rate warmup and decay schedule.NTv3 uses lower phase-two learning rates than DNABERT-2 and HyenaDNA; the target encoder tracks the context encoder through EMA updates.Source: §A.3 · pp. 22–23
Objective and Predictor
- The objective combines sequence-level cosine latent prediction with MLM or NTP, variance regularization, and covariance loss; the predictor is lightweight and backbone-dependent.The latent, generative, variance, and covariance loss weights are 1.0, 0.1, 25.0, and 0.5, respectively; Transformer backbones use a four-layer predictor, while mean-pooled representations use an MLP.Source: §A.3 · pp. 20, 23