Attention with Linear Biases (ALiBi)方法介绍：提升Transformer外推能力的创新方法

最新推荐文章于 2025-05-29 11:43:08 发布

阿正的梦工坊

最新推荐文章于 2025-05-29 11:43:08 发布

阅读量1.1k

点赞数 11

分类专栏： Deep Learning DL Papers 文章标签： transformer 深度学习人工智能

本文链接：https://blog.youkuaiyun.com/shizheng_Li/article/details/147171686

版权

Deep Learning 同时被 2 个专栏收录

290 篇文章

订阅专栏

DL Papers

139 篇文章

订阅专栏

ALiBi方法介绍：提升Transformer外推能力的创新方法

在自然语言处理（NLP）领域，Transformer模型因其强大的建模能力而广受欢迎。然而，一个长期存在的问题是：Transformer在处理比训练时更长的序列时，往往表现不佳，即缺乏外推能力（extrapolation）。2022年ICLR会议上发表的论文《Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation》提出了一种简单而高效的方法——Attention with Linear Biases (ALiBi)，显著提升了Transformer在长序列上的外推能力。本文将详细介绍ALiBi的核心方法、其为何能提高外推能力，以及后续工作如何进一步优化这一能力。

下文中图片来自于原论文：https://arxiv.org/pdf/2108.12409

ALiBi的核心方法

传统位置编码的局限性

传统的Transformer模型（如Vaswani等人在2017年提出的模型）通常使用正弦位置编码（sinusoidal position embeddings）或学习位置编码来为输入序列中的每个词添加位置信息。这些编码被加到词嵌入上，帮助模型区分序列中不同位置的词。然而，这些方法存在一个问题：模型在训练时只见过固定长度（例如 $L$ ）的序列，当推理时遇到更长的序列（ $L_{\text{valid}} > L$ ）时，性能会显著下降，尤其是在正弦位置编码中，模型仅能外推到略长于训练长度的序列（例如 $L + 50$ ）。

ALiBi的创新设计

ALiBi提出了一种完全不同的位置编码方式，摒弃了传统的嵌入式位置编码，转而通过线性偏置直接修改注意力机制的计算方式。具体方法如下：

去除位置嵌入：ALiBi不再在输入词嵌入上添加位置编码，而是直接在注意力机制中引入位置信息。这种设计避免了传统方法因位置嵌入固定长度而导致的外推限制。
线性偏置的注意力计算：在Transformer的注意力机制中，查询向量（query, $\mathbf{q}_i$ ）与键向量（key, $\mathbf{k}_j$ ）的点积用于计算注意力分数。ALiBi在点积后添加了一个与距离成比例的负偏置，公式如下：

$\text{softmax}\left(\mathbf{q}_i \mathbf{K}^{\top} + m \cdot [-(i-1), \ldots, -2, -1, 0]\right)$

其中：
- $\mathbf{q}_i \mathbf{K}^{\top}$ 是标准的查询-键点积。
- $m$ 是一个针对每个注意力头（head）固定的标量，称为斜率（slope）。
- $\ldots, -2, -1, 0]$ 表示偏置项，与查询和键之间的距离成线性关系。距离越远，负偏置越大，从而降低远距离键的注意力分数。

在这里插入图片描述

斜率的选择：对于具有 $n$ 个注意力头的模型，ALiBi使用几何序列来设定斜率。例如，对于8个头的模型，斜率为 $\frac{1}{2^1}, \frac{1}{2^2}, \ldots, \frac{1}{2^8}$ 。这些斜率无需训练，而是预先固定，类似于正弦位置编码的波长设置。这种设计不仅简单，还能泛化到不同数据集和模型规模。
实现简便：ALiBi只需修改注意力掩码矩阵（mask matrix），将线性偏置添加到查询-键点积中即可。这种改动仅需几行代码，且不增加运行时开销。

ALiBi的特性

无额外参数：与需要学习大量参数的T5偏置方法不同，ALiBi的偏置是静态的，无需训练。
内存效率：虽然ALiBi的掩码矩阵因多头而略有增加（例如增加约100MB），但其允许在短序列上训练，从而大幅降低整体内存需求。
倾向近期信息：ALiBi的负偏置使模型更关注靠近查询的键，表现出对近期信息的归纳偏置（inductive bias towards recency）。

为什么ALiBi能提高外推能力？

ALiBi的外推能力提升主要源于以下几个方面：

