Fast and Simplex: 2-Simplicial Attention in Triton

近期研究表明，训练损失会随着模型规模和所处理的词元（token）数量呈现出幂律关系，而要实现计算最优的模型，就需要同时扩展模型规模和词元数量。然而，这些缩放定律假设数据是无限供应的，并且主要适用于计算资源受限的情况。随着现代大语言模型越来越多地依赖于大规模的互联网数据集，它们受限于计算资源这一假设正逐渐失效。这种变化突显了需要优先考虑词元效率的模型架构。 在本研究中，我们探索了2-单纯形Transformer（2-simplicial Transformer）的应用。该架构通过高效的Triton内核实现，将标准的点积注意力机制推广到三线性函数。我们证明，与标准Transformer相比，2-单纯形Transformer具有更高的词元效率：在词元预算固定的情况下，规模相近的模型在数学、编程、推理和逻辑等任务上的表现优于传统的点积注意力模型。我们通过实验证明，在知识和推理任务中，2-单纯形注意力改变了缩放定律中的指数，从而量化了其性能增益。

论文旨在解决当前大规模语言模型中训练损失随模型大小和token数量变化的scaling laws假设不适用于数据受限情况的问题。随着模型越来越大，它们对互联网规模数据集的依赖使得计算限制（compute-bound）的假设不再成立，因此需要更注重token效率的架构。

论文提出了一种新的注意力机制——2-simplicial attention，它将标准点积注意力推广到三线性函数，并通过高效的Triton内核实现。这种架构在固定token预算下展现出比传统Transformer更强的token效率，在数学、编码、推理和逻辑任务上表现更优。此外，该方法改变了知识和推理任务中scaling laws的指数。

1. 2-simplicial Transformer在多个复杂任务上优于标准Transformer，尤其是在token效率方面 2. 实验验证了新架构在相同token预算下的性能优势 3. 作者量化了其对scaling laws的影响，指出其指数变化的意义 4. 使用了高效Triton kernel进行实现，提升了计算效率

1. Scaling Laws for Neural Language Models (arXiv:2001.08361) 2. Language Models (Mostly) Know What They Know (arXiv:2107.08981) 3. Efficient Large-Scale Language Model Training via Sketching-Based Acceleration (arXiv:2204.06683) 4. Training Compute-Optimal Large Language Models (arXiv:2203.03545)