RPiAE: A Representation-Pivoted Autoencoder Enhancing Both Image Generation and Editing

Yue Gong; Hongyu Li; Shanyuan Liu; Bo Cheng; Yuhang Ma; Liebucha Wu; Xiaoyu Wu; Manyuan Zhang

cs.CVMar 19, 2026

RPiAE: A Representation-Pivoted Autoencoder Enhancing Both Image Generation and Editing

Authors: Yue Gong, Hongyu Li, Shanyuan Liu, Bo Cheng, Yuhang Ma, Liebucha Wu, Xiaoyu Wu, Manyuan Zhang,

Organizations: Beihang University · AI Research · The Chinese University of Hong Kong

Abstract

Diffusion models have become the dominant paradigm for image generation and editing, with latent diffusion models shifting denoising to a compact latent space for efficiency and scalability. Recent attempts to leverage pretrained visual representation models as tokenizer priors either align diffusion features to representation features or directly reuse representation encoders as frozen tokenizers. Although such approaches can improve generation metrics, they often suffer from limited reconstruction fidelity due to frozen encoders, which in turn degrades editing quality, as well as overly high-dimensional latents that make diffusion modeling difficult. To address these limitations, We propose Representation-Pivoted AutoEncoder, a representation-based tokenizer that improves both generation and editing. We introduce Representation-Pivot Regularization, a training strategy that enables a representation-initialized encoder to be fine-tuned for reconstruction while preserving the semantic structure of the pretrained representation space, followed by a variational bridge which compress latent space into a compact one for better diffusion modeling. We adopt an objective-decoupled stage-wise training strategy that sequentially optimizes generative tractability and reconstruction-fidelity objectives. Together, these components yield a tokenizer that preserves strong semantics, reconstructs faithfully, and produces latents with reduced diffusion modeling complexity. Experiments demonstrate that RPiAE outperforms other visual tokenizers on text-to-image generation and image editing, while delivering the best reconstruction fidelity among representation-based tokenizers.

Figures & tables

Figure 1: Motivation. A practical tokenizer for diffusion must simultaneously achieve high reconstruction fidelity for editing and strong generative tractability for diffusion training, while preserving the semantic structure of pretrained representation models.

Figure 2: Overview of RPiAE. A pretrained RM encoder extracts representation features, which are compressed by a variational bridge into diffusion-friendly latents and decoded back for pixel-space reconstruction; a frozen pivot replica provides semantic supervision during training.

Table 3

Figure 3: Three-stage training of RPiAE: (I) pivot-regularized encoder tuning, (II) variational bridge training with KL regularization, and (III) decoder specialization under fixed latents.

Method	Tokenizer	Image Reconstruction				#Params	Epochs	Generation w/o CFG				Generation w/ CFG
Method	Tokenizer	rFID $\downarrow$	PSNR $\uparrow$	LPIPS $\downarrow$	SSIM $\uparrow$	#Params	Epochs	gFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$	gFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
Generation Models w/o RM
MaskGIT [ 3 ]	VQGAN [ 40 ]	2.23	17.9	0.202	0.422	227M	555	6.18	182.1	0.80	0.51	–	–	–	–
LlamaGen [ 32 ]	VQGAN [ 40 ]	0.59	24.5	–	0.813	3.1B	300	9.38	112.9	0.69	0.67	2.18	263.3	0.81	0.58
MaskDiT-XL [ 46 ]	SD-VAE [ 29 ]	0.61	26.9	0.130	0.736	675M	1600	5.69	177.9	0.74	0.60	2.28	276.6	0.80	0.61
DiT-XL [ 26 ]	SD-VAE [ 29 ]	0.61	26.9	0.130	0.736	675M	1400	9.62	121.5	0.67	0.67	2.27	278.2	0.83	0.57

Table 3: Reconstruction and class-conditional generation on ImageNet-1K at

256^{2}

. We report reconstruction metrics for the tokenizer and generation metrics for the diffusion model trained on the corresponding latents, evaluated both without and with CFG. † indecates reproducing results.

