H-JEPA-LM

Hierarchical Joint-Embedding Predictive Language Model
A fundamentally different approach to language modeling — predicting latent representations, not tokens.

---

Two models in one package

This repo contains two models. Both predict latent representations instead of tokens, but they differ in complexity:

	JEPA-LM	H-JEPA-LM
Class	JEPELM	HJEPELM
Module	jepalm.model	jepalm.hjepa
Architecture	Single-level encoder + predictor	Multi-level hierarchical encoder
Action conditioning	No	Yes
World model planning	No	Yes (latent-space rollouts)
Best for	Research, baselines, simpler experiments	Production, action-aware tasks
Parameters	~5.6M (tiny config)	~5.8M (tiny config)

When to use which

• Use JEPA-LM if you want a simple, clean baseline. Single-level prediction with EMA target encoder. Good for understanding JEPA fundamentals.
• Use H-JEPA-LM if you need hierarchical prediction, action conditioning, or world model planning. This is the full-featured model.

Both models share the same core idea: predict latent representations (not tokens) via cosine similarity loss, with an EMA target encoder to prevent collapse.

Install

pip install git+https://github.com/Griffith-7/H-JEPA-LM.git

Or clone and install locally:

git clone https://github.com/Griffith-7/H-JEPA-LM.git
cd H-JEPA-LM
pip install -e ".[dev]"

Quick start

JEPA-LM (base model)

from jepalm import JEPELM, JEPAConfig

config = JEPAConfig(
    enc_hidden_dim=256,
    enc_num_layers=4,
    enc_num_heads=4,
    pred_hidden_dim=128,
    pred_num_layers=2,
)

model = JEPELM(config)
print(f"Parameters: {model.count_parameters()['total']:,}")

import torch
x = torch.randint(0, 1000, (2, 32))
out = model(x)
print(f"JEPA loss: {out['jepa_loss']:.4f}")
print(f"Cosine similarity: {out['cosine_similarity']:.4f}")

# Generate text from the latent space (latent -> decoder -> tokens)
prompt = torch.randint(0, 1000, (2, 16))
generated = model.generate(prompt, max_new_tokens=20, temperature=0.8)
print(f"Generated tokens: {generated.shape}")

H-JEPA-LM (hierarchical + action conditioning)

from jepalm import HJEPELM, HConfig

config = HConfig(
    vocab_size=30522,
    dim=256,
    layers=4,
    heads=4,
    ff_dim=1024,
    max_len=128,
    num_levels=2,
    action_dim=64,
)

model = HJEPELM(config)
print(f"Parameters: {model.count_parameters()['total']:,}")

import torch
x = torch.randint(0, 30522, (2, 64))

# Basic forward pass (returns multi-level latent outputs)
outputs = model(x)
print(f"Highest level shape: {outputs[-1].shape}")

# Compute JEPA loss (with optional action conditioning)
loss, val = model.compute_loss(x)
print(f"Loss: {loss.item():.4f}")

# With actions
actions = torch.randn(2, 64)
loss, val = model.compute_loss(x, actions=actions)
print(f"Loss (with action): {loss.item():.4f}")

# World model: predict next state
next_state = model.predict_next_state(x, actions)
print(f"Next state: {next_state.shape}")

# Plan actions to reach a goal
goal = torch.randn(2, 256)
planned = model.plan_actions(x, goal, num_steps=5, num_candidates=10)
print(f"Planned actions: {planned.shape}")

# Generate text from the hierarchical latent space
generated = model.generate(x, max_new_tokens=20, temperature=0.8)
print(f"Generated tokens: {generated.shape}")

CLI

# Show both models and parameter counts
hjepa info

# Train base JEPA-LM
hjepa train --preset tiny --epochs 5

# Run 5-way benchmark
hjepa bench

# Run all tests
hjepa test

How JEPA-LM works

Traditional LLMs predict the next token left-to-right. JEPA-LM predicts meaning — the latent representation of masked text — in embedding space.

Input: "The cat sat on the [MASK] because it was [MASK]"
       |
