Train the best ever transformer-VAE

[Fraser]

Hi Patrick, thanks for the feedback!

I’ve linked a revised plan bellow.

The idea here is just to take an existing flax-T5 model and stick an autoencoder between the encoder & decoder. As seen here.

I’ve currently had calls with 3 other team members and we’re really exited to see what this produces!

Train a VAE to interpolate on English sentences

Transformer-VAE Convert a T5 model into a variational autoencoder for text. I have already made a project that does this in PyTorch but its never been trained at scale. This project is to convert the autoencoder into Flax so it can be trained efficiently on a TPU to train the largest ever Transformer-VAE! Language The model will be trained in english. Model Build on T5-base, this will match with the Optimus model. Only additional parameters come from a small Autoencoder module that will …