[NeurIPS 2020] Official Implementation: "SMYRF: Efficient Attention using Asymmetric Clustering".

Last update: Dec 22, 2022

Related tags

Deep Learning smyrf

Overview

SMYRF: Efficient attention using asymmetric clustering

Get started:

Abstract

We propose a novel type of balanced clustering algorithm to approximate attention. Attention complexity is reduced from O(N^2) to O(NlogN), where N is the sequence length. Our algorithm, SMYRF, uses Locality Sensitive Hashing (LSH) in a novel way by defining new Asymmetric transformations and an adaptive scheme that produces balanced clusters. The biggest advantage of SMYRF is that it can be used as a drop-in replacement for dense attention layers without any retraining. On the contrary, prior fast attention methods impose constraints (e.g. tight queries and keys) and require re-training from scratch. We apply our method to pre-trained state-of-the-art Natural Language Processing and Computer Vision models and we report significant memory and speed benefits. Notably, SMYRF-BERT outperforms (slightly) BERT on GLUE, while using $50%$ less memory. We also show that SMYRF can be used interchangeably with dense attention before and after training. Finally, we use SMYRF to train GANs with attention in high resolutions. Using a single TPU, we train BigGAN on Celeba-HQ, with attention at resolution 128x128 and 256x256, capable of generating realistic human faces.

Authors: Giannis Daras, Nikita Kitaev, Augustus Odena, Alexandros G. Dimakis

Results

Memory-quality trade-off

GLUE benchmark

	Avg.	#	C	CoLA	MNLI-m/mm	MRPC	QNLI	QQP	RTE	SST-2	STS-B
BERT₁₂₈	82.69	1	1	57.83	84.43/84.68	88.41	91.31	89.70	65.70	93.46	88.73
SMYRF-BERT_2x32	82.98	2	32	58.79	83.76/84.27	87.69	91.14	89.72	68.59	93.23	89.65
SMYRF-BERT_2x16	81.74	2	16	58.90	82.86/83.49	85.72	89.53	89.33	64.98	93.12	87.75
BERT₆₄	81.57	1	64	58.80	82.34/82.47	87.02	90.48	89.69	61.73	93.00	88.64
BERT₃₂	73.56	1	32	56.40	64.51/63.41	77.89	79.81	88.59	55.23	92.66	83.53

Interchangeability of SMYRF and dense attention

Results on IMDB dataset. Using dense attention on inference consistently improves results, nearly matching dense attention perf.

	Memory	SMYRF Inference	Accuracy
RoBERTa	100%	☒	94.96%
SMYRF-RoBERTa	50%	☒	93.72%
SMYRF-RoBERTa	50%	☑	94.62%
BERT	100%	☒	94.12%
SMYRF-BERT	50%	☒	92.64%
SMYRF-BERT	50%	☑	93.54%

Smyrf-BigGAN training on Celeba-HQ-128

Generated faces by a Smyrf-BigGAN trained on 128x128 resolution with attention at 128x128, using 50% of dense memory.

Results after 120k iterations:

	Resolution	Attention	#	C	FID
BigGAN	128x128	64x64	1	4096	26.06
Smyrf-BigGAN	128x128	128x128	4	2048	25.03

where # denotes number of hashes and C number of queries per cluster.

What's here

The code hosted in this repository is the one we used to run all the experiments in the paper. Get started:

For a deeper dive, look at the examples/ folder where we have code for pre-training SMYRF-BigGAN, sampling from a pre-trained BigGAN with SMYRF, finetuning state-of-the-art NLP models with SMYRF and a lot more.

Acknowledgments

We would like to wholeheartedly thank the TensorFlow Research Cloud (TFRC) program that gave us access to Cloud TPUs and GCP credits to train our models.

The code for the NLP experiments is exclusively based on the HuggingFace transformers library. We are very grateful to the authors of the library for their work.

The code for the CV experiments is based on the PyTorch implementation of BigGAN available in this url. The code has been expanded to support training on TPUs. Again, we want to thank the author for open-sourcing this implementation.

Comments

Auto-regressive

Hi Giannis!

