v0.6.0 — JEPA comes to Apple Silicon

v0.6.0 — JEPA comes to Apple Silicon

A big feature release. mlx-tune now supports the JEPA family — Yann LeCun's Joint-Embedding Predictive Architecture — end to end on a Mac, across vision, language, and world models:

• Vision SSL — train LeJEPA from scratch, or load Meta's pretrained I-JEPA (images) and V-JEPA 2 (video) and fine-tune them with LoRA.
• LLM-JEPA — bring the JEPA objective to LLM fine-tuning (the JEPA paper's own LLM follow-up).
• LeWM — train a latent world model from pixels and plan with it (CEM/MPC).
• V-JEPA 2 world model surface — Meta's masked-latent predictor is ported too: anticipate the rest of a clip in representation space and score how surprising the actual future was. Plus Meta's fine-tuned SSv2 action classifiers (174 classes) — zero training, just .predict().

As far as I can tell, this is the first solid train-and-fine-tune JEPA stack on MLX — and the first time LLM-JEPA and LeWorldModel exist on MLX at all.

Both pretrained vision ports are converted HuggingFace→MLX and numerically verified against the reference PyTorch models (cosine similarity 1.000000), so the features you get out are the real thing.

Upgrade

uv pip install -U mlx-tune

Purely additive — no existing API changes. New entry points: FastJEPAModel, FastVideoJEPAModel, LLMJEPATrainer, FastWorldModel (+ predict_latents / latent_energy / VJEPA2PretrainedVideoClassifier on the V-JEPA 2 track).

What is JEPA (and is it an LLM)?

JEPA is a predict-in-representation-space paradigm. A model is trained to predict one view's representation from another, in latent space — capturing predictable structure (semantics, dynamics) instead of reconstructing pixels/tokens. The hard problem is representation collapse; each variant fights it differently. Most JEPA work is vision, but the same idea now extends to LLMs (LLM-JEPA) and to world models (LeWM).

For more details about JEPA in mlx-tune: JEPA Docs

New examples

• examples/58_lejepa_pretraining.py — LeJEPA from-scratch + linear probe
• examples/59_ijepa_feature_extraction.py — I-JEPA: probes + LoRA fine-tune + save/load/predict
• examples/60_vjepa2_video.py — V-JEPA 2: video probe + LoRA fine-tune + save/load/predict
• examples/61_llm_jepa_finetuning.py — LLM-JEPA: JEPA objective for LLM fine-tuning
• examples/62_lewm_world_model.py — LeWM: train a latent world model + CEM/MPC planning
• examples/63_jepa_dense_regression.py — counting / depth / segmentation heads
• examples/64_vjepa2_predictor_classifier.py — V-JEPA 2 predictor (anticipation + surprise energy) + Meta's pretrained SSv2 classifier

Full Changelog: v0.5.1...v0.6.0