Rápido e Simples: Atenção 2-Simpléxica em Triton

Fast and Simplex: 2-Simplicial Attention in Triton

July 3, 2025

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

cs.AI

Resumo

Trabalhos recentes mostraram que a perda de treinamento escala como uma lei de potência tanto com o tamanho do modelo quanto com o número de tokens, e que alcançar modelos computacionalmente ótimos requer escalar o tamanho do modelo e a contagem de tokens em conjunto. No entanto, essas leis de escala assumem um suprimento infinito de dados e se aplicam principalmente em cenários limitados por computação. À medida que os modelos de linguagem grandes modernos dependem cada vez mais de conjuntos de dados massivos em escala da internet, a suposição de que eles são limitados por computação está se tornando menos válida. Essa mudança destaca a necessidade de arquiteturas que priorizem a eficiência de tokens. Neste trabalho, investigamos o uso do Transformer 2-simplicial, uma arquitetura que generaliza a atenção de produto escalar padrão para funções trilineares por meio de uma implementação eficiente de kernel Triton. Demonstramos que o Transformer 2-simplicial alcança uma eficiência de tokens melhor do que os Transformers padrão: para um orçamento fixo de tokens, modelos de tamanho semelhante superam suas contrapartes de produto escalar em tarefas envolvendo matemática, codificação, raciocínio e lógica. Quantificamos esses ganhos demonstrando que a atenção 2-simplicial altera o expoente nas leis de escala para tarefas de conhecimento e raciocínio em comparação com a atenção de produto escalar.

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.