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 8.
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 Flax/JAX Projects
  
  
    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 …