I’ve made a more clear project post here, be sure to give it a like if your interested in helping out!

  
    
    
    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 …