DINO-SAE: DINO Spherical Autoencoder for High-Fidelity Image Reconstruction and Generation

TL;DR

DINO-SAE addresses the trade-off between semantic fidelity and pixel-level reconstruction when using pretrained Vision Foundation Models as encoders. It introduces a Hierarchical Convolutional Patch Embedding and Cosine Similarity Alignment to preserve semantic direction while improving reconstruction quality, and models the latent space as a spherical manifold with Riemannian Flow Matching to enhance generation. On ImageNet-1K, it achieves state-of-the-art reconstruction metrics (0.37 rFID, 26.2 dB PSNR) and demonstrates faster, more stable diffusion training with latents derived from DINO-SAE (gFID 3.47 at 80 epochs). These results indicate that semantic-aligned, high-fidelity tokenizers can substantially improve both reconstruction and downstream generative performance for VFM-based autoencoders.

Abstract

Recent studies have explored using pretrained Vision Foundation Models (VFMs) such as DINO for generative autoencoders, showing strong generative performance. Unfortunately, existing approaches often suffer from limited reconstruction fidelity due to the loss of high-frequency details. In this work, we present the DINO Spherical Autoencoder (DINO-SAE), a framework that bridges semantic representation and pixel-level reconstruction. Our key insight is that semantic information in contrastive representations is primarily encoded in the direction of feature vectors, while forcing strict magnitude matching can hinder the encoder from preserving fine-grained details. To address this, we introduce Hierarchical Convolutional Patch Embedding module that enhances local structure and texture preservation, and Cosine Similarity Alignment objective that enforces semantic consistency while allowing flexible feature magnitudes for detail retention. Furthermore, leveraging the observation that SSL-based foundation model representations intrinsically lie on a hypersphere, we employ Riemannian Flow Matching to train a Diffusion Transformer (DiT) directly on this spherical latent manifold. Experiments on ImageNet-1K demonstrate that our approach achieves state-of-the-art reconstruction quality, reaching 0.37 rFID and 26.2 dB PSNR, while maintaining strong semantic alignment to the pretrained VFM. Notably, our Riemannian Flow Matching-based DiT exhibits efficient convergence, achieving a gFID of 3.47 at 80 epochs.

Figures (11)

• Figure 1: Generated images from DiT$^{DH}$ model trained on DINO-SAE latents
• Figure 2: DINO-SAE architecture and training loss.
• Figure 3: Linear Probing result of DINOv3 and DINO-SAE encoder. We compare the linear probing of feature from DINOv3 and our DINO-SAE encoder. With our cosine similarity loss, DINO-SAE achieves high linear probing results, deteriorating only within 3% on Top-1 accuracy compared to DINOv3's linear probing results
• Figure 4: PCA visualization of feature representations. We compare the principal components of feature maps from the Original Image, the Frozen Teacher (DINOv3), and the Student trained with our Cosine Similarity. Notably, our results demonstrate that the Cosine Similarity loss effectively preserves global semantic information, even as the patch embeddings are updated during training. This validates that directional alignment faithfully maintains the teacher's semantic structure and object boundaries.
• Figure 5: Visual comparison of reconstruction quality. Left: GT, Middle: RAE, Right: DINO-SAE.