Nano-JEPA

Nano-JEPA is a lightweight, efficient implementation of the Joint Embedding Predictive Architecture (JEPA) designed for community-driven research.

Performance

Trained on a dual-GPU setup achieving ~49% MFU (Model Flops Utilization).

Variant	Parameters	Configuration
200M	200M	12-layer, embed_dim 768
1B	1.1B	Coming soon
2B	2.2B	Coming soon

Quick Start

You can load the weights using standard PyTorch:

import torch
from model import NanoJEPA

# Initialize architecture (ensure this matches your training config)
model = NanoJEPA(
    img_size=224, 
    patch_size=16, 
    embed_dim=768, 
    context_depth=12, 
    predictor_depth=4
)

# Load weights
state_dict = torch.load("weights/nano-jepa-200M.pth")
model.load_state_dict(state_dict)
model.eval()