2-单纯形注意力提升模型效率,相同算力下表现更优
Fast and Simplex: 2-Simplicial Attention in Triton
- 用三线性函数替代标准注意力,通过Triton高效实现
- 在数学、编程等任务上,相同token数下性能显著提升
- 改变知识与推理任务的缩放指数,适合资源受限场景
近期研究表明,训练损失随模型规模和词元数量呈幂律增长,且达到计算最优需同步扩大模型与数据量。但这些规律假设数据无限,主要适用于计算瓶颈场景。随着大语言模型越来越多依赖互联网级数据集,计算瓶颈假设逐渐失效。这凸显了提升词元效率架构的重要性。本文研究2-单纯形Transformer,该架构通过高效Triton内核实现,将标准点积注意力推广至三线性函数。实验表明,在固定词元预算下,同等规模的2-单纯形模型在数学、编程、推理与逻辑任务上优于标准Transformer。我们量化了这一优势:2-单纯形注意力改变了知识与推理任务的缩放律指数,相较于点积注意力更具优势。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。