README.md

---
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`) |