摘要

Transformer 在自然语言处理与计算机视觉领域均取得了显著成功。然而，由于自注意力机制（self-attention）的时间与内存复杂度随输入序列长度呈二次增长，将其扩展至长序列任务（如长文档或高分辨率图像）的成本极高，难以实际应用。为此，本文提出 Long-Short Transformer（Transformer-LS），一种针对长序列建模的高效自注意力机制，可在语言与视觉任务中实现线性时间与内存复杂度。该方法通过融合一种基于动态投影的新颖长程注意力机制以捕捉远距离依赖关系，同时结合短程注意力机制以捕获细粒度的局部相关性。为解决两种注意力机制间存在的尺度差异问题，我们进一步提出一种双归一化策略。Transformer-LS 可无缝应用于自回归与双向模型，且无需引入额外计算开销。在多个语言与视觉任务中，该方法均显著优于现有最先进模型，涵盖 Long Range Arena 基准测试、自回归语言建模以及 ImageNet 图像分类任务。例如，在 enwik8 数据集上，Transformer-LS 仅使用比先前方法一半的参数量，即可实现 0.97 的测试比特每字符（BPC）性能，同时运行速度更快，并在相同硬件条件下支持长达三倍的序列长度，相较其全注意力版本具有显著优势。在 ImageNet 分类任务中，该模型亦取得了当前最优结果——例如，一个参数量仅为 5580 万的中等规模模型，仅在 224×224 的 ImageNet-1K 数据集上训练，即可达到 84.1% 的 Top-1 准确率；同时在高分辨率图像处理方面展现出更强的可扩展性。相关源代码与预训练模型已开源，地址为：https://github.com/NVIDIA/transformer-ls。