Update README.md

README.md

@@ -13,109 +13,129 @@ tags:
   - image-feature-extraction
 ---
 
-# STELLAR — Sparse Visual Representations via Spatial–Semantic Factorization
+# STELLAR: Learning Sparse Visual Representations via Spatial–Semantic Factorization
 
-**STELLAR** learns a **unified sparse visual representation** that supports both
-**reconstruction** and **semantics** using as few as **16 tokens**. By factorizing
-*"what"* (semantics) from *"where"* (spatial layout), each image is encoded as the
-low-rank product of a **localization** matrix and a **semantics** matrix.
+**How many tokens are needed for one image?**
+
+We show that with a **spatial–semantic factorized representation**, **16 semantic
+tokens** paired with explicit spatial maps are enough for strong **image recognition
+and reconstruction**. STELLAR separates **what** an image contains from **where** it
+appears, representing the image as the low-rank product of a localization matrix
+and a semantics matrix.
+
+STELLAR-H achieves **2.60 reconstruction FID** and **79.10% ImageNet linear-probing
+accuracy**, using approximately **90% fewer latent values** than a dense grid when
+counting both factors.
 
 <p align="center">
   <img src="factorization.svg" alt="Spatial–semantic factorization" width="720">
 </p>
 
-- 📄 **Paper:** [arXiv:2602.01905](https://arxiv.org/abs/2602.01905) (ICML 2026)
-- 💻 **Code:** [github.com/microsoft/STELLAR](https://github.com/microsoft/STELLAR)
+- **ICML 2026:** [Paper](https://openreview.net/pdf?id=ysOOfySED6) / [arXiv](https://arxiv.org/abs/2602.01905)
+- **Code and configs:** [microsoft/STELLAR](https://github.com/microsoft/STELLAR)
+- **Training and evaluation:** [Usage guide](https://github.com/microsoft/STELLAR/blob/main/docs/usage.md)
 
 These checkpoints contain the **full set of trained STELLAR modules** (encoder, sparse
 tokens, projections, reconstruction decoder, and clustering heads), so a single file
 supports feature extraction, image reconstruction, and continued pretraining. All
 models are self-supervised on **ImageNet-1K** at 224×224.
 
-## Highlights
-
-- **Sparse & unified** — one small set of tokens serves both high-level semantics
-  and pixel-level reconstruction.
-- **Factorized latents** — each token captures a concept (*what*) together with a
-  spatial map of *where* it appears.
-- **Strong on both axes** — STELLAR-H reaches **2.60 FID** (reconstruction) and
-  **79.1%** ImageNet linear-probing accuracy with just **16 tokens**.
+The headline results use the evaluation protocols in our ICML work. Reconstruction
+FID is **not generation FID**; loading a checkpoint or running a new probe does not
+automatically reproduce those numbers. The released files contain pretraining
+weights, not the separately finetuned B/L reconstruction probes.
 
 ## Available models
 
-| Model | Backbone | Tokens | Params | Type | File |
+| Model | Backbone | Semantic tokens | Feature width | Type | Weights |
 | :--- | :--- | :---: | :---: | :--- | :--- |
-| `stellar-b16` | ViT-B/16 | 16 | 88M | main | [`stellar-b16.safetensors`](stellar-b16.safetensors) |
-| `stellar-l16` | ViT-L/16 | 16 | 307M | main | [`stellar-l16.safetensors`](stellar-l16.safetensors) |
-| `stellar-h16` | ViT-H/14 | 16 | 636M | main | [`stellar-h16.safetensors`](stellar-h16.safetensors) |
-| `stellar-b8`  | ViT-B/16 | 8  | 88M | ablation | [`stellar-b8.safetensors`](stellar-b8.safetensors) |
-| `stellar-b24` | ViT-B/16 | 24 | 88M | ablation | [`stellar-b24.safetensors`](stellar-b24.safetensors) |
+| `stellar-b16` | ViT-B/16 | 16 | 768 | main | [safetensors](stellar-b16.safetensors) |
+| `stellar-l16` | ViT-L/16 | 16 | 1024 | main | [safetensors](stellar-l16.safetensors) |
+| `stellar-h16` | ViT-H/14 | 16 | 1280 | main | [safetensors](stellar-h16.safetensors) |
+| `stellar-b8` | ViT-B/16 | 8 | 768 | ablation | [safetensors](stellar-b8.safetensors) |
+| `stellar-b24` | ViT-B/16 | 24 | 768 | ablation | [safetensors](stellar-b24.safetensors) |
 
-The main models (`b16`, `l16`, `h16`) are recommended for downstream use; the 8- and
-24-token base models are ablations on the number of sparse tokens.
+Start with **B16** for a smaller backbone; use L16 or H16 when memory permits.
+The B8 and B24 variants study the number of sparse tokens. The token count in the
+model name is independent of the backbone patch size: H16 uses a ViT-H/14 encoder.
 
 ## Usage
 
-Install the STELLAR code and the Hub helpers:
+Use Python 3.10 and run from the GitHub code directory:
 
 ```bash
-pip install huggingface_hub safetensors
-git clone https://github.com/microsoft/STELLAR && cd STELLAR
-pip install -r requirements.txt
+git clone https://github.com/microsoft/STELLAR.git
+cd STELLAR
+python -m pip install -r requirements-inference.txt
 ```
 
+Feature extraction needs no Azure account, Olympus trainer, VQGAN weights, or
+separate MAE download. This is a custom PyTorch model; use the repository loader,
+not `transformers.AutoModel` or `pipeline()`.
+
 ### Quick start
 
-From the STELLAR code directory, use the [`load_stellar.py`](load_stellar.py) helper
-(it downloads the weights from the Hub for you):
+The [GitHub loader](https://github.com/microsoft/STELLAR/blob/main/load_stellar.py)
+downloads and validates the required weights. Replace `your_image.jpg` with a local image:
 
 ```python
 import torch
-from load_stellar import load_stellar, list_models
-
-print(list_models())                   # ['stellar-b16', 'stellar-l16', ...]
-model = load_stellar("stellar-b16")     # purpose="encode" (default)
+from load_stellar import load_stellar
+from examples.common import read_image
 
-# RGB image in [0, 1], resized to 224×224 (ImageNet normalization is applied internally)
-image = torch.rand(1, 3, 224, 224)
+revision = "6794be9a20fb9d3944c5bcc6003512a347fa518f"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = load_stellar("stellar-b16", revision=revision, device=device)
+image = read_image("your_image.jpg").unsqueeze(0).to(device)
 with torch.no_grad():
     out = model.encode(image)
 
-out["sparse"]    # (1, K, D)   sparse concept tokens   ("what")
-out["spatial"]   # (1, P, K)   per-token spatial maps  ("where")
-out["dense"]     # (1, P, D)   dense per-patch features
-out["cls"]       # (1, 1, D)   global image token
+print(out["sparse"].shape)
+print(out["spatial"].shape)
 ```
 
-### Reconstruction & continued pretraining
+For B16, these shapes are `(1, 16, 768)` and `(1, 196, 16)`. The helper converts
+to RGB, resizes the short side to 256 and center-crops to 224. Inputs must be
+floating-point RGB tensors in `[0, 1]`; **do not apply ImageNet normalization**,
+which is already performed inside the model.
+
+Pin the Hub revision for repeatable downloads. After caching the weights,
+`local_files_only=True` enables offline loading. Set `HF_HUB_CACHE` before starting
+Python to choose a cache directory. Allow roughly 0.5 GB per B checkpoint, 1.4 GB
+for L16, and 2.7 GB for H16, plus the tokenizer if reconstructing.
 
-The same checkpoint can be loaded for other purposes via the `purpose` argument. Image
-reconstruction and continued pretraining use the decoder, which predicts
-[MaskGIT-VQGAN](https://huggingface.co/fun-research/TiTok) tokens — pass the tokenizer
-path as `vq_model`:
+### Reconstruction
+
+Continuing from the quick start, download the external
+[MaskGIT-VQGAN tokenizer](https://huggingface.co/fun-research/TiTok) and load the decoder:
 
 ```python
-# 1. encode -> factorized features (sparse concept tokens + spatial maps)
-model = load_stellar("stellar-b16", purpose="reconstruct", vq_model=VQGAN_PATH)
-features = model.encode(image)          # dict: sparse (B,K,D), spatial (B,P,K), ...
-
-# 2. decode the factorized features -> VQGAN decoder -> pixels
-out = model.reconstruct(features)       # or model.reconstruct(features["sparse"], features["spatial"])
-pixels = out["reconstruction"]          # (B, 3, H, W) RGB in [0, 1]
-#                                         224x224 for /16 models, 256x256 for the /14 H model
-#   out["tokens"] : (B, P) predicted VQGAN token ids
-#   out["logits"] : (B, P, 1024) raw codebook logits
-
-# continued pretraining (all modules, gradients enabled)
-model = load_stellar("stellar-b16", purpose="pretrain", vq_model=VQGAN_PATH)
-losses = model({"image": image, "labels": labels, ...})["predictions"]
-```
+from huggingface_hub import hf_hub_download
+from torchvision.transforms.functional import to_pil_image
 
-`reconstruct` is the **decoder half** of STELLAR: it takes the factorized features and
-runs low-rank dense map → ViT decoder → VQGAN decoder to return RGB pixels. See
-[`examples/reconstruction.ipynb`](examples/reconstruction.ipynb) for an end-to-end demo
-that loads an image and displays the reconstruction.
+vq_path = hf_hub_download(
+  "fun-research/TiTok",
+  "maskgit-vqgan-imagenet-f16-256.bin",
+  revision="ab646ed225080a3acb7c78440a574d7f67f16fa7",
+)
+model = load_stellar(
+  "stellar-b16", purpose="reconstruct", vq_model=vq_path,
+  revision=revision, device=device,
+)
+with torch.no_grad():
+  features = model.encode(image)
+  reconstruction = model.reconstruct(features)
+pixels = reconstruction["reconstruction"]
+to_pil_image(pixels[0].cpu()).save("reconstruction.png")
+```
 
+`reconstruct` takes **features, not an image**. You can also pass
+`model.reconstruct(features["sparse"], features["spatial"])`. Both factors are
+needed: semantic tokens alone do not specify spatial layout. B/L return RGB
+224x224 pixels; H returns 256x256 pixels. Outputs are in `[0, 1]`, with VQ token
+IDs in `reconstruction["tokens"]` and logits in `reconstruction["logits"]`.
+VQ decoding uses argmax and a frozen tokenizer, so this is not a differentiable
+RGB decoder for end-to-end pixel losses.
 
 ### What the model returns
 
@@ -130,39 +150,38 @@ that loads an image and displays the reconstruction.
 `B` = batch, `K` = number of sparse tokens, `P` = number of patches (196 for /16 at
 224², 256 for /14), `D` = embedding dim (768 / 1024 / 1280 for B / L / H).
 
-### Loading the weights manually
+## Training and evaluation
 
-```python
-import json, torch
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
-from src.models.stellar_model import STELLARModel
-
-repo = "microsoft/STELLAR"
-cfg = json.load(open(hf_hub_download(repo, "config.json")))["models"]["stellar-b16"]
-state = load_file(hf_hub_download(repo, cfg["weights"]))
-
-model = STELLARModel(
-    num_sparse_tokens=cfg["num_sparse_tokens"],
-    num_decoder_layers=cfg["num_decoder_layers"],
-    spatial_temp=cfg["spatial_temp"],
-    vit_pretrained=cfg["backbone"],
-    do_recon=False, do_clustering=False, vq_model=None,
-)
-model.load_state_dict(state, strict=False)   # encoder-only build ignores decoder/heads
-model.eval()
-features = model.encode(torch.rand(1, 3, 224, 224))
-```
+The GitHub repository provides image-folder training, checkpoint resume, and
+classification, segmentation and reconstruction probes. Install the training
+dependencies and prepare your data using the
+[usage guide](https://github.com/microsoft/STELLAR/blob/main/docs/usage.md).
 
-> **Tip:** download with `huggingface_hub` (as above) rather than `git clone` so that
-> downloads are registered on the Hub — `git clone` is not counted in download stats.
+| Goal | Configuration or command |
+| :--- | :--- |
+| Pretrain on ImageNet or custom images | [configs/stellar.yaml](https://github.com/microsoft/STELLAR/blob/main/configs/stellar.yaml) |
+| Train a classification probe | [configs/eval_cls.yaml](https://github.com/microsoft/STELLAR/blob/main/configs/eval_cls.yaml) |
+| Train a segmentation probe | [configs/eval_seg.yaml](https://github.com/microsoft/STELLAR/blob/main/configs/eval_seg.yaml) |
+| Train a reconstruction probe | [configs/eval_recon.yaml](https://github.com/microsoft/STELLAR/blob/main/configs/eval_recon.yaml) |
+| Measure released-decoder rFID/LPIPS | [Reconstruction metrics](https://github.com/microsoft/STELLAR/blob/main/docs/usage.md#reconstruction-metrics) |
+
+See the [configuration guide](https://github.com/microsoft/STELLAR#configuration-guide)
+for data paths, batch size, devices, learning rate and model settings. Evaluation
+recipes **train new heads on frozen features**; testing requires a trained head
+checkpoint. The default pretraining recipe starts from MAE, not these Hub weights.
+
+For a custom continued-pretraining loop, `load_stellar(..., purpose="pretrain",
+vq_model=...)` loads all trained modules in train mode. Supply the multi-crop batch
+contract described in the usage guide. This is **weights-only initialization**,
+not an optimizer resume; use a full Lightning checkpoint with `scratch.resume`
+to resume an interrupted training run.
 
 ## Model details
 
 - **Architecture:** ViT encoder (MAE-initialized) + learned sparse latent queries with
   spatial–semantic factorization.
 - **Pretraining data:** ImageNet-1K (self-supervised; labels not used).
-- **Input:** RGB images in `[0, 1]`, resized to 224×224 (bicubic). ImageNet mean/std
+- **Input:** RGB images in `[0, 1]` at 224x224. ImageNet mean/std
   normalization is applied **inside** the model — pass raw `[0, 1]` images.
 - **Weights:** the complete set of trained STELLAR modules (encoder, sparse tokens,
   projections, reconstruction decoder, and clustering heads), stored in `safetensors`.
@@ -172,13 +191,28 @@ features = model.encode(torch.rand(1, 3, 224, 224))
 
 ## Intended uses & limitations
 
-- **Intended use:** extracting compact sparse/dense visual features for downstream
-  recognition, segmentation, retrieval, reconstruction, and analysis.
-- **Limitations:** pretrained on ImageNet-1K at 224×224, so features reflect that
-  distribution; performance on very different domains (e.g. medical, satellite) may
-  require fine-tuning. The models are research artifacts and are not safety-tested for
+- **Supported research uses:** compact visual features for recognition,
+  segmentation, reconstruction, retrieval experiments and representation analysis.
+- **Compression is not encoder acceleration:** the ViT still processes dense
+  patches. Fewer output latent values do not imply a 90% encoder or VLM speedup.
+- **Tokens are learned concepts:** they are not guaranteed objects, class IDs,
+  temporal identities or pretrained language-aligned embeddings.
+- **Domain shift:** ImageNet-pretrained features may require adaptation and
+  held-out evaluation on medical, satellite or other substantially different data.
+- **Deployment:** these are research artifacts, not safety-tested models for
   production decision-making.
 
+### Connections to VLMs and latent generation
+
+Sparse semantic tokens provide a compact input to a learned language-model
+adapter. Appendix B.3 evaluates alignment to a CLIP text tower using a trained
+probe; **no pretrained VLM adapter or captioning system is released**.
+
+The paired semantic and spatial factors may also be studied as latents for
+representation autoencoders (RAEs) or latent diffusion. **We have not evaluated
+RAE-style generation and make no generation-quality claims.** Reconstruction
+results do not establish unconditional or text-to-image generation performance.
+
 ## Citation
 
 ```bibtex
@@ -187,10 +221,10 @@ features = model.encode(torch.rand(1, 3, 224, 224))
   author    = {Zhao, Theodore Zhengde and Kiblawi, Sid and Yang, Jianwei and Usuyama, Naoto and Tan, Reuben and Codella, Noel C and Naumann, Tristan and Poon, Hoifung and Wei, Mu},
   booktitle = {International Conference on Machine Learning (ICML)},
   year      = {2026},
-  url       = {https://arxiv.org/abs/2602.01905},
+  url       = {https://openreview.net/pdf?id=ysOOfySED6},
 }
 ```
 
 ## License
 
-Released under the [MIT License](LICENSE).
+Released under the [MIT License](LICENSE).