JEPA-DNA: Genomic Pre-training Framework

• JEPA-DNA is a genomic pre-training framework that augments traditional masked language modeling with a latent prediction task to overcome token-level limitations.
• It employs dual encoders and a predictor head to fuse local nucleotide recovery with global functional embeddings, promoting biologically robust representations.
• Empirical evaluations report consistent AUROC improvements across multiple genomics benchmarks, indicating stronger performance on tasks requiring functional understanding.

Searching arXiv for the primary JEPA-DNA paper and closely related JEPA sequence-modeling work to ground the article with current citations. JEPA-DNA is a pre-training framework for genomic foundation models that combines a standard language-model style objective on DNA—masked language modeling (MLM) or next-token prediction (NTP)—with a Joint-Embedding Predictive Architecture (JEPA) objective that predicts latent representations of masked sequence segments rather than only raw nucleotides. Its central claim is that token-level recovery captures local genomic syntax and fine-grained motif patterns but often fails to capture broader functional context, whereas latent grounding via a supervised global token forces the model to predict high-level functional embeddings of masked genomic segments and thereby promotes a more robust, biologically grounded representation space (Larey et al., 19 Feb 2026).

1. Definition and conceptual rationale

JEPA-DNA was introduced to address what the authors describe as a “granularity trap” in genomic foundation models. In this view, MLM and NTP provide supervision purely in token space, rewarding the recovery of exact nucleotides and short-range sequence statistics, including biologically unimportant or noisy patterns such as repetitive elements or neutral polymorphisms. JEPA-DNA augments this regime with a latent-space prediction task so that the model is also required to predict a sequence-level embedding of masked content, not just its token identities (Larey et al., 19 Feb 2026).

The framework is explicitly model-agnostic. It is presented as a training recipe that can be deployed either as a standalone from-scratch objective or as a continual pre-training enhancement for existing genomic foundation models. In the main instantiation, JEPA-DNA is added to DNABERT-2, but the paper states that the method extends both MLM and NTP paradigms and only requires a global aggregation token such as a prepended $[CLS]$ or an appended $[EOS]$ (Larey et al., 19 Feb 2026).

The operative concept is “latent grounding.” The masked context is processed by a context encoder, while the full unmasked sequence is processed by a target encoder. The model is then trained so that a predictor, conditioned on the masked-context representation, matches the target encoder’s global embedding. A plausible implication is that the representation is encouraged to preserve functional properties that remain predictable from context while deemphasizing lower-level details that are less relevant to biological function (Larey et al., 19 Feb 2026).

2. Objective, architecture, and latent grounding

JEPA-DNA adds three components to a standard sequence backbone: a context encoder $\mathcal{E}_\theta$, a target encoder $\mathcal{E}_{\bar{\theta}}$, and a predictor head $\mathcal{P}_\phi$. The context encoder receives a masked sequence and produces token embeddings including a context $[CLS]$ embedding. The target encoder is a structural duplicate of the context encoder, updated by exponential moving average, and receives the unmasked sequence to produce a target $[CLS]$ embedding. The predictor consumes the context-side representation and outputs a predicted target $[CLS]$ embedding (Larey et al., 19 Feb 2026).

The latent target is defined as

$$\mathbf{z}_{target} = \mathcal{E}_{\bar{\theta}}(\mathbf{x})_{cls},$$

while the predictor output is

$$\hat{\mathbf{z}} = \mathcal{P}_\phi(\mathbf{h}_{cls}; \mathbf{h}_{1:L}).$$

The JEPA loss is a cosine-alignment objective: $[EOS]$0 In parallel, the generative branch preserves the usual token-level objective. For MLM, if $[EOS]$1 denotes masked indices, the loss is

$[EOS]$2

The total training objective is

$[EOS]$3

where $[EOS]$4 and $[EOS]$5 are VICReg-style variance and covariance regularizers used to prevent representation collapse and to encourage diverse, decorrelated latent features (Larey et al., 19 Feb 2026).

The predictor is a 3-layer Transformer encoder of dimension 384 with 3 heads, pre-norm, GELU activations, and frozen sinusoidal positional embeddings. Its $[EOS]$6 output is projected back to the 768-dimensional target space. In the DNABERT-2 implementation, the backbone is a 12-layer Transformer encoder with hidden dimension 768, 12 attention heads, and approximately 117M parameters, using the DNABERT-2 BPE tokenizer and a maximum pre-training length of 512 tokens (Larey et al., 19 Feb 2026).

