---
license: mit
tags:
  - jepa
  - vicreg
  - vit3d
  - physics
  - self-supervised
  - representation-learning
datasets:
  - polymathic-ai/active_matter
library_name: pytorch
---

# ViT3D-d6 / VICReg / FFT — `active_matter` (epoch 29)

A 6-block 3D Vision Transformer pretrained with VICReg in a JEPA-style setup
on the [`active_matter`](https://polymathic-ai.org/the_well/datasets/active_matter/)
dataset from The Well. This checkpoint is the encoder weights at pretrain
epoch 29 — the best-validation epoch in our sweep.

The encoder produces a frozen `(B, 384, 16, 16)` feature map from a
`(B, 11, 16, 256, 256)` input. Linear and k-NN probes on top of those frozen
features regress the active-matter parameters $\alpha$ (alignment strength)
and $\zeta$ (active stress).

## Architecture

| Component | Spec |
|---|---|
| Input | $(B, 11, 16, 256, 256)$ |
| 3D PatchEmbed | `Conv3d(11 → 384, kernel=stride=4×16×16)` |
| Tokens | $T'{\times}H'{\times}W' = 4{\times}16{\times}16 = 1024$ |
| Transformer blocks | 6 × pre-norm, $h{=}6$, MLP ratio 4, QKV bias |
| Pos. embedding | learnable, $1024 \times 384$ |
| Output | $(B, 384, 16, 16)$ — time collapsed by mean over $T'$ |
| **Total params** | **≈ 15.4 M** |

The companion `ConvPredictor` used during JEPA pretraining (≈ 7.1 M) is
**not** included — only the encoder is published.

## Training recipe

| | |
|---|---|
| Objective | JEPA + VICReg ($\lambda_{\text{sim}}{=}2,\ \lambda_{\text{std}}{=}40,\ \lambda_{\text{cov}}{=}2$) |
| Preprocessing | band-limited FFT resize (preserves periodic BCs) |
| Optimiser | AdamW, lr $5\!\times\!10^{-4}$, wd $0.05$, cosine, 3 warmup epochs |
| Precision | bf16 |
| Batch size | 4 |
| Epochs | 30 (this checkpoint = epoch 29) |
| Hardware | single A100 |

## Evaluation — frozen-encoder probes on held-out test split

Lower MSE is better.

| Probe | mean MSE $\downarrow$ | $\alpha$ MSE $\downarrow$ | $\zeta$ MSE $\downarrow$ | $k$ / metric |
|---|---|---|---|---|
| Linear | **0.107** | **0.016** | **0.197** | — |
| k-NN   | **0.120** | **0.009** | **0.231** | $k=20$, cosine |

Every CNN variant we tried (VICReg baseline, VICReg+FFT, Conv+Attn,
Conv+Attn×6) lands at linear MSE 0.22–0.27 on the same data; the 3D-patch
tokeniser is the change that unlocks the ~3× improvement.

## Usage

```python
import torch
from huggingface_hub import hf_hub_download
from physics_jepa.utils.model_utils import ViT3DEncoder

ckpt_path = hf_hub_download(
    repo_id="szcharlesji/vit3d-d6-active-matter",
    filename="ViT3DEncoder_29.pth",
)
encoder = ViT3DEncoder(
    in_chans=11, num_frames=16, img_size=(256, 256),
    patch_size=(4, 16, 16), embed_dim=384, depth=6, num_heads=6,
    mlp_ratio=4.0,
)
encoder.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
encoder.eval()

# x: (B, 11, 16, 256, 256), float
with torch.no_grad():
    feat = encoder(x)   # (B, 384, 16, 16)
```


## Citation

```bibtex
@misc{ji2026jepa-active-matter,
  author = {Charles Cheng Ji, Zhanhe Shi, Richard Wang, Romina Yalovetzky},
  title  = {Physics-Aware Representation Learning for Physical Systems},
  year   = {2026},
}
```