Method	GenEval							DPG-Bench	GEdit-Bench-EN
Method	Sin.Obj.	Two.Obj.	Counting	Colors	Pos	Color.Attr.	Overall	Average	G_SC	G_PQ	G_O
Open-source Models
Flux-VAE [ 19 ]	0.93	0.68	0.48	0.76	0.51	0.44	0.63	77.44	3.71	8.42	3.86
RAE-DINOv2-S [ 45 ]	0.97	0.78	0.52	0.81	0.60	0.52	0.70	80.81	2.86	7.04	3.09
RAE-DINOv2-B [ 45 ]	0.93	0.65	0.46	0.73	0.47	0.37	0.60	78.58	2.71	8.03	3.12
+ DDT Head	0.96	0.81	0.51	0.82	0.59	0.44	0.69	80.82	3.25	7.76	3.51

Table 4: Comparison of different Method for Text-to-Image Generation and Image Editing on GenEval, DPG-Bench and GEdit-Bench-EN.

Figure 4: Performance of GenEval, DPG-Bench, and GEdit over training for different encoders. Our method achieves both a higher performance ceiling and faster convergence.

Figure 5: Visualizations of Text to Image Generation in (a) and Image Editing in (b).

Table 5: Ablation studies on GEdit-Bench-EN / ImgEdit-Bench.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Field	$\mathcal{E}_{\theta}$	$\mathcal{D}_{\phi}$	$\mathcal{E}^{p}$	$\mathcal{B}_{e}$	$\mathcal{B}_{d}$	LightningDiT
Architecture	Input dim.	$224\times 224\times 3$	$16\times 16\times 768$	$224\times 224\times 3$	$16\times 16\times 768$	$16\times 16\times 64$	$16\times 16\times 64$
	Output dim.	$16\times 16\times 768$	$256\times 256\times 3$	$16\times 16\times 768$	$16\times 16\times 64$	$16\times 16\times 768$	$16\times 16\times 64$
	Hidden dim.	768	1536	768	768	768	1152
	Num. layers	12	24	12	1	6	28
	MLP Ratio	4	4	4	1	4	4
	Dim. per head	64	64	64	96	96	72

Appendix

Table 6: Implementation details of RPiAE and LightningDiT used in our experiments.

Category	Field	$\mathcal{E}_{\theta}$	$\mathcal{D}_{\phi}$	$\mathcal{E}^{p}$	DiT DH
Architecture	Input dim.	$224\times 224\times 3$	$16\times 16\times 768$	$224\times 224\times 3$	$16\times 16\times 768$
	Output dim.	$16\times 16\times 768$	$256\times 256\times 3$	$16\times 16\times 768$	$16\times 16\times 768$
	Hidden dim.	768	128	768	[1152, 2048]
	Num. layers	12	–	12	[28, 2]
	MLP Ratio	4	–	4	4
	Dim. per head	64	–	64	[72, 128]

Appendix

Table 7: Implementation details of RPiAE and DiT DH used in our preliminary experiments.

Generation Models w/o RM
Method	Tokenizer	#Params	Epochs	Generation w/o CFG				Generation w/ CFG
Method	Tokenizer	#Params	Epochs	gFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$	gFID $\downarrow$	IS $\uparrow$	Prec. $\uparrow$	Rec. $\uparrow$
MaskGIT [ 3 ]	VQGAN [ 40 ]	227M	555	6.18	182.1	0.80	0.51	–	–	–	–
LlamaGen [ 32 ]	VQGAN [ 40 ]	3.1B	300	9.38	112.9	0.69	0.67	2.18	263.3	0.81	0.58
MaskDiT-XL [ 46 ]	SD-VAE [ 29 ]	675M	1600	5.69	177.9	0.74	0.60	2.28	276.6	0.80	0.61
DiT-XL [ 26 ]	SD-VAE [ 29 ]	675M	1400	9.62	121.5	0.67	0.67	2.27	278.2	0.83	0.57

Appendix

Table 8: Class-conditional generation on ImageNet-1K at

