几十行PyTorch代码让GPU利用率猛涨，MIT团队打造新模型架构，渲染质量优于3D高斯泼溅

正文

请到「今天看啥」查看全文

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

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

（来源： arXiv ）

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

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

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

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

（来源： arXiv ）

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

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

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

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

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

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

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

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

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

请到「今天看啥」查看全文