https://github.com/lucidrains/peer-pytorch
Pytorch implementation of the PEER block from the paper, Mixture of A Million Experts, by Xu Owen He at Deepmind
https://github.com/lucidrains/peer-pytorch
artificial-intelligence deep-learning mixture-of-experts product-key
Last synced: about 1 year ago
JSON representation
Pytorch implementation of the PEER block from the paper, Mixture of A Million Experts, by Xu Owen He at Deepmind
- Host: GitHub
- URL: https://github.com/lucidrains/peer-pytorch
- Owner: lucidrains
- License: mit
- Created: 2024-07-09T15:46:42.000Z (almost 2 years ago)
- Default Branch: main
- Last Pushed: 2024-08-23T13:19:56.000Z (almost 2 years ago)
- Last Synced: 2025-04-10T01:15:14.439Z (about 1 year ago)
- Topics: artificial-intelligence, deep-learning, mixture-of-experts, product-key
- Language: Python
- Homepage:
- Size: 271 KB
- Stars: 123
- Watchers: 4
- Forks: 3
- Open Issues: 4
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
## PEER - Pytorch
Pytorch implementation of the PEER block from the Deepmind paper, Mixture of A Million Experts, by Xu Owen He.
## Install
```bash
$ pip install PEER-pytorch
```
## Usage
```python
import torch
from PEER_pytorch import PEER
peer = PEER(
dim = 512,
heads = 8, # tested up to 32 - (hk = heads * num_experts_per_head (16))
num_experts = 1_000_000, # he chose 1 million
num_experts_per_head = 16, # he settled on 16, but was 32 in PKM paper
dim_key = 128,
pre_rmsnorm = True
).cuda()
x = torch.randn(2, 1024, 512).cuda()
out = peer(x) + x
assert x.shape == out.shape
```
## Citations
```bibtex
@inproceedings{He2024MixtureOA,
title = {Mixture of A Million Experts},
author = {Xu Owen He},
year = {2024},
url = {https://api.semanticscholar.org/CorpusID:271038610}
}
```
```bibtex
@article{Csordas2023ApproximatingTF,
title = {Approximating Two-Layer Feedforward Networks for Efficient Transformers},
author = {R'obert Csord'as and Kazuki Irie and J{\"u}rgen Schmidhuber},
journal = {ArXiv},
year = {2023},
volume = {abs/2310.10837},
url = {https://api.semanticscholar.org/CorpusID:264172384}
}
```