正文

斯坦福、UCSD、UC伯克利和Meta了推出TTT大语言模型，一种全新大语言模型（LLM）架构，或将取代最火的Transformer和Mamba。

TTT通过对输入token进行梯度下降，压缩上下文，这种方法被称为“测试时间训练层（Test-Time-Training layers，TTT）”。TTT层直接替代了注意力机制，实现了具有表现力记忆的线性复杂度架构。

来自斯坦福、UCSD、UC伯克利和Meta的研究团队设计了TTT，用机器学习模型取代了RNN的隐藏状态。这个研究了一年多的项目，或将改变我们的语言模型方法。

结果显示，TTT-Linear和TTT-MLP赶超或击败了Transformer和Mamba。研究人员Xiaolong Wang表示：“不敢相信，我们真的做到了。”

未来，TTT不仅适用于语言建模，甚至还可用于长视频建模，对帧进行密集采样，这对于TTT层来说前景广阔。

TTT研发历史：五年构想，一朝实现

2020年，OpenAI的缩放定律显示LSTM无法有效使用长上下文，Mamba虽有改进，但仍存在相似问题。

TTT的构想已超过5年，但团队在过去1.5年中才开始全力研发，目标是开发具有线性复杂度和更强隐藏状态的新LLM架构。

传统序列模型存储历史上下文于隐藏状态中，像Mamba的RNN层会压缩为固定大小状态，尽管效率高，但受限于表达能力。

团队想法是，将上下文压缩到模型权重中，就像LLM处理互联网数据一样，增强表达能力，保持隐藏状态的固定大小。

研究采用自监督学习，更新隐藏状态权重，对每个token进行梯度下降。隐藏状态只存在于端到端架构中的一层，其他组件如QKV投影矩阵通过标准交叉熵目标函数学习。

TTT层作为信息压缩和模型记忆机制，直接替代Transformer中的自注意力层，通过表达性记忆解锁线性复杂性架构。

在125M到1.3B参数规模的大模型对比中，TTT-Linear和TTT-MLP击败了Transformer和Mamba，展现了更低困惑度和更少FLOP，并且更好利用长上下文。

为提升TTT层的高效性，研究团队采取了两项创新措施。

首先，类似于常规训练中的mini-batch，他们在TTT中使用mini-batch token进行梯度步进，以实现更好的并行性。

其次，团队为每个TTT mini-batch操作开发了一种对偶形式，充分利用现代GPU和TPU。该方法输出与原始实现相当，但训练速度提升了5倍以上。

TTT层在实际运行时间中表现优异。下图显示，TTT-Linear在8k上下文中比Transformer更快，并与Mamba相当，且FLOP更少。TTT层不仅克服了RNN的局限，还实现了线性复杂度与高效运行时间的结合。

TTT的研究方法：突破与创新

自注意力机制的隐藏状态是一个随时间线性增长的Key-Value缓存列表。其更新规则将当前的KV元组添加到列表中，输出规则扫描所有元组，形成注意力矩阵。这虽然不需要压缩历史上下文，但处理时间会随上下文长度线性增长。为保持长上下文的高效和表现力，研究人员提出了一种新的压缩启发式方法。

TTT层的隐藏状态是一个模型，更新规则是自监督学习的一步梯度下降。具体来说，隐藏状态相当于模型的权重，这个模型可以是线性模型、小型神经网络等。输出规则则是由更新后的模型对输入数据进行预测。更新规则在自监督损失上进行梯度下降，选择重构图片作为自监督损失，使模型在部分信息图片中重构出图片。

TTT层与RNN层和自注意力机制类似，可以将输入序列映射到输出序列。每个输入序列训练一个不同的权重序列，因此被称为测试时间训练层（TTT）。TTT层与现有的序列建模层一样，可以在任何更大的神经网络架构中替换它们。

在TTT层的训练中，前向传播和后向传播过程相似。训练方式与任何其他Transformer模型相同，可以使用相同的数据、方法和目标（如下一个token预测）来优化网络的其他部分参数。训练更大的神经网络称为外循环，而在每个TTT层内训练权重称为内循环。内循环优化的是模型f的参数，外循环优化的是网络其他部分的参数。

TTT的关键是自监督任务，它决定了从测试序列中学习的特征类型。研究人员直接优化自监督任务以实现下一个token预测的最终目标。通过添加外循环参数，使这个任务可学习。最新的自监督损失包含了内循环优化的参数，外循环则优化损失函数的超参数。

