Poly Attention
Implementation of Poly-Attention, a general scheme for higher-order self-attention
Install
$ pip install poly-attention
Usage
import torch
from poly_attention import PolyAttention
attn = PolyAttention(
dim = 512,
heads = 8,
dim_head = 64,
causal = False
)
tokens = torch.randn(1, 1024, 512)
out = attn(tokens) # (1, 1024, 512)
A Vision Transformer based on Poly-Attention
import torch
from poly_attention import PolyViT
vit = PolyViT(
image_size = 256,
patch_size = 32,
num_classes = 1000,
dim = 1024,
depth = 6,
heads = 16,
mlp_dim = 2048,
order = 2 # standard poly attention order 2
)
images = torch.randn(1, 3, 256, 256)
preds = vit(images) # (1, 1000)
Quick test
python train_function_composition.py --poly_layers=1 --base_layers=2
Appreciation
-
@dillfrescott for submitting a stability fix
-
@pranoyr for more efficient caching of the GQA heads!
Citations
@inproceedings{chakrabarti2026poly,
title = {Poly-attention: a general scheme for higher-order self-attention},
author = {Chakrabarti, Sayak and Pitassi, Toniann and Alman, Josh},
booktitle = {International Conference on Learning Representations (ICLR)},
year = {2026}
}
@misc{kayyam2026transformersneedprojectionssystematic,
title = {Do Transformers Need Three Projections? Systematic Study of QKV Variants},
author = {Ali Kayyam and Anusha Madan Gopal and M Anthony Lewis},
year = {2026},
eprint = {2606.04032},
archivePrefix = {arXiv},
primaryClass = {cs.LG},
url = {https://arxiv.org/abs/2606.04032},
}
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
poly_attention-0.2.5.tar.gz
(11.7 kB
view details)
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file poly_attention-0.2.5.tar.gz.
File metadata
- Download URL: poly_attention-0.2.5.tar.gz
- Upload date:
- Size: 11.7 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.8.17
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
44c251ea35e738fd9d6902c3da5e396ef67fa79b7df1086baf704299c246d606
|
|
| MD5 |
2497430f4c9c1accabccc9c7bb1b877b
|
|
| BLAKE2b-256 |
6d86ee001c0acc0d5333d2c0dbd63b9ed4180e290fcb5cfdc761249745007ec4
|
File details
Details for the file poly_attention-0.2.5-py3-none-any.whl.
File metadata
- Download URL: poly_attention-0.2.5-py3-none-any.whl
- Upload date:
- Size: 13.1 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via:
uv/0.8.17
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
9eab10c5b258057f3a148a2675a1e4c986595710ad04dba483abec6c09f5d1e4
|
|
| MD5 |
38292253c13a2570859b8bfe266d9332
|
|
| BLAKE2b-256 |
0faf4b1a288cb661ee579160f05caf9e5cbfe57bb1a6bc4b5cd24e4f383fa27a
|
