jbilcke/jepa-dna-mlx

JEPA-DNA MLX

Three genomic foundation models, packaged together for local inference on Apple silicon.

This repository provides float32 Safetensors conversions of NVIDIA's JEPA-DNA DNABERT-2, HyenaDNA and NTv3 target-encoder checkpoints. It also includes variable-length Core AI model assets for DNABERT-2 and NTv3.

JEPA-DNA adds a joint-embedding prediction objective to DNA language-model pretraining. These checkpoints produce contextual representations of genomic sequences for research in feature extraction, embedding comparison, linear probing and sequence-perturbation analysis. See the JEPA-DNA paper for the training method and original biological evaluations.

This is a community conversion of the NVIDIA releases. The converted tensors preserve their original values: no additional training, quantization or reduced-precision conversion was performed.

Models at a glance

Variant	Backbone	Approx. source size	Sequence embedding	Float32 weights	Validated input window
DNABERT-2	Transformer with ALiBi and DNA BPE	117M parameters	768 values	458.62 MiB	512 tokens, including CLS/SEP
HyenaDNA	Long convolution with learned filters	0.6M parameters	128 values	2.43 MiB	8192 tokens, including SEP
NTv3	Convolutional U-Net and Transformer	100M parameters	768 values	406.25 MiB	8192 nucleotides

Serialized sizes include stored buffers and any duplicated tied tensors; they are not peak-memory estimates or exact unique-parameter counts. Each variant contains a full tensor inventory in manifest.json.

Download

Download the complete collection with the Hugging Face CLI:

hf download jbilcke/jepa-dna-mlx --local-dir ./jepa-dna-mlx

To download only one variant's Safetensors and supporting metadata:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="jbilcke/jepa-dna-mlx",
    local_dir="./jepa-dna-mlx",
    allow_patterns=["dnabert2/*", "README.md", "LICENSE.pdf", "NOTICE.md", "licenses/*"],
    ignore_patterns=["dnabert2/coreai/*"],
)

Replace dnabert2 with hyenadna or ntv3 as needed. For reproducible experiments, pass a specific Hub commit as revision.

Weight format and compatibility

Each model.safetensors contains the target encoder and its language-model head in float32, retaining the source tensor names and PyTorch storage layouts. MLX can read the arrays directly:

import mlx.core as mx

weights = mx.load("./jepa-dna-mlx/dnabert2/model.safetensors")
print(len(weights))  # 142 stored tensors

Array loading is not a complete forward pass. Inference requires an implementation of the corresponding DNA architecture, including the correct attention/filter operations, convolution layout conversions, tokenizer and pooling. These are genomic encoders and are not supported by a generic chat-model or mlx_lm loader. The configurations retain their upstream architecture metadata; this multi-model repository is not a root-level Transformers AutoModel package.

Outputs consist of contextual hidden states, vocabulary logits and sequence representations. The checkpoint heads do not directly produce supervised task labels or calibrated pathogenicity scores.

Input and pooling conventions

Variant	Tokenization	Padding	Sequence representation
DNABERT-2	Case-preserving DNA BPE, CLS then sequence then SEP	Right, to the batch maximum	CLS hidden state
HyenaDNA	Uppercase single bases, followed by SEP	Left, to the batch maximum	Last valid hidden state
NTv3	Uppercase single bases, no specials for ordinary encoding	Right, to the next multiple of 128	Mask-weighted mean of full-resolution hidden states

DNABERT-2 token positions can cover multiple nucleotides. Preserve its BPE merge ranks and case behavior; lower-case bases are not implicitly uppercased. HyenaDNA and NTv3 use their own single-base vocabularies, so their token IDs are not interchangeable.

The validated preprocessing truncates DNABERT-2 and HyenaDNA to the windows above, including their special tokens, and rejects NTv3 inputs beyond 8192 bases. Those are the evaluated sequence policies; a larger value in an original backbone configuration does not establish numerical validation at that larger context.

Padding affects HyenaDNA and NTv3 backbone computations. Running a sequence alone or with a longer batch companion can yield different representations. Keep the padding, windowing and batch policy consistent when comparing embeddings. DNABERT-2 restores zero hidden states at masked padding positions before its language-model head.

Core AI assets

