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

TL;DR

DINO-SAE is a generative autoencoder that uses a frozen DINOv2 encoder as a semantic tokenizer and aligns to it with a cosine-similarity loss — releasing the magnitude axis so the encoder can preserve high-frequency detail — then trains a DiT on the resulting spherical latent manifold via Riemannian Flow Matching. On ImageNet-1K, the autoencoder hits 0.37 rFID / 26.2 dB PSNR and the manifold-aware DiT reaches gFID 3.47 in 80 epochs. The thesis is that VFM features naturally live on a hypersphere and generation should respect that geometry instead of working in ambient Euclidean space.

Key claims

Forcing the encoder to match VFM features in both direction and magnitude (standard L2 distillation) caps reconstruction fidelity because magnitude matching destroys high-frequency detail; a directional (cosine-similarity) objective preserves semantics while freeing magnitudes for detail retention [Abstract].
A Hierarchical Convolutional Patch Embedding stem replacing the standard ViT linear patchifier mitigates the patch-embedding information bottleneck and improves local structure / texture reconstruction [§Contributions, Abstract].
SSL/contrastive VFM representations (DINO, CLIP) empirically lie on a hyperspherical manifold rather than ℝ^d, so the natural generative dynamics are geodesic flows on the sphere rather than Euclidean velocity fields [Intro / geometric perspective].
Modeling the diffusion process via Riemannian Flow Matching (Chen & Lipman 2023) on the hypersphere — eliminating redundant radial variation and modeling only directional dynamics — yields more efficient generative training than Euclidean flow matching over the same latent [§Contributions].
On ImageNet-1K, DINO-SAE achieves rFID 0.37 and 26.2 dB PSNR, reported as state-of-the-art reconstruction at this scale while retaining semantic alignment to the pretrained VFM [Abstract].
The Riemannian-FM-trained DiT reaches gFID 3.47 in 80 epochs, which the authors frame as efficient convergence relative to comparable Euclidean DiT baselines on VFM latents [Abstract, Fig. 1].

Method

The autoencoder is built around a frozen DINOv2 backbone. The encoder is a small adaptor on top of DINO patch features whose output is supervised by a Cosine Similarity Alignment loss against the DINO target — only the direction of the latent vector needs to match, leaving magnitude free for the decoder to use as a detail-carrying channel. The patch-embedding stage that feeds DINO is replaced with a Hierarchical Convolutional Patch Embedding (HCPE) stem: a small convolutional pyramid that does multi-scale feature aggregation before tokenization, designed to retain local structure and texture that a single linear patch projection loses.

For the generative model, the authors observe that contrastive/SSL features cluster on the unit sphere (and that L2-normalizing the latent has a negligible effect on reconstruction). They therefore train a DiT on the spherical latent manifold using Riemannian Flow Matching: the prior is a distribution on S^{d-1}, the interpolation between noise and data follows the geodesic (great-circle) path rather than a straight line in ℝ^d, and the network regresses a tangent-space velocity field. At sampling time the ODE is integrated on the sphere via exponential maps. The combined system is referred to as DiT^{DH} (DiT with DINO-Hyperspherical latent).

Results

On ImageNet-1K reconstruction, DINO-SAE reports rFID 0.37 and PSNR 26.2 dB, positioned by the authors as SOTA among VFM-based generative autoencoders — addressing the standard complaint that DINO-tokenizer pipelines (RAE-style) trade reconstruction fidelity for semantic alignment.

For generation, the headline number is gFID 3.47 at 80 training epochs for the Riemannian-FM DiT on DINO-SAE latents. The framing is convergence efficiency: hitting that gFID quickly because the model is no longer wasting capacity on radial directions that the data doesn’t actually populate. The paper attributes the speed-up jointly to (a) cleaner latents from the cosine-similarity-aligned, HCPE-stemmed autoencoder, and (b) the geometry-matched RFM objective.

Quantitative comparisons against Euclidean-FM baselines and against RAE-style VFM autoencoders are reported in the full tables (only partially visible in the fetched excerpt); the qualitative claim is “matches or beats prior VFM-AE+DiT pipelines at a fraction of the epochs.”

Why it’s interesting

This sits at the exact intersection Jiaming flagged: it’s a representation autoencoder (in the RAE / DINO-tokenizer lineage that Improved Baselines with Representation Autoencoders just generalized) and a manifold flow-matching paper. The shared diagnosis with RAEv2 is the same — “matching the full VFM feature, magnitudes included, is the wrong objective” — but the proposed fixes are orthogonal: RAEv2 sums multiple encoder layers and bolts on REPA, while DINO-SAE drops magnitude matching entirely and changes the geometry of the generative process to match. If both diagnoses are correct, a natural follow-up is a multi-layer + cosine-aligned + RFM stack. For Luma the manifold-FM side is the more transferable bit: the same trick (geodesic flow on a constrained latent manifold rather than ambient Euclidean) should apply anywhere a learned latent is known to live on a lower-dimensional surface — e.g. normalized world-model state embeddings, camera-pose manifolds, or any tokenizer that already L2-normalizes its codes.