+-----------------------------------------------------+
|  Bidirectional Encoder                              |
|  -> Latent representations for all tokens           |
+-----------------------------------------------------+
|  Predictor (narrow bottleneck)                      |
|  -> Predicts latents for masked spans               |
+-----------------------------------------------------+
|  EMA Target Encoder (stop-gradient)                 |
|  -> Stable targets, prevents embedding collapse     |
+-----------------------------------------------------+
|  JEPA Loss (cosine similarity in latent space)      |
+-----------------------------------------------------+

What makes H-JEPA-LM different

H-JEPA-LM extends JEPA-LM with three additions:

1. Hierarchical prediction — Encodes at multiple levels (token details → semantic meaning). The per-level loss weights are learned during training (softmax over a free level_weights vector), instead of hardcoded constants, so the model can discover how much each abstraction level matters.
2. Action conditioning — Actions are encoded and fused into the highest-level predictions. The model learns to predict what happens when you take an action.
3. World model planning — Plans sequences of actions by rolling out predictions in latent space and selecting the best trajectory.
4. Real decoder — An autoregressive LightweightDecoder is trained (weak NTP loss, default lambda_ntp=0.1) alongside the multi-level JEPA loss, giving a working latent → text path via model.generate().

JEPA-LM:                              H-JEPA-LM:
  Encoder                                Hierarchical Encoder
    |                                      Level 0 (tokens) ──┐
    |                                      Level 1 (semantic) ─┤
  Predictor                               Hierarchical Predictor
    |                                      + Action Conditioning
  Target Encoder (EMA)                    Target Encoders (EMA, per level)
    |                                      + World Model Planning
  JEPA Loss                               Multi-level JEPA Loss

Key innovation: JEPA as primary objective

┌──────────────────────────────────────────────────┐
│           Why this is different                   │
├──────────────────────────────────────────────────┤
│                                                  │
│  LLM-JEPA (ICLR 2026):                          │
│    Existing LLM + bolted-on JEPA loss            │
│    -> JEPA is secondary (optional)               │
│    -> Causal attention (left-to-right)           │
│    -> Needs paired Text<->Code data              │
│                                                  │
│  JEPA-LM / H-JEPA-LM (This repo):               │
│    New architecture with JEPA as core             │
│    -> JEPA is PRIMARY objective                   │
│    -> Bidirectional attention                     │
│    -> Self-supervised via span masking            │
│    -> EMA target encoder prevents collapse        │
│                                                  │
└──────────────────────────────────────────────────┘

Benchmarks

Honesty note: All numbers below come from the 5-way comparison in benchmark_hjepa.py, which trains tiny models (dim=128, 4 layers, 3 epochs, ~2000 Wikitext articles, frozen-encoder linear probes trained for 3 epochs). These results are illustrative of the mechanism, not evidence of scaling. The embedding "diversity" metrics below are not the headline: diverse embeddings that are wrong are easy to produce. The meaningful task metrics — SST-2 classification accuracy, next-token reasoning top-1/top-5, and next-token perplexity (identical linear-probe head for every model) — are printed by the benchmark and saved to benchmark_hjepa_results/results.json. The core open question from the Reddit discussion — whether JEPA-style models can match an autoregressive LM in generative perplexity at scale — is not yet answered; the perplexity column is the honest baseline to beat.

Model	Cosine Sim ↓	Embed Std ↑	SV Ratio ↓	Params
GPT (NTP)	0.998	0.040	0.898	4.7M
BERT (MLM)	0.850	0.297	0.576	8.7M
LLM-JEPA	0.998	0.040	0.898	4.9M
JEPA-LM	0.857	0.231	0.514	5.6M
H-JEPA-LM	0.774	0.308	0.457	5.8M

At this tiny scale H-JEPA-LM produces embeddings with ~23% lower pairwise cosine similarity than LLM-JEPA. Treat this as illustrative only — it measures embedding diversity, not task quality.

Full results from a real run

The 5-way comparison above was re-run on a single GPU and reproduces the diversity numbers. The complete task metrics from that run:

