File size: 2,956 Bytes
b3b7d5f 1f75ef4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | ---
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.
## Stage-1 artifacts (for pixel-space decoding)
Mirrors of the stage-1 release linked from the DSR README (Google Drive), needed only
for the notebook's whole-pipeline figure (encode → noise → DiT → decode):
| File | What |
|---|---|
| `rae_decoder_vitxl_siglip2b_256.pt` | RAE ViT-XL decoder for the SigLIP2-B latent space (ImageNet-256) |
| `rae_stats_siglip2b_imagenet1k.pt` | Per-channel latent normalization statistics (`mean` / `var`) |
|