Modality Distribution Mapping (MODIST MAP)

This repository creates a generative mapping between joint embeddings spaces.

Structure

The repo is structured as follows:

1. config - contains the lightning config files
2. checkpoints - the directory of model checkpoints. Only the rvq and transformer checkpoints are included.

Inside the src/modist_map we have the source code:

1. rvq - contains the code to train the rvq models. This is not implemented with lightning as it uses an external python library.
2. vis - contains visualisation code
3. data - contains code for making text and audio embeddings using CLAP as well as the lightning data module. Also code for creating a balanced dataset.
4. callbacks - contains callbacks usied in the Lightning configuration
5. checkpoints - the directory of model checkpoints. Only the rvq and transformer checkpoints are included.
6. networks - pytorch/lightning modules
7. lightning_utils - utils used by lightning
8. models - contains the mmap model which is a LightningModule
9. examples - examples to better understand model usage and the codebase

Approach

This is achieved by first creating discrete code sequences independently for both modalitites using vector quantisation and then training two auto-regressive transformers over these sequences in either direction.

Residual Vector Quantization

Give two RVQ with 4 code books each for both audio and text modalities, the two RVQs can transform a text and audio embedding into sequences:

audio sequence:

a_1, a_2, a_3, a_4

text sequence:

t_1, t_2, t_3, t_4

each element can be one from a fixed code_book size and would likely differ between modalities.

Generative models are then trained on sequences as follows:

audio > text

<bos>, a_1, a_2, a_3, a_4, <eos>, <bos>, t_1, t_2, t_3, t_4, <eos>

text > audio

<bos>, t_1, t_2, t_3, t_4, <eos>, <bos>, a_1, a_2, a_3, a_4, <end>

Causal and Masked Models

Given causal masking, the task is to predict the next sequence. E.g.

input: <bos>, a_1, a_2, a_3, a_4, <eos>, <bos>, t_1, t_2, t_3, t_4
output: a_1, a_2, a_3, a_4, <eos>, <bos>,, t_1, t_2, t_3, t_4, <eos>

When both causal transformers are created this allows for unconstrained generation (starting with the token only) or completing partial codes.

Not implemented yet

Given masked language modelling, the task is to predict masked tokens E.g.

input: <bos>, a_1, a_2, a_3, a_4, <eos>, <bos>, t_1, t_2, t_3, t_4, <bos>
output: <bos>, <mask>, a_2, <mask>, a_4, <eos>, <bos>, t_1, t_2, <mask>, t_4, <bos>

And this allows for any part of a sequence to be regenerated.