A notable technical detail is “re-masking” in latent space. After the context encoder processes the masked input, hidden states corresponding to masked positions are replaced by a learnable mask embedding before entering the predictor. This prevents the predictor from using leaked information from masked positions and makes the JEPA branch depend on visible context rather than trivial identity recovery (Larey et al., 19 Feb 2026).

3. Training regime and continual pre-training configuration

In the reported experiments, JEPA-DNA is used primarily as a continual pre-training procedure rather than from-scratch training. The setup begins from a pre-trained DNABERT-2 checkpoint. Both the context and target encoders are initialized from this checkpoint, while the predictor is initialized from scratch. Pre-training is then resumed with the combined MLM+JEPA objective for five epochs on a subset of the original DNABERT-2 corpus (Larey et al., 19 Feb 2026).

The pre-training corpus consists of the human genome GRCh38/hg38 together with five model organisms: mouse, zebrafish, fruit fly, C. elegans, and Arabidopsis. Sequences are filtered to valid nucleotides, chunked into fixed-length segments with 50% overlap, and yield approximately 4.76M training sequences covering approximately 7.6B base pairs (Larey et al., 19 Feb 2026).

Masking is span-based rather than purely token-independent. JEPA-DNA samples 1–3 contiguous spans per sequence, covering 20–40% of the sequence length, and uses the same masking pattern for the MLM loss and the JEPA branch. This makes the prediction problem more contextual and more strongly tied to long-range sequence information than standard low-rate random masking (Larey et al., 19 Feb 2026).

Optimization proceeds in two phases. First, the predictor is warmed up for 1000 steps with the encoder frozen and a learning rate of $[EOS]$7. Then full training begins with the encoder unfrozen, using SGD with momentum 0.9, weight decay 0.01, batch size 32, gradient accumulation 4 for an effective batch size of 128, a linear warmup from $[EOS]$8 to $[EOS]$9, and cosine decay to $\mathcal{E}_\theta$0. The target encoder EMA momentum is increased from 0.996 toward 1.0 across training. The implementation is in PyTorch, uses FlashAttention, and is run on two NVIDIA GPUs with DataParallel (Larey et al., 19 Feb 2026).

The variance and covariance regularizers are computed with additional care. The paper states that variance is measured with an extra forward pass in evaluation mode and with uniform sequence length within a batch, to avoid artificially inflated variance from dropout or padding effects. This detail indicates that the JEPA branch is treated as a representation-learning signal whose geometry must be explicitly stabilized, not merely as an auxiliary reconstruction loss (Larey et al., 19 Feb 2026).

4. Empirical performance across supervised and zero-shot genomics benchmarks

JEPA-DNA is evaluated in two regimes: frozen linear probing and zero-shot variant effect prediction. In linear probing, the encoder is frozen and a linear classifier is trained on top of the $\mathcal{E}_\theta$1 embedding. For variant tasks, the classifier uses the concatenation of reference and variant $\mathcal{E}_\theta$2 embeddings. In zero-shot evaluation, cosine similarity between reference and mutant sequence embeddings is used directly as a surrogate for functional effect (Larey et al., 19 Feb 2026).

On GUE and VariantBenchmarks tasks, JEPA-DNA generally improves over MLM-only DNABERT-2. The reported AUROC gains include $\mathcal{E}_\theta$3 on GUE TF Binding, $\mathcal{E}_\theta$4 on GUE Promoter, $\mathcal{E}_\theta$5 on GUE Splice Site, $\mathcal{E}_\theta$6 on VariantBenchmarks Coding Pathogenicity, $\mathcal{E}_\theta$7 on VariantBenchmarks Non-coding Pathogenicity, $\mathcal{E}_\theta$8 on VariantBenchmarks Expression Effect, $\mathcal{E}_\theta$9 on VariantBenchmarks meQTL, and $\mathcal{E}_{\bar{\theta}}$0 on LRB Causal eQTL. The only reported AUROC decrease in this table is $\mathcal{E}_{\bar{\theta}}$1 on VariantBenchmarks sQTL (Larey et al., 19 Feb 2026).

In zero-shot tasks, the strongest gains appear on BEND Expression Effect, TraitGym Mendelian, and ClinVar. AUROC improves from 0.490 to 0.524 on BEND Expression Effect, from 0.507 to 0.544 on TraitGym Mendelian, and from 0.528 to 0.544 on Songlab ClinVar. TraitGym Complex and LRB Pathogenic OMIM do not improve in the reported setup, which the paper associates with long-context difficulty under a 512-token limit (Larey et al., 19 Feb 2026).

