仅需几十行 PyTorch 代码即可大幅提高 GPU 利用率，在英伟达 A100 上的 GPU 利用率高达 70%。这一能力由一款名为 LaCT 的新模型架构实现，它由北京大学本科校友、美国麻省理工学院博士生张天远和所在团队打造。

图丨张天远（来源：https://tianyuanzhang.com/）

研究团队通过在不同模态的任务中采用范围从 2000 到 100 万 tokens 的大块更新打造了这种新模型架构。该架构集成了用于捕捉长上下文的大块测试时训练，并集成了用于针对局部结构进行建模的窗口注意力机制。

由于每个大块内的 tokens 都被视为一个无序集，因此研究团队将窗口注意力集成到 LaCT 中，以便捕获块内的局部依赖关系，这让 LaCT 能够显著提高并行性。

这能够扩展非线性快速权重，从而能够提高存储容量。而且，这种简单的实现方式能够轻松地集成更有效的测试时优化器（比如 Muon）。

此外，LaCT 的大块设计也能够很自然地用于针对各种 N 维数据进行建模，原因在于它可以将块大小与数据的内部结构对齐，例如将图像或连续视频帧内的 tokens 分组为一个块。

为了验证本次方法的有效性，研究团队在不同数据模态和任务中，包括从图像集、语言模型和自回归视频扩散模型中进行了新视图合成。

实验结果显示：研究团队的模型能够处理多达 128 张分辨率为 960×536 的输入图像，形成最多 100 万 tokens 的序列，并且在此输入规模下，在渲染质量方面优于 3D 高斯泼溅（3D Gaussian Splatting）技术。

尽管语言数据本身并不显式包含块状结构，但是与 DeltaNet 等 SOTA 方法相比，研究团队的模型实现了大致相当的性能表现。

研究团队还通过将 LaCT 与滑动窗口注意力相结合，将一个 140 亿参数的双向视频扩散 Transformer 适配为自回归模型。这种适配后的模型可以生成包含多达 56000 个视觉 tokens 的连贯视频。

与此同时，在最长的序列实验中，他们进行了超过 100 万个上下文长度的新视图合成。

目前，研究团队已经开源了代码和权重，预计 LaCT 将能推动人们对于更高效长上下文建模架构的探索（https://tianyuanzhang.com/projects/ttt-done-right/）。

吞吐量开销小至 1%-3%

当前，处理长上下文的需求正在迅速增长。虽然 softmax 注意力已成为建模各类数据的解决方案，但其计算成本随序列长度呈二次方增长，这推动了人们对更高效长上下文建模的广泛研究。

最近，测试时间训练（TTT，Test-Time Training）已成为一种富有前景的高效二次序列建模方法。测试时间训练能将循环神经网络中的循环状态概念扩展到一个小型的、在线自适应的子网络。这个子网络的参数也被称为快速权重，它们通过自监督目标在线快速适配，以便记忆上下文中的信息。

近期，多个团队均探索了快速权重网络的各种在线目标、优化器和架构。尽管如此，已有的测试时训练方法仍然难以有效扩展到长上下文场景，根本原因在于测试时训练层的硬件利用率极低，在当前 GPU 上硬件利用率通常低于峰值算力的 5%。

这种低效性是由于使用了小批量规模，即每隔一个 token 或每 16 个到 64 个 tokens 更新一次快速权重，之所以这样做是因为传统观点认为这种方式对于上下文学习更加有效。

但是，这种小批量处理方式