Update README.md

README.md

@@ -1,3 +1,100 @@
----
-license: mit
----
+---
+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},
+}
+```