Schnell und Einfach: 2-Simpliziale Aufmerksamkeit in Triton

Fast and Simplex: 2-Simplicial Attention in Triton

July 3, 2025

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

cs.AI

Zusammenfassung

Aktuelle Arbeiten haben gezeigt, dass der Trainingsverlust als Potenzgesetz sowohl mit der Modellgröße als auch mit der Anzahl der Tokens skaliert und dass die Erreichung rechenoptimaler Modelle eine gemeinsame Skalierung von Modellgröße und Tokenanzahl erfordert. Diese Skalierungsgesetze gehen jedoch von einer unbegrenzten Datenmenge aus und gelten hauptsächlich in rechenbeschränkten Szenarien. Da moderne große Sprachmodelle zunehmend auf massive, internetgroße Datensätze angewiesen sind, wird die Annahme, dass sie rechenbeschränkt sind, immer weniger zutreffend. Diese Verschiebung unterstreicht die Notwendigkeit von Architekturen, die die Token-Effizienz priorisieren. In dieser Arbeit untersuchen wir die Verwendung des 2-simplizialen Transformers, einer Architektur, die die Standard-Dot-Produkt-Attention auf trilineare Funktionen verallgemeinert und durch eine effiziente Triton-Kernel-Implementierung realisiert wird. Wir zeigen, dass der 2-simpliziale Transformer eine bessere Token-Effizienz als Standard-Transformer erreicht: Bei einem festen Token-Budget übertreffen gleich große Modelle ihre Dot-Produkt-Pendants bei Aufgaben in den Bereichen Mathematik, Programmierung, logisches Denken und Logik. Wir quantifizieren diese Verbesserungen, indem wir demonstrieren, dass die 2-simpliziale Attention den Exponenten in den Skalierungsgesetzen für Wissens- und Denkaufgaben im Vergleich zur Dot-Produkt-Attention verändert.

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.