在实际应用中，TTT层在浮点运算（FLOP）方面已经非常高效。为了进一步提升效率，研究人员提出了mini-batch梯度下降和对偶形式。mini-batch梯度下降通过一次并行计算多个梯度，提升了并行化效率。对偶形式在mini-batch结束时计算更新规则，避免了大量内存占用和I/O成本。

TTT层的两种变体分别是TTT-Linear和TTT-MLP。TTT-Linear采用线性模型，而TTT-MLP则采用两层MLP，类似于Transformer的MLP结构。TTT-MLP在隐藏维度上是输入维度的四倍，并包含层归一化和残差连接，以提升稳定性。

研究表明，具有线性模型和batch GD的TTT层等同于线性注意力，证实了TTT层在不同实例化下的广泛适用性。研究人员通过创新设计，实现了TTT层在长上下文处理中的卓越性能，提供了大语言模型发展的新方向。

TTT的实验结果：卓越性能验证

研究人员通过与Transformer和Mamba进行比较，评估了TTT-Linear和TTT-MLP的性能。实验使用了The Pile数据集，执行了2k和8k上下文长度的标准实验。

在2k上下文中，TTT-Linear、Mamba和Transformer表现相当，TTT-MLP在较大的FLOP预算下稍差。尽管TTT-MLP在每个模型大小上复杂度更好，但额外的FLOP成本抵消了优势。在8k上下文中，TTT-Linear和TTT-MLP明显优于Mamba，即使是使用Transformer架构的TTT-MLP在1.3B参数规模下也略胜一筹。

为了评估长上下文处理能力，研究人员使用了Pile的子集Books，对1k到32k上下文长度进行了实验。在2k上下文中，Mamba表现略好于TTT-Linear。而在32k上下文中，TTT-Linear和TTT-MLP均优于Mamba，即使采用Transformer架构的TTT-MLP在32k上下文中也比Mamba表现稍好。

研究人员将上下文长度作为一个可以选择的超参数，并选择了困惑度中的argmin。在10^20 FLOP后，TTT-Linear和TTT-MLP的性能几乎完全重叠，而Mamba和TF Finetune的性能也大部分重叠。TF Finetune的性能明显优于TF Pretrain，因为它受益于长上下文，而不会在训练FLOP中产生极大的成本。

LLM训练和推理可以分解为前向、后向和生成。由于前向和后向都可以并行化，研究人员使用了对偶形式。生成新token（解码）是顺序的，因此使用了原始形式。实验在TPU上运行，结果显示，在上下文为2k的情况下，TTT-Linear每次迭代训练需要0.27秒，比Transformer快10%。

为了公平比较，研究人员还将方法在GPU上运行，证明了mini-batch TTT和对偶形式的有效性。前向内核批大小为16的延迟实验显示，所有模型参数均为1.3B（Mamba为1.4B）。随着上下文长度增加，Transformer每个token的时间线性增长，但TTT和Mamba大致保持不变。

作者介绍：TTT项目背后的团队

UCSD助理教授Xiaolong Wang是TTT研究的主要推动者之一。他在研究提交后发推表示祝贺，并回顾了这一项目的历程，提到TTT的构想已有五年时间，但正式研究持续了一年半。

Yu Sun是斯坦福大学的博士后，师从Carlos Guestrin、Tatsu Hashimoto和Sanmi Koyejo。他在UC Berkeley获得了电子工程科学博士学位，导师是Alyosha Efros和Moritz Hardt，并在康奈尔大学获得学士学位。Yu Sun专注于测试时间训练（TTT）的算法框架，他在2022年11月与Xinhao Li共同启动了这一项目，并在2023年6月起全职负责该项目。

Xinhao Li是UCSD的博士生，导师是Xiaolong Wang教授。他的研究兴趣主要是深度学习和计算机视觉。在斯坦福大学Tatsunori Hashimoto教授团队中担任访问学生期间，与Yu Sun博士合作。在2024年3月之前，他是TTT早期代码库的主要贡献者。

Karan Dalal是UC Berkeley的本科生，他于2023年6月全职加入TTT项目，与Xinhao Li合作共同领导当前代码库的开发工作。

结语：展望未来的TTT

尽管TTT的原理复杂，但其核心在于提升长文本处理的效率和效果。当前，Transformer是主流架构，但在处理长文本时效率低下，而TTT有效解决了这一问题。

AI技术日新月异，正因有科学家的不断研究，我们才能在这些基础上构建更好的应用。TTT的出现，预示着大语言模型发展的新方向，为未来的技术进步打下坚实基础。