Snel en Simpel: 2-Simpliciale Aandacht in Triton

Fast and Simplex: 2-Simplicial Attention in Triton

July 3, 2025

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

cs.AI

Samenvatting

Recent werk heeft aangetoond dat het trainingsverlies schaalt volgens een machtswet met zowel de modelgrootte als het aantal tokens, en dat het bereiken van compute-optimale modellen vereist dat de modelgrootte en het tokenaantal samen worden geschaald. Deze schaalwetten gaan echter uit van een oneindige hoeveelheid data en zijn voornamelijk van toepassing in compute-gebonden situaties. Aangezien moderne grote taalmodellen steeds meer vertrouwen op enorme internet-schaal datasets, wordt de aanname dat ze compute-gebonden zijn steeds minder geldig. Deze verschuiving benadrukt de noodzaak voor architecturen die token-efficiëntie prioriteren. In dit werk onderzoeken we het gebruik van de 2-simpliciale Transformer, een architectuur die standaard dot-product aandacht generaliseert naar trilineaire functies via een efficiënte Triton kernel-implementatie. We tonen aan dat de 2-simpliciale Transformer betere token-efficiëntie bereikt dan standaard Transformers: voor een vast tokenbudget presteren vergelijkbaar grote modellen beter dan hun dot-product tegenhangers bij taken die wiskunde, coderen, redeneren en logica betreffen. We kwantificeren deze winst door aan te tonen dat 2-simpliciale aandacht de exponent in de schaalwetten verandert voor kennis- en redeneertaken in vergelijking met dot-product aandacht.

English

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.