Fast and Simplex: 2-Simplicial Attention in Triton
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.
686732e19db35afc9c304cc9
Aurko Roy
686732e19db35afc9c304cca
Timothy Chou
686732e19db35afc9c304ccb
Sai Surya Duvvuri
686732e19db35afc9c304ccc
Sijia Chen
686732e19db35afc9c304ccd
Jiecao Yu
686732e19db35afc9c304cce
Xiaodong Wang
686732e19db35afc9c304ccf
Manzil Zaheer
686732e19db35afc9c304cd0
Rohan Anil
Fast and Simplex: 2-Simplicial Attention in Triton
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.
The 2-simplicial Transformer improves token efficiency over standard Transformers, offering better performance on knowledge and reasoning tasks with a fixed token budget.
2-simplicial Transformer trilinear functions Triton kernel dot-product attention token efficiency scaling laws knowledge tasks reasoning tasks