Toeplitz MLP Mixer：低复杂度且信息丰富的序列模型

基于 Transformer 的大语言模型在某些方面受限于注意力机制二次方的时间和空间计算复杂度。我们引入了 Toeplitz MLP Mixer (TMM)，一种类 Transformer 架构，它将注意力替换为在序列维度上的三角掩码 Toeplitz 矩阵乘法，在训练期间产生 $\mathcal{O}(dn \log n)$ 的时间和 $\mathcal{O}(dn)$ 的空间复杂度，并在推理预填充阶段产生 $\mathcal{O}(dn)$ 的时间和空间复杂度。尽管缺乏其他次二次方架构中复杂的输入调制或状态维护机制，TMM 在每单位计算量和设备内存所达到的损失方面产生了更高的训练效率。我们证明 TMM 能够保留更多输入信息，从而提升了复制能力，我们认为这归因于其缺乏架构偏置。与更高的输入信息保留率相一致，与同类架构相比，TMM 在信息检索和上下文学习基准准确率上表现出优越的性能。最后，我们从算子指标理论的角度进行了分析，并表明，与直觉相反，因果不可逆模型的已训练 Toeplitz 层比在其输入上实际可逆的模型更有可能是可逆的或近乎可逆的。