Would be great if this was set up with the basics of a Python package and any conveniences that make sense. Since there are a couple of different subcomponents in this project, we probably also want to figure out a directory structure so that people working on different aspects aren't stepping on toes.

I've been told JAX plays nice with other frameworks, so a PyTorch dataloader is honestly fine. Should load in text (possibly using https://github.com/leogao2/lm_dataformat) and spectrograms from disk in order to construct batches of paired data.

Copy the Transformer-in-Transformer architecture, with some modifications. First, the patches should always cover the full vertical (Mel filters) range. So the outer patch size might be 80 x 8. The inner patch size should be 80 x 1, which will give us frame-level information. Second, both should use rotary embeddings for relative positional encoding.

Feel free to take inspiration from: https://github.com/NZ99/transformer_in_transformer_flax https://github.com/lucidrains/transformer-in-transformer

Char or unigram-level transformer (though from the perspective of the transformer, it's just a sequence of token IDs). Likely does not need to have a context greater than, say, 128 tokens as the dataset clips are typically less than 10 seconds of audio, so we'll be cropping to that size or less anyway. Use rotary relative positional embeddings.

Tokenize the text either via chars or unigrams. Figure out most appropriate method here: ideally we want to be able to accommodate all English text, punctuation, emojis, and potentially other text.

An optax training loop with an input pipeline will let us take the model code and run training with datasets.

https://github.com/deepmind/optax

Shouldn't be too terrible, especially once we've got a dataloading class set up.

It would be good to have config system for the model. This should specify stuff around model shapes/sizes, microbatch size, hyperparams, and anything we'd want stored in the Weights and Biases config for a run.

Switch dataset formatting to TFRecords, from current version which just uses a .pt files within a directory. This should make streaming possible, which may important for real training. Slightly awkward because the preprocessing is in PyTorch and the training is in JAX. I'm told the best way to read off TFRecords during training is through tf.data

https://github.com/kingoflolz/mesh-transformer-jax

Use the pod orchestation code from here. Effectively, we should borrow everything and modify the transformer_shard and tfrecord_loader files.

Should write datasets to TFRecords. Probably some format where the spectrogram tensor and its padded, tokenized text tensor are stored as a single example.

Both preprocess.py and train.py now use TFRecords as the format for pre-processed data. Would close #25. Currently untested on GCP, but it works with TFRecords at local paths.

We can start thinking about how to evaluate our models once trained. The simplest would be contrastive loss over some held out set with a fixed batch size. Another would be to evaluate how well CLAP score correlates with MOS (Mean Opinion Score), which is the gold standard subjective eval in the audio NN literature. We could probably also try linear probe training on the Google Speech Commands dataset.

An alternative to the symmetric cross entropy loss, we could optimize for alignment and uniformity on the unit hypersphere as was done in this paper:

https://arxiv.org/abs/2005.10242

The code itself is like 4 lines of PyTorch, with the caveat that we'd need to write a pdist JAX function in order to use it.