Add DSR reference checkpoints (baseline + in_context_36 ep80 EMA) and model card

README.md

@@ -0,0 +1,57 @@
+---
+pipeline_tag: unconditional-image-generation
+tags:
+- diffusion-transformer
+- imagenet
+- registers
+- outlier-tokens
+---
+
+# DSR — RAE-DiT-XL ImageNet-256 reference checkpoints
+
+Stage-2 diffusion-transformer checkpoints accompanying
+**[DSR: Dual-Stage Registers](https://github.com/a-little-hoof/DSR)**
+(*Taming Outlier Tokens in Diffusion Transformers*), released so the outlier-token
+visualizations can be reproduced without retraining.
+
+Both are EMA weights of a **DiT-XL** (depth 28, hidden 1152, 16 heads, patch 1,
+`in_channels` 768) operating in the SigLIP2-B RAE latent space (16×16×768,
+ImageNet-256 class-conditional), matching the DSR repo's
+`configs/stage2/sampling/ImageNet256/DiT-XL_SigLIP2-B_data_lognormal.yaml`.
+
+| File | Model | Training config (DSR repo) | Notes |
+|---|---|---|---|
+| `dit_xl_siglip2b_lognormal_ema.pt` | baseline, no registers | `DiT-XL_SigLIP2-B_data_lognormal` | full training run |
+| `dit_xl_siglip2b_lognormal_in_context_36_ep80_ema.pt` | + 36 in-context diffusion registers at block 8 | `*_data_lognormal-in_context_36` | **epoch-80 snapshot** (early): samples look under-trained; intended for comparing token-norm structure, where the register effect is already fully developed |
+
+## Usage
+
+The DSR notebook [`reproduce_outlier_tokens.ipynb`](https://github.com/a-little-hoof/DSR)
+downloads these automatically when no local checkpoint is found. Manual loading:
+
+```python
+import torch
+from huggingface_hub import hf_hub_download
+
+path = hf_hub_download("smallAI/DSR-DiT-XL-ImageNet256",
+                       "dit_xl_siglip2b_lognormal_ema.pt")
+state_dict = torch.load(path, map_location="cpu", weights_only=True)
+# build LightningDiT from the DSR repo's sampling config, then:
+# model.load_state_dict(state_dict, strict=True)
+```
+
+Both files are plain `state_dict`s (no optimizer state) and load with
+`torch.load(..., weights_only=True)`.
+
+## What they demonstrate
+
+Sampling the baseline while recording per-token L2 norms at every block shows outlier
+tokens — max/median patch-token norm rising to ~20× in blocks ~14–24, at fixed spatial
+positions shared across samples. In the register model the patch tokens stay near-flat
+and the 36 register tokens carry the outlier mass instead. See the DSR repository for
+the full method and the notebook that produces these figures.
+
+## License / intended use
+
+Released for research reproducibility. The models were trained on ImageNet-1k; use of
+these weights must comply with the ImageNet terms of access.

dit_xl_siglip2b_lognormal_ema.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:572276ef69ec48e1fb910289c8b48ac4b49a3d4db0f9c32bc4afc97f1d363261
+size 2709223719

dit_xl_siglip2b_lognormal_in_context_36_ep80_ema.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5062a0a801133d00bdb5156a73a61c59436217dbf52fc0a2e56bb4d4a5c020d
+size 2709763825