Gaussian-JEPA

Joint-Embedding Predictive Learning for 3D Gaussian Splats

Bin Ren¹, Qi Ma², Yue Li³, Zongyan Han¹, Yidi Li⁴, Yuqian Fu⁵
Rao Muhammad Anwer¹, Theo Gevers³, Fahad Shahbaz Khan¹, Salman Khan¹

¹Mohamed bin Zayed University of Artificial Intelligence (MBZUAI)    ²ETH Zürich    ³University of Amsterdam
⁴Taiyuan University of Technology    ⁵King Abdullah University of Science and Technology (KAUST)

Overview

Gaussian-JEPA is a self-supervised framework for learning object-level representations directly from 3D Gaussian Splatting (3DGS) assets. It replaces masked-attribute reconstruction with latent prediction: an online encoder represents visible Gaussian tokens, while an exponential-moving-average target encoder provides stop-gradient supervision for held-out spatial blocks.

The reference model uses all 14 Gaussian attributes, four non-overlapping target blocks with heterogeneous spatial support, complementary latent projections, and feature-space grounding. The repository provides the pretraining implementation, downstream transfer protocols, frozen representation diagnostics, shape completion, and part segmentation.

Method

A 1K-Gaussian observation is grouped into 64 local tokens. The target sampler selects four non-overlapping blocks of sizes [11, 9, 7, 5]; their complement forms a shared 32-token context. The context encoder runs once, after which a shared predictor estimates two latent projections for each target block from features supplied by the EMA encoder. Training remains entirely in representation space and does not decode coordinates, opacity, covariance, or spherical-harmonic coefficients.

Results

Gaussian classification results below use 1K Gaussians for both pretraining and transfer. Detailed settings and complete metrics are available in docs/results.

Evaluation	Gaussian-MAE	Gaussian-JEPA
ModelNet10 linear probing (%)	93.50	93.72
ModelNet40 linear probing (%)	88.97	90.47
ShapeNet-Part class mIoU (%)	84.2	84.5
Resampling drift ↓	0.1202	0.0850
Partial-observation R@1 AUC ↑	41.82	52.74
Shape-completion Chamfer distance ↓	0.0732	0.0678

Installation

The reference environment uses Python 3.9, PyTorch 2.0.1, and CUDA 11.8. pointnet2_ops and knn_cuda are compiled during installation, so the CUDA toolkit exposed by nvcc must be compatible with the installed PyTorch build.

git clone https://github.com/Amazingren/Gaussian-JEPA.git
cd Gaussian-JEPA
conda env create -f environment.yml
conda activate gaussian_jepa

See docs/getting_started.md for dataset layouts, configuration details, and the complete evaluation workflow.

Data

Pretraining uses ShapeSplatsV1, while classification uses ModelNet Splats. Request access and accept the corresponding dataset terms before downloading. Dataset locations can be configured without modifying tracked source files:

export SHAPENET55GS_PLY_ROOT=/path/to/shapesplat_ply
export MODELNETGS_PLY_ROOT=/path/to/modelsplat_ply

The ShapeNet55 and ModelNet split files used by the loaders are included under datasets/.

Checkpoints

Pretrained weights are not included in the current public release. The commands below use checkpoints/gaussian_jepa_ep300.pth as the local path for a compatible 300-epoch pretraining checkpoint. See CHECKPOINTS.md for the expected layout and state-dict structure.

mkdir -p checkpoints
sha256sum checkpoints/gaussian_jepa_ep300.pth
python tools/smoke_test_release.py \
  --checkpoint checkpoints/gaussian_jepa_ep300.pth

The smoke test verifies CUDA extensions, strict checkpoint loading, finite pretraining losses, and one backward pass without requiring a dataset.

Pretraining

python main.py \
  --config cfgs/pretrain/gaussian_jepa.yaml \
  --exp_name gaussian_jepa \
  --launcher none \
  --seed 0

