Geometric Autoencoder for Diffusion Models

Hangyu Liu; Jianyong Wang; Yutao Sun

cs.CVMar 11, 2026

Geometric Autoencoder for Diffusion Models

Authors: Hangyu Liu, Jianyong Wang, Yutao Sun

Organizations: Shanghai Innovation Institute · Tsinghua University

Abstract

Latent diffusion models have established a new state-of-the-art in high-resolution visual generation. Integrating Vision Foundation Model priors improves generative efficiency, yet existing latent designs remain largely heuristic. These approaches often struggle to unify semantic discriminability, reconstruction fidelity, and latent compactness. In this paper, we propose Geometric Autoencoder (GAE), a principled framework that systematically addresses these challenges. By analyzing various alignment paradigms, GAE constructs an optimized low-dimensional semantic supervision target from VFMs to provide guidance for the autoencoder. Furthermore, we leverage latent normalization that replaces the restrictive KL-divergence of standard VAEs, enabling a more stable latent manifold specifically optimized for diffusion learning. To ensure robust reconstruction under high-intensity noise, GAE incorporates a dynamic noise sampling mechanism. Empirically, GAE achieves compelling performance on the ImageNet-1K $256 \times 256$ benchmark, reaching a gFID of 1.82 at only 80 epochs and 1.31 at 800 epochs without Classifier-Free Guidance, significantly surpassing existing state-of-the-art methods. Beyond generative quality, GAE establishes a superior equilibrium between compression, semantic depth and robust reconstruction stability. These results validate our design considerations, offering a promising paradigm for latent diffusion modeling. Code and models are publicly available at https://github.com/sii-research/GAE.

Figures & tables

Panel of Figure 1 — Figure 1: Overview of GAE performance. Left: GAE establishes a superior Pareto curve between linear probing accuracy and latent dimension. Right: GAE with 32 latent dimension significantly accelerates convergence and delivers superior generation results at both 80 and 800 training epochs.

Figure 2: Overview of the GAE architecture. The input image is processed through a pixel-level branch ( $E_{p}$ , $A_{p}$ ) and a frozen semantic branch (VFM, $E_{sp}$ ). A Semantic Preservation loss $L_{sp}$ aligns the latent mean with the features from the decoupled Semantic Teacher.

Figure 3: Illustration of the three latent alignment paradigms: The term $L_{sp}$ aligns the AE representations with those of the VFM. Pre Alignment aligns high-dimensional encoder features directly; Post Alignment projects AE latents back to a high-dimensional space for supervision; Latent Alignment operates within the compressed latent bottleneck via a projection of VFM features.

Metric	Autoencoder			Teacher
Metric	Pre	Post	Latent	SVD	DinoL
rFID $\downarrow$	0.40	0.48	0.51	–	–
LP $\uparrow$	20.9	60.8	63.2	60.9	83.7

Table 1: Pilot Study of alignment paradigms on ImageNet.

Arch	LP $\uparrow$
Single Attn	62.8
Attn + Linear	63.4
Attn + Patch Conv	75.6

Table 2: Ablation of Downsampler in Semantic Teacher.

Method	Epochs	#Params	Generation w/o CFG				Generation w/ CFG
			gFID	IS	Pre.	Rec.	gFID	IS	Pre.	Rec.
AutoRegressive
MaskGIT ( Chang et al., 2022 )	555	227M	6.18	182.1	0.80	0.51	–	–	–	–
LlamaGen ( Sun et al., 2024b )	300	3.1B	9.38	112.9	0.69	0.67	2.18	263.3	0.81	0.58
VAR ( Tian et al., 2024 )	350	2.0B	–	–	–	–	1.80	365.4	0.83	0.57
MagViT-v2 ( Yu et al., 2023 )	1080	307M	3.65	200.5	–	–	1.78	319.4	–	–

Table 3: Quantitative comparison of class-conditional Image Generation on ImageNet 256

\times

256. GAE achieves 1.31 FID without guidance and 1.13 with Classifier-Free Guidance, outperforming all prior methods by a large margin. ∗ denotes RAE utilizes the AutoGuidance ( Karras et al., 2024 ) instead of standard Classifier-Free Guidance.

Figure 4: Decoder stability against latent noise injection. We evaluate rFID by adding varying levels of Gaussian noise to the latent representations before decoding. The results demonstrate that models trained with higher

\sigma

exhibit superior tolerance to latent distribution shifts, ensuring stable performance during generative sampling.

Dim	$C_{\sigma}$	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	gFID $\downarrow$
	0.05	0.37	27.78	0.101	0.784	2.58
