Train a VAE to interpolate on English sentences
Transformer-VAEConvert 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!LanguageThe model will be trained in english.ModelBuild on T5-base, this will match with the Optimus model.Only additional parameters come from a small Autoencoder module that will go between the encoder and decoder.DatasetsUse the wikipedia sentences dataset from OPTIMUS.This comes tokenized so we’ll need to use its tokenizer with T5.github.comChunyuanLI/Optimus/blob/master/download_datasets.md# Download/Pre-process Datasets

## Wikipedia

Download processed files (11.78G) below, and unzip it (298 files)

https://textae.blob.core.windows.net/optimus/data/datasets/wikipedia_json_64_filtered.zip

Download raw file (11.79G):

https://textae.blob.core.windows.net/optimus/data/datasets/wikipedia.segmented.nltk.txt

Our pre-processing protocal: We split the original wiki text into 298 files, and loop over files in one epoch.

We filter each sentence in wiki based on two constraints: (1) The sentence length is smaller than 64. (2) The tokenized sentence length is smaller than 256 (so that the encoder can take the entire sentence).

To filter the sentence, please change the data folders and run the script:

    sh scripts/scripts_local/run_data_filtering_wiki.shThis file has been truncated.show originalTraining scriptsThe originalPyTorch training scriptis adapted from the old Huggingface clm training script so using theflax clm scriptshould be a good base to build on.ChallengesThe original model was made with PyTorch so there will be some features that can’t be ported over. E.g. I added a prism layer to the PyTorch code which requires FFTs.Desired project outcomeA colab notebook where people can explore the Transformer-VAE’s latent space.Interpolate between sentences.Transfer the style/content of one sentence to another.Do gradient descent on a sentences latent code to get the desired sentiment, classification score, etc.ReadsHere are some background links to understand the context behind this project:Initial Transformer-VAE postand theimprovements post.MMD-VAE, the MMD loss that Transformer-VAE uses.OPTIMUSCurrent SOTA text-vae will give a sense of the outputs would should expet.
2021-06-28T17:51:59Z
How does this work?Currenlty the takes a transformer encoder and decoder and puts a VAE between them..1554×742 43.5 KBThe VAE forms a compressed latent code which allows interpolating on the training data..1540×926 69 KBFor regulatisation an MMD loss is used instead of KL Divergence which reduces posterior collapse..1528×916 60.8 KB
2021-06-28T17:52:49Z
More Detailed ToDosSign in to GitHub · GitHub
2021-06-28T17:57:01Z
old post:Train the best ever transformer-VAE - #14 by Fraser
2021-06-30T12:56:44Z
I really like this project and all the info provided here! Thanks a lotLet’s try to make this happen! Will ask some people as well
2021-07-16T09:11:03Z
thanks for the awesome information.
2021-11-16T14:57:54Z
alexsunny123:thanks for the awesome information.https://ometv.onlhttps://chatroulette.topthanks my issue has been fixed.