Update README.md

README.md

@@ -1,3 +1,74 @@
----
-license: mit
----
+---
+license: mit
+tags:
+  - self-supervised-learning
+  - world-models
+  - equivariance
+  - vision
+  - pytorch
+datasets:
+  - 3DIEBench
+  - STL10
+---
+
+# seq-JEPA: Autoregressive Predictive Learning of Invariant-Equivariant World Models
+
+<p align="center">
+  <a href="https://openreview.net/forum?id=GKt3VRaCU1"><img src="https://img.shields.io/badge/NeurIPS%202025-Paper-blue" alt="Paper"></a>
+  <a href="https://hafezgh.github.io/seq-jepa/"><img src="https://img.shields.io/badge/Project-Page-green" alt="Project Page"></a>
+  <a href="https://github.com/hafezgh/seq-jepa"><img src="https://img.shields.io/badge/GitHub-Code-black" alt="Code"></a>
+</p>
+
+## Model Description
+
+By processing views sequentially with action conditioning, seq-JEPA naturally segregates representations for equivariance- and invariance-demanding tasks.
+
+## Available Checkpoints
+
+| Checkpoint | Dataset | Training | Download |
+|------------|---------|----------|----------|
+| `3diebench_rot_seqlen3.pth` | 3DIEBench | seq-len=3, rotation conditioning | [Download](https://huggingface.co/Hafez/seq-JEPA/resolve/main/3diebench_rot_seqlen3.pth) |
+| `stl10_pls.pth` | STL10 | PLS (predictive learning across saccades) | [Download](https://huggingface.co/Hafez/seq-JEPA/resolve/main/stl10_pls.pth) |
+
+## Usage
+
+First, clone the repository to access model definitions:
+
+git clone https://github.com/hafezgh/seq-jepa.git
+cd seq-jepaThen load the checkpoints:
+
+import torch
+from models import SeqJEPA_Transforms, SeqJEPA_PLS
+
+# 3DIEBench checkpoint
+kwargs = {
+    "num_heads": 4, "n_channels": 3, "num_enc_layers": 3,
+    "num_classes": 55, "act_cond": True, "pred_hidden": 1024,
+    "act_projdim": 128, "act_latentdim": 4, "cifar_resnet": False,
+    "learn_act_emb": True
+}
+model = SeqJEPA_Transforms(img_size=128, ema=True, ema_decay=0.996, **kwargs)
+ckpt = torch.load('3diebench_rot_seqlen3.pth')
+model.load_state_dict(ckpt['model_state_dict'])
+
+# STL10 PLS checkpoint  
+kwargs = {
+    "num_heads": 4, "n_channels": 3, "num_enc_layers": 3,
+    "num_classes": 10, "act_cond": True, "pred_hidden": 1024,
+    "act_projdim": 128, "act_latentdim": 2, "cifar_resnet": True,
+    "learn_act_emb": True, "pos_dim": 2
+}
+model = SeqJEPA_PLS(fovea_size=32, img_size=96, ema=True, ema_decay=0.996, **kwargs)
+ckpt = torch.load('stl10_pls.pth')
+model.load_state_dict(ckpt['model_state_dict'])## Citation
+
+## Citation
+
+@inproceedings{
+ghaemi2025seqjepa,
+title={seq-{JEPA}: Autoregressive Predictive Learning of Invariant-Equivariant World Models},
+author={Hafez Ghaemi and Eilif Benjamin Muller and Shahab Bakhtiari},
+booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
+year={2025},
+url={https://openreview.net/forum?id=GKt3VRaCU1}
+}