快速且单纯：基于Triton的2-单纯注意力机制

近期工作表明，训练损失随模型大小和token数量呈幂律标度，且实现计算最优模型需要同时缩放模型大小和token数量。然而，这些标度律假设数据无限供应，且主要适用于计算受限场景。随着现代大型语言模型日益依赖大规模互联网数据集，计算受限的假设正变得不再成立。这一转变凸显了对优先考虑token效率的架构的需求。在本工作中，我们研究了2-单纯Transformer，该架构通过高效的Triton内核实现，将标准点积注意力泛化为三线性函数。我们证明，2-单纯Transformer比标准Transformer具有更好的token效率：在固定token预算下，同等大小的模型在涉及数学、编码、推理和逻辑的任务上优于其点积对应模型。我们通过证明与点积注意力相比，2-单纯注意力改变了知识和推理任务标度律中的指数，来量化这些增益。