High order Moment Models

We propose an alternative to classical attention that scales linearly with the number of tokens and is based on high order moments.

The HoMM scheme is as follows: Having a query token $x_q$ and a set of context tokens $x_c$, we first use a projection $ho$ to map each token $x_c$ to a high dimension space, where the high-order moments are computed recursively (by chunking and performing element-wise product, and then averaging over the tokens). $x_q$ is projected into the same high dimensional space with a projection $s$. The element-wise product of the two corresponds to $x_q$ selecting the information it needs in the high-order moments of $x_c$. The results is then projected back to the same space as $x_q$ and added to the original tokens via a residual connection.

/!\ Help welcome: DM me on twitter (https://twitter.com/david_picard), or submit an issue, or email me!

Changelog

• 20240122: diffusion branch is merged. Imagenet models are still training and improving.
• 20240120: metrics are fixed. Diffusion branch started. imagenet classification progress (53%->60%).
• 20240119: support for lightning and hydra added! Welcome to the multigpu world!

Fix me

Easy targets if you want to contribute

• Fix the MAE training with lightning+hydra
• Make an evaluation script for MAE: it loads the encoder from a MAE checkpoint and trains a classifier on top of it on imagenet. Add the fine-tune all model option
• fix the diffusion samplers, they don't really work.
• Search a good architecture for diffusion, maybe inspired from RIN
• Make a script that leverages a search tool (like https://docs.ray.io) to search for good hyper params (lr, wd, order, order_expand and ffw_expand mainly)

Currently testing on

• Vision: ImageNet classification (best 224x224 model score so far: 61.7% top-1 // 20230122)
• Vision: Masked Auto Encoder pretraining
• Probabilistic Time Series Forecasting: Running comparisons against AutoML Forecasting evaluations

Launching a classification training run

This repo supports hydra for handling configs. Look at src/configs to edit them. Here is an example of a training run:

python src/train.py data.dataset_builder.data_dir=path_to_imagenet seed=3407 model.network.dim=128  data.size=224 model.network.kernel_size=32 model.network.nb_layers=12 model.network.order=2 model.network.order_expand=4 model.network.ffw_expand=4  model.network.dropout=0.0 model.optimizer.optim.weight_decay=0.01 model.optimizer.optim.lr=1e-3 data.full_batch_size=1024 trainer.max_steps=300000 model.lr_scheduler.warmup_steps=10000 computer.num_workers=8 computer.precision=bf16-mixed data/additional_train_transforms=randaugment data.additional_train_transforms.randaugment_p=0.1 data.additional_train_transforms.randaugment_magnitude=6 model.train_batch_preprocess.apply_transform_prob=1.0 checkpoint_dir="./checkpoints/"

Launching MAE training run

python src/train.py --config-name train_mae data.dataset_builder.data_dir=path_to_dataset seed=3407 model.network.dim=128  data.size=256 model.network.kernel_size=16 model.network.nb_layers=8 model.network.order=4 model.network.order_expand=8 model.network.ffw_expand=4  model.network.dropout=0.0 model.optimizer.optim.weight_decay=0.01 model.optimizer.optim.lr=1e-3 data.full_batch_size=256 trainer.max_steps=300000 model.lr_scheduler.warmup_steps=10000 computer.num_workers=8 computer.precision=bf16-mixed data/additional_train_transforms=randaugment data.additional_train_transforms.randaugment_p=0.1 data.additional_train_transforms.randaugment_magnitude=6 model.train_batch_preprocess.apply_transform_prob=1.0 checkpoint_dir="./checkpoints/"

GAT-HoMM: a Graph Neural Network with HoMM Attention

• Results: accuracy on the 1000 test nodes of the Cora dataset (https://arxiv.org/pdf/1710.10903.pdf) of 0.805
• To reproduce, execute:

  • python src/train_gnn.py
    

• To run hyperparameter optimization, execute:

  • python src/optimize_hps_gnn.py
    

• Illustrative notebook: src/gnn_homm_nb.ipynb
• Default configuration file for train_gnn.py: src/configs/train_gnn.yml
• Default configuration file for optimize_hps_gnn.py: src/configs/hp_opt_gnn.yml

TODO:

• Vision: diffusion model
• NLP: sentence embedding
• NLP: next token prediction
• Graphs?

Ablation

On imagenet, with the following parameters:

• image size: 160
• patch size: 16
• # of layers: 8
• batch size: 512
• weight decay: 0.01
• # of training steps: 150k
• optimizer: AdamW
• rand-augment + cutmix/mixup

| dim | o | oe | acc | Flops | # params | |-----|---|----|------|-------|----------| | 320 | 1 | 8 | 43.6 | 2.6G | 26M | | 320 | 2 | 4 | 47.6 | 2.6G | 26M | | 320 | 4 | 2 | 46.1 | 2.6G | 26M | | 256 | 2 | 8 | 47.9 | 2.9G | 29M | | 256 | 4 | 4 | 46.1 | 2.9G | 29M |

Clearly, having the second order makes a big difference. Having the fourth order not so much. It's better to have a higher dimension and lower expansion than the contrary.