3.3 注意力机制的高级变体

3.3 注意力机制的高级变体

读者读完这节，能够深入理解跨模态注意力、层次化注意力、动态注意力等高级变体，掌握复杂的注意力架构设计。

3.3.1 跨模态注意力 (Cross-Modal Attention)

跨模态注意力是连接不同类型数据（如文本、图像、音频）的关键技术，它允许模型在多种模态之间建立语义联系。

基本原理

跨模态注意力通过计算不同模态特征之间的相似度，实现跨模态的信息融合：

class CrossModalAttention(nn.Module):
    def __init__(self, text_dim, image_dim, hidden_dim, num_heads=8):
        super().__init__()
        self.text_dim = text_dim
        self.image_dim = image_dim
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        
        # 模态特定的线性变换
        self.text_projection = nn.Linear(text_dim, hidden_dim)
        self.image_projection = nn.Linear(image_dim, hidden_dim)
        
        # 跨模态注意力
        self.cross_attention = nn.MultiheadAttention(hidden_dim, num_heads)
        
        # 输出投影
        self.output_projection = nn.Linear(hidden_dim, text_dim)
        
    def forward(self, text_features, image_features):
        """
        Args:
            text_features: [batch_size, seq_len, text_dim]
            image_features: [batch_size, img_seq_len, image_dim]
        Returns:
            cross_modal_output: [batch_size, seq_len, text_dim]
        """
        batch_size, seq_len, text_dim = text_features.shape
        _, img_seq_len, image_dim = image_features.shape
        
        # 模态投影
        text_proj = self.text_projection(text_features)  # [batch, seq_len, hidden]
        image_proj = self.image_projection(image_features)  # [batch, img_seq_len, hidden]
        
        # 跨模态注意力：文本关注图像
        cross_output, _ = self.cross_attention(
            query=text_proj,
            key=image_proj,
            value=image_proj
        )
        
        # 输出投影
        output = self.output_projection(cross_output)
        
        return output
​

多模态融合策略

跨模态注意力有多种融合策略，每种策略适用于不同的场景：

class MultiModalFusion(nn.Module):
    def __init__(self, modal_dims, fusion_type='cross_attention', hidden_dim=512):
        super().__init__()
        self.modal_dims = modal_dims
        self.fusion_type = fusion_type
        self.hidden_dim = hidden_dim
        self.num_modalities = len(modal_dims)
        
        # 每个模态的投影层
        self.projections = nn.ModuleList([
            nn.Linear(dim, hidden_dim) for dim in modal_dims
        ])
        
        if fusion_type == 'cross_attention':
            self.fusion_layer = self._cross_attention_fusion
        elif fusion_type == 'co_attention':
            self.fusion_layer = self._co_attention_fusion
        elif fusion_type == 'hierarchical':
            self.fusion_layer = self._hierarchical_fusion
            
    def _cross_attention_fusion(self, features):
        """交叉注意力融合"""
        # 选择第一个模态作为查询
        query = features[0]
        
        # 其他模态作为键和值
        keys = features[1:]
        values = features[1:]
        
        # 跨模态注意力
        fused_features = []
        for key, value in zip(keys, values):
            attention_output, _ = nn.MultiheadAttention(
                self.hidden_dim, 8
            )(query, key, value)
            fused_features.append(attention_output)
        
        # 融合所有模态
        fused = torch.mean(torch.stack(fused_features), dim=0)
        return fused
    
    def _co_attention_fusion(self, features):
        """协同注意力融合"""
        batch_size = features[0].shape[0]
        
        # 计算模态间的注意力矩阵
        attention_matrices = []
        for i in range(self.num_modalities):
            for j in range(self.num_modalities):
                if i != j:
                    # 模态i到模态j的注意力
                    attn_matrix = torch.matmul(
                        features[i], features[j].transpose(-2, -1)
                    )
                    attention_matrices.append(attn_matrix)
        
        # 使用注意力矩阵进行融合
        fused_features = []
        for i, feature in enumerate(features):
            modality_attention = torch.stack(attention_matrices[i*self.num_modalities:(i+1)*self.num_modalities])
            weighted_feature = torch.matmul(modality_attention.mean(dim=0), feature)
            fused_features.append(weighted_feature)
        
        return torch.mean(torch.stack(fused_features), dim=0)
    
    def _hierarchical_fusion(self, features):
        """层次化融合"""
        # 首先进行两两模态融合
        pairwise_fused = []
        for i in range(len(features)):
            for j in range(i+1, len(features)):
                pairwise = torch.mean(torch.stack([features[i], features[j]]), dim=0)
                pairwise_fused.append(pairwise)
        
        # 然后进行全局融合
        global_fused = torch.mean(torch.stack(pairwise_fused), dim=0)
        return global_fused
    
    def forward(self, features_dict):
        """
        Args:
            features_dict: dict containing modal features
                         e.g., {'text': tensor, 'image': tensor, 'audio': tensor}
        Returns:
            fused_output: tensor of shape [batch_size, seq_len, hidden_dim]
        """
        # 提取并投影各模态特征
        features = []
        for modality, feature in features_dict.items():
            projected = self.projections[len(features)](feature)
            features.append(projected)
        
        # 应用融合策略
        fused = self.fusion_layer(features)
        
        return fused
