快速 Transformer

Transformer 是非常成功的模型，在许多自然语言处理任务中达到了最先进 的性能。然而，由于自注意力机制的时间复杂度呈二次方增长，将其扩展到 长序列上非常困难。

本库是我们针对 Transformer 的快速注意力机制研究而开发的。您可以在文 档中找到我们的论文列表 https://fast-transformers.github.io_，以及相关 论文和我们已实现的论文。

快速入门

以下代码构建了一个使用 softmax 注意力的 Transformer 和一个使用线性注 意力的 Transformer，并比较了它们对包含 1000 个元素的序列进行编码所 需的时间。

.. code:: python

import torch
from fast_transformers.builders import TransformerEncoderBuilder

# 创建 Transformer 构建器
builder = TransformerEncoderBuilder.from_kwargs(
    n_layers=8,
    n_heads=8,
    query_dimensions=64,
    value_dimensions=64,
    feed_forward_dimensions=1024
)

# 构建一个使用 softmax 注意力的 Transformer
builder.attention_type = "full"
softmax_model = builder.get()

# 构建一个使用线性注意力的 Transformer
builder.attention_type = "linear"
linear_model = builder.get()

# 构造虚拟输入
X = torch.rand(10, 1000, 8*64)

# 为 CUDA 准备一切
X = X.cuda()
softmax_model.cuda()
softmax_model.eval()
linear_model.cuda()
linear_model.eval()

# 预热 GPU
with torch.no_grad():
    softmax_model(X)
    linear_model(X)
torch.cuda.synchronize()

# 测量执行时间
softmax_start = torch.cuda.Event(enable_timing=True)
softmax_end = torch.cuda.Event(enable_timing=True)
linear_start = torch.cuda.Event(enable_timing=True)
linear_end = torch.cuda.Event(enable_timing=True)

with torch.no_grad():
    softmax_start.record()
    y = softmax_model(X)
    softmax_end.record()
    torch.cuda.synchronize()
    print("Softmax: ", softmax_start.elapsed_time(softmax_end), "ms")
    # Softmax: 144 ms（在 GTX1080Ti 上）

with torch.no_grad():
    linear_start.record()
    y = linear_model(X)
    linear_end.record()
    torch.cuda.synchronize()
    print("Linear: ", linear_start.elapsed_time(linear_end), "ms")
    # Linear: 68 ms（在 GTX1080Ti 上）

依赖与安装

Fast Transformers 库有以下依赖项：

• PyTorch
• C++ 工具链
• CUDA 工具链（如果您希望编译用于 GPU）

对于大多数机器来说，安装非常简单：

.. code:: bash

pip install --user pytorch-fast-transformers

注意：macOS 用户应确保已安装 llvm 和 libomp。可以使用 homebrew <https://brew.sh>_ 包管理器运行 brew install llvm libomp 来完成安装。

文档

我们有一个专门的 文档网站 <https://fast-transformers.github.io/>，但也鼓 励您阅读 源代码 <https://github.com/idiap/fast-transformers>。

研究

我们的研究

要了解本库中一些注意力机制实现背后的理论，我们鼓励您关注我们的研 究成果。

* Transformer 就是 RNN：具有线性注意力的快速自回归 Transformer (`2006.16236 <https://arxiv.org/abs/2006.16236>`_)
* 具有聚类注意力的快速 Transformer (`2007.04825 <https://arxiv.org/abs/2007.04825>`_)

如果您觉得我们的研究有所帮助或产生了影响，请考虑引用以下文献：

.. code::

    @inproceedings{katharopoulos_et_al_2020,
        author = {Katharopoulos, A. and Vyas, A. and Pappas, N. and Fleuret, F.},
        title = {Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention},
        booktitle = {国际机器学习大会（ICML）论文集},
        year = {2020}
    }

    @article{vyas_et_al_2020,
        author={Vyas, A. and Katharopoulos, A. and Fleuret, F.},
        title={Fast Transformers with Clustered Attention},
        booktitle = {神经信息处理系统国际会议（NeurIPS）论文集},
        year={2020}
    }

其他人的研究

