深入解析cabm自注意力机制：从原理到新手实践指南

1次阅读

共计 3062 个字符，预计需要花费 8 分钟才能阅读完成。

自注意力机制 (Self-Attention) 是 Transformer 架构的核心组件，它通过计算序列中每个元素与其他元素的相关性，动态地分配注意力权重。cabm 自注意力是其中一种高效实现方式，特别适合处理长序列数据。

cabm 自注意力的计算过程可以分解为以下几个步骤：

输入嵌入转换：给定输入序列 $X \in \mathbb{R}^{n\times d_{model}}$，首先通过线性变换得到 Query、Key 和 Value 矩阵：
$$
Q = XW_Q, \quad K = XW_K, \quad V = XW_V
$$
其中 $W_Q, W_K, W_V \in \mathbb{R}^{d_{model}\times d_k}$ 是可学习参数矩阵
注意力分数计算：
$$
\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V
$$
这里除以 $\sqrt{d_k}$ 是为了防止点积结果过大导致 softmax 梯度消失
多头注意力：将上述过程并行执行 $h$ 次，然后将结果拼接：
$$
\text{MultiHead}(Q,K,V) = \text{Concat}(head_1,…,head_h)W_O
$$

以下是 cabm 自注意力的完整 PyTorch 实现，包含 masked attention 功能：

import torch
import torch.nn as nn
import math

class CABMAttention(nn.Module):
    """
    CABM 自注意力实现
    Args:
        embed_dim: 输入嵌入维度
        num_heads: 注意力头数
        dropout: dropout 概率
        bias: 是否使用偏置项
    """
    def __init__(self, embed_dim, num_heads, dropout=0.1, bias=True):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert (self.head_dim * num_heads == embed_dim), "embed_dim 必须能被 num_heads 整除"

        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, key_padding_mask=None):
        """
        Args:
            x: 输入序列 [batch_size, seq_len, embed_dim]
            key_padding_mask: 用于忽略 padding 的 mask [batch_size, seq_len]
        Returns:
            注意力输出 [batch_size, seq_len, embed_dim]
        """
        batch_size, seq_len, embed_dim = x.size()

        # 线性变换得到 Q,K,V
        q = self.q_proj(x)  # [batch_size, seq_len, embed_dim]
        k = self.k_proj(x)
        v = self.v_proj(x)

        # 重塑为多头形式
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

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

        # 应用 mask
        if key_padding_mask is not None:
            attn_scores = attn_scores.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2),
                float('-inf')
            )

        # softmax 归一化
        attn_weights = torch.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 加权求和
        output = torch.matmul(attn_weights, v)

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)

        # 最终线性变换
        output = self.out_proj(output)

        return output

主要内存消耗来自 Q、K、V 矩阵，大小为 $O(n^2)$，其中 n 是序列长度
多头注意力将内存消耗降低为 $O(n^2/h)$，h 是头数

使用 Flash Attention 技术减少内存访问次数
采用稀疏注意力模式处理超长序列
实现 KV 缓存避免重复计算

# 使用 PyTorch 的 scaled_dot_product_attention 高效实现
from torch.nn.functional import scaled_dot_product_attention

class ParallelCABMAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.inner_attn = nn.MultiheadAttention(embed_dim, num_heads)

    def forward(self, x, mask=None):
        return self.inner_attn(x, x, x, key_padding_mask=mask)[0]