Metric 3D understanding emerges from self-supervised RGB-D pretraining — and transfers zero-shot from bedrooms to highways to colonoscopies, for $4 of compute.
Spatial-JEPA learns the geometry of 3D scenes the way I-JEPA learns images: by predicting masked latent representations — never regressing depth, never seeing a 3D label. A small RGB-D encoder trained this way develops an internal sense of metric distance that holds up across domains it was never trained on, and rivals foundation models trained on >100M images at a fraction of the cost.
Metric understanding is measured as Spearman ρ between latent feature distances and true metric depth — higher means the encoder's representation tracks real 3D geometry.
| Encoder | Trained on | Evaluated on | Spearman ρ |
|---|---|---|---|
| Spatial-JEPA | SUN RGB-D (6.7K indoor) | NYUv2 (indoor) | 0.625 |
| Spatial-JEPA | SUN RGB-D (6.7K indoor) | KITTI (outdoor, zero-shot) | 0.911 |
| Spatial-JEPA | SUN RGB-D (6.7K indoor) | C3VD (colonoscopy, zero-shot) | 0.337 |
| I-JEPA (RGB only) | SUN RGB-D | NYUv2 | 0.256 |
| DINOv2 ViT-B/14 | ImageNet-22K (142M) | NYUv2 | 0.222 |
| DepthAnything v2 (upper bound) | — | NYUv2 | 0.984 |
The takeaways:
- Geometry is emergent. No depth regression head, no 3D supervision — metric structure falls out of the self-supervised objective alone.
- It transfers zero-shot. An indoor-only encoder scores ρ = 0.911 on outdoor KITTI, nearly matching an encoder trained on KITTI itself (0.903). Metric understanding outlasts visual appearance.
- Depth conditioning is the lever. Adding depth to the input moves ρ from 0.256 → 0.555; the asymmetric cross-attention architecture adds the final +0.07 to 0.625.
- It's absurdly cheap. ~$4 and 2 hours on a single A100 — ~125,000× cheaper than DINOv2, with 2.8× better geometry. Emergence saturates at ~1,000 training images.
Full experimental detail, ablations, and reproduction commands are in docs/RESULTS.md.
RGB ─► RGB ViT encoder ─┐
├─► asymmetric cross-attention fusion ─► mixer ─► latent
Depth ─► Depth ViT enc. ─┘ │
▼
I-JEPA objective: predict masked-region latents
from visible context (Smooth-L1 in latent space)
- Paired transforms — crop/flip parameters are sampled once and applied identically to RGB and depth, preserving pixel correspondence (independent transforms silently corrupt it).
- Fixed depth normalization — divide by a global constant, not per-sample stats, so metric scale is preserved across the dataset.
- EMA target encoder — cosine momentum 0.996 → 1.0; the target stabilizes early and is effectively frozen by the end of training.
See Key design decisions below and the source in models/ and training/.
spatial-jepa/
├── models/ # ViT blocks, masking, I-JEPA baseline, Spatial-JEPA (RGBD)
├── training/ # training loops, losses, schedulers, ablation variants
├── data/ # RGB-D dataset loaders (NYUv2, SUN RGB-D, KITTI, C3VD, …)
├── evaluation/ # latent-geometry ρ, probes, transfer & robustness evals
├── demo/ # Gradio "3D Oracle" app + inference + sample images
├── nyuv2/ # vendored pytorch-nyuv2 loader (no upstream setup.py)
├── scripts/
│ ├── modal/ # cloud experiment drivers (training + evaluation on Modal)
│ └── *.sh/*.py # local runners, Colab setup, figure generation
├── results/ # SpatialQA benchmark (26,838 questions, JSONL)
├── figures/ # qualitative results & visualizations
├── docs/ # full results report, reproducibility guide, experiment log
└── config.py # centralized training config
pip install -r requirements.txt
# Sanity-check the data pipeline (loads a batch, plots RGB+depth, checks alignment)
python data/verify.py
# Forward passes — prints tensor shapes
python models/ijepa_baseline.py
python models/spatial_jepa.py
# Smoke-test training (2 epochs)
python training/train_baseline.py --epochs 2 --batch_size 8Apple Silicon (M-series): ./scripts/run_local_m4.sh all runs the full verify → forward → train smoke test. Device selection (CUDA → MPS → CPU) is handled automatically in utils/device.py.
Cloud (Modal): training and evaluation at scale run via the drivers in scripts/modal/, e.g. modal run scripts/modal/modal_train_sunrgbd.py. These reference a Modal secret named wandb-api-key for experiment logging — create your own with modal secret create wandb-api-key WANDB_API_KEY=….
results/spatialqa_nyuv2_val.jsonl contains 26,838 metric spatial-reasoning questions auto-generated from NYUv2 with segmentation + depth ground truth — distance, comparison, direction, height, and proximity. Using only RGB at inference (RGB → monocular depth → Spatial-JEPA → reasoning), the system answers height questions at 97% and comparison at 71% accuracy.
| Decision | Why |
|---|---|
| Paired RGB/depth transforms | Independent augmentation breaks depth-pixel correspondence |
| Global (not per-sample) depth norm | Preserves metric scale for the geometry evaluation |
| Learned positional embeddings | Predictor queries "what lives at position p?" — learned PE beats fixed sine/cosine here |
| Cosine EMA momentum 0.996 → 1.0 | Slow early target updates; near-frozen target at convergence |
Datasets are not bundled (see .gitignore). Loaders expect standard releases of NYUv2, SUN RGB-D, KITTI, and C3VD — see data/ and the download drivers in scripts/modal/. The nyuv2/ package is vendored from xapharius/pytorch-nyuv2 (no upstream setup.py).
@software{patil_spatial_jepa_2026,
author = {Patil, Heramb},
title = {Spatial-JEPA: Emergent Metric 3D Understanding from Self-Supervised RGB-D Pretraining},
year = {2026},
url = {https://github.com/herambpatilofficial/Spatial-JEPA}
}Released under the MIT License.