PEER-pytorch

PEER - Pytorch

Pytorch implementation of the PEER block from the Deepmind paper, Mixture of A Million Experts, by Xu Owen He.

Install

$ pip install PEER-pytorch

Usage

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

@inproceedings{He2024MixtureOA,
    title   = {Mixture of A Million Experts},
    author  = {Xu Owen He},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:271038610}
}

@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}
}

@inproceedings{anonymous2025continual,
    title   = {Continual Learning via Sparse Memory Finetuning},
    author  = {Anonymous},
    booktitle = {Submitted to The Fourteenth International Conference on Learning Representations},
    year    = {2025},
    url     = {https://openreview.net/forum?id=LGo7U1m24L},
    note    = {under review}
}

Metadata

Release files for PEER-pytorch 0.2.2

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for PEER-pytorch 0.2.2
File	Size	Uploaded
peer_pytorch-0.2.2.tar.gz	268.3 kB	Nov 1, 2025	Details

Built distribution (wheel)

Table of built distributions (wheels) for PEER-pytorch 0.2.2
File	Interpreter	ABI	Platform	Reset
peer_pytorch-0.2.2-py3-none-any.whl 11.2 kB Nov 1, 2025	Python 3	none	any	Details

Total release size: 279.6 kB

Release files / peer_pytorch-0.2.2.tar.gz

Download URL	peer_pytorch-0.2.2.tar.gz
Size	268.3 kB
Tags	Source
SHA-256 checksum How to use checksums	`54e9794c3e8df635e8b4d6a1f5dd8c24c65c85e9ca5c8c3ea6a03ce1525c7963`
BLAKE2b-256 checksum How to use checksums	`575102cb3eb051abe7ad12ad80ed196943d6969fb8bc6d74430588b9f6e2dc45`
Upload date	Nov 1, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.9.25`

Release files / peer_pytorch-0.2.2-py3-none-any.whl

Download URL	peer_pytorch-0.2.2-py3-none-any.whl
Size	11.2 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`e9bc1a251bb428ab6ce4043f28ecaacd9a3de39b2b13b1133eacbc5c767e0d8d`
BLAKE2b-256 checksum How to use checksums	`1dc7beb72a8f683baf3a3ffa1f0b76974854b4dac4a5fdf4f661190cdae71143`
Upload date	Nov 1, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.9.25`