JEPA on Apple Silicon

Three ways to use JEPA

mlx-tune exposes JEPA through an Unsloth-flavoured API. Pick the track that matches your data and goal.

What is JEPA (and is it an LLM)?

JEPA is a vision/video architecture, not a language model. Every JEPA here is a Vision Transformer that operates on image patches or video tubelets. The "JEPA" idea is how it's trained: take two views of an input, encode both, and train the model so one view's representation predicts the other's — in latent space, never reconstructing raw pixels. This captures predictable, abstract structure (semantics, dynamics) and discards unpredictable detail.

The enemy of every JEPA is representation collapse (the encoder outputs a constant). Different variants fight it differently:

Supported models

LeJEPA presets (random init — train from scratch)

Pretrained checkpoints (downloaded & converted on first use)

LeJEPA — self-supervised pretraining

Train a ViT on your own images, no labels required. The loss has a single hyperparameter lam (the SIGReg weight, ~0.05).

from mlx_tune import FastJEPAModel, JEPATrainer, JEPAConfig, linear_probe

# 1. Build a randomly-initialised ViT (LeJEPA trains from scratch).
model, _ = FastJEPAModel.from_pretrained("vit-tiny", img_size=128)

# 2. Self-supervised pretraining — `images` is just a list of HWC arrays / PIL images.
trainer = JEPATrainer(
    model,
    args=JEPAConfig(num_epochs=5, batch_size=64, lam=0.05),
    train_dataset=images,
)
trainer.train()

# 3. Use the learned encoder: extract features, or evaluate with a linear probe.
feats = model.encode(images)                       # (N, dim) frozen features
acc = linear_probe(model, tr_x, tr_y, te_x, te_y)  # quick representation-quality check

# 4. Save / reload the encoder.
model.save_pretrained("my_lejepa")
model, _ = FastJEPAModel.from_pretrained("my_lejepa")

I-JEPA — pretrained image encoder

Load Meta's pretrained I-JEPA and put its features to work. Inputs default to 224×224 (the resolution it was trained at); other sizes work too via position-embedding interpolation — see Scale-up.

from mlx_tune import FastJEPAModel, linear_probe, knn_probe, attentive_probe

# Downloads ~2.5 GB on first run, then converts HF -> MLX (no torch needed).
model, _ = FastJEPAModel.from_pretrained("facebook/ijepa_vith14_1k")

# Three frozen-feature probes — none of them backprop through the 630M encoder:
linear    = linear_probe(model, tr_x, tr_y, te_x, te_y)     # logistic regression on mean-pooled features
knn       = knn_probe(model, tr_x, tr_y, te_x, te_y, k=20)  # label-light, no training
attentive = attentive_probe(model, tr_x, tr_y, te_x, te_y)  # attention-pooling head — the paper's eval

# Or just grab features for retrieval / clustering.
feats   = model.encode(my_images)          # (N, 1280)  mean-pooled
tokens  = model.encode_tokens(my_images)   # (N, T, 1280)  per-patch tokens

Domain-adaptive continued pretraining (warm-start LeJEPA)

Because LeJEPA and the pretrained encoders are all ViTs, you can continue self-supervised pretraining from a pretrained checkpoint — adapt Meta's ImageNet-trained I-JEPA to your own domain with the LeJEPA objective, no labels. Just pass the loaded model straight to JEPATrainer:

from mlx_tune import FastJEPAModel, JEPATrainer, JEPAConfig

# Start from pretrained I-JEPA instead of random init...
model, _ = FastJEPAModel.from_pretrained("facebook/ijepa_vith14_1k")

# ...and keep training it on your unlabeled domain images with LeJEPA / SIGReg.
JEPATrainer(model, args=JEPAConfig(img_size=224, num_epochs=3, learning_rate=1e-4),
            train_dataset=my_domain_images).train()

model.save_pretrained("ijepa_my_domain")   # domain-adapted encoder

V-JEPA 2 — pretrained video world model

Load Meta's V-JEPA 2 video encoder. Clips are (T, 256, 256, 3) arrays; T must be a multiple of the tubelet size (2). Frames are resized to 256×256 automatically.

from mlx_tune import (
    FastVideoJEPAModel, video_linear_probe, video_knn_probe, video_attentive_probe,
)

model, _ = FastVideoJEPAModel.from_pretrained("facebook/vjepa2-vitl-fpc64-256")

# Frozen clip features, or any of the three video probes.
clip_feats = model.encode(my_clips)                            # (N, 1024)
linear     = video_linear_probe(model, tr_x, tr_y, te_x, te_y)
knn        = video_knn_probe(model, tr_x, tr_y, te_x, te_y, k=20)
attentive  = video_attentive_probe(model, tr_x, tr_y, te_x, te_y)  # paper's eval