Thanks for the great paper! I am interested in your asymmetric LSH, as I think having separate query / key space (as opposed to shared QK as in Reformer) will bring performance improvements in LSH-based attention.

I saw that you recommended to a previous user to use this form of clustering for the auto-regressive case, and just wanted to probe if you had considered the scenario where a bucket of queries do not get matched with any keys from the past at all. This was an issue I had with trying to make separate QK space work with routing transformer, but just wondering if you had identified and found a solution to this problem.

Phil

opened by lucidrains 2
Logging and scoring

Currently logging and scoring is disabled for TPU BigGAN for maximum efficiency. We can probably re-write the logger and scorer to lower their performance bottleneck by converting most cpu materializations to XLA ops.
bug example

opened by giannisdaras 0
Ema not working on TPU

Exponential moving average on weights of G is not working on TPUs. The problem is related to the loading of the state dict: https://github.com/ajbrock/BigGAN-PyTorch/blob/master/utils.py#L614

For now, we disable ema.
bug example

opened by giannisdaras 0

[NeurIPS 2020] Official Implementation: "SMYRF: Efficient Attention using Asymmetric Clustering".

Related tags

Overview

SMYRF: Efficient attention using asymmetric clustering

Abstract

Results

Memory-quality trade-off

GLUE benchmark

Interchangeability of SMYRF and dense attention

Smyrf-BigGAN training on Celeba-HQ-128

What's here

Acknowledgments

You might also like...

Code for ICE-BeeM paper - NeurIPS 2020

Code for Discriminative Sounding Objects Localization (NeurIPS 2020)

Advances in Neural Information Processing Systems (NeurIPS), 2020.

Neuron Merging: Compensating for Pruned Neurons (NeurIPS 2020)

Multi-Task Temporal Shift Attention Networks for On-Device Contactless Vitals Measurement (NeurIPS 2020)

Defending graph neural networks against adversarial attacks (NeurIPS 2020)

Code for the Population-Based Bandits Algorithm, presented at NeurIPS 2020.

Code release for NeurIPS 2020 paper "Co-Tuning for Transfer Learning"

Discovering Interpretable GAN Controls [NeurIPS 2020]

Comments

Auto-regressive

Logging and scoring

Ema not working on TPU

Releases(1.0)

1.0(Apr 30, 2020)

Owner

Giannis Daras

《LXMERT: Learning Cross-Modality Encoder Representations from Transformers》(EMNLP 2020)

A simple tutoral for error correction task, based on Pytorch

The official implementation of our CVPR 2021 paper - Hybrid Rotation Averaging: A Fast and Robust Rotation Averaging Approach

Clairvoyance: a Unified, End-to-End AutoML Pipeline for Medical Time Series

A Broad Study on the Transferability of Visual Representations with Contrastive Learning

Multi-Agent Reinforcement Learning (MARL) method to learn scalable control polices for multi-agent target tracking.

Densely Connected Convolutional Networks, In CVPR 2017 (Best Paper Award).

Source for the paper "Universal Activation Function for machine learning"

PyToch implementation of A Novel Self-supervised Learning Task Designed for Anomaly Segmentation

I created My own Virtual Artificial Intelligence named genesis, He can assist with my Tasks and also perform some analysis,,

Rethinking Transformer-based Set Prediction for Object Detection

potpourri3d - An invigorating blend of 3D geometry tools in Python.

SparseML is a libraries for applying sparsification recipes to neural networks with a few lines of code, enabling faster and smaller models

Liver segmentation using MONAI and pytorch

A Deep Learning based project for creating line art portraits.

Pytorch implementation of Distributed Proximal Policy Optimization: https://arxiv.org/abs/1707.02286

Memory Defense: More Robust Classificationvia a Memory-Masking Autoencoder

[CVPR 2022] Pytorch implementation of "Templates for 3D Object Pose Estimation Revisited: Generalization to New objects and Robustness to Occlusions" paper

Code for Referring Image Segmentation via Cross-Modal Progressive Comprehension, CVPR2020.

A transformer-based method for Healthcare Image Captioning in Vietnamese