REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion
  Transformers
In this paper we tackle a fundamental question: "Can we train latent
diffusion models together with the variational auto-encoder (VAE) tokenizer in
an end-to-end manner?" Traditional deep-learning wisdom dictates that
end-to-end training is often preferable when possible. However, for latent
diffusion transformers, it is observed that end-to-end training both VAE and
diffusion-model using standard diffusion-loss is ineffective, even causing a
degradation in final performance. We show that while diffusion loss is
ineffective, end-to-end training can be unlocked through the
representation-alignment (REPA) loss -- allowing both VAE and diffusion model
to be jointly tuned during the training process. Despite its simplicity, the
proposed training recipe (REPA-E) shows remarkable performance; speeding up
diffusion model training by over 17x and 45x over REPA and vanilla training
recipes, respectively. Interestingly, we observe that end-to-end tuning with
REPA-E also improves the VAE itself; leading to improved latent space structure
and downstream generation performance. In terms of final performance, our
approach sets a new state-of-the-art; achieving FID of 1.26 and 1.83 with and
without classifier-free guidance on ImageNet 256 x 256. Code is available at
https://end2end-diffusion.github.io.
67fdd3e3913c97aa32f94e9b
Xingjian Leng
claimed_verified
2025-04-16T09:23:36.449Z
641480554b1701c01cdb36c4
/avatars/f1f6b294e0236d76a68c099164c81f36.svg
Xingjian Leng
user
xingjianleng
67fdd3e3913c97aa32f94e9c
Jaskirat Singh
claimed_verified
2025-04-25T08:35:52.907Z
60cc389a0844fb1605fef405
/avatars/ec11f85735e0525439e8821cf6d12e53.svg
Jaskirat Singh
user
jsingh
67fdd3e3913c97aa32f94e9d
Yunzhong Hou
admin_assigned
2025-04-17T08:34:11.946Z
6752870ec63bc5b670b1b27e
https://cdn-avatars.huggingface.co/v1/production/uploads/6752870ec63bc5b670b1b27e/3CdHxnyTKbGup-1V67nEV.jpeg
Yunzhong Hou
user
yunzhong-hou
67fdd3e3913c97aa32f94e9e
Zhenchang Xing
admin_assigned
2025-04-17T08:34:21.190Z
63bf533f4a2beec65568f813
/avatars/bba583653b5cada8dd4a3ff2281e9dec.svg
Xing
user
Zhenchang
67fdd3e3913c97aa32f94e9f
Saining Xie
admin_assigned
2025-04-17T08:34:39.700Z
6596422646624a86ff3b3bda
/avatars/216e12b77e45ac5f1fa20932f5745411.svg
Saining Xie
user
sainx
67fdd3e3913c97aa32f94ea0
Liang Zheng
admin_assigned
2025-04-17T08:34:56.041Z
666351ebd86c026caa135e5c
/avatars/50a37f7e999f660c69f518b71577eb7d.svg
Liang Zheng
user
liangzheng06
REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion
  Transformers
In this paper we tackle a fundamental question: "Can we train latent
diffusion models together with the variational auto-encoder (VAE) tokenizer in
an end-to-end manner?" Traditional deep-learning wisdom dictates that
end-to-end training is often preferable when possible. However, for latent
diffusion transformers, it is observed that end-to-end training both VAE and
diffusion-model using standard diffusion-loss is ineffective, even causing a
degradation in final performance. We show that while diffusion loss is
ineffective, end-to-end training can be unlocked through the
representation-alignment (REPA) loss -- allowing both VAE and diffusion model
to be jointly tuned during the training process. Despite its simplicity, the
proposed training recipe (REPA-E) shows remarkable performance; speeding up
diffusion model training by over 17x and 45x over REPA and vanilla training
recipes, respectively. Interestingly, we observe that end-to-end tuning with
REPA-E also improves the VAE itself; leading to improved latent space structure
and downstream generation performance. In terms of final performance, our
approach sets a new state-of-the-art; achieving FID of 1.26 and 1.83 with and
without classifier-free guidance on ImageNet 256 x 256. Code is available at
https://end2end-diffusion.github.io.
End-to-end training of latent diffusion models and VAE tokenizers using representation-alignment loss enhances diffusion model training speed and VAE performance.
latent diffusion models variational auto-encoder VAE diffusion-loss representation-alignment loss REPA loss REPA-E FID ImageNet classifier-free guidance