The benchmark coverage is broad. Supervised tasks include GUE Promoter, GUE TF Binding, GUE Splice Site, multiple VariantBenchmarks tasks, and LRB Causal eQTL. Zero-shot tasks include BEND Expression Effect, BEND Disease Variant, TraitGym Complex, TraitGym Mendelian, Songlab ClinVar, and LRB Pathogenic OMIM (Larey et al., 19 Feb 2026).

The empirical pattern supports the claim that JEPA-DNA is especially useful for tasks where broader regulatory or functional semantics matter. This suggests that the added latent objective most strongly benefits sequence-level abstractions such as pathogenicity, expression effect, and Mendelian variant interpretation, while offering smaller gains when the downstream signal is less aligned with the learned global embedding (Larey et al., 19 Feb 2026).

5. Relation to broader JEPA sequence modeling

JEPA-DNA belongs to a broader family of sequence models that combine token-level modeling with latent-space prediction. In proteins, ProteinJEPA found that the strongest recipe was not all-position latent prediction but “masked-position MLM+JEPA,” in which the latent loss is applied only at masked positions while retaining MLM cross-entropy. Under matched wall-clock budgets, that recipe yields more downstream wins than losses against MLM-only continuation for pretrained ESM2-35M and ESM2-150M, whereas JEPA-only collapses in nearly every experiment (Ofer et al., 8 May 2026). A plausible implication is that JEPA-DNA’s decision to complement rather than replace generative supervision is consistent with a broader pattern across biological sequence modeling.

In single-cell transcriptomics, Cell-JEPA makes a related argument in a different modality: reconstruction objectives on highly sparse counts encourage models to encode measurement artifacts, whereas predicting cell-level embeddings from partial observations forces the model to learn dropout-robust features. Cell-JEPA reports 0.72 AvgBIO in zero-shot transfer versus 0.53 for scGPT and frames latent prediction as a way to learn a stable cell manifold rather than to reconstruct every noisy input component (ElSheikh et al., 2 Feb 2026). JEPA-DNA’s emphasis on latent grounding rather than pure nucleotide recovery is aligned with that same shift from raw reconstruction toward predictive representation learning.

Outside biology, LLM-JEPA extends JEPA to language by pairing the usual next-token objective with an embedding-space regression term that predicts one view from another, such as text from code, and reports improvements across NL-RX, GSM8K, Spider, and other datasets (Huang et al., 11 Sep 2025). That work argues that JEPA acts as a regularizer rather than a replacement for generative modeling, echoing the biological-sequence evidence from ProteinJEPA and the hybrid objective used in JEPA-DNA. Taken together, these results position JEPA-DNA as the genomic instance of a wider design pattern: retain the generative loss needed for local syntax and exact token recovery, but add a latent objective to shape the geometry of higher-level semantics (Huang et al., 11 Sep 2025).

6. Limitations, open questions, and future directions

The JEPA-DNA paper is explicit about several limitations. First, the experimental evidence is restricted to continual pre-training of DNABERT-2; from-scratch JEPA-DNA training is proposed but not reported. Second, the implementation is demonstrated only on a Transformer encoder, although the authors state that the framework should extend to state-space models, long-convolution models, and decoder-style NTP architectures (Larey et al., 19 Feb 2026).

Sequence length is another clear constraint. The backbone uses a 512-token context window, and longer benchmark inputs are centrally truncated. This likely limits performance on long-range tasks such as TraitGym Complex and LRB Pathogenic OMIM, where the reported zero-shot results do not improve over the baseline (Larey et al., 19 Feb 2026).

The paper also notes that biologically informed masking remains largely unexplored. The reported system uses generic span masking, but future versions are proposed to incorporate multi-block or biology-aware masking patterns that more directly probe regulatory logic and long-range dependencies. The authors likewise mention next-sequence prediction-style auxiliary tasks, alternative global aggregation strategies beyond $\mathcal{E}_{\bar{\theta}}$2/$\mathcal{E}_{\bar{\theta}}$3, and broader comparisons with contrastive self-supervision and MSA-based models as open directions (Larey et al., 19 Feb 2026).

A final conceptual limitation is that the “functional embeddings” are self-supervised rather than externally grounded in assay labels. The target $\mathcal{E}_{\bar{\theta}}$4 is produced by an EMA copy of the same model, not by epigenomic or perturbational supervision. The reported downstream gains indicate that these embeddings are biologically useful, but they do not identify a unique biological ontology. This suggests that JEPA-DNA should be understood not as supervised functional annotation, but as a way to bias genomic foundation models toward a representation space in which functional distinctions become more linearly accessible and more stable under sequence perturbation (Larey et al., 19 Feb 2026).