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
rFID ↓
PSNR ↑
LPIPS ↓
SSIM ↑
gFID ↓
IS ↑
Prec. ↑
Rec. ↑
gFID ↓
IS ↑
Prec. ↑
Rec. ↑
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 2562 . 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
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
Eθ
Dϕ
Ep
Be
Bd
LightningDiT
Architecture
Input dim.
224×224×3
16×16×768
224×224×3
16×16×768
16×16×64
16×16×64
Output dim.
16×16×768
256×256×3
16×16×768
16×16×64
16×16×768
16×16×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
Eθ
Dϕ
Ep
DiT DH
Architecture
Input dim.
224×224×3
16×16×768
224×224×3
16×16×768
Output dim.
16×16×768
256×256×3
16×16×768
16×16×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.
Method
Tokenizer
#Params
Epochs
Generation w/o CFG
Generation w/ CFG
gFID ↓
IS ↑
Prec. ↑
Rec. ↑
gFID ↓
IS ↑
Prec. ↑
Rec. ↑
Generation Models w/o RM
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 2562 . 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 2562 .
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.
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.
Representation Autoencoders (RAE) replace traditional VAE with pretrained vision encoders. In this paper, we systematically investigate several design choices and find three insights which simplify and improve RAE. First, we study a generalized formulation where the representation is defined as sum of the last k encoder layers rather than solely the final layer. This simple change greatly improves reconstruction without encoder finetuning or specialized data (e.g., text, faces). Second, we study the prevalent assumption that RAE (using pretrained representation as encoder) replaces representation alignment (REPA), which distills the same representation to intermediate layers instead. Through large-scale empirical analysis, we uncover a surprising finding: RAE and REPA exhibit complementary working mechanisms, allowing the same representation to be used as both encoder and target for intermediate diffusion layers. Finally, the original RAE struggles with classifier-free guidance (CFG) and requires training a second, weaker diffusion model for AutoGuidance (AG). We show that REPA itself can be viewed as x-prediction in RAE latent space. By simply re-parameterizing the output of the DiT model, it can provide guidance for "free". Overall, RAEv2 leads to more than 10x faster convergence over the original RAE, achieving a state-of-the-art gFID of 1.06 in just 80 epochs on ImageNet-256. On FDr6, RAEv2 achieves a state-of-the-art 2.17 at just 80 epochs compared to the previous best 3.26 (800 epochs) without any post-training. This motivates EPFID@k (epochs to reach unguided gFID < k) as a measure of training efficiency. RAEv2 attains an EPFID@2 of 35 epochs, versus 177 for the original RAE. We also validate our approach across diverse settings for text-to-image generation and navigation world models, showing consistent improvements. The code is available at https://raev2.github.io.
Text-guided diffusion image editing aims to modify semantic attributes of an image while preserving its identity, layout, and background. However, naïvely switching the text condition during sampling often causes global drift, as denoising dynamics propagate changes across tokens and can disrupt unedited regions. To address this issue, we propose Asynchronous Token Decoding Edit (ATDEdit), an inference-time framework that views each sampler step as a parallel update of a globally coupled token matrix and enables token-indexed condition switching with differentiated update policies. Instead of applying synchronous target-conditioned updates to all tokens, ATDEdit estimates editable locations using token-wise conditional surprisal and applies target-conditioned corrections to the selected token set. It supplies source key/value memory at keep-token positions and projects selected keep-token latent rows back to their source values; these operations promote background preservation but do not constitute a pixel-level invariance guarantee. This approach combines local editing and background preservation without external or user-provided spatial masks and without model fine-tuning. On PIE-Bench, ATDEdit achieves the strongest reported preservation metrics, including 27.44~dB PSNR and 0.055 LPIPS, while retaining competitive semantic alignment.
Yang Shi, Liangsi Lu, Minzhe Guo +4
Guangdong University of Technology Guangzhou, China · Hong Kong Baptist University Hong Kong, China · Peking University Beijing, China +1