Rápido y Simple: Atención 2-Simpléctica en 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

Resumen

Trabajos recientes han demostrado que la pérdida durante el entrenamiento escala como una ley de potencia tanto con el tamaño del modelo como con el número de tokens, y que lograr modelos óptimos en términos de cómputo requiere escalar conjuntamente el tamaño del modelo y el recuento de tokens. Sin embargo, estas leyes de escalado asumen un suministro infinito de datos y se aplican principalmente en entornos limitados por el cómputo. A medida que los modelos de lenguaje modernos y de gran escala dependen cada vez más de conjuntos de datos masivos a escala de internet, la suposición de que están limitados por el cómputo se está volviendo menos válida. Este cambio resalta la necesidad de arquitecturas que prioricen la eficiencia en el uso de tokens. En este trabajo, investigamos el uso del Transformer 2-simplicial, una arquitectura que generaliza la atención estándar de producto punto a funciones trilineales mediante una implementación eficiente de un kernel Triton. Demostramos que el Transformer 2-simplicial logra una mayor eficiencia en el uso de tokens que los Transformers estándar: para un presupuesto fijo de tokens, modelos de tamaño similar superan a sus contrapartes de producto punto en tareas que involucran matemáticas, programación, razonamiento y lógica. Cuantificamos estas mejoras al demostrar que la atención 2-simplicial modifica el exponente en las leyes de escalado para tareas de conocimiento y razonamiento en comparación con la atención de producto punto.

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.