快速且简单：Triton 中的 2-单纯形注意力

最近的工作表明，训练损失随着模型大小和 token 数量的增加而呈现幂律关系，并且实现计算最佳模型需要同时扩展模型大小和 token 数量。然而，这些缩放定律假设数据无限供应，并且主要适用于计算受限的设置。随着现代大型语言模型越来越依赖于大规模的互联网规模数据集，它们是计算受限的假设正变得越来越不有效。这种转变凸显了对优先考虑 token 效率的架构的需求。 在这项工作中，我们研究了 2-单面 Transformer 的使用，这是一种通过高效的 Triton 内核实现将标准点积注意力推广到三线性函数的架构。我们证明了 2-单面 Transformer 比标准 Transformer 实现了更好的 token 效率：对于固定的 token 预算，类似大小的模型在涉及数学、编码、推理和逻辑的任务中优于它们的点积对应模型。我们通过证明与点积注意力相比，2-单面注意力改变了知识和推理任务缩放定律中的指数来量化这些收益。