One Layer Is Enough: Adapting Pretrained Visual Encoders for Image Generation

TL;DR

FAE (Feature Auto-Encoder) compresses frozen pretrained visual representations (DINOv2, SigLIP) into a low-dimensional generation-friendly latent using a deliberately minimal encoder: a single self-attention layer followed by a linear projection. A “double decoder” then separates feature reconstruction (lightweight 6-layer Transformer reconstructing the original embedding) from pixel synthesis (ViT-L pixel decoder trained on noise-augmented embeddings). On ImageNet 256×256, FAE-SiT-XL hits FID 1.29 with CFG / 1.48 without CFG at 800 epochs, and FID 1.70 / 2.08 at only 80 epochs — near-SOTA against RAE while keeping the generator architecture unchanged.

Key claims

A single self-attention layer + linear projection is the right encoder for adapting frozen visual representations to generation: a deep Transformer encoder overfits the (easy) feature-reconstruction objective and discards information, while a pure linear projection cannot de-redundantize patch-wise global signal [§4.1, Table 5].
The adaptation task (feature reconstruction from clean unmasked inputs) is substantially weaker than the original masked-prediction SSL task, so keeping the encoder shallow and the latents close to the pre-trained embeddings yields higher reconstruction quality [§4.1].
Separating feature reconstruction from pixel synthesis (“double decoder” design) preserves the frozen encoder’s semantics while giving the pixel decoder freedom for high-fidelity generation; the pixel decoder operates entirely in embedding space and can be reused across FAE variants without architectural changes [§4.2].
The Gaussian embedding decoder (pixel decoder trained on noise-injected frozen embeddings, σ ≈ scaled to embedding norm) is already a strong image decoder before being fine-tuned on FAE’s reconstructed embeddings, indicating the FAE latent preserves most of the pre-trained embedding’s information [§4.2 Pixel Fine-Tuning].
ImageNet 256×256 without CFG: FAE w/ timestep shift reaches FID 2.08 at 80 epochs and 1.48 at 800 epochs — SOTA at this no-CFG regime, beating RAE (DiT-XL) 1.87 and VA-VAE 2.17 [Table 1].
ImageNet 256×256 with CFG: FAE w/ timestep shift reaches FID 1.70 at 80 epochs and 1.29 at 800 epochs — near-SOTA, behind RAE (DiTDH-XL, 839M, 800ep) 1.13 but at smaller backbone (675M) [Table 1].
The recipe is generator-agnostic: FAE drops into both SiT (diffusion) and STARFlow (normalizing flow) with default parameterizations, simply replacing their native latent representation [§4.3].
FAE’s latent preserves DINOv2’s downstream understanding: ImageNet linear probe drops only from 87.00% (DINOv2-g/14) to 86.17% (FAE on DINOv2-g/14), and cross-image part-level patch matching is retained [§4.4, Table in §4.4].
On CC12M T2I (256², 400 epochs, T5-XL text encoder), FAE+SigLIP reaches FID 7.47 (no-CFG) / 6.90 (CFG) on MS-COCO — competitive with much larger / web-data models [§5.2].
The encoder design has implicit theoretical support: across DINOv2 feature-centering strategies (moving-average, Sinkhorn–Knopp) and unified tokenizers like UniTok, a recurring principle is removing redundant global information across patches — which is exactly what a single attention layer does and what a linear projection cannot [§4.1].

Method

FAE has three components: a single-layer feature encoder, a feature decoder, and a pixel decoder.

Encoder. A frozen pretrained ViT (DINOv2 or SigLIP) produces patch embeddings $x \in \mathbb{R}^{N \times D}$ . The encoder is one self-attention layer followed by a linear projection mapping to $z \in \mathbb{R}^{N \times d}$ with $d \ll D$ (e.g., 16×16×32 for ImageNet). The self-attention layer’s role is to remove redundant global information shared across patches; a pure linear projection cannot do this.