$d=32$	0.1	0.45	27.37	0.107	0.774	2.36
	0.2	0.57	26.99	0.116	0.759	2.32
	0.1	0.31	29.60	0.076	0.845	3.15
$d=64$	0.2	0.43	28.92	0.086	0.827	3.12
	0.3	0.49	28.30	0.094	0.811	3.11

Table 4: Impact of standard deviation

C_{\sigma}

and latent dimension

d

on ImageNet-1K. gFID is reported without CFG. Autoencoders are trained on ImageNet-1K for 100 epochs.

Method	Dim	LP $\uparrow$	Reconstruction				gFID (w/o CFG) $\downarrow$
Method	Dim	LP $\uparrow$	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	80 Ep	160 Ep	800 Ep
VTP-L	64	73.9	0.36	26.83	0.126	0.750	2.81	–	–
FAE	64	–	0.66	–	–	–	–	2.02	–
GAE	64	78.3	0.38	29.09	0.081	0.834	1.96	1.53	1.29

Table 5: Quantitative comparison in the 64-dim latent space on ImageNet

256\times 256

. GAE consistently outperforms existing 64-dim baselines in semantic alignment, reconstruction quality, and generative convergence speed.

$\lambda_{sp}$	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	LP $\uparrow$	gFID $\downarrow$
0.0	0.76	27.16	0.157	0.77	5.74	12.55
0.5	0.40	28.11	0.095	0.796	63.5	2.35
1.0	0.45	27.37	0.107	0.774	69.2	2.36
2.0	0.59	26.50	0.122	0.749	71.4	2.45

Table 6: Ablation of Semantic loss weight

\lambda_{sp}

on ImageNet-1K. gFID is reported without CFG. Autoencoders are trained on ImageNet-1K for 100 epochs using a consistent 32-dimensional latent bottleneck.

Method	rFID $\downarrow$	gFID $\downarrow$
KL constraint	0.977	16.72
Latent Norm	0.764	12.55

Table 7: Comparison between vanilla VAE and AE with Latent Normalization on ImageNet.

Type	Dim	Layers	Heads	LP $\uparrow$	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	gFID $\downarrow$
ViT-L	1024	24	16	69.2	0.45	27.37	0.107	0.774	2.36
ViT-B	768	12	12	62.5	0.41	26.70	0.122	0.754	2.43
ViT-S	384	12	6	–	1.55	23.85	0.212	0.625	–

Table 8: Scalability analysis of the Autoencoder backbone. gFID is reported without CFG. Autoencoders are trained on ImageNet-1K for 100 epochs using a consistent 32-dimensional latent bottleneck.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Qualitative results. Samples are generated by GAE ( $d=32$ ) after 800 epochs of training. A Classifier-Free Guidance (CFG) scale of $w=3.3$ is utilized to enhance visual fidelity.

Epochs	80	160	320	480	640	800
$d=32$	1.82	1.51	1.37	1.33	1.31	1.31
$d=64$	1.96	1.53	1.39	1.36	1.31	1.29

Appendix

Table 9: Detailed gFID scores without Classifier-Free Guidance for GAE with different latent dimension

d

at different training epochs on ImageNet

256\times 256

Method	Dim ( $d$ )	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$
SD-VAE	4	0.610	26.260	0.134	0.722
MAR-VAE	16	0.530	26.184	0.135	0.716
VAVAE	32	0.280	27.710	0.097	0.779
AlignTok	32	0.260	25.830	0.117	–
VTP-L	64	0.360	26.830	0.126	0.750
RAE (ViT-XL)	768	0.570	18.860	0.225	0.620

Appendix