The predictor: anticipate the rest of a clip

The masked-latent predictor — the "world model" half of V-JEPA 2 — is ported too (HF parity cosine 1.000000) and loads by default. Give it the latents of the first frames and it predicts the latents of the rest; latent_energy scores how surprising the actual future was (anomaly / discontinuity detection in representation space):

from mlx_tune import FastVideoJEPAModel, latent_energy

model, _ = FastVideoJEPAModel.from_pretrained("facebook/vjepa2-vitl-fpc64-256")

# Context = first 4 frames; predict the latents of everything after.
predicted, actual = model.predict_latents(clip, context_frames=4)

energy = latent_energy(predicted, actual)                  # scalar surprise
e_map  = latent_energy(predicted, actual, per_token=True)  # localise it
# A clip with a hard cut scores higher energy than a coherent one.

Meta's fine-tuned action classifiers (zero training)

Meta's VJEPA2ForVideoClassification checkpoints (attentive pooler + head, fine-tuned on Something-Something-v2) load straight into MLX — from_pretrained auto-detects them. Logits match HuggingFace (cosine 1.000000, same top-1):

clf, _ = FastVideoJEPAModel.from_pretrained("facebook/vjepa2-vitl-fpc16-256-ssv2")

results = clf.predict([clip], top_k=5)      # 174 SSv2 action classes
for r in results[0]:
    print(f"{r['prob']:.3f}  {r['label']}")

clf.save_pretrained("ssv2_mlx")             # MLX-native save/reload

Downstream classification (frozen / LoRA / full)

Both pretrained tracks share the same downstream API. Attach a classification head and pick how much of the encoder to train:

Image classification (I-JEPA or LeJEPA encoder)

from mlx_tune import (
    FastJEPAModel, JEPAClassifierTrainer, JEPAClassifierConfig,
)

model, _ = FastJEPAModel.from_pretrained("facebook/ijepa_vith14_1k")

# Attach a head; LoRA-adapt the encoder. (frozen base is frozen FIRST, then wrapped)
clf = FastJEPAModel.for_image_classification(
    model, num_classes=10, finetune="lora", r=8,
)

trainer = JEPAClassifierTrainer(
    clf,
    JEPAClassifierConfig(img_size=224, batch_size=6, num_epochs=5),
    train_images, train_labels,
    eval_images=val_images, eval_labels=val_labels,
)
trainer.train()
print(f"accuracy: {trainer.evaluate():.3f}")

Video classification (V-JEPA 2 encoder)

from mlx_tune import (
    FastVideoJEPAModel, VideoClassifierTrainer, VideoClassifierConfig,
)

model, _ = FastVideoJEPAModel.from_pretrained("facebook/vjepa2-vitl-fpc64-256")
clf = FastVideoJEPAModel.for_video_classification(model, num_classes=2, finetune="lora", r=8)

trainer = VideoClassifierTrainer(
    clf, VideoClassifierConfig(batch_size=2, num_epochs=5),
    train_clips, train_labels, eval_videos=val_clips, eval_labels=val_labels,
)
trainer.train()

Save, reload, and predict

Trained classifiers save the encoder, LoRA adapters, and head together. Reloading reconstructs the architecture (re-applying the LoRA structure) and produces identical predictions.

# after trainer.train() ...
clf.save_pretrained("my_classifier")

# Reload in a fresh process and run inference.
clf = FastJEPAModel.load_classifier("my_classifier")        # video: FastVideoJEPAModel.load_classifier
preds = clf.predict(new_images)                             # class ids   (N,)
probs = clf.predict(new_images, return_probs=True)          # softmax probs (N, num_classes)

LLM-JEPA — the JEPA objective for LLM fine-tuning

Every track above is vision. LLM-JEPA (arXiv 2509.14252, Huang, LeCun & Balestriero) brings the same Joint-Embedding Predictive idea to language-model fine-tuning. It is a training objective, not a model: it augments standard next-token prediction (NTP) with a JEPA term that aligns the embeddings of two “views” of the same item — e.g. a natural-language description (text) and its regex / SQL / code (code). It is implemented here for the first time on MLX.

The fine-tuned artifact is a normal LoRA-fine-tuned LLM — save / merge / generate work exactly as with SFT. The JEPA term only shapes the representation during training (the paper reports accuracy gains and resistance to overfitting). The dataset is a list of dicts with the two views under text_field / code_field (default "text" / "code"), falling back to prompt / completion.

from mlx_tune import FastLanguageModel, LLMJEPATrainer, LLMJEPAConfig

model, tokenizer = FastLanguageModel.from_pretrained("mlx-community/Qwen3.5-0.8B-MLX-4bit")
model = FastLanguageModel.get_peft_model(model, r=16, lora_alpha=16)