The optional dnabert2/coreai and ntv3/coreai directories contain self-contained float32 .aimodel assets, tokenizer metadata, a deployment manifest, file checksums and numerical validation results. They target macOS 27 / iOS 27 and use Core AI runtime specialization. No architecture-specific AOT .aimodelc files are included.

Variant	Batch size	Tokenized/padded length	Function
DNABERT-2	1–4	3–512, every integer	main
NTv3	1–4	128	short
NTv3	1–4	256–8192, multiples of 128	main

Inputs are int32 input_ids and attention_mask, each [batch, length]. Outputs are float32 hidden ([batch, length, 768]), logits ([batch, length, vocabulary]) and pooled ([batch, 768]). The output vocabulary width is 4096 for DNABERT-2 and 11 for NTv3. Masks must preserve each model's right-padding and special-token conventions.

Keep each Core AI directory intact when installing it: deployment.json specifies the exact file inventory and hashes. HyenaDNA is provided as MLX-compatible weights because the evaluated Core AI exporter does not support its FFT lowering.

Conversion validation

• All 421 stored tensors across the three Safetensors checkpoints were checked against the source tensors, preserving values exactly.
• MLX inference was compared with upstream PyTorch on short, mixed-length and maximum-context fixtures. Token IDs and masks matched exactly; hidden states, logits and pooled representations passed numerical comparisons. See each variant's validation.json.
• Core AI passed 36 DNABERT-2 and 24 NTv3 batch/length cases, plus the upstream short, long and context-boundary fixtures. The gates require finite float32 outputs, matching shapes, maximum absolute error below 0.005, RMSE below 0.0005, and pooled cosine similarity of at least 0.99999. Full reports are in the respective coreai/validation.json files.

Core AI export validation used an Apple M4 Mac with 24 GiB unified memory, macOS 27, PyTorch 2.9.0 and coreai-torch 0.4.2. Model preparation, memory use, latency and thermal behavior depend on device and workload. Physical iPhone performance has not been measured.

These checks measure conversion fidelity. The original paper's downstream biological benchmark scores have not been independently reproduced for this converted release.

Repository layout

README.md
LICENSE.pdf
NOTICE.md
SHA256SUMS
models.json
licenses/
dnabert2/
  model.safetensors
  config.json
  tokenizer.json
  manifest.json
  provenance.json
  validation.json
  MODEL_CARD.md
  coreai/
hyenadna/
  model.safetensors
  config.json
  tokenizer_config.json
  manifest.json
  provenance.json
  validation.json
  MODEL_CARD.md
ntv3/
  model.safetensors
  config.json
  vocab.json
  manifest.json
  provenance.json
  validation.json
  MODEL_CARD.md
  coreai/

SHA256SUMS covers the release files. Per-model provenance records the exact NVIDIA checkpoint revision, source hash, backbone configuration revision and converted-weight hash. Original NVIDIA model cards are retained unchanged for attribution and reference.

Sources and license

The source checkpoints are NVIDIA DNABERT-2, NVIDIA HyenaDNA and NVIDIA NTv3. Their underlying architectures come from DNABERT-2, HyenaDNA and NTv3.

The converted JEPA-DNA weights remain under the NVIDIA Non-Commercial License. Its use limitation defines non-commercial use as research or evaluation, with the exception specified there for NVIDIA and its affiliates. Redistribution must preserve the license and notices. This is a custom license, not CC BY-NC-4.0; format conversion does not grant broader rights. See NOTICE.md and licenses/ for additional backbone attribution and terms, including the InstaDeep model license for NTv3.

Intended use is non-commercial genomics research and evaluation under the applicable terms. The models are not validated medical devices or clinical decision systems. This conversion is not an official NVIDIA, InstaDeep or Apple release, and no endorsement is implied.

Citation

Please cite the original JEPA-DNA research and the relevant backbone when using these weights:

@article{larey2026jepa,
  title={JEPA-DNA: Grounding Genomic Foundation Models through Joint-Embedding Predictive Architectures},
  author={Ariel Larey and Elay Dahan and Amit Bleiweiss and Raizy Kellerman and Guy Leib and Omri Nayshool and Dan Ofer and Tal Zinger and Dan Dominissini and Gideon Rechavi and Nicole Bussola and Simon Lee and Shane O'Connell and Dung Hoang and Marissa Wirth and Alexander W. Charney and Yoli Shavit and Nati Daniel},
  journal={arXiv preprint arXiv:2602.17162},
  year={2026}
}