相对位置信息的引入：
- 传统正弦位置编码为每个位置分配一个固定的绝对位置向量，这种方式在训练时限定了序列长度，导致模型难以处理未见过的长序列。
- ALiBi通过相对距离的线性偏置，将位置信息融入注意力计算中。无论序列多长，偏置始终基于查询和键的相对距离生成，因此天然支持任意长度的序列。
避免值向量中的位置信息：
- 在ALiBi中，位置信息仅影响查询和键的注意力分数，而不直接融入值向量（values）。这与T5偏置和旋转位置编码（rotary position embeddings）类似，减少了位置信息对模型输出的直接干扰，从而增强了外推时的稳定性。
缓解早期标记诅咒（Early Token Curse）：
- 在非重叠推理（nonoverlapping inference）中，序列被分成固定长度的子序列，子序列开头的预测缺乏足够的上下文，导致性能下降（即早期标记诅咒）。
- ALiBi通过支持更长的序列输入，减少了子序列划分的频率，从而让更多预测拥有充足的上下文，显著降低了早期标记诅咒的影响。实验表明，ALiBi的性能提升主要源于这一机制，而非利用更长的上下文模式。
高效的归纳偏置：
- ALiBi的线性偏置鼓励模型优先关注近期的词，这种偏置与语言的局部依赖特性相符。即使在长序列中，模型也能有效聚焦于相关信息，避免因序列过长而导致的注意力分散。

实验验证

论文在多个数据集上验证了ALiBi的优越性：

WikiText-103：ALiBi模型在 $L = 512$ 训练时，可外推到 $L_{\text{valid}}=3072$ ，获得比正弦模型（ $L = 3072$ ）更好的困惑度（perplexity），且训练速度快1.84倍。
Toronto BookCorpus：无需调整斜率参数，ALiBi在书籍领域的数据集上同样表现出色，验证了方法的泛化能力。
CC100+RoBERTa：在1.3亿参数模型上，ALiBi以 $L = 1024$ 训练，外推到 $L_{\text{valid}}=2048$ 时，困惑度优于正弦模型（ $L = 2048$ ），内存使用量减少3.1GB，训练速度提升11%。

这些结果表明，ALiBi不仅在外推时表现优异，即使在训练长度等于推理长度时，也能超越传统方法。

后续工作如何进一步提高外推能力？

ALiBi的提出为Transformer的外推能力研究开辟了新方向，后续工作在以下几个方面进行了探索和改进：

优化相对位置偏置：
- Wennberg & Henter (2021)提出了基于径向基函数（radial-basis function）的相对位置偏置方法，虽然增加了可训练参数，但为探索非线性偏置提供了思路。
- Distance Aware Transformer (Wu et al., 2021)尝试通过乘法而非加法引入距离偏置，尽管实验表明加法（如ALiBi）效果更佳，但此类研究丰富了偏置设计的可能性。
结合缓存机制：
- Transformer-XL (Dai et al., 2019)通过缓存机制扩展推理时的上下文长度。ALiBi的相对位置偏置可与缓存结合，进一步提升长序列处理能力，同时保持高效性。
针对特定任务的优化：
- 在机器翻译、文本分类等任务中，研究者尝试将ALiBi的思路应用于特定领域。例如，Kiyono等人 (2021) 探索了位置编码在翻译任务中的外推性，提出移位绝对位置编码（Shifted Absolute Position Embedding），为ALiBi的泛化提供了借鉴。
- Longformer (Beltagy et al., 2020)通过稀疏注意力机制处理长文档，ALiBi的线性偏置可与其结合，减少对长序列训练的依赖。
利用更长上下文：
- ALiBi的分析表明，其性能提升主要源于减少早期标记诅咒，而非充分利用长上下文。未来工作可探索如何让模型更高效地挖掘长序列中的模式，例如通过动态调整斜率或引入层次化注意力机制。
与其他方法的融合：
- ALiBi与Routing Transformer (Roy et al., 2020) 或 kNN-LM (Khandelwal et al., 2020) 等方法正交，结合这些技术可能进一步提升性能，尤其是在低资源场景下。

结论

ALiBi通过简单而优雅的线性偏置设计，解决了Transformer模型在长序列外推上的难题。其核心在于去除传统位置嵌入、引入相对位置偏置，并通过固定斜率实现高效的注意力计算。ALiBi不仅提升了外推能力，还在训练效率和内存使用上展现出显著优势。后续工作在偏置设计、缓存机制和任务优化等方面进一步扩展了ALiBi的潜力，为NLP模型处理超长序列奠定了基础。