• Efficient Attention：具有线性复杂度的注意力机制 (1812.01243 <https://arxiv.org/abs/1812.01243>_)
• Linformer：线性复杂度的自注意力机制 (2006.04768 <https://arxiv.org/abs/2006.04768>_)
• Reformer：高效的 Transformer (2001.04451 <https://arxiv.org/abs/2001.04451>_)

支持、许可与版权

本软件采用 MIT 许可证发布，这意味着您可以按照自己的意愿和目的使 用它。有关支持、版权和许可证的所有信息都可以在仓库中的 LICENSE <https://github.com/idiap/fast-transformers/blob/master/LICENSE>_ 文件中找到。

fast-transformers 快速上手指南

fast-transformers 是一个专注于加速 Transformer 模型注意力机制的开源库，旨在解决标准自注意力机制在处理长序列时计算量呈二次方增长的问题。它提供了线性注意力（Linear Attention）等多种高效实现。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统：Linux 或 macOS（Windows 支持可能受限，建议优先使用 Linux）。
• 核心依赖：
  • Python 3.6+
  • PyTorch (需与您的 CUDA 版本匹配)
• 编译工具链（用于构建自定义 CUDA/C++ 算子）：
  • Linux: GCC, CUDA Toolkit。
  • macOS: 必须安装 llvm 和 libomp。
    • 使用 Homebrew 安装命令：brew install llvm libomp

提示：如果您在中国大陆，建议在安装 PyTorch 时使用国内镜像源以加速下载。

安装步骤

方法一：通过 PyPI 安装（推荐）

对于大多数用户，直接使用 pip 安装是最简单的方式：

pip install --user pytorch-fast-transformers

国内加速方案： 如果使用官方源下载缓慢，推荐使用清华大学或阿里云镜像源：

pip install --user pytorch-fast-transformers -i https://pypi.tuna.tsinghua.edu.cn/simple

方法二：从源码安装（适用于开发者）

如果您需要最新的功能或进行二次开发：

git clone https://github.com/idiap/fast-transformers.git
cd fast-transformers
pip install --user .

基本使用

以下示例展示了如何使用 TransformerEncoderBuilder 快速构建一个标准 Softmax 注意力模型和一个线性注意力（Linear Attention）模型，并对比它们在处理长度为 1000 的序列时的推理速度。

import torch
from fast_transformers.builders import TransformerEncoderBuilder

# 1. 创建 Transformer 构建器
builder = TransformerEncoderBuilder.from_kwargs(
    n_layers=8,
    n_heads=8,
    query_dimensions=64,
    value_dimensions=64,
    feed_forward_dimensions=1024
)

# 2. 构建标准 Softmax 注意力模型
builder.attention_type = "full"
softmax_model = builder.get()

# 3. 构建线性注意力模型 (Fast Transformer)
builder.attention_type = "linear"
linear_model = builder.get()

# 4. 构造虚拟输入数据 (Batch=10, Seq_Len=1000, Features=8*64)
X = torch.rand(10, 1000, 8*64)

# 5. 将数据和模型移至 GPU 并设置为评估模式
X = X.cuda()
softmax_model.cuda()
softmax_model.eval()
linear_model.cuda()
linear_model.eval()

# 6. GPU 预热 (Warmup)
with torch.no_grad():
    softmax_model(X)
    linear_model(X)
torch.cuda.synchronize()

# 7. 测量执行时间
softmax_start = torch.cuda.Event(enable_timing=True)
softmax_end = torch.cuda.Event(enable_timing=True)
linear_start = torch.cuda.Event(enable_timing=True)
linear_end = torch.cuda.Event(enable_timing=True)

# 测试 Softmax 注意力
with torch.no_grad():
    softmax_start.record()
    y = softmax_model(X)
    softmax_end.record()
    torch.cuda.synchronize()
    print("Softmax: ", softmax_start.elapsed_time(softmax_end), "ms")

# 测试线性注意力
with torch.no_grad():
    linear_start.record()
    y = linear_model(X)
    linear_end.record()
    torch.cuda.synchronize()
    print("Linear: ", linear_start.elapsed_time(linear_end), "ms")

预期结果： 在相同的硬件环境下（如 GTX 1080 Ti），线性注意力模型的推理时间通常显著低于标准 Softmax 模型（例如：Softmax 约 144ms vs Linear 约 68ms），尤其在序列长度增加时优势更加明显。