Metric	GPT (NTP)	BERT (MLM)	JEPA-LM	LLM-JEPA	H-JEPA-LM
SST-2 Accuracy	0.5092	0.5092	0.5092	0.5092	0.5092
Reasoning Top-1	0.0000	0.0000	0.0000	0.0000	0.0000
Next-token PPL (probe, ↓)	22554	33570	32545	31669	20592

Read honestly: at this scale (dim=128, 3 epochs, ~2000 articles, 3 probe epochs) every model is at chance on the task metrics — SST-2 accuracy is the random baseline (0.5), reasoning top-1 is 0, and all perplexities are near-uniform for a 30k-token vocabulary (random ≈ 30k). The benchmark therefore cannot yet demonstrate task-quality gains; it only demonstrates that the embedding-diversity claim reproduces. Proving JEPA-LM can match an autoregressive baseline in perplexity/accuracy requires scaling training well beyond this toy setup — this is the open research question.

Why these metrics matter

• Cosine Similarity — How similar embeddings are to each other. Lower = more diverse. GPT embeddings are nearly identical (0.998).
• Embedding Std Dev — Variation in embedding magnitudes. Higher = more information encoded.
• SV Ratio — Balance of embedding dimensions used. Lower = more balanced representations.

The full metric set (SST-2 accuracy, reasoning top-1/top-5, and a fair next-token perplexity via an identical linear probe on every model's frozen embeddings) is produced by:

python benchmark_hjepa.py   # writes benchmark_hjepa_results/results.json

Limitations

• No task-quality result yet. A full benchmark run reproduces the diversity claims but shows every model at chance on SST-2 accuracy, reasoning, and near-uniform next-token perplexity (dim=128, 3 epochs, ~2000 articles, 3 probe epochs). Whether JEPA-LM/H-JEPA-LM can match autoregressive next-token perplexity at scale is the open question — see the "Next-token PPL" metric.
• Generation re-encodes each step, so generate() is O(S²) per sequence; it is a working latent→text path, not an optimized inference engine.
• Hierarchy is still structural. Per-level loss weights are now learned, but the number of levels is still a hyperparameter (HConfig.num_levels).
• Future work: comparing against an image-JEPA trained on rendered text (as suggested in the Reddit thread) is an interesting experiment but is not implemented here.

Project structure

H-JEPA-LM/
├── jepalm/                    # Core package (pip installable)
│   ├── __init__.py           # Public API: JEPELM, HJEPELM, configs
│   ├── cli.py                # CLI entry points
│   │
│   ├── model.py              # JEPA-LM (base model)
│   ├── config.py             # JEPAConfig
│   ├── encoder.py            # Bidirectional encoder
│   ├── target_encoder.py     # EMA target encoder
│   ├── predictor.py          # Narrow predictor
│   ├── decoder.py            # Lightweight autoregressive decoder
│   ├── generation.py         # Sampling + generation utilities
│   ├── loss.py               # JEPA + NTP loss
│   ├── masking.py            # Span masking
│   │
│   ├── hjepa.py              # H-JEPA-LM (hierarchical model)
│   │
│   ├── train.py              # Training loop
│   ├── dataset.py            # Dataset loading
│   └── eval.py               # Evaluation
├── tests/                     # 21 unit tests
├── benchmarks/               # Benchmark charts & scripts
├── .github/workflows/        # CI/CD
├── demo.ipynb                # Jupyter notebook demo
├── benchmark_hjepa.py        # 5-way comparison benchmark
├── train.py                  # Training entry point
├── test_model.py             # Quick smoke test
├── pyproject.toml            # Package configuration
├── LICENSE
└── README.md

Testing

pip install -e ".[dev]"
pytest tests/ -v

References

• I-JEPA — Image-based JEPA (CVPR 2023)
• V-JEPA — Video-based JEPA (2024)
• LLM-JEPA — Text JEPA bolted onto existing LLMs (ICLR 2026)
• JEPA-LM / H-JEPA-LM (Ours) — JEPA as primary objective for text

License

MIT — contributions welcome.