注意力机制从零实现

注意力机制是 Transformer 的核心。从零实现注意力机制有助于理解其数学原理和工程细节。

多头注意力实现

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, Q, K, V, mask=None):
        B = Q.shape[0]
        Q = self.W_q(Q).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(K).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(V).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores.masked_fill_(mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        return self.W_o(out.transpose(1, 2).contiguous().view(B, -1, -1))

注意力计算复杂度

标准注意力的复杂度为 O(n²)，n 为序列长度。这就是为什么长上下文推理如此昂贵——序列翻倍，计算量翻四倍。

相关资源

• Attention Is All You Need
• Attention 综述
• PyTorch Attention API