超越 Transformer 与 Mamba，Meta 联合斯坦福等高校推出最强架构 TTT

近日，斯坦福、UCSD、UC 伯克利和 Meta 的研究人员提出了一种全新架构，用机器学习模型取代 RNN 的隐藏状态。

这个模型通过对输入 token 进行梯度下降来压缩上下文，这种方法被称为「测试时间训练层（Test-Time-Training layers，TTT）」。该研究作者之一 Karan Dalal 表示，他相信这将根本性地改变语言模型方法。

自注意力机制在处理长上下文时表现良好，但其复杂度是二次的。现有的 RNN 层具有线性复杂度，但其在长上下文中的表现受限于其隐藏状态的表达能力。随着上下文长度的增加，成本也会越来越高。

作者提出了一种具有线性复杂度和表达能力强的隐藏状态的新型序列建模层。关键思路是让隐藏状态本身成为一个机器学习模型，并将更新规则设为自监督学习的一步。

论文中提出了两种实例：TTT-Linear 和 TTT-MLP，它们的隐藏状态分别是线性模型和两层 MLP。团队在 125M 到 1.3B 参数规模上评估了实例，并与强大的 Transformer 和现代 RNN Mamba 进行了比较。结果显示，与 Mamba 相比，TTT-Linear 的困惑度更低，FLOP 更少（左），对长上下文的利用更好（右）：

这个结果代表了现有 RNN 的尴尬现实。一方面，RNN（与 Transformer 相比）的主要优点是其线性（与二次型）复杂性。这种渐近优势只有在长上下文的实践中才能实现，根据下图，这个长度是 8k。另一方面，一旦上下文足够长，现有的 RNN（如 Mamba）就很难真正利用所依赖的额外信息。

并且，大量的实验结果表明：TTT-Linear 和 TTT-MLP 都匹配或超过基线。与 Transformer 类似，它们可以通过限制更多的代币来不断减少困惑，而 Mamba 在 16k 上下文后则不能。经过初步的系统优化，TTT Linear 在 8k 环境下已经比 Transformer 更快，并且在 wall-clock 时间上与 Mamba 相匹配。

TTT 层在理论上和实验评估中表现出色，尤其是在长上下文处理和硬件效率方面。如果在实际应用中能够解决一些潜在的工程挑战，如大规模部署和集成问题，工业界对 TTT 层的接受度也将逐步提升。

论文链接：https://arxiv.org/pdf/2407.04620v1