对于研究者和开发者来说，ALiBi提供了一个易于实现且效果显著的工具，值得在各种NLP任务中尝试。未来，结合更智能的上下文利用策略，ALiBi的理念有望推动Transformer在更广泛场景下的应用。

参考文献

Press, O., Smith, N. A., & Lewis, M. (2022). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022.
Vaswani, A., et al. (2017). Attention is All You Need. NeurIPS 2017.
Dai, Z., et al. (2019). Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. ACL 2019.

为什么是负偏置？

在ALiBi方法中，选择负偏置（ $\ldots, -2, -1, 0]$ ）而非正偏置，是为了在注意力机制中引入一种偏向近期信息的归纳偏置（inductive bias towards recency）。这种设计并非仅仅为了表示距离大小，而是有意通过负偏置来降低远距离键的注意力分数，从而让模型更关注靠近查询的上下文。以下是详细的解释：

1. 负偏置的核心作用：抑制远距离注意力

在Transformer的注意力机制中，注意力分数是通过查询向量（ $\mathbf{q}_i$ ）与键向量（ $\mathbf{k}_j$ ）的点积计算的，经过softmax函数后转化为概率分布：
$\text{softmax}\left(\mathbf{q}_i \mathbf{K}^{\top}\right)$
ALiBi在点积后添加了一个线性偏置项：
$\text{softmax}\left(\mathbf{q}_i \mathbf{K}^{\top} + m \cdot [-(i-1), \ldots, -2, -1, 0]\right)$
这里的偏置项是负数（例如 $- (i - 1), - 2, - 1, 0$ ），并且随着查询和键的距离（ $i - j$ ）增加，负值变得更大（更负）。由于softmax函数对输入的指数化处理（ $\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$ ），负偏置会降低远距离键的注意力分数：

当距离较远时，例如 $i - j = 10$ ，偏置为 $- 10 m$ （ $m$ 为正斜率），这会使 $\mathbf{q}_i \cdot \mathbf{k}_j - 10m$ 变小，进而使 $e^{\mathbf{q}_i \cdot \mathbf{k}_j - 10m}$ 在softmax中贡献较小的概率。
当距离很近时，例如 $i - j = 1$ ，偏置为 $- m$ ，对注意力分数的减小作用较弱，概率较高。

因此，负偏置的设计使得模型更倾向于关注靠近当前位置的词，符合语言处理中近期上下文通常更重要的规律。

2. 正偏置会有什么效果？

如果使用正偏置（例如 $\ldots, 2, 1, 0]$ ），效果会完全相反：

随着距离增加，正偏置值变大（例如 $i - j = 10$ 时偏置为 $10 m$ ），这会增加远距离键的注意力分数（因为 $e^{\mathbf{q}_i \cdot \mathbf{k}_j + 10m}$ 更大）。
结果是，模型会更倾向于关注远距离的词，而忽略近期的上下文。

在语言建模中，远距离的信息通常不如近期信息重要。例如，预测下一个词时，前面几个词的上下文往往比几十个词之前的上下文更有预测价值。正偏置会导致模型将注意力分散到较远的、不太相关的词上，可能降低模型的性能，尤其是在长序列中更容易导致注意力分散。

3. 负偏置与外推能力的关系

ALiBi的目标之一是提升Transformer在长序列上的外推能力，即在推理时处理比训练时更长的序列。负偏置在这一场景下有以下优势：

局部性偏好：负偏置鼓励模型聚焦于局部上下文，即使序列长度增加，模型也能保持对近期信息的关注。这种局部性偏好减少了模型对全局序列结构的依赖，使得外推到任意长度的序列时，注意力机制的行为仍然可控。
稳定性：正偏置会导致远距离键的注意力分数随距离指数级增长（因softmax的指数化），在长序列中可能引发数值不稳定或注意力过度集中于某些远距离位置。负偏置则通过抑制远距离分数，保持了注意力的稳定性。

4. 正偏置是否完全不可行？

理论上，正偏置并非完全不可用，但需要额外的机制来控制其影响。例如，可以通过调整斜率 $m$ 或引入归一化来限制远距离注意力的过度增强。然而，这会增加复杂性，而ALiBi的设计目标是简单高效。负偏置天然符合语言的局部依赖特性，且无需额外调整即可实现良好的外推效果。

5. 实验支持

