从堆叠、分层到双线性：手把手带你复现注意力机制的几次关键进化

从堆叠、分层到双线性：手把手带你复现注意力机制的几次关键进化

注意力机制在深度学习领域的崛起，彻底改变了我们处理序列数据的方式。从最初的简单堆叠到如今复杂的双线性结构，每一次进化都带来了性能的显著提升。本文将带你穿越这段技术发展史，通过代码实践亲身体验每一次关键突破背后的设计哲学。

1. 注意力机制的起源与堆叠实现

2014年，Bahdanau等人首次将注意力机制引入机器翻译任务，开启了这一领域的先河。堆叠注意力（Stacked Attention）作为早期代表，通过多层注意力堆叠实现对输入数据的逐步聚焦。

1.1 基础注意力模块实现

我们先从最基础的注意力模块开始。以下是一个使用PyTorch实现的基础注意力层：

import torch import torch.nn as nn import torch.nn.functional as F class BasicAttention(nn.Module): def __init__(self, hidden_dim): super().__init__() self.attn = nn.Linear(hidden_dim * 2, hidden_dim) self.v = nn.Linear(hidden_dim, 1, bias=False) def forward(self, hidden, encoder_outputs): # hidden: (batch_size, hidden_dim) # encoder_outputs: (batch_size, seq_len, hidden_dim) seq_len = encoder_outputs.size(1) hidden = hidden.unsqueeze(1).repeat(1, seq_len, 1) # (batch_size, seq_len, hidden_dim) energy = torch.tanh(self.attn(torch.cat([hidden, encoder_outputs], dim=2))) # (batch_size, seq_len, hidden_dim) attention = self.v(energy).squeeze(2) # (batch_size, seq_len) return F.softmax(attention, dim=1)

注意：基础注意力模块的计算复杂度为O(n^2)，这在处理长序列时会成为性能瓶颈。

1.2 堆叠注意力的实现技巧

堆叠注意力的核心思想是通过多层注意力网络逐步细化关注区域。在视觉问答(VQA)任务中，这种结构特别有效：

class StackedAttention(nn.Module): def __init__(self, hidden_dim, num_layers=2): super().__init__() self.layers = nn.ModuleList([BasicAttention(hidden_dim) for _ in range(num_layers)]) def forward(self, question_embed, image_features): # question_embed: (batch_size, hidden_dim) # image_features: (batch_size, num_regions, hidden_dim) u = question_embed for layer in self.layers: attn_weights = layer(u, image_features) # (batch_size, num_regions) attn_applied = torch.bmm(attn_weights.unsqueeze(1), image_features).squeeze(1) # (batch_size, hidden_dim) u = u + attn_applied # 残差连接 return u

堆叠注意力的优势在于：

• 逐步细化关注区域
• 通过残差连接保持信息流动
• 相对简单的实现复杂度

但在实际应用中，我们发现它存在以下局限：

1. 各层注意力独立计算，缺乏全局协调
2. 对多层次语义理解能力有限
3. 计算资源消耗随层数线性增长

2. 分层注意力网络(HAM)的突破

分层注意力网络(Hierarchical Attention Model)的提出，解决了堆叠注意力在处理多层次语义信息时的不足。它通过构建层次化的注意力机制，实现了从局部到全局的信息整合。

2.1 HAM的核心架构

HAM通常包含两个主要层次：

1. 词级注意力：处理原始输入序列
2. 句级注意力：聚合词级表示

class HierarchicalAttention(nn.Module): def __init__(self, word_dim, sentence_dim): super().__init__() self.word_attention = BasicAttention(word_dim) self.sentence_attention = BasicAttention(sentence_dim) def forward(self, document): # document: (batch_size, num_sentences, num_words, word_dim) batch_size, num_sentences, num_words, _ = document.size() # 词级注意力 word_encoded = [] for i in range(num_sentences): sentence = document[:, i, :, :] # (batch_size, num_words, word_dim) word_attn = self.word_attention(sentence.mean(dim=1), sentence) # (batch_size, num_words) word_encoded.append(torch.bmm(word_attn.unsqueeze(1), sentence).squeeze(1)) # (batch_size, word_dim) word_encoded = torch.stack(word_encoded, dim=1) # (batch_size, num_sentences, word_dim) # 句级注意力 sentence_attn = self.sentence_attention(word_encoded.mean(dim=1), word_encoded) # (batch_size, num_sentences) doc_encoded = torch.bmm(sentence_attn.unsqueeze(1), word_encoded).squeeze(1) # (batch_size, word_dim) return doc_encoded

2.2 HAM的优化技巧

在实践中，我们发现以下技巧可以显著提升HAM性能：

