https://github.com/lucidrains/g-mlp-pytorch
Implementation of gMLP, an all-MLP replacement for Transformers, in Pytorch
https://github.com/lucidrains/g-mlp-pytorch
artificial-intelligence deep-learning multilayer-perceptron
Last synced: 9 months ago
JSON representation
Implementation of gMLP, an all-MLP replacement for Transformers, in Pytorch
- Host: GitHub
- URL: https://github.com/lucidrains/g-mlp-pytorch
- Owner: lucidrains
- License: mit
- Created: 2021-05-18T02:21:09.000Z (about 5 years ago)
- Default Branch: main
- Last Pushed: 2021-08-14T18:54:41.000Z (almost 5 years ago)
- Last Synced: 2025-09-25T16:07:32.610Z (10 months ago)
- Topics: artificial-intelligence, deep-learning, multilayer-perceptron
- Language: Python
- Homepage:
- Size: 34.2 MB
- Stars: 430
- Watchers: 8
- Forks: 58
- Open Issues: 4
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
## gMLP - Pytorch
Implementation of gMLP, an all-MLP replacement for Transformers, in Pytorch
## Install
```bash
$ pip install g-mlp-pytorch
```
## Usage
For masked language modelling
```python
import torch
from torch import nn
from g_mlp_pytorch import gMLP
model = gMLP(
num_tokens = 20000,
dim = 512,
depth = 6,
seq_len = 256,
circulant_matrix = True, # use circulant weight matrix for linear increase in parameters in respect to sequence length
act = nn.Tanh() # activation for spatial gate (defaults to identity)
)
x = torch.randint(0, 20000, (1, 256))
logits = model(x) # (1, 256, 20000)
```
For image classification
```python
import torch
from g_mlp_pytorch import gMLPVision
model = gMLPVision(
image_size = 256,
patch_size = 16,
num_classes = 1000,
dim = 512,
depth = 6
)
img = torch.randn(1, 3, 256, 256)
logits = model(img) # (1, 1000)
```
You can also add a tiny amount of attention (one-headed) to boost performance, as mentioned in the paper as `aMLP`, with the addition of one extra keyword `attn_dim`. This applies to both `gMLPVision` and `gMLP`
```python
import torch
from g_mlp_pytorch import gMLPVision
model = gMLPVision(
image_size = 256,
patch_size = 16,
num_classes = 1000,
dim = 512,
depth = 6,
attn_dim = 64
)
img = torch.randn(1, 3, 256, 256)
pred = model(img) # (1, 1000)
```
Non-square images and patch sizes
```python
import torch
from g_mlp_pytorch import gMLPVision
model = gMLPVision(
image_size = (256, 128),
patch_size = (16, 8),
num_classes = 1000,
dim = 512,
depth = 6,
attn_dim = 64
)
img = torch.randn(1, 3, 256, 128)
pred = model(img) # (1, 1000)
```
## Experimental
A independent researcher proposes using a multi-headed approach for gMLPs in a blogpost on Zhihu. To do so, just set `heads` to be greater than `1`
```python
import torch
from torch import nn
from g_mlp_pytorch import gMLP
model = gMLP(
num_tokens = 20000,
dim = 512,
depth = 6,
seq_len = 256,
causal = True,
circulant_matrix = True,
heads = 4 # 4 heads
)
x = torch.randint(0, 20000, (1, 256))
logits = model(x) # (1, 256, 20000)
```
## Citations
```bibtex
@misc{liu2021pay,
title = {Pay Attention to MLPs},
author = {Hanxiao Liu and Zihang Dai and David R. So and Quoc V. Le},
year = {2021},
eprint = {2105.08050},
archivePrefix = {arXiv},
primaryClass = {cs.LG}
}
```
```bibtex
@software{peng_bo_2021_5196578,
author = {PENG Bo},
title = {BlinkDL/RWKV-LM: 0.01},
month = aug,
year = 2021,
publisher = {Zenodo},
version = {0.01},
doi = {10.5281/zenodo.5196578},
url = {https://doi.org/10.5281/zenodo.5196578%7D
}
```