Joint training of VAE and Transformer

[minghuix98]

Thanks for your excellent work.
Could you clarify whether the VAE and the transformer that operates on the latent features can be jointly optimized—that is, trained end-to-end?
Thanks for your time!

[lucidrains]

@minghuix98 hey Xu, at the moment not, but i could make it customizable next week if you have that need

[minghuix98]

@minghuix98 hey Xu, at the moment not, but i could make it customizable next week if you have that need

Hi Phil Wang,

Thanks so much for your prompt reply!
I'd appreciate that very much and didn't expect you would work on that before.
My question arise from that I gone through some previous work: REPA and REPA-E for image generation. In REPA-E, they tested and reported the naïve end-to-end training, which is shown as below.

I have two related questions and would be really grateful for your insights:

1. Do you think the core requirement for enabling end-to-end training is the use of stop-gradient, or is it necessary to combine stop-gradient with L_REPA (alignment loss)?
2. I believe that the currently developed unified models (for both image generation and understanding) could also be extended to support end-to-end training in conjunction with a VAE. What are your thoughts on this?

Best,
Minghui

[lucidrains]

ohh, will check out the REPA paper and their proposal, thanks!

[lucidrains]

@minghuix98 hmm, correct me if i'm wrong, but REPA depends on a pretrained DINO? i'm not sure if it can be applied here if so

[minghuix98]

@minghuix98 hmm, correct me if i'm wrong, but REPA depends on a pretrained DINO? i'm not sure if it can be applied here if so

Yeah, it's correct! And REPA-E uses some tricks (including stop gradient, as well as still using a alignment loss that relies on DINO [similar with REPA] ) to make VAE and diffusion transformer can be trained end-to-end.

So I was just curious about which trick is sufficient and critical to make it end-to-end. I don't have to investigate them too much deep, just my curiosity. Maybe I will test them more if I needed in the future.

Thanks for your time, it's totally ok for me. And thanks again for your big contribution to the community.