Быстро и просто: 2-симплициальное внимание в Triton

Fast and Simplex: 2-Simplicial Attention in Triton

July 3, 2025

Авторы: Aurko Roy, Timothy Chou, Sai Surya Duvvuri, Sijia Chen, Jiecao Yu, Xiaodong Wang, Manzil Zaheer, Rohan Anil

cs.AI

Аннотация

Недавние исследования показали, что ошибка обучения масштабируется по степенному закону как с размером модели, так и с количеством токенов, и что достижение оптимальных по вычислительным ресурсам моделей требует совместного масштабирования размера модели и количества токенов. Однако эти законы масштабирования предполагают бесконечный объем данных и применимы в основном в условиях, ограниченных вычислительными ресурсами. Поскольку современные крупные языковые модели всё больше полагаются на огромные наборы данных интернет-масштаба, предположение о том, что они ограничены вычислительными ресурсами, становится менее справедливым. Этот сдвиг подчеркивает необходимость в архитектурах, которые уделяют приоритетное внимание эффективности использования токенов. В данной работе мы исследуем использование 2-симплициального Transformer — архитектуры, которая обобщает стандартное скалярное произведение внимания на трилинейные функции с помощью эффективной реализации ядра Triton. Мы демонстрируем, что 2-симплициальный Transformer достигает лучшей эффективности использования токенов по сравнению со стандартными Transformer: при фиксированном бюджете токенов модели аналогичного размера превосходят свои аналоги с использованием скалярного произведения в задачах, связанных с математикой, программированием, рассуждениями и логикой. Мы количественно оцениваем эти преимущества, показывая, что 2-симплициальное внимание изменяет показатель в законах масштабирования для задач, связанных с знаниями и рассуждениями, по сравнению со скалярным произведением внимания.

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.