strassen-attention 0.1.5

Strassen Attention

Strassen Attention

Implementation of Strassen attention, from Kozachinskiy et al. of National Center of AI in Chile 🇨🇱

Install

$ pip install strassen-attention

Usage

import torch
from strassen_attention import strassen_attend

q = torch.randn(1, 8, 32, 16)
k = torch.randn(1, 8, 32, 16)
v = torch.randn(1, 8, 32, 16)

attended = strassen_attend(
    q,
    k,
    k.clone(),
    v,
    v.clone()
)

assert attended.shape == q.shape

For the multi-head attention module

import torch
from strassen_attention.strassen_mha import StrassenMHA

mha = StrassenMHA(dim = 512, causal = True)

tokens = torch.randn(1, 256, 512)

assert mha(tokens).shape == tokens.shape

Strassen attention transformer

import torch
import torch
from strassen_attention.strassen_transformer import StrassenTransformer

transformer = StrassenTransformer(dim = 512, depth = 4)

x = torch.randn(1, 16 * 16, 512)
assert transformer(x).shape == x.shape

Citations

@misc{kozachinskiy2025strassenattentionunlockingcompositional,
    title   = {Strassen Attention: Unlocking Compositional Abilities in Transformers Based on a New Lower Bound Method}, 
    author  = {Alexander Kozachinskiy and Felipe Urrutia and Hector Jimenez and Tomasz Steifer and Germán Pizarro and Matías Fuentes and Francisco Meza and Cristian B. Calderon and Cristóbal Rojas},
    year    = {2025},
    eprint  = {2501.19215},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2501.19215}, 
}

@article{Peng2024OnLO,
    title     = {On Limitations of the Transformer Architecture},
    author    = {Binghui Peng and Srini Narayanan and Christos Papadimitriou},
    journal   = {ArXiv},
    year      = {2024},
    volume    = {abs/2402.08164},
    url       = {https://api.semanticscholar.org/CorpusID:267636545}
}