Fast and Simplex: 2-Simplicial Attention in Triton

Aurko Roy, Timothy Chou, Sai Surya Duvvuri, Sijia Chen, Jiecao Yu, Xiaodong Wang, Ying Zhang, Manzil Zaheer, Rohan Anil

Keywords: higher order attention, sliding window attention, triton

TL;DR: Higher order sliding window attention improves scaling law exponents on reasoning tasks

Abstract:

Recent work has shown that training loss scales as a power law with both model size and the
number of tokens, and that achieving compute-optimal models requires scaling model size and
token count together. However, these scaling laws assume an infinite supply of data and
apply primarily in compute-bound settings. As modern large language models increasingly rely
on massive internet-scale datasets, the assumption that they are compute-bound is becoming
less valid. This shift highlights the need for architectures that prioritize token
efficiency.

In this work, we investigate the use of the 2-simplicial Transformer, an architecture that
generalizes standard dot-product attention to trilinear functions through an efficient
Triton kernel implementation. We demonstrate that the 2-simplicial Transformer achieves
better token efficiency than standard Transformers: for a fixed token budget, similarly
sized models outperform their dot-product counterparts on tasks involving mathematics,
coding, reasoning, and logic. We quantify these gains by demonstrating
that $2$-simplicial attention changes the exponent
in the scaling laws for knowledge and reasoning tasks compared to dot product attention.

Primary Area: foundation or frontier models, including LLMs

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2026/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 10293