Table 10: Quantitative comparison of reconstruction quality on ImageNet-1K. We compare our GAE with various generative and alignment-based methods across different latent dimensions (

d

Teacher	Teacher LP $\uparrow$	AE Reconstruction				AE Latent LP $\uparrow$	Gen.
Model	LP	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	LP	gFID $\downarrow$
MAE-L	33.9	0.36	28.71	0.084	0.813	29.8	4.20
DINOv2-L	76.0	0.45	27.37	0.107	0.774	69.2	2.36

Appendix

Table 11: Comparison of SP Teachers (MAE-L vs. DINOv2-L) on ImageNet-1K. gFID is reported without CFG.

Epochs	Timeshift	CFG Interval	CFG Weight	gFID $\downarrow$
800	0.4	0.27	2.3	1.15
	0.4	0.27	2.5	1.14
	0.4	0.27	3.1	1.14
	0.4	0.27	3.5	1.15
	0.4	0.3	2.7	1.15
	0.4	0.3	3.3	1.13

Appendix

Table 12: Ablation study of CFG parameters for GAE on ImageNet

256\times 256

. We report the gFID scores across different guidance intervals and weights.

Figure 6: Comparison of downsampler architectures. Left: The FAE downsampler. Right: Our downsampler with patch convolution.

Category	Field	$E_{p}$	$A_{p}$	$D_{p}$	$E^{*}_{sp}$	$D_{sp}$	LDM
Architecture	Input dim.	16 $\times$ 16 $\times$ 3	16 $\times$ 16 $\times$ 1024	16 $\times$ 16 $\times$ 32	16 $\times$ 16 $\times$ 3	16 $\times$ 16 $\times$ 32	16 $\times$ 16 $\times$ 32
	Output dim.	16 $\times$ 16 $\times$ 1024	16 $\times$ 16 $\times$ 32	16 $\times$ 16 $\times$ 3	16 $\times$ 16 $\times$ 32	16 $\times$ 16 $\times$ 3	16 $\times$ 16 $\times$ 32
	Hidden dim.	1024	–	1024	1024	1024	1152
	Num. layers	24	–	24	24/1/1	4	28
	Num. heads	16	–	16	16/16/–	16	16
	MLP Ratio	4	–	4	4/–/–	4	4

Appendix

Table 13: Detailed architectural and optimization configurations for main generation.

E_{p}

and

D_{p}

denote the pixel encoder and decoder, while

A_{p}

represents the linear projector. The semantic encoder

E^{*}_{sp}

integrates a DINOv2-L backbone with downsampler, and

D_{sp}

serves as an auxiliary decoder for semantic training.

Method	Dim	Flatten $\uparrow$	GAP $\uparrow$
Part 1: Linear Probing results for Fig. 1
SD-VAE	16	4.99	1.16
MAR-VAE	16	13.7	4.29
VAVAE	32	43.1	29.0
Aligntok	32	–	35.1
VTP-L	64	73.9	67.6

Appendix

Table 14: Comprehensive LP results.

Latent Diffusion Models

Latent diffusion models denoise compact representations to generate or transform images, video, audio, and other complex data efficiently.

See all 207 papers

What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion
May 8, 2026Zhengrong Yue, Taihang Hu, Mengting Chen +8Latent Diffusion Models Autoencoders
Shanghai Jiao Tong University · Alibaba Group · Beihang University +4
Distribution Matching Variational AutoEncoder
Dec 8, 2025Sen Ye, Jianning Pei, Mengde Xu +4Variational Autoencoders Generative Modeling
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · University of Chinese Academy of Sciences · Tencent +1
Repurposing Geometric Foundation Models for Multi-view Diffusion
Mar 23, 2026Wooseok Jang, Seonghu Jeon, Jisang Han +5Multi-View Diffusion Models Novel View Synthesis
KAIST AI · New York University · Intel Labs

May 8, 2026J/K move · Enter open · S save

What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion

Zhengrong Yue, Taihang Hu, Mengting Chen, Haiyu Zhang+7

Shanghai Jiao Tong University · Alibaba Group · Beihang University+4

Latent Diffusion Models Autoencoders Diffusion-Based Image Generation Diffusion Models

Open paper AskPDF arXiv

Tokenizers are a crucial component of latent diffusion models, as they define the latent space in which diffusion models operate. However, existing tokenizers are primarily designed to improve reconstruction fidelity or inherit pretrained representations, leaving unclear what kind of latent space is truly friendly for generative modeling. In this paper, we study this question from the perspective of latent manifold organization. By constructing controlled tokenizer variants, we identify three key properties of a diffusion-friendly latent manifold: coherent spatial structure, local manifold continuity, and global manifold semantics. We find that these properties are more consistent with downstream generation quality than reconstruction fidelity. Motivated by this finding, we propose the Prior-Aligned AutoEncoder (PAE), which explicitly shapes the latent manifold instead of leaving diffusion-friendly manifold to emerge indirectly from reconstruction or inheritance. Specifically, PAE leverages refined priors derived from VFMs and perturbation-based regularization to turn spatial structure, local continuity, and global semantics into explicit training objectives. On ImageNet 256x256, PAE improves both training efficiency and generation quality over existing tokenizers, reaching performance comparable to RAE with up to 13x faster convergence under the same training setup and achieving a new state-of-the-art gFID of 1.03. These results highlight the importance of organizing the latent manifold for latent diffusion models.

Abstract

Latent Diffusion Models

Explore similar work

What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion

Distribution Matching Variational AutoEncoder

Repurposing Geometric Foundation Models for Multi-view Diffusion