​

实际应用案例

视觉问答 (VQA) 中的跨模态注意力：

class VQAModel(nn.Module):
    def __init__(self, text_dim, image_dim, hidden_dim, num_classes):
        super().__init__()
        self.text_dim = text_dim
        self.image_dim = image_dim
        self.hidden_dim = hidden_dim
        
        # 文本编码器
        self.text_encoder = nn.LSTM(text_dim, hidden_dim, batch_first=True)
        
        # 图像编码器
        self.image_encoder = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((14, 14)),
            nn.Flatten(),
            nn.Linear(64*14*14, hidden_dim)
        )
        
        # 跨模态注意力
        self.cross_modal_attention = CrossModalAttention(
            text_dim, image_dim, hidden_dim
        )
        
        # 分类器
        self.classifier = nn.Linear(hidden_dim, num_classes)
        
    def forward(self, text, image):
        # 编码文本
        text_features = self.text_encoder(text)[0]  # [batch, seq_len, hidden]
        
        # 编码图像
        image_features = self.image_encoder(image)  # [batch, hidden]
        image_features = image_features.unsqueeze(1).expand(-1, text_features.shape[1], -1)
        
        # 跨模态注意力
        fused_features = self.cross_modal_attention(text_features, image_features)
        
        # 全局池化和分类
        pooled = torch.mean(fused_features, dim=1)
        logits = self.classifier(pooled)
        
        return logits
​

3.3.2 层次化注意力 (Hierarchical Attention)

层次化注意力通过在多个粒度上应用注意力机制，捕获不同层次的信息模式。

多粒度注意力架构

class HierarchicalAttention(nn.Module):
    def __init__(self, input_dim, hidden_dims, num_heads_list):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dims = hidden_dims
        self.num_layers = len(hidden_dims)
        self.num_heads_list = num_heads_list
        
        # 多层注意力
        self.attention_layers = nn.ModuleList([
            nn.MultiheadAttention(hidden_dims[i], num_heads_list[i])
            for i in range(self.num_layers)
        ])
        
        # 层次间的跳跃连接
        self.skip_connections = nn.ModuleList([
            nn.Linear(input_dim if i == 0 else hidden_dims[i-1], hidden_dims[i])
            for i in range(self.num_layers)
        ])
        
        # 输出层
        self.output_layer = nn.Linear(hidden_dims[-1], input_dim)
        
    def forward(self, x):
        """
        Args:
            x: input tensor [batch_size, seq_len, input_dim]
        Returns:
            output: tensor [batch_size, seq_len, input_dim]
        """
        batch_size, seq_len, input_dim = x.shape
        
        current_input = x
        
        for i, (attention_layer, skip_conn) in enumerate(
            zip(self.attention_layers, self.skip_connections)
        ):
            # 跳跃连接
            skip_input = skip_conn(current_input)
            
            # 注意力计算
            attention_output, _ = attention_layer(
                current_input, current_input, current_input
            )
            
            # 残差连接
            if i == 0:
                current_input = attention_output + skip_input
            else:
                current_input = attention_output + current_input + skip_input
            
            # 层归一化
            current_input = nn.LayerNorm(self.hidden_dims[i])(current_input)
        
        # 最终输出
        output = self.output_layer(current_input)
        
        return output
​

树状注意力结构

class TreeAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim, tree_structure):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.tree_structure = tree_structure  # 定义树的层次结构
        
        # 为每个树节点创建注意力层
        self.node_attentions = nn.ModuleDict()
        for node_id in tree_structure.nodes:
            self.node_attentions[str(node_id)] = nn.MultiheadAttention(
                input_dim if tree_structure.nodes[node_id]['level'] == 0 else hidden_dim,
                8
            )
        
        # 节点间的信息传递
        self.node_connections = nn.ModuleDict()
        for edge in tree_structure.edges:
            src, dst = edge
            self.node_connections[f"{src}->{dst}"] = nn.Linear(
                input_dim if tree_structure.nodes[src]['level'] == 0 else hidden_dim,
                hidden_dim
            )
    
    def forward(self, x, tree_positions):
        """
        Args:
            x: input features [batch_size, seq_len, input_dim]
            tree_positions: positions in tree structure [batch_size, seq_len]
        Returns:
            tree_output: output features [batch_size, seq_len, hidden_dim]
        """
        # 按树层次处理
        for level in range(max(self.tree_structure.nodes[node]['level'] for node in self.tree_structure.nodes)):
            level_nodes = [node for node in self.tree_structure.nodes 
                          if self.tree_structure.nodes[node]['level'] == level]
            
            # 处理同一层的节点
            for node in level_nodes:
                node_mask = (tree_positions == node)
                if node_mask.any():
                    node_features = x[node_mask]
                    
                    # 获取父节点的信息
                    parent = self.tree_structure.nodes[node].get('parent')
                    if parent is not None:
                        parent_mask = (tree_positions == parent)
                        if parent_mask.any():
                            parent_features = x[parent_mask]
                            
                            # 节点间信息传递
                            connection = self.node_connections[f"{parent}->{node}"]
                            parent_info = connection(parent_features.mean(dim=1, keepdim=True))
                            
                            # 节点自注意力
                            node_output, _ = self.node_attentions[str(node)](
                                node_features, node_features, node_features
                            )
                            
                            # 融合父节点信息
                            x[node_mask] = node_output + parent_info
        
        return x
​

实际应用：文档层次化理解

class DocumentHierarchicalAttention(nn.Module):
    def __init__(self, word_dim, sentence_dim, doc_dim, num_heads):
        super().__init__()
        self.word_dim = word_dim
        self.sentence_dim = sentence_dim
        self.doc_dim = doc_dim
        
        # 词级注意力
        self.word_attention = nn.MultiheadAttention(word_dim, num_heads)
        
        # 句级注意力
        self.sentence_attention = nn.MultiheadAttention(sentence_dim, num_heads)
        
        # 文档级注意力
        self.doc_attention = nn.MultiheadAttention(doc_dim, num_heads)
        
        # 跨层级信息传递
        self.word_to_sentence = nn.Linear(word_dim, sentence_dim)
        self.sentence_to_doc = nn.Linear(sentence_dim, doc_dim)
        self.doc_to_sentence = nn.Linear(doc_dim, sentence_dim)
        self.sentence_to_word = nn.Linear(sentence_dim, word_dim)
        
    def forward(self, word_embeddings, sentence_masks):
        """
        Args:
            word_embeddings: [batch_size, doc_len, sent_len, word_dim]
            sentence_masks: [batch_size, doc_len, sent_len]
        Returns:
            document_representation: [batch_size, doc_len, sent_len, word_dim]
        """
        batch_size, doc_len, sent_len, word_dim = word_embeddings.shape
        
        # 1. 词级注意力
        word_features = word_embeddings.view(batch_size * doc_len, sent_len, word_dim)
        word_output, _ = self.word_attention(word_features, word_features, word_features)
        word_output = word_output.view(batch_size, doc_len, sent_len, word_dim)
        
        # 2. 句级注意力
        sentence_features = torch.mean(word_output, dim=2)  # [batch, doc_len, word_dim]
        sentence_info = self.word_to_sentence(sentence_features)  # [batch, doc_len, sentence_dim]
        
        sentence_output, _ = self.sentence_attention(
            sentence_info, sentence_info, sentence_info
        )
        
        # 3. 文档级注意力
        doc_features = torch.mean(sentence_output, dim=1)  # [batch, sentence_dim]
        doc_info = self.sentence_to_doc(doc_features)  # [batch, doc_dim]
        
        doc_output, _ = self.doc_attention(
            doc_info.unsqueeze(1), doc_info.unsqueeze(1), doc_info.unsqueeze(1)
        )
        doc_output = doc_output.squeeze(1)  # [batch, doc_dim]
        
        # 4. 自上而下的信息传递
        doc_to_sentence = self.doc_to_sentence(doc_output.unsqueeze(1)).unsqueeze(1)  # [batch, 1, 1, doc_dim]
        sentence_output = sentence_output + doc_to_sentence.expand(-1, doc_len, -1)
        
        sentence_to_word = self.sentence_to_word(sentence_output.unsqueeze(2)).unsqueeze(2)  # [batch, doc_len, 1, sentence_dim]
        word_output = word_output + sentence_to_word.expand(-1, -1, sent_len, -1)
        
        return word_output, sentence_output, doc_output
​

总结

本节介绍了注意力机制的高级变体：

1. 跨模态注意力: 实现不同模态间的信息融合，适用于多模态任务
2. 层次化注意力: 在多个粒度上应用注意力，捕获层次化信息
3. 实际应用: 提供了VQA和文档理解的实例

这些高级变体为解决复杂任务提供了强大的工具。