# 改进版的多头分层注意力实现 class MultiHeadHierarchicalAttention(nn.Module): def __init__(self, word_dim, sentence_dim, num_heads=4): super().__init__() self.word_attentions = nn.ModuleList([BasicAttention(word_dim) for _ in range(num_heads)]) self.sentence_attentions = nn.ModuleList([BasicAttention(sentence_dim) for _ in range(num_heads)]) self.word_proj = nn.Linear(word_dim * num_heads, word_dim) self.sentence_proj = nn.Linear(sentence_dim * num_heads, sentence_dim) def forward(self, document): # 实现类似上述但包含多头注意力 # 详细实现略 pass

3. 分层共同注意力网络

分层共同注意力(Hierarchical Co-Attention)进一步扩展了HAM的概念，引入了双向注意力机制，特别适合视觉-语言多模态任务。

3.1 共同注意力的实现原理

共同注意力的核心是同时计算：

• 图像到问题的注意力
• 问题到图像的注意力

class CoAttention(nn.Module): def __init__(self, visual_dim, question_dim): super().__init__() self.W = nn.Linear(question_dim, visual_dim) def forward(self, visual_feats, question_feats): # visual_feats: (batch_size, num_regions, visual_dim) # question_feats: (batch_size, seq_len, question_dim) # 计算亲和矩阵 transformed_question = self.W(question_feats) # (batch_size, seq_len, visual_dim) affinity = torch.bmm(visual_feats, transformed_question.transpose(1, 2)) # (batch_size, num_regions, seq_len) # 图像到问题的注意力 v2q_attn = F.softmax(affinity.max(dim=1)[0], dim=-1) # (batch_size, seq_len) q_attended = torch.bmm(v2q_attn.unsqueeze(1), question_feats).squeeze(1) # (batch_size, question_dim) # 问题到图像的注意力 q2v_attn = F.softmax(affinity.max(dim=2)[0], dim=-1) # (batch_size, num_regions) v_attended = torch.bmm(q2v_attn.unsqueeze(1), visual_feats).squeeze(1) # (batch_size, visual_dim) return q_attended, v_attended

3.2 分层共同注意力的完整实现

将共同注意力与分层结构结合，我们得到完整的Hierarchical Co-Attention网络：

class HierarchicalCoAttention(nn.Module): def __init__(self, visual_dim, word_dim, sentence_dim): super().__init__() self.word_level_coattn = CoAttention(visual_dim, word_dim) self.sentence_level_coattn = CoAttention(visual_dim, sentence_dim) def forward(self, visual_feats, question_words, question_sentences): # 词级共同注意力 word_q, word_v = self.word_level_coattn(visual_feats, question_words) # 句级共同注意力 sentence_q, sentence_v = self.sentence_level_coattn(visual_feats, question_sentences) # 融合各层级表示 fused_visual = word_v + sentence_v fused_question = word_q + sentence_q return fused_visual, fused_question

这种结构在VQA任务中表现出色，因为它能够：

• 同时捕捉视觉和语言线索
• 在不同粒度上建立跨模态关联
• 通过层级结构处理复杂语义关系

4. 双线性注意力网络的创新

双线性注意力网络(Bilinear Attention Networks)代表了注意力机制的最新进展，通过双线性交互实现了更精细的特征融合。

4.1 双线性注意力的数学基础

传统的注意力机制通常使用点积或加性注意力，而双线性注意力引入了更丰富的交互方式：

注意力分数 = softmax(Q^T W K)

其中W是可学习的双线性权重矩阵，允许查询(Query)和键(Key)之间进行更复杂的交互。

4.2 PyTorch实现细节

class BilinearAttention(nn.Module): def __init__(self, query_dim, key_dim, attn_dim): super().__init__() self.W = nn.Parameter(torch.randn(query_dim, attn_dim, key_dim) * 0.01) self.query_proj = nn.Linear(query_dim, query_dim) self.key_proj = nn.Linear(key_dim, key_dim) def forward(self, query, keys): # query: (batch_size, query_dim) # keys: (batch_size, seq_len, key_dim) query = self.query_proj(query) # (batch_size, query_dim) keys = self.key_proj(keys) # (batch_size, seq_len, key_dim) # 双线性交互 query = query.unsqueeze(1) # (batch_size, 1, query_dim) intermediate = torch.matmul(query, self.W) # (batch_size, 1, attn_dim, key_dim) intermediate = intermediate.squeeze(1).transpose(1, 2) # (batch_size, key_dim, attn_dim) scores = torch.matmul(keys, intermediate) # (batch_size, seq_len, attn_dim) scores = scores.mean(dim=2) # (batch_size, seq_len) attn_weights = F.softmax(scores, dim=1) context = torch.bmm(attn_weights.unsqueeze(1), keys).squeeze(1) return context, attn_weights

