Source-linked AI summary
One Layer Is Enough: Adapting Pretrained Visual Encoders for Image Generation
Yuan Gao, Chen Chen, Tianrong Chen, Jiatao Gu
TL;DR
Adapting high-dimensional pretrained visual representations to generation is difficult because understanding and generative models favor different latent properties. FAE uses a single-attention compression module with separate reconstruction and generation decoders, achieving strong results across generative settings, including ImageNet FID 1.48 without CFG after 800 epochs. Its reconstruction fidelity remains below methods that explicitly optimize image reconstruction.
Problem
Pretrained representation encoders favor high-dimensional features, while generative models require compact latents that preserve noise and fine-scale image information.
Method
FAE compresses pretrained features with a single attention layer and uses separate decoders for feature reconstruction and image generation.
Results
Across class-conditional and text-to-image benchmarks, FAE performs strongly; without CFG it achieves state-of-the-art ImageNet FID 1.48 after 800 epochs and 2.08 after 80 epochs.
Takeaways & Limitations
FAE provides a simple, general mechanism for adapting pretrained vision encoders to diffusion and normalizing-flow generation.
Takeaways & Limitations
Because its encoder lacks an explicit image reconstruction loss, FAE has worse rFID and tokenizer fidelity than methods such as VA-VAE.
Abstract
from arXiv · showhide
Visual generative models (e.g., diffusion models) typically operate in compressed latent spaces to balance training efficiency and sample quality. In parallel, there has been growing interest in leveraging high-quality pre-trained visual representations, either by aligning them inside VAEs or directly within the generative model. However, adapting such representations remains challenging due to fundamental mismatches between understanding-oriented features and generation-friendly latent spaces. Representation encoders benefit from high-dimensional latents that capture diverse hypotheses for masked regions, whereas generative models favor low-dimensional latents that must faithfully preserve injected noise. This discrepancy has led prior work to rely on complex objectives and architectures. In this work, we propose FAE (Feature Auto-Encoder), a simple yet effective framework that adapts pre-trained visual representations into low-dimensional latents suitable for generation using as little as a single attention layer, while retaining sufficient information for both reconstruction and understanding. The key is to couple two separate deep decoders: one trained to reconstruct the original feature space, and a second that takes the reconstructed features as input for image generation. FAE is generic; it can be instantiated with a variety of self-supervised encoders (e.g., DINO, SigLIP) and plugged into two distinct generative families: diffusion models and normalizing flows. Across class-conditional and text-to-image benchmarks, FAE achieves strong performance. For example, on ImageNet 256x256, our diffusion model with CFG attains a near state-of-the-art FID of 1.29 (800 epochs) and 1.70 (80 epochs). Without CFG, FAE reaches the state-of-the-art FID of 1.48 (800 epochs) and 2.08 (80 epochs), demonstrating both high quality and fast learning.
1 Introduction
Pre-trained visual representations improve generative modeling but are difficult to adapt because understanding and generation favor different latent structures. FAE addresses this mismatch with a minimal feature-compression design and achieves strong ImageNet generation quality with fast convergence.
- Motivation: Pre-trained representations improve training efficiency and generative fidelity when incorporated into diffusion pipelines or VAEs.
- Motivation: Understanding-oriented encoders require high-dimensional latents for masked-region prediction, whereas generative models favor spaces that preserve injected noise during denoising.
- Method: FAE compresses pre-trained embeddings with a single attention layer and linear projection, then reconstructs the original features with a lightweight decoder.
- Evaluation: FAE is integrated with diffusion frameworks including SiT and LightningDiT, as well as normalizing-flow models such as STARFlow.
- Results: FAE reaches FID 1.48 after 800 epochs and 2.08 after 80 epochs without CFG on ImageNet 256×256, achieving state-of-the-art results.With CFG, it reaches FID 1.29 after 800 epochs and 1.70 after 80 epochs.
2 Related Work
Prior work adapts visual representations mainly for discriminative tasks or bridges them to generation through alignment losses and architectural changes. FAE instead reuses pretrained visual encoders through lightweight internal adaptation for generative modeling.
- Visual Representation Learning: Self-supervised visual adaptation has predominantly targeted classification, segmentation, and retrieval rather than generative modeling.
- Visual Generative Models: Diffusion models commonly generate in compressed VAE latent spaces, while transformer and normalizing-flow approaches provide alternative generative architectures.
- Representation Alignment: REPA and related methods align diffusion features with pretrained image embeddings, while VA-VAE performs alignment within a VAE.
- Representation Alignment: FAE internalizes alignment with a single attention layer rather than relying on external alignment losses or auxiliary projection heads.
- Representation Alignment: Unlike approaches that directly modify the latent space, FAE repurposes pretrained ViTs as the generative backbone through lightweight adaptation.
3 Motivation
Pre-trained representation features are high-dimensional and semantically rich, whereas generative latents must support stable denoising and fine-detail recovery. FAE therefore compresses them into a compact latent space while keeping the representation and generation interfaces compatible.
- 3 Motivation: High-dimensional self-supervised features can destabilize generative denoising because generation must recover fine details while preserving injected noise.Dino-V2-G, for example, has feature dimension 1,536.
- 3 Motivation: Directly using pretrained embeddings as diffusion latents requires wider channels and additional generator heads, coupling architecture design to encoder dimensionality.
- 3 Motivation: FAE seeks a low-dimensional generative latent space that can reuse existing architectures while remaining close to the pretrained feature space.
- 3 Motivation: FAE introduces a lightweight encoder that compresses high-dimensional representation features into a compact latent space tailored for generation.
4 Method
FAE compresses pretrained visual embeddings into compact latents with a minimal attention-based encoder, then separates feature reconstruction from pixel synthesis. This design preserves semantic structure while supporting efficient generation.
- Feature Encoder: FAE uses a single self-attention layer and linear projection to compress pretrained patch embeddings into low-dimensional latents.The minimal encoder keeps the mapping close to the original feature space while reducing parameters.
- Ablation: The single-attention encoder outperforms both a purely linear encoder and a deep Transformer encoder in reconstruction quality.This ablation supports attention as an effective way to remove redundant information across patches.
- Double Decoder: The feature decoder reconstructs pretrained embeddings, while a separate pixel decoder converts those reconstructed features into images.This double-decoder design preserves encoder semantics while giving pixel synthesis a dedicated pathway.
- Feature Decoder: Feature reconstruction uses a standard VAE objective with L2 reconstruction and KL regularization to keep compact latents near a simple prior.The reconstructed embeddings remain compatible with downstream tasks trained on the original representations.
- Pixel Decoder: The pixel decoder is trained with adversarial, perceptual, and reconstruction losses that target realism, semantic fidelity, and low-level detail.It is first trained on noisy pretrained embeddings and then fine-tuned on reconstructed embeddings.
- Semantic Preservation: FAE preserves pretrained semantic geometry: similar patches remain close, and corresponding parts across images remain matchable in the latent space.The framework applies to representations such as DINOv2 and SigLIP without architectural changes.
5 Experiments
FAE performs strongly across class-conditional, text-to-image, and latent-flow generation settings. It reaches strong ImageNet and COCO FID scores while learning efficiently and using relatively limited training data for text-to-image generation.
- Evaluation Setup: FAE is evaluated on ImageNet class-conditional generation and CC12M-trained text-to-image generation, with additional tests of semantic understanding.The study also examines linear probing, text-image retrieval, and applicability to normalizing flows.
- Class-conditional Image Generation: FID 1.48 without CFG after 800 epochs and 2.08 after 80 epochs establish state-of-the-art ImageNet generation performance.With CFG, FID reaches 1.29 after 800 epochs and 1.70 after 80 epochs.
- Text-to-Image Generation: FID 6.90 with CFG and 7.47 without CFG are achieved on MS-COCO after 400 epochs using only CC12M for pretraining.The results are described as near state of the art despite substantially less training data than typical web-scale systems.
- Text-to-Image Generation: FAE produces visually coherent text-to-image samples that follow short prompts, indicating good text–image alignment.These samples use a SigLIP2-FAE backbone and a 2B-parameter MMDiT decoder at 384 × 384 resolution.
- Latent Normalizing Flows: The FAE-based STARFlow variant achieves FID 2.67 versus 4.51 for the SD-VAE baseline after 400 training epochs.The FAE variant also converges substantially faster in guided and unguided settings.
6 Ablation Study
The ablations show that a single attention layer offers the best balance of generation quality, convergence, and feature reconstruction, while diffusion-side design choices further improve performance.
- FAE Model Structure: The single-attention encoder outperforms linear and 4/6-layer transformer variants in FID while matching deeper transformers on embedding reconstruction.It also substantially exceeds the linear baseline in reconstruction similarity.
- LDM Model Structure: Combining SwiGLU, ROPE, and RMSNorm in the latent diffusion model improves convergence speed and generation quality, with the largest gains when all three are included.The resulting structure is equivalent to LightningDiT.
- Token Dimension: 32-dimensional tokens achieve the best final FID and IS scores and the fastest convergence despite 64 dimensions yielding lower rFID.Time-shift substantially narrows the performance gap between token dimensions.
- Time Shift: Time-shift accelerates convergence and bridges quality differences across latent dimensions, reaching state-of-the-art convergence speed and generation quality within 64 epochs.The ablation changes loss weighting and the diffusion trajectory.
7 Conclusion
FAE adapts pretrained visual representations for generation with a single attention layer and two lightweight decoders, achieving strong and efficient results across benchmarks. Its main limitation is weaker reconstruction fidelity than methods explicitly optimized for image reconstruction.
- 7 Conclusion: FAE uses a single attention layer with two lightweight decoders and achieves strong, efficient performance across class-conditional and text-to-image benchmarks.On ImageNet 256×256, it reports FID 1.29 with CFG after 800 epochs and 1.48 without CFG.
- 7 Conclusion: Without CFG, FAE achieves state-of-the-art FID of 1.48 after 800 epochs and 2.08 after 80 epochs.These results indicate both high sample quality and fast learning behavior.
- 7 Conclusion: FAE’s rFID and tokenizer fidelity lag behind VA-VAE because its encoder lacks an explicit image reconstruction loss.VA-VAE directly optimizes reconstruction quality.
A FAE Encoder Structure
The FAE encoder merges linear layers in its attention module and uses a larger per-head dimension; its structure is presented in Figure 7.
- A FAE Encoder Structure: The encoder merges linear layers in the attention module and uses a larger per-head dimension.The resulting encoder structure is shown in Figure 7.
B Ablation on FlowMatching Timesteps Shift
The flow-matching timestep-shift ablation shows faster convergence, reduced differences across latent token dimensions, and a modest final FID improvement.
- B Ablation on FlowMatching Timesteps Shift: Timestep shift accelerates convergence, reduces discrepancies across latent token dimensions, and modestly improves final FID.The result is reported for the flow-matching timestep-shift ablation.
C rFID
FAE’s disentangled encoder training preserves its representation-learning setup but yields weaker reconstruction fidelity than methods directly optimizing reconstruction. The accompanying material also describes how patch-similarity visualizations compare DINOv2 embeddings with FAE latents.
- C rFID: FAE’s rFID and tokenizer reconstruction fidelity lag behind VA-VAE because encoder training is disentangled from image reconstruction loss.VA-VAE directly optimizes reconstruction quality, unlike FAE’s training setup.
- C rFID: Similarity visualizations compare DINOv2 embedding maps with corresponding FAE-latent maps using the original query patch as reference.Darker colors indicate higher similarity, and query patches are marked with red rectangles.
- C rFID: The visual analysis includes similarity examples for cat and impala images.These examples are presented as figures accompanying the reconstruction-fidelity discussion.
F Matching most Similar patch pair across two images
FAE retains DINOv2-like semantic patch matching across different images. The evaluation selects animal-related patches, matches them by cosine similarity, and visualizes 16 patch pairs per example.
- F Matching most Similar patch pair across two images: FAE latents retain DINOv2’s cross-image patch-matching behavior, reliably matching semantically corresponding regions.The result suggests preservation of fine-grained, part-level semantics rather than only coarse global information.
- F Matching most Similar patch pair across two images: Animal-related patches are identified with K-Means, then each selected patch is matched to the highest-cosine-similarity patch in the second image.Each example visualizes 16 selected patch pairs.
- F Matching most Similar patch pair across two images: The bird and elephant figures illustrate matching the most similar patch pairs across two photographs.These figures provide cross-image examples of the patch-matching procedure.