The canonical configuration trains for 300 epochs with 1,024 Gaussians, 64 groups of 32 neighbors, and a total batch size of 256. For distributed training, launch main.py with torchrun and set --launcher pytorch.

Downstream evaluation

ModelNet classification

Six configurations cover full fine-tuning, linear probing, and MLP-3 probing on ModelNet10 and ModelNet40. For example:

python main.py \
  --config cfgs/finetune/modelnet10_linear.yaml \
  --finetune_model \
  --ckpts checkpoints/gaussian_jepa_ep300.pth \
  --exp_name modelnet10_linear \
  --seed 0

Select another protocol by changing the configuration filename. The transfer loader extracts the released JEPA_encoder.* weights automatically.

Frozen Gaussian diagnostics

The resampling and partial-observation evaluators give Gaussian-JEPA and a compatible Gaussian-MAE E(All) checkpoint identical primitives, groupings, and observation masks.

python tools/eval_frozen_embeddings.py \
  --jepa-ckpt checkpoints/gaussian_jepa_ep300.pth \
  --mae-ckpt /path/to/gaussian_mae_ep300.pth \
  --gs-root "$MODELNETGS_PLY_ROOT" \
  --num-objects 0 \
  --output-dir outputs/resampling

python tools/eval_partial_retrieval.py \
  --jepa-ckpt checkpoints/gaussian_jepa_ep300.pth \
  --mae-ckpt /path/to/gaussian_mae_ep300.pth \
  --gs-root "$MODELNETGS_PLY_ROOT" \
  --num-objects 0 \
  --output-dir outputs/partial_observation

--num-objects 0 evaluates all 2,467 ModelNet40-GS test objects. A positive value runs a category-balanced subset for a faster diagnostic.

Shape completion and part segmentation

• completion_gs/ contains frozen-encoder, partial-to-complete Gaussian prediction and render-space evaluation.
• segmentation_gs/ contains ShapeNet-Part fine-tuning and evaluation.
• viz/ contains the rendering and qualitative-analysis utilities used by the release.

Feature extraction

The input PLY must follow the standard 3DGS property convention used by ShapeSplats.

python tools/extract_features.py \
  --ply /path/to/point_cloud.ply \
  --checkpoint checkpoints/gaussian_jepa_ep300.pth \
  --output outputs/object_features.npz \
  --seed 0

The output stores a normalized 768-D object embedding, 64 local 384-D token features, group centers, sampled Gaussians, and sampling metadata. The exact input contract and validated scope are specified in MODEL_CARD.md.

Repository structure

assets/             figures used by this README
cfgs/               pretraining and ModelNet transfer configurations
models/             Gaussian-JEPA and downstream model definitions
datasets/           Gaussian loaders and dataset splits
tools/              training, feature extraction, and frozen evaluation
completion_gs/      Gaussian shape completion
segmentation_gs/    ShapeNet-Part segmentation
viz/                visualization and rendering utilities
docs/               setup guides and evaluation protocols

Acknowledgements

This implementation builds on ShapeSplat / Gaussian-MAE and Point-MAE. We thank their authors for releasing code and data. Third-party provenance and license terms are recorded in THIRD_PARTY_NOTICES.md.

Citation

@article{ren2026gaussian,
  title   = {Gaussian-JEPA: Joint-Embedding Predictive Learning for 3D Gaussian Splats},
  author  = {Ren, Bin and Ma, Qi and Li, Yue and Han, Zongyan and Li, Yidi and Fu, Yuqian and Anwer, Rao Muhammad and Gevers, Theo and Khan, Fahad Shahbaz and Khan, Salman},
  journal = {arXiv preprint arXiv:2608.15651},
  year    = {2026}
}

License

This repository is released under the Creative Commons Attribution-ShareAlike 4.0 International License. Third-party components remain subject to their original terms.

Please use CONTRIBUTING.md for bug reports and proposed extensions.