快速与单纯：Triton中的2-单纯形注意力机制

摘要

近期研究表明，訓練損失隨模型規模與標記數量呈冪律關係增長，且實現計算最優模型需同步擴展模型規模與標記數量。然而，這些擴展定律基於數據無限供應的假設，主要適用於計算受限的場景。隨著現代大型語言模型日益依賴於海量的互聯網規模數據集，它們處於計算受限的假設正逐漸失效。這一轉變凸顯了對優先考慮標記效率的架構之需求。在本研究中，我們探討了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.

快速与单纯：Triton中的2-单纯形注意力机制

Fast and Simplex: 2-Simplicial Attention in Triton

摘要

Support