multiscreen 0.1.8

Implementation of MultiScreen

MultiScreen

Implementation of Multiscreen proposed by Ken Nakanishi for Screening is Enough

Basically it is a non-softmax attention with ReLU squared activation, content similarity thresholding, and aggressive normalization of the values.

Install

$ pip install multiscreen

Usage

import torch
from multiscreen import MultiScreen

multi_screen = MultiScreen(
    num_tokens = 256,
    dim = 512,
    depth = 6,
    heads = 8,
    dim_keys = 16,     # paper says 16 or 32
    dim_values = 64    # paper says 64 or 128
)

token_ids = torch.randint(0, 256, (1, 1024))

logits = multi_screen(token_ids)
assert logits.shape == (1, 1024, 256)

Enwik8

$ uv run train.py

Citations

@misc{nakanishi2026screening,
    title   = {Screening Is Enough},
    author  = {Ken M. Nakanishi},
    year    = {2026},
    eprint  = {2604.01178},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2604.01178},
}

@misc{zhang2021sparse,
    title   = {Sparse Attention with Linear Units},
    author  = {Biao Zhang and Ivan Titov and Rico Sennrich},
    year    = {2021},
    eprint  = {2104.07012},
    archivePrefix = {arXiv},
    primaryClass = {cs.CL}
}

@misc{horuz2025resurrectionrelu,
    title   = {The Resurrection of the ReLU},
    author  = {Coşku Can Horuz and Geoffrey Kasenbacher and Saya Higuchi and Sebastian Kairat and Jendrik Stoltz and Moritz Pesl and Bernhard A. Moser and Christoph Linse and Thomas Martinetz and Sebastian Otte},
    year    = {2025},
    eprint  = {2505.22074},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2505.22074},
}