Feature decoder. A 6-layer Transformer with hidden dimension matching the pretrained backbone, RoPE + RMSNorm + SwiGLU, trained with simple L2 reconstruction + KL regularization to recover $\hat{x}$ from $z$ . The simple L2 loss (vs. VA-VAE’s more complex objectives) is the load-bearing choice that lets FAE’s reconstructed features zero-shot drop into downstream understanding pipelines.

Pixel decoder. A ViT-L decoder trained in two stages, both in embedding space: (1) inject Gaussian noise into the frozen pre-trained embeddings and train the decoder to reconstruct images (adversarial + LPIPS + L1 losses) — this “Gaussian embedding decoder” is robust to embedding perturbations; (2) fine-tune the same decoder on the feature decoder’s reconstructed embeddings $\hat{x}$ . Stage (1) alone already produces strong generation, evidence that FAE’s compressed latent preserves most pre-trained information.

Generative model. Standard SiT/LightningDiT or STARFlow trained directly on $z$ , with default hyperparameters. The decoupling of latent design from generator architecture is the explicit contrast against RAE, which requires wider DiT backbones, additional heads, and noise-augmented decoders to handle high-dimensional latents.

Results

ImageNet-256 class-conditional (675M backbone, batch 512):

Without CFG: 2.08 FID at 80 epochs, 1.48 at 800 epochs — beats RAE (DiT-XL, 676M) 1.87, VA-VAE 2.17, REPA 5.90 [Table 1].
With CFG: 1.70 at 80 epochs, 1.29 at 800 epochs — behind RAE (DiTDH-XL, 839M) 1.13, ahead of RAE (DiT-XL) 1.41 and VA-VAE 1.35 [Table 1].
Timestep shift gives consistent ~0.2 FID improvement at both 80ep and 800ep regimes.

T2I on CC12M (no web data, 604M, MS-COCO eval): FID 7.47 no-CFG / 6.90 CFG at 400 epochs, competitive with much larger web-data models like SDVAE+T5 (21.25) and Re-Imagen (6.88, 2.5B+750M). DINOv2 and SigLIP variants both work; SigLIP slightly better at this scale.

Semantic preservation: ImageNet linear probe 86.17% (FAE on DINOv2-g/14) vs 87.00% (DINOv2-g/14), 80.80% (DINOv2-S/14). Patch-wise similarity structure and cross-image part matching qualitatively preserved [§4.4, Figs. 4, 8-11].

Why it’s interesting

FAE is the third entry in the rapidly-coalescing 2025-late/2026-early “RAE family” of recipes that put frozen DINOv2/SigLIP at the heart of the diffusion latent space, and it sharpens the engineering tradeoff against the original RAE recipe in a concrete way. RAE (Diffusion Transformers with Representation Autoencoders) keeps the high-dimensional pretrained features and pays for it with architectural changes — wider DiT, the DiTDH head, dimension-aware noise shift, noise-augmented decoder. RAE T2I scaling (Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders) reports that two of those four ingredients become optional at scale. RAEv2 (Improved Baselines with Representation Autoencoders) keeps the same direction and stacks multi-layer aggregation + REPA-complementarity for 10× speedup. FAE takes the opposite engineering bet: compress to a low-dimensional latent so the existing generator runs unmodified, and put the burden on a careful double-decoder. The resulting numbers (FID 1.29 with CFG on a 675M backbone) are competitive with RAE’s 839M DiTDH-XL while keeping SiT/LightningDiT unchanged — this is the cleanest evidence yet that the RAE family’s architectural exotica may be substitutable by latent-side engineering. For Luma’s DiT recipes this matters because (a) FAE’s latent is shape-compatible with existing pipelines (16×16×32 ≈ SD-VAE shape) and (b) the same trained pixel decoder is reusable across encoder choices, lowering the swap-cost of trying DINOv2 vs SigLIP vs future SSL encoders. The negative-result framing in §4.1 — that a deep encoder hurts reconstruction because it overfits the easy adaptation task — is the kind of mechanistic-clarity result that should propagate downstream.