RAE: decode a frozen representation

An interactive field guide

RAE reuses a strong visual encoder and learns a decoder for its features. Its reconstruction decoder is a single-pass ViT.

Before this article: Semantic representations · DiT

What was already possible

A low-dimensional VAE learns its encoder mainly for reconstruction and regularization. Modern pretrained encoders already contain useful visual structure. RAE asks whether their features can serve as the image representation for diffusion generation, while keeping that pretrained structure intact.

What changes, and why

Freeze a DINOv2, SigLIP2, or MAE encoder. Train a ViT decoder to reconstruct RGB from its patch features. A second-stage DiT learns to generate those features. This separates three roles: feature extraction, feature-to-pixel reconstruction, and latent generation. The reconstruction decoder itself does not run diffusion.

Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.

1. Frozen feature extraction256² → 224² → [B,256,768]2. Single-pass ViT decoder[B,256,768] → [B,3,256,256]3. Separate latent generatornoise → generated feature grid

Frozen feature extraction

Keep patch features; discard class/register outputs.

with no_grad():
    z = encoder(resize(image,224)).patch_features

Single-pass ViT decoder

Project features, add fixed positions, decode patches, and unpatchify.

patches = decoder(project(z)+fixed_position)
image = unpatchify(pixel_head(patches))

Separate latent generator

Adapt model width and time schedule to the representation dimension.

generated_z = diffusion_sampler(latent_dit)
image = rae_decoder(generated_z)

Inspect the pinned source

$$\alpha=\sqrt{d/d_{\mathrm{ref}}},\qquad t^\prime=\frac{\alpha t}{1+(\alpha-1)t}$$

Work through the mechanism

In the DINOv2-B recipe, resize 256² input to 224² before the patch-14 encoder. The 16 × 16 feature grid has 768 channels. A patch-16 decoder returns 256² pixels. This is 196608 latent scalars, compared with 4096 in a 32 × 32 × 4 VAE grid. Equal Transformer token counts do not mean equal information budgets.

The schedule demo fixes 256 latent positions and varies channels C. It illustrates a dimension-dependent time shift relative to a 4096-scalar reference. At C=16 the dimensions match and the shift is identity. A time shift is one adaptation; it cannot by itself solve the modeling difficulty of a larger feature space.

Change a parameter, inspect the computed values, then compare this small example with the model described above.

Analytical teaching model · controls change the example, not a trained research model.

Latent channels C

768

Base time t × 100

50

Reset

The straight path has known velocity. A learned model must estimate velocity; this demo does not measure its error.

Computed values for the current example

Quantity	Value
α = √(256C / 4096)	6.928
base t	0.500
t	0.874
x_t = (1 − t) × 2 + t × (−1)	-0.622
v = ε − x	-3

Training is not sampling

The paper decoder recipe uses L1, LPIPS, and an adversarial term. Noise-augmented decoder training makes it tolerate imperfect generated latents. In the reported ablation this improves generation while worsening exact reconstruction. The released code was later refactored: its noise distribution and optimizer recipe must be recorded separately from the paper.

Reported evidence, not a local reproduction

Setting	Metric	Value	Interpretation
ImageNet-256; DINOv2-B / ViT-XL decoder	rFID ↓ / linear top-1 ↑	0.49 / 84.5%	Reconstruction and frozen-feature classification
Same decoder size; MAE-B encoder	rFID ↓ / linear top-1 ↑	0.16 / 68.0%	Better reconstruction does not rank semantics
Noise augmentation: off → on	gFID ↓; rFID ↓	4.81 → 4.28; 0.49 → 0.57	Section 4 controlled ablation, distinct from main-table ranking

What this does not establish

RAE inherits encoder pretraining cost and semantic behavior. Its high-dimensional continuous values are not a compact discrete bitstream. The lowest rFID in the table does not correspond to the highest linear-probe accuracy. Also distinguish fixed 2D sine-cosine positions in the reconstruction decoder from RoPE in the later generative backbone.

Check your understanding

Which encoder weights change during RAE decoder training?

None in the frozen-encoder recipe. Improvements in reconstruction come from learning the decoder, while encoder features remain fixed.

Sources and next steps

• RAE · 2510.11690v1
• Method / pinned implementation source
• AToken
• RAE-T2I

Return to the learning map

Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.