论文中提到，作者尝试了不同的偏置设计，包括可训练的斜率，但发现固定负偏置（以几何序列设置斜率）效果最佳。他们还探索了其他斜率选择，发现负偏置在 $(0, 1)$ 范围内、密度靠近0时表现稳健。这种设计在WikiText-103、Toronto BookCorpus和CC100+RoBERTa等数据集上均验证了其优越性，尤其是在外推到长序列（如从 $L = 512$ 到 $L_{\text{valid}}=3072$ ）时，负偏置的ALiBi模型显著优于正弦位置编码和其他方法。

总结

ALiBi选择负偏置而非正偏置，不是因为距离大小的表示，而是为了通过降低远距离键的注意力分数，引入对近期信息的偏好。这种设计不仅符合语言建模的局部依赖特性，还增强了模型在长序列上的外推能力，同时保持了计算的简单性和稳定性。正偏置会优先关注远距离信息，可能导致性能下降和注意力分散，因此不适合ALiBi的目标场景。

代码实现

以下是一个可运行的ALiBi（Attention with Linear Biases）实现，包括训练和推理代码，使用PyTorch框架。代码实现了Transformer语言模型，并按照论文中描述的斜率（slopes）设计支持8个和16个注意力头的情况，同时泛化到任意头数 $n$ 。为了简化演示，我们使用一个小型数据集进行训练和推理，并提供清晰的注释。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator

# 生成斜率（Slopes）函数
def get_slopes(n_heads):
    """
    根据头数n_heads生成ALiBi的斜率。
    - 对于8个头：几何序列 [1/2^1, 1/2^2, ..., 1/2^8]
    - 对于16个头：从1/sqrt(2)开始，比例为1/sqrt(2)
    - 通用情况：从2^(-8/n)开始，比例为2^(-8/n)
    """
    if n_heads == 8:
        return [1 / (2 ** i) for i in range(1, 9)]  # [1/2, 1/4, ..., 1/256]
    elif n_heads == 16:
        # 插值生成16个斜率，从1/2^0.5开始，比例为1/2^0.5
        slopes = []
        for i in range(16):
            slopes.append(1 / (2 ** (0.5 + 0.5 * i)))
        return slopes
    else:
        # 通用公式：从2^(-8/n)开始，比例为2^(-8/n)
        start = 2 ** (-8 / n_heads)
        return [start * (start ** i) for i in range(n_heads)]

