Linformer: Self-Attention with Linear Complexity

下载 PDF

Linformer: Self-Attention with Linear Complexity - 中文验证版

英文原文卡片:linformer_2020.md

状态:已翻译。

元数据

计算设置

论文明确报告了加速器硬件。所有预训练实验使用 BookCorpus 加英文 Wikipedia,约 3300M 词,采用 MLM 目标,并"在 64 块 Tesla V100 GPU 上并行化训练 250k 步更新"。论文还指出默认使用混合精度训练,但这与所提议的注意力机制正交。对于部署式测量,它在单块 16GB Tesla V100 GPU 卡上基准测试推理速度和内存。这些推理测试使用随机生成的输入,序列长度最高达 65536,并根据内存中能容纳的最大 batch 选择 batch size,因此内存数字实际上是设备容量结果,而非仅理论计数。

瓶颈

Linformer 针对与其他 2020 年长上下文论文相同的二次自注意力瓶颈,但诊断不同。标准自注意力形成 R^(n x n) 中的上下文映射矩阵 P,需要 n x d 矩阵乘法和 O(n^2) 时间和空间。论文同时强调训练和部署成本:大型 Transformer 训练缓慢,而实际推理通常需要蒸馏或压缩。它还论证常见的节省内存技术无法解决全部问题。微批处理和梯度检查点可以将训练装入内存,但以时间换空间且不加速推理。蒸馏加速较小的学生模型,但不加速教师训练。稀疏和 LSH 方法降低渐近成本,但可能引入顺序操作、常量因子或性能下降。

方法适配

该方法通过将注意力矩阵视为低秩来适配注意力。论文首先对预训练 RoBERTa 注意力矩阵进行谱分析,发现长尾奇异值分布,意味着 P 的大部分可以用少量方向近似。Linformer 没有为每个注意力矩阵计算 SVD,而是插入学习的投影矩阵 E_i 和 F_i(属于 R^(n x k))用于键和值。键和值从 n 个位置投影到 k 个投影位置,注意力计算为 n x k 上下文矩阵。当 k 远小于 n 时,这将注意力成本从 O(n^2) 变为 O(nk),定理部分论证 k 的选择无需随序列长度线性增长。

该设计是设备友好的,因为它保持密集线性代数。与任意稀疏注意力不同,Linformer 将计算保留为投影和密集矩阵乘法,V100 GPU 执行效率高。论文还探索了投影共享以减少参数和内存开销:headwise 共享、key-value 共享和 layerwise 共享。Layerwise 共享在整个模型中使用单一投影矩阵,实验报告退化很少,这对部署内存和训练内存都重要。

证据

预训练证据是低秩投影不会使语言建模崩溃。在最大序列长度 n = 512 和 n = 1024 下,验证困惑度随 k 增加接近标准 Transformer;论文强调 n = 512 时 k = 128 和 n = 1024 时 k = 256 几乎持平。在更长序列预训练比较中,n 取 512、1024、2048 和 4096 而 k 固定为 256,最终困惑度在收敛后大致相同。这是计算受 k 而非比例 n/k 控制这一主张的核心经验支持。

下游结果延续该叙述。在 SST-2、IMDB、QNLI 和 QQP 上,同语料 RoBERTa-base 基线平均 92.25。n = 512、k = 256 且具有共享 key-value 和 layerwise 投影的 Linformer 平均 92.30,n = 1024、k = 256 变体保持在约 91.83-92.18。论文指出所有基线和 Linformer 变体均以相同目标、语料和多达 250K 步更新预训练,而 Linformer 以更短的挂钟时间达到这些更新。

单 GPU 推理表给出最清晰的计算设备证据。在 16GB V100 上,Linformer(k = 128)在 n = 512 时已快 1.5 倍且允许 1.7 倍更大最大 batch size。在 n = 4096 时,k = 128 提供 3.4 倍时间节省和 14 倍内存节省;k = 256 提供 3.2 倍和 13 倍。在 n = 65536 时,k = 128 提供 20 倍时间节省和 60 倍内存节省,k = 256 提供 18 倍和 52 倍。这些是在相同 V100 容量约束下对比标准 Transformer 报告的。

历史影响

Linformer 成为用密集低秩替代品而非稀疏模式替换注意力矩阵的关键早期示例。其计算贡献是主张长上下文注意力可以保持 GPU 友好的密集线性代数,同时仍避免二次内存。它还通过将理论复杂度与命名 GPU 上的具体推理内存和最大 batch 测量配对,影响了后来的评估习惯。

局限

局限很重要。最大序列长度证据(最高 65536)是在随机生成数据上的推理基准,而非经该长度长文档训练的应用。低秩假设可能在对需要许多细粒度全局交互的任务上失败,固定投影可能丢失稀疏或精确注意力能保留的信息。论文在选定的 GLUE 和 IMDB 任务上报告强对等性,但这些任务本身并不证明在非常长的推理或检索密集型文档上的质量。最后,预训练设置仍使用 64 块 V100 GPU,因此该方法降低了每序列注意力成本,但并未降低预训练本身的规模。

链接