PAE: What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion [](https://github.com/ZhengrongYue/PAE) [](https://huggingface.co/yuezhengrong/PAE-collections) [](https://www.modelscope.cn/models/ZhengrongYue/PAE-Collections) This project presents **PAE** (Prior-Aligned AutoEncoder), a tokenizer framework that explicitly shapes a **diffusion-friendly latent manifold** for latent diffusion models. Instead of relying solely on reconstruction fidelity or passively inheriting pretrained representations, PAE identifies and optimizes three key properties of a diffusion-friendly latent space — **spatial structure coherence**, **local manifold continuity**, and **global manifold semantics** — through targeted prior-alignment regularizations. On ImageNet 256×256, PAE achieves a new **state-of-the-art gFID of 1.03** with up to **13× faster convergence** than RAE under the same LightningDiT setup. Prior alignment constructs a diffusion-friendly latent manifold. Left: Compared with reconstruction-oriented counterparts, the prior-aligned latent manifold is more structurally coherent, locally continuous, and semantically organized. Right: PAE yields faster convergence, better generation quality, and robust few-step sampling performance. Class-conditional samples generated by PAE with LightningDiT-XL/1 on ImageNet 256×256. ## 🔥 Updates * **[2026.05.09]** 🚀 🚀 🚀 We release **PAE**. Code and pretrained models are now available! ## ✨ Highlights - 🎯 **New Perspective**: We study what makes a latent manifold diffusion-friendly, identifying three key properties: spatial structure coherence, local manifold continuity, and global manifold semantics. - 🏗️ **Explicit Manifold Shaping**: PAE turns these properties into explicit training objectives via three prior-alignment regularizations (SSR, MCR, SCR), rather than leaving them to emerge indirectly. - ⚡ **13× Faster Convergence**: PAE reaches performance comparable to RAE with up to 13× fewer training epochs under the same LightningDiT setup. - 🏆 **State-of-the-Art**: Achieves gFID **1.03** on ImageNet 256×256, the best result among all compared methods. - 🔄 **Encoder-Agnostic**: Compatible with multiple VFM backbones including DINOv2, SigLIP2, DINOv3, and MAE. ## 🏛️ Architecture Overview of the PAE framework. A frozen VFM provides stable representation features. DAM injects pixel detail while preserving the VFM as the dominant semantic source. Three prior-alignment objectives explicitly shape the latent manifold. ## ❤️ Acknowledgement Our work builds upon the foundations laid by many excellent projects in the field. We would like to thank the authors of [LightningDiT](https://github.com/hustvl/LightningDiT), [RAE](https://github.com/bytetriper/RAE), [GAE](https://github.com/sii-research/GAE), [ADM](https://github.com/openai/guided-diffusion). We are grateful for their contributions to the community.