# ALiBi注意力层
class ALiBiAttention(nn.Module):
    def __init__(self, d_model, n_heads, max_seq_len):
        super(ALiBiAttention, self).__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.max_seq_len = max_seq_len

        self.query = nn.Linear(d_model, d_model)
        self.key = nn.Linear(d_model, d_model)
        self.value = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)

        # 获取ALiBi斜率
        self.slopes = torch.tensor(get_slopes(n_heads), device='cuda' if torch.cuda.is_available() else 'cpu')

        # 预计算偏置矩阵
        self.register_buffer('bias', self._create_alibi_bias())

    def _create_alibi_bias(self):
        # 创建ALiBi的偏置矩阵，形状为 (n_heads, max_seq_len, max_seq_len)
        positions = torch.arange(self.max_seq_len).unsqueeze(0).unsqueeze(0)  # (1, 1, max_seq_len)
        distances = positions - torch.arange(self.max_seq_len).unsqueeze(1)  # (1, max_seq_len, max_seq_len)
        bias = distances * self.slopes.view(-1, 1, 1)  # (n_heads, max_seq_len, max_seq_len)
        return bias

    def forward(self, x, mask=None):
        B, L, _ = x.size()  # 批大小、序列长度、模型维度
        q = self.query(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)  # (B, n_heads, L, head_dim)
        k = self.key(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.value(x).view(B, L, self.n_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)  # (B, n_heads, L, L)

        # 添加ALiBi偏置
        bias = self.bias[:, :L, :L]  # 裁剪到当前序列长度
        scores = scores + bias

        # 因果掩码（仅关注之前的词）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Softmax归一化
        attn = F.softmax(scores, dim=-1)
        context = torch.matmul(attn, v)  # (B, n_heads, L, head_dim)
        context = context.transpose(1, 2).contiguous().view(B, L, self.d_model)  # (B, L, d_model)
        return self.out(context)

# Transformer模型
class ALiBiTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, n_layers, max_seq_len, dropout=0.1):
        super(ALiBiTransformer, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_dropout = nn.Dropout(dropout)
        self.layers = nn.ModuleList([
            nn.ModuleDict({
                'attn': ALiBiAttention(d_model, n_heads, max_seq_len),
                'norm1': nn.LayerNorm(d_model),
                'ffn': nn.Sequential(
                    nn.Linear(d_model, d_model * 4),
                    nn.ReLU(),
                    nn.Linear(d_model * 4, d_model),
                    nn.Dropout(dropout)
                ),
                'norm2': nn.LayerNorm(d_model)
            }) for _ in range(n_layers)
        ])
        self.fc_out = nn.Linear(d_model, vocab_size)
        self.max_seq_len = max_seq_len

    def forward(self, x, mask=None):
        x = self.embedding(x)
        x = self.pos_dropout(x)

        for layer in self.layers:
            attn_out = layer['attn'](x, mask)
            x = layer['norm1'](x + attn_out)
            ffn_out = layer['ffn'](x)
            x = layer['norm2'](x + ffn_out)

        logits = self.fc_out(x)
        return logits

# 简单文本数据集
class TextDataset(Dataset):
    def __init__(self, texts, vocab, max_len):
        self.texts = texts
        self.vocab = vocab
        self.max_len = max_len
        self.tokenizer = get_tokenizer('basic_english')

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        tokens = self.tokenizer(text)[:self.max_len]
        indices = [self.vocab[token] for token in tokens]
        if len(indices) < self.max_len:
            indices += [self.vocab['<pad>']] * (self.max_len - len(indices))
        return torch.tensor(indices, dtype=torch.long)

# 训练函数
def train_model(model, train_loader, optimizer, criterion, device, epochs=5):
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for batch in train_loader:
            optimizer.zero_grad()
            input_ids = batch[:, :-1].to(device)  # 输入序列
            target_ids = batch[:, 1:].to(device)  # 目标序列（移位一位）
            
            # 创建因果掩码
            seq_len = input_ids.size(1)
            mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).to(device)
            
            outputs = model(input_ids, mask)
            loss = criterion(outputs.view(-1, outputs.size(-1)), target_ids.view(-1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f'Epoch {epoch+1}, Loss: {total_loss / len(train_loader):.4f}')

# 推理函数
def generate_text(model, vocab, start_text, max_len, device):
    model.eval()
    tokenizer = get_tokenizer('basic_english')
    tokens = tokenizer(start_text)
    input_ids = [vocab[token] for token in tokens]
    input_ids = input_ids[:max_len-1]
    input_ids = torch.tensor([input_ids], dtype=torch.long).to(device)
    
    with torch.no_grad():
        for _ in range(max_len - len(input_ids[0])):
            seq_len = input_ids.size(1)
            mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).to(device)
            outputs = model(input_ids, mask)
            probs = F.softmax(outputs[:, -1, :], dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=1)
    
    # 解码生成的token
    itos = vocab.get_itos()
    generated = [itos[idx] for idx in input_ids[0].cpu().numpy()]
    return ' '.join(generated)

# 主函数
def main():
    # 超参数
    vocab_size = 10000
    d_model = 256
    n_heads = 8  # 可改为16或其他值
    n_layers = 4
    max_seq_len = 64
    batch_size = 32
    epochs = 5
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 创建简单数据集
    texts = [
        "The quick brown fox jumps over the lazy dog.",
        "A cat sits on the mat and watches the world.",
        "Sunny hills bloom with vivid colors daily.",
        "Stars shine brightly in the clear night sky."
    ] * 1000  # 增加数据量

    # 构建词汇表
    tokenizer = get_tokenizer('basic_english')
    def yield_tokens(texts):
        for text in texts:
            yield tokenizer(text)
    vocab = build_vocab_from_iterator(yield_tokens(texts), specials=['<pad>', '<unk>'], max_tokens=vocab_size)
    vocab.set_default_index(vocab['<unk>'])

    # 数据加载
    dataset = TextDataset(texts, vocab, max_seq_len)
    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # 初始化模型
    model = ALiBiTransformer(vocab_size, d_model, n_heads, n_layers, max_seq_len).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss(ignore_index=vocab['<pad>'])

    # 训练
    print("开始训练...")
    train_model(model, train_loader, optimizer, criterion, device, epochs)

    # 推理
    print("\n生成文本...")
    start_text = "The quick brown"
    generated = generate_text(model, vocab, start_text, max_len=20, device=device)
    print(f"生成的文本: {generated}")

if __name__ == "__main__":
    main()