4.3 双线性注意力的优势分析

与传统注意力机制相比，双线性注意力具有以下优势：

1. 更丰富的特征交互：通过双线性变换允许查询和键之间进行更复杂的交互
2. 更强的表达能力：可学习权重矩阵W能够捕捉更复杂的关联模式
3. 灵活的参数控制：attn_dim参数可以调节模型的容量和计算复杂度

在实际应用中，我们通常会在双线性注意力基础上添加以下改进：

• 多头机制
• 层归一化
• 残差连接
• 低秩分解（减少参数数量）

class MultiHeadBilinearAttention(nn.Module): def __init__(self, query_dim, key_dim, attn_dim, num_heads=8): super().__init__() self.heads = nn.ModuleList([ BilinearAttention(query_dim // num_heads, key_dim // num_heads, attn_dim // num_heads) for _ in range(num_heads) ]) self.proj = nn.Linear((key_dim // num_heads) * num_heads, key_dim) def forward(self, query, keys): # 分割输入到各个头 query = query.chunk(len(self.heads), dim=-1) keys = keys.chunk(len(self.heads), dim=-1) outputs = [] for head, q, k in zip(self.heads, query, keys): context, _ = head(q, k) outputs.append(context) # 合并多头输出 combined = torch.cat(outputs, dim=-1) return self.proj(combined)

5. 实战：在VQA任务中比较各代注意力

为了直观展示不同注意力机制的差异，我们在VQA v2数据集上进行了对比实验。以下是关键实现步骤：

5.1 数据准备

from torchvision import transforms from PIL import Image # 图像预处理 transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 文本预处理 def preprocess_question(question, word2idx): tokens = question.lower().split() return [word2idx.get(token, word2idx['<unk>']) for token in tokens]

5.2 模型训练框架

import torch.optim as optim from torch.utils.data import DataLoader def train_model(model, train_loader, val_loader, epochs=10): criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) for epoch in range(epochs): model.train() train_loss = 0.0 for images, questions, answers in train_loader: optimizer.zero_grad() outputs = model(images, questions) loss = criterion(outputs, answers) loss.backward() optimizer.step() train_loss += loss.item() # 验证阶段 model.eval() val_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for images, questions, answers in val_loader: outputs = model(images, questions) loss = criterion(outputs, answers) val_loss += loss.item() _, predicted = outputs.max(1) correct += predicted.eq(answers).sum().item() total += answers.size(0) print(f'Epoch {epoch+1}: Train Loss {train_loss/len(train_loader):.4f}, ' f'Val Loss {val_loss/len(val_loader):.4f}, ' f'Accuracy {100.*correct/total:.2f}%')

5.3 性能对比结果

我们在相同条件下训练了四种注意力模型，结果如下：

从实验结果可以看出：

1. 模型复杂度与性能呈正相关
2. 双线性注意力虽然计算成本较高，但性能优势明显
3. 共同注意力在多模态任务中表现突出

6. 注意力机制的未来发展方向

尽管双线性注意力已经取得了显著成果，但这一领域仍在快速发展。以下是一些值得关注的方向：

1. 稀疏注意力：通过限制注意力范围降低计算复杂度

   • 局部窗口注意力
   • 轴向注意力
   • 稀疏变换器

2. 内存高效的注意力：

   • 线性注意力变体
   • 核化注意力
   • 低秩近似

3. 内容感知的动态注意力：

   • 根据输入动态调整注意力机制
   • 混合专家注意力
   • 可微架构搜索

# 一个简单的稀疏注意力实现示例 class SparseBilinearAttention(nn.Module): def __init__(self, query_dim, key_dim, attn_dim, sparsity=0.5): super().__init__() self.bilinear = BilinearAttention(query_dim, key_dim, attn_dim) self.sparsity = sparsity def forward(self, query, keys): context, attn_weights = self.bilinear(query, keys) # 应用稀疏化 if self.sparsity > 0: k = int(attn_weights.size(-1) * (1 - self.sparsity)) values, indices = torch.topk(attn_weights, k, dim=-1) sparse_weights = torch.zeros_like(attn_weights).scatter_(-1, indices, values) sparse_weights = F.normalize(sparse_weights, p=1, dim=-1) context = torch.bmm(sparse_weights.unsqueeze(1), keys).squeeze(1) return context

在实际项目中，选择哪种注意力机制取决于具体需求：

• 计算资源有限：堆叠或分层注意力
• 多模态任务：共同注意力
• 追求最佳性能：双线性注意力
• 超长序列：稀疏注意力变体