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

仅需几十行 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 更新一次快速权重，之所以这样做是因为传统观点认为这种方式对于上下文学习更加有效。

但是，这种小批量处理方式会导致并行效率低下以及计算密度不足，尤其在使用大型非线性快速权重时，会给硬件高效实现带来重大挑战，以至于实际算力利用率难以突破 10% 的有效阈值。基于此，本次研究团队采用相反的策略并引入了 LaCT。

如下图所示，LaCT 块由三种类型的层组成：窗口注意力层、大块测试时训练层和前馈层。

每一层都配备了残差连接，这一设计也遵循了 Transformer 架构中的标准做法。窗口注意力层通过执行局部自注意力，来捕捉局部依赖关系。而在测试时训练层，研究团队则将序列分割成了大块。

研究团队表示，历史上下文通过“更新”操作逐渐被压缩到快速权重中，最新的权重被“应用”到当前的查询向量（Q）上，以便计算其对应的输出。前馈层则执行与 Transformer 中类似的通道混合操作。

由于测试时训练的“更新”操作和“应用”操作是解耦的，因此可以自适应地设置块大小，并以不同的顺序应用这些操作，进而能够模拟不同类型的数据依赖关系。

当分块大小等于完整序列长度时，会先执行“应用”操作再执行“更新”操作，这在概念上与全注意力机制相似。通过交替使用“更新”操作和“应用”操作，能够形成分块因果掩码，其中分块大小与块大小互相对应。在两个操作之间切换顺序会导致掩码发生偏移，偏移掩码不会在块内泄露未来信息，这在语言建模中构建完整因果掩码时非常重要。

大块测试时训练层会将数据视为集合序列，因为其快速的权重更新会忽略每个块内的 tokens 顺序和空间局部性。然而，许多数据模态比如视频、图像集合或文本，并不完全符合这种基于集合的视角。对于这些模态而言，块内结构和局部性对于捕获整体数据结构至关重要。

因此，研究团队将局部窗口注意力层与测试时训练层集成在一起，以便处理块内的数据结构。此外，窗口注意力机制能有效捕捉数据中的局部特征。对于测试时训练层来说，这让它能够将其固定大小的快速权重容量集中用于建模非局部依赖关系。

总的来说，LaCT 是一种混合架构，它采用二次计算注意力机制来处理局部结构，针对非局部上下文采用线性计算的测试时训练机制。上下文并行（CP，Context Parallelism）沿着上下文长度维度针对序列进行分区，并将分片分布在多个设备上来进行并行计算。

前馈层和窗口注意力均属于局部操作算子，因此天然地支持上下文并行。对于测试时训练层，小块难以支持上下文并行，因此更倾向于使用张量并行。

研究团队的大块测试时训练层通过在块内分片 tokens 来实现上下文并行。在训练新视图合成时，他们采用了这种并行方法，并观察到 1% 至 3% 的极小吞吐量开销。与此同时，LaCT 架构可以与数据并行、流水线并行和张量并行等其他并行策略兼容。

实验涵盖：新视图合成、语言建模和自回归视频生成

如前所述，研究团队开展了关于新视图合成、语言建模和自回归视频生成的实验。在与线性成本基线方法的对比实验中，研究团队为其增加了相同的窗口注意力模块，以便确保能够进行公平的比较。

表丨对每个实验中关键因素的总结（来源：arXiv）

在新视图合成上，研究团队在场景级和物体级数据集上对本次方法进行评估。他们使用 Objaverse 数据集进行物体级训练，并遵循 LVSM 和 GS - LRM 的设置。

训练完成之后，研究团队在 Google Scanned Objects（GSO）数据集上进行评估，该数据集的分辨率分别为 256×256 和 512×512。每次评估涉及 4 到 48 个输入视图，且每个物体有 8 个新视图。

对于场景级评估，研究团队采用挑战性较高的 DL3DV 场景数据集，其中包含超过 11000 个训练场景和 140 个测试场景，每个场景大约有 300 个视图，评估的分辨率为 960 × 536。

对于物体级评估，研究团队使用了如下两个基线模型：全注意力模型和寄存器注意力模型。

全注意力基线模型将测试时训练层替换为逐块因果注意力层，实现了输入 tokens 之间的双向交互和来自新视图的交叉注意力。

寄存器注意力模型将输入 tokens 压缩到 4096 个寄存器中，并通过与这些寄存器的交叉注意力解码新视图。

在场景级评估中，研究团队与 LongLRM 进行对比，LongLRM 是一种结合了 Mamba 和全注意力机制的模型，可用于 3D 高斯溅射（3D Gaussian splat）预测。此外，他们还与纯基于优化的 3D 高斯溅射方法进行了对比。

表丨对所有模型计算复杂性的总结（来源：arXiv）

在性能评估上，研究团队采用每 tokens 损失度量来评估模型有效使用完整上下文的能力。出现单调递减的损失表示上下文利用成功，而处于平稳状态则表示上下文使用有限。

另据悉，他们从原始 LaCT 块中移除了窗口注意力层，将滑动窗口注意力（SWA，sliding window-attention）层直接集成到大块测试时训练层中，并将模型与全注意力模型、门控线性注意力（GLA，Gated Linear Attention）和 DeltaNet 进行了比较。

为了确保公平性，他们为 GLA 和 DeltaNet 都增强了相同的滑动窗口注意力层，并采用 100 万的 RoPE 库进行 32K tokens 上下文的培训。

表丨对所有方法机制和训练吞吐量的总结（来源：arXiv）

为了比较块递归和逐 tokens 递归，在条件受控的实验中，研究团队的线性大块递归策略在相同状态大小下优于线性逐 tokens 递归策略在。

由于语言本身并不天然存在块状结构，研究团队提出的线性大块递归变体在初始阶段性能不如 GLA 和 DeltaNet 等逐 token 方法。然而，当将其与大规模非线性状态以及 Muon 优化器相结合时，该变体的表现将超越这些逐 token 方法。

总的来说，本次成果凸显了大块测试时训练在计算效率和性能上的优势，为更高效且可扩展的长上下文序列建模铺平了道路。

通过消除对于低级硬件特定实现的依赖，LaCT 使人们能够更广泛地探索架构设计空间。未来，研究团队希望这项工作能够启发并加速长上下文建模和测试时训练领域的新研究。

参考资料：

来源：DeepTech深科技一点号