高速かつ単純：Tritonにおける2-単体的アテンション

要旨

最近の研究では、訓練損失がモデルサイズとトークン数の両方に対してべき乗則に従ってスケールすること、そして計算最適なモデルを達成するためにはモデルサイズとトークン数を共にスケールさせる必要があることが示されています。しかし、これらのスケーリング則は無限のデータ供給を前提としており、主に計算資源が制約となる状況に適用されます。現代の大規模言語モデルがますます大規模なインターネット規模のデータセットに依存するようになるにつれ、それらが計算資源に制約されているという前提は次第に妥当性を失っています。この変化は、トークン効率を優先するアーキテクチャの必要性を浮き彫りにしています。本研究では、標準的なドット積アテンションを効率的なTritonカーネル実装を通じて三線形関数に一般化する2-シンプリシャルトランスフォーマーの使用を調査します。2-シンプリシャルトランスフォーマーが標準的なトランスフォーマーよりも優れたトークン効率を達成することを実証します：固定されたトークン予算のもとで、同程度のサイズのモデルが数学、コーディング、推論、論理を含むタスクにおいてドット積ベースのモデルを上回ります。これらの利点を定量化するために、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