Rapide et Simple : Attention 2-Simplicale dans 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

Résumé

Des travaux récents ont montré que la perte d'entraînement évolue selon une loi de puissance à la fois avec la taille du modèle et le nombre de tokens, et que l'obtention de modèles optimaux en termes de calcul nécessite de faire évoluer conjointement la taille du modèle et le nombre de tokens. Cependant, ces lois de mise à l'échelle supposent une quantité infinie de données et s'appliquent principalement dans des contextes limités par le calcul. Alors que les modèles de langage modernes reposent de plus en plus sur des ensembles de données massifs à l'échelle d'Internet, l'hypothèse selon laquelle ils sont limités par le calcul devient moins valable. Ce changement met en évidence la nécessité d'architectures qui privilégient l'efficacité en termes de tokens. Dans ce travail, nous étudions l'utilisation du Transformer 2-simplicial, une architecture qui généralise l'attention par produit scalaire standard à des fonctions trilinéaires grâce à une implémentation efficace via un noyau Triton. Nous démontrons que le Transformer 2-simplicial atteint une meilleure efficacité en tokens que les Transformers standards : pour un budget de tokens fixe, des modèles de taille similaire surpassent leurs homologues à produit scalaire sur des tâches impliquant les mathématiques, le codage, le raisonnement et la logique. Nous quantifions ces gains en montrant que l'attention 2-simpliciale modifie l'exposant dans les lois de mise à l'échelle pour les tâches de connaissance et de raisonnement par rapport à l'attention par produit scalaire.

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.