# Two "views" of the same item: a description and its regex
data = [
    {"text": "match one or more digits", "code": r"\d+"},
    {"text": "match a hex color code",   "code": r"#[0-9a-fA-F]{6}"},
    # ...
]

trainer = LLMJEPATrainer(
    model, data, tokenizer=tokenizer,
    args=LLMJEPAConfig(jepa_lambda=0.1, jepa_distance="cosine", max_steps=40),
)
trainer.train()

# The result is a normal LoRA model — generate / merge as usual:
print(model.generate(prompt="match one or more digits", max_tokens=16))
# model.save_pretrained_merged("./llm_jepa_merged")

Dense & regression heads (counting, depth, segmentation)

I-JEPA's own paper-headline tasks are depth prediction and object counting — i.e. "more than classify". Three non-classification heads attach to any JEPA encoder and reuse the same frozen / LoRA / full fine-tuning machinery:

from mlx_tune import (FastJEPAModel, JEPAClassifierConfig,
                      JEPARegressionTrainer, JEPADenseTrainer)

model, _ = FastJEPAModel.from_pretrained("facebook/ijepa_vith14_1k")

# object counting (scalar regression)
reg = FastJEPAModel.for_image_regression(model, out_dim=1, finetune="lora")
JEPARegressionTrainer(reg, JEPAClassifierConfig(img_size=224),
                      images, counts).train()
print(reg.evaluate())                 # {"mae": ..., "rmse": ...}

# depth map (dense regression) — targets are per-pixel maps at img_size
dense = FastJEPAModel.for_dense_prediction(model, out_channels=1, task="regression")
# segmentation — out_channels = num classes, integer mask targets
seg = FastJEPAModel.for_dense_prediction(model, out_channels=21, task="segmentation")

Dense heads predict at the patch grid and bilinearly upsample to img_size. Reload with FastJEPAModel.load_regressor(dir) / load_dense(dir).

Native resolution & training over a real corpus

Non-224 inputs. Pretrained I-JEPA position embeddings were tied to 224²; now you can load at any (patch-divisible) size — the learned grid is bicubically interpolated at load (and on-the-fly for off-grid inputs). The native-resolution path is bit-identical.

# run a pretrained encoder at 384x384
model, _ = FastJEPAModel.from_pretrained("facebook/ijepa_vith14_1k", img_size=384)

Streaming + resumable checkpoints. ImageFolderDataset decodes images lazily from a folder (recursively) on the prefetch thread, so the corpus never has to fit in memory. Set save_steps + resume to checkpoint the encoder + optimizer state and continue an interrupted run.

from mlx_tune import FastJEPAModel, JEPAConfig, JEPATrainer, ImageFolderDataset

model, _ = FastJEPAModel.from_pretrained("vit-base", img_size=224)
ds = ImageFolderDataset("/path/to/images")     # lazy, recursive
cfg = JEPAConfig(img_size=224, batch_size=64, num_epochs=10,
                 save_steps=500, resume=True)   # resumes from /checkpoint
JEPATrainer(model, cfg, ds).train()

LeWM — a latent world model you can train (and plan with)

LeWM (LeWorldModel, arXiv 2603.19312, Maes/Le Lidec/Scieur/LeCun/Balestriero) trains a latent world model end-to-end from pixels with only two loss terms and one tunable hyperparameter:

This is the trainable, on-device world model — the counterpart to V-JEPA 2-AC's load-and-plan (which is a design writeup for a future release, see jepa.md §9).

from mlx_tune import FastWorldModel, LeWMConfig, LeWMTrainer, plan_cem, PointMassEnv

env = PointMassEnv(size=48)                       # toy 2-D control env
data = env.collect(n_episodes=80, ep_len=10)      # {"frames", "actions"} trajectories

model = FastWorldModel.from_pretrained("lewm-tiny", img_size=48, action_dim=2)
LeWMTrainer(model, LeWMConfig(img_size=48, action_dim=2, sigreg_lambda=0.05,
                              max_steps=250), data).train()

# plan: pick actions whose predicted latent matches a goal image's latent (MPC)
goal_z = model.encode([env.render([0.8, 0.2])])[0]
z = model.encode([env.render()])[0]
action = plan_cem(model, z, goal_z, horizon=3, action_dim=2)

Runnable examples

python examples/58_lejepa_pretraining.py
python examples/59_ijepa_feature_extraction.py
python examples/60_vjepa2_video.py
python examples/61_llm_jepa_finetuning.py
python examples/62_lewm_world_model.py
python examples/63_jepa_dense_regression.py
python examples/64_vjepa2_predictor_classifier.py

Tips & gotchas