对角注意力稀疏：从N²到N×d的自回归图像生成效率革命
如果最近在折腾自回归图像生成，你多半会被两个数字搞得头疼：一个是序列长度到了几千甚至上万，另一个是注意力复杂度跟着序列长度平方往上翻。这篇论文要处理的正是第二件事。它没有堆一堆花哨模块，而是把注意力矩阵中“离对角线不远的一小条”单独提出来，做成可解释、可落地的对角注意力稀疏结构。换句话说，它假设图像 token 的有效依赖大多集中在时间上相邻的上下文里，远处信息少量建模就够。这篇解读我会顺着它的设计动机、掩码定义、实现细节、实验观察，再延伸到化学图像生成和跨模态协同生成这些实际场景，尽量把每步为什么这么干都讲透，适合刚接触自回归图像生成、或者正在被注意力显存账单逼疯的同学做参考。1. 自回归图像生成的两座山：长上下文与大矩阵
1.1 从 PixelCNN 到 VQGAN + Transformer 的路线演进
自回归图像生成并不是一个新概念。早期 PixelCNN 这一类模型是直接在原始像素上逐点预测，条件概率写成 (p(x)=\prod_t p(x_t|x_{<t}))。它思路很直接，但问题也明显：像素级生成速度太慢，感受野被卷积核限制得比较厉害，图像细节越多，计算越吃紧。
后来 VQGAN 这类方法把图像压成一个离散 token 序列。VQGAN 先用 CNN 把图像编码进一个尺寸更小的特征图，然后在 codebook 里找一个最接近的离散编码。这样一张 256×256 的图像可以变成 16×16=256 个视觉 token，比直接预测像素门槛低得多。紧接着的 Transformer 模块再把这些 token 当作像文本一样的序列来建模，用标准的 causal mask 做自回归预测。整体上跟 GPT 处理文本的过程非常相似，只是输入不是词而是视觉 token。
这种“VQGAN 压缩 + Transformer 序列建模”的路线，强项在于能借用语言模型那边的工程沉淀，包括大规模预训练技巧、采样策略、各种分布式训练工具。但你一旦把分辨率拉高，图像 token 数就从 256 涨到 1024、4096，甚至更多。于是原来在文本里已经很头疼的二次复杂度问题，在图像这块被放得更大了。
1.2 因果注意力为什么会成为瓶颈
Transformer 做序列建模时，每个位置都要和所有位置算一个 attention score。由于自回归生成不能看未来，所以还需要加因果 mask，把未来位置全部遮掉。这样一来，注意力矩阵就变成了一个下三角矩阵：query 只能看到排在它前面的 key。
问题就出在这里。假设序列长度是 N，每一层的注意力和前馈网络计算中，仅 attention score 矩阵就有 N×N 这么多个数值需要计算和存储。这个复杂度是平方级别的，图像 token 越多，N 越大，计算量和显存占用都以二次方速度增长。拿一张 1024×1024 的图像来说，如果 patch 是 16×16，那序列长度就是 4096，attention 矩阵的单层尺寸就已经到了 1600 万级别的元素，这还只是一个 head、一层 Transformer。
更麻烦的是，这种全量注意力在图像上未必有价值。图像里的绝大多数信息都是局部相关的：纹理连续、边缘连贯、颜色过渡平缓，真正需要跨很远位置的语义依赖往往出现在少量结构节点上。但全注意力不管三七二十一，把远处 token 也拿来算一遍点积，最终很多权重趋近于零，等于白白浪费了算力。如果我们能在“不显著损失生成质量”的前提下，把注意力限制到最可能产生有效依赖的范围内，问题就解决了。
2. 对角注意力稀疏：把 N² 直接变成 N×d
2.1 一个简单且可解释的掩码：对角带窗口
论文给出的方案非常朴素：在因果 mask 的基础上，再限制 query 与 key 之间的距离。假设序列长度为 N，定义一个带宽 d，那么位置 i 的 query 只能关注满足 (j < i) 且 (i - j \le d) 的 key。把注意力矩阵画出来，允许计算的位置并不是整个下三角，而是靠近主对角线的一条带状区域。这就是“对角稀疏”这个名字的由来。
用 PyTorch 写这个 mask 几乎就是几
页:
[1]