256^{2}

. We report generation metrics for the diffusion model trained on the corresponding latents, evaluated both without and with CFG. † indicates reproduced results. Best and second-best results are highlighted in bold and underlined, respectively.

Figure 6: Qualitative comparison of image reconstruction results between our method and RAE. Our method achieves better reconstruction quality, particularly in preserving regular structural patterns and color fidelity. As shown in the examples, it reconstructs clearer and more plausible textures for structured regions such as wall tiles, mesh grids, fences, and honeycomb, while also reducing the color shifts that are visible in RAE reconstructions.

Figure 7: Selected samples for class-conditional generation on ImageNet-1K at $256^{2}$ .

Figure 8: More qualitative comparison on GenEval. Our method generates images with higher visual fidelity, richer details, and more accurate structural composition.

Figure 9: More qualitative comparison on image editing. RPiAE achieves the strongest instruction-following ability while maintaining high reconstruction quality.

Model	Top-1 Acc. $\uparrow$	Top-5 Acc. $\uparrow$
DINOv2-B	84.56	97.04
RPiAE	84.18	96.91

Appendix

Table 9: ImageNet linear probing results.

T2I Diffusion Models

Text-to-image (T2I) diffusion models generate images from prompts by learning to reverse noise under text conditioning.

See all 393 papers

RoAR-VAE: Improving Latent Diffusability through Robust Prediction
Oct 19, 2026Yang Yue, Fangyun Wei, Yitong Wang +4Image Tokenization Variational Autoencoders
Tsinghua University · Microsoft Research
Improved Baselines with Representation Autoencoders
May 18, 2026Jaskirat Singh, Boyang Zheng, Zongze Wu +3Diffusion Model Guidance Diffusion Transformer
Adobe Research · ANU · New York University
Diffusion Image Editing via Asynchronous Token Decoding
Aug 10, 2026Yang Shi, Liangsi Lu, Minzhe Guo +4Text-Guided Image Editing Diffusion-Based Image Editing
Guangdong University of Technology Guangzhou, China · Hong Kong Baptist University Hong Kong, China · Peking University Beijing, China +1

Oct 19, 2026J/K move · Enter open · S save

RoAR-VAE: Improving Latent Diffusability through Robust Prediction

Yang Yue, Fangyun Wei, Yitong Wang, Jinjing Zhao+3

Tsinghua University · Microsoft Research

Image Tokenization Variational Autoencoders Latent Diffusion Models Diffusion-Based Image Generation

Open paper AskPDF arXiv

An effective visual tokenizer must preserve image fidelity while producing latents that downstream diffusion models can learn efficiently. Prior work improves latent diffusability through semantic alignment or joint tokenizer--diffusion training. Revisiting REPA-E, a representative approach, we find that removing its diffusion objective and dedicated generative blocks while retaining semantic prediction from corrupted latents yields comparable downstream generation performance under a matched setting. This observation motivates a general principle of robust prediction: generation-relevant information should remain recoverable from corrupted latents. We instantiate this principle with two complementary objectives. RobustAlign predicts pretrained semantic features from corrupted latents to preserve high-level structure, while RobustRecon reconstructs the clean image to retain fine-grained spatial and appearance information. Together, they form Robust Alignment and Reconstruction VAE (RoAR-VAE). Our analysis suggests that robust prediction improves diffusability by making useful information more accessible in high-variance latent directions, which diffusion models tend to learn earlier. Empirically, RoAR-VAE improves average generation benchmark performance by 2.19 points over the strongest baseline tokenizer under the same training budget, while improving small-text reconstruction accuracy by 16.93 percentage points over FLUX.1-VAE. These results demonstrate that RoAR-VAE achieves a favorable balance between reconstruction fidelity and downstream generative learnability.

Abstract

T2I Diffusion Models

Explore similar work

RoAR-VAE: Improving Latent Diffusability through Robust Prediction

Improved Baselines with Representation Autoencoders

Diffusion Image Editing via Asynchronous Token Decoding