3.3 注意力机制的高级变体(续)

3.3 注意力机制的高级变体（续）

3.3.3 动态注意力 (Dynamic Attention)

动态注意力根据输入内容动态调整注意力模式，实现自适应的信息选择和整合。

基于内容的动态注意力

class DynamicContentAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_heads=8, max_dynamic_heads=16):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.max_dynamic_heads = max_dynamic_heads
        
        # 内容分析网络
        self.content_analyzer = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, max_dynamic_heads * 2)
        )
        
        # 动态注意力权重计算
        self.dynamic_attention = nn.MultiheadAttention(input_dim, num_heads)
        
        # 动态头选择
        self.head_selector = nn.Sequential(
            nn.Linear(max_dynamic_heads * 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, max_dynamic_heads),
            nn.Sigmoid()
        )
        
    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
        
        # 分析内容，确定动态注意力配置
        content_analysis = self.content_analyzer(x)  # [batch, seq_len, max_dynamic_heads * 2]
        
        # 选择合适的注意力头
        head_selection = self.head_selector(content_analysis)  # [batch, seq_len, max_dynamic_heads]
        
        # 动态注意力计算
        dynamic_output, _ = self.dynamic_attention(x, x, x)
        
        # 根据内容分析调整输出
        adjusted_output = dynamic_output * head_selection.mean(dim=1, keepdim=True)
        
        return adjusted_output
​

基于位置的动态注意力

class DynamicPositionAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim, max_seq_len=1024):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.max_seq_len = max_seq_len
        
        # 位置编码
        self.positional_encoding = nn.Parameter(
            torch.randn(max_seq_len, hidden_dim)
        )
        
        # 动态注意力掩码生成
        self.mask_generator = nn.Sequential(
            nn.Linear(input_dim + hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, max_seq_len),
            nn.Sigmoid()
        )
        
        # 注意力计算
        self.attention = nn.MultiheadAttention(input_dim, 8)
        
    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
        
        # 位置编码
        pos_encoding = self.positional_encoding[:seq_len, :].unsqueeze(0).expand(batch_size, -1, -1)
        
        # 位置特征
        position_features = torch.cat([x, pos_encoding], dim=-1)
        
        # 生成动态注意力掩码
        dynamic_mask = self.mask_generator(position_features)  # [batch, seq_len, max_seq_len]
        dynamic_mask = dynamic_mask[:, :, :seq_len]  # 截断到实际长度
        
        # 应用掩码
        masked_attention = self.attention(
            x, x, x,
            attn_mask=(dynamic_mask < 0.5)
        )
        
        return masked_attention[0]
​

自适应注意力窗口

class AdaptiveWindowAttention(nn.Module):
    def __init__(self, input_dim, hidden_dim, min_window=4, max_window=128):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.min_window = min_window
        self.max_window = max_window
        
        # 窗口大小预测器
        self.window_predictor = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )
        
        # 多窗口注意力
        self.window_attentions = nn.ModuleList([
            nn.MultiheadAttention(input_dim, 8)
            for _ in range(int(math.log2(max_window // min_window)) + 1)
        ])
        
        # 窗口融合
        self.window_fusion = nn.Linear(len(self.window_attentions), 1)
        
    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
        
        # 预测每个位置的最优窗口大小
        window_predictions = self.window_predictor(x).squeeze(-1)  # [batch, seq_len]
        
        # 多尺度窗口注意力
        window_outputs = []
        for i, window_attention in enumerate(self.window_attentions):
            window_size = self.min_window * (2 ** i)
            
            # 应用窗口注意力
            window_output = []
            for pos in range(seq_len):
                start = max(0, pos - window_size // 2)
                end = min(seq_len, pos + window_size // 2 + 1)
                
                if end - start > 1:
                    window_input = x[:, start:end, :]
                    window_attn, _ = window_attention(
                        window_input, window_input, window_input
                    )
                    window_output.append(window_attn[:, pos - start, :].unsqueeze(1))
            
            if window_output:
                window_output = torch.cat(window_output, dim=1)
                window_outputs.append(window_output)
        
        # 根据预测权重融合不同窗口的输出
        if window_outputs:
            window_outputs = torch.stack(window_outputs, dim=-1)  # [batch, seq_len, num_windows]
            
            # 计算融合权重
            window_weights = self.window_fusion(window_outputs.transpose(-2, -1)).squeeze(-1)
            window_weights = F.softmax(window_weights, dim=-1)
            
            # 加权融合
            final_output = torch.sum(window_outputs * window_weights.unsqueeze(-1), dim=-1)
        else:
            final_output = x
        
        return final_output
​

3.3.4 注意力机制的组合与演化

多种注意力机制的组合

class HybridAttentionMechanism(nn.Module):
    def __init__(self, input_dim, hidden_dim, attention_types):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.attention_types = attention_types
        
        # 初始化不同类型的注意力
        self.attention_components = nn.ModuleDict()
        
        if 'self' in attention_types:
            self.attention_components['self'] = nn.MultiheadAttention(input_dim, 8)
        
        if 'cross' in attention_types:
            self.attention_components['cross'] = nn.MultiheadAttention(input_dim, 8)
        
        if 'local' in attention_types:
            self.attention_components['local'] = LocalAttention(input_dim, window_size=64)
        
        if 'global' in attention_types:
            self.attention_components['global'] = nn.MultiheadAttention(input_dim, 8)
        
        # 注意力融合权重
        self.fusion_weights = nn.Linear(len(attention_types) * input_dim, input_dim)
        
        # 输出投影
        self.output_projection = nn.Linear(input_dim, input_dim)
        
    def forward(self, query, key=None, value=None):
        """
        Args:
            query: query tensor [batch_size, seq_len, input_dim]
            key: key tensor (optional) [batch_size, seq_len, input_dim]
            value: value tensor (optional) [batch_size, seq_len, input_dim]
        Returns:
            output: tensor [batch_size, seq_len, input_dim]
        """
        attention_outputs = []
        
        # 计算各种注意力的输出
        for attn_type, attn_module in self.attention_components.items():
            if attn_type == 'self':
                output, _ = attn_module(query, query, query)
            elif attn_type == 'cross' and key is not None and value is not None:
                output, _ = attn_module(query, key, value)
            elif attn_type == 'local':
                output = attn_module(query)
            elif attn_type == 'global':
                output, _ = attn_module(query, query, query)
            
            attention_outputs.append(output)
        
        # 融合不同注意力的输出
        concatenated = torch.cat(attention_outputs, dim=-1)
        fused = self.fusion_weights(concatenated)
        
        # 输出投影
        output = self.output_projection(fused)
        
        return output
​

注意力机制的演化策略

class AttentionEvolution(nn.Module):
    def __init__(self, input_dim, hidden_dim, evolution_stages):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.evolution_stages = evolution_stages
        
        # 演化阶段的注意力组件
        self.stage_attentions = nn.ModuleList()
        for stage_config in evolution_stages:
            stage_attention = self._create_attention_stage(stage_config)
            self.stage_attentions.append(stage_attention)
        
        # 演化控制网络
        self.evolution_controller = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, len(evolution_stages)),
            nn.Softmax(dim=-1)
        )
        
    def _create_attention_stage(self, config):
        """根据配置创建注意力阶段"""
        stage_type = config['type']
        
        if stage_type == 'simple':
            return nn.MultiheadAttention(self.input_dim, 8)
        elif stage_type == 'local':
            return LocalAttention(self.input_dim, config.get('window_size', 64))
        elif stage_type == 'global':
            return nn.MultiheadAttention(self.input_dim, 8)
        elif stage_type == 'adaptive':
            return AdaptiveWindowAttention(self.input_dim, self.hidden_dim)
        else:
            raise ValueError(f"Unknown attention stage type: {stage_type}")
    
    def forward(self, x):
        """
        Args:
            x: input tensor [batch_size, seq_len, input_dim]
        Returns:
            output: tensor [batch_size, seq_len, input_dim]
            evolution_weights: weights for each stage [batch_size, num_stages]
        """
        # 计算演化权重
        evolution_weights = self.evolution_controller(x)
        
        # 计算各阶段的输出
        stage_outputs = []
        for i, stage_attention in enumerate(self.stage_attentions):
            stage_output = stage_attention(x)
            stage_outputs.append(stage_output)
        
        # 根据演化权重融合各阶段输出
        weighted_output = torch.zeros_like(x)
        for i, output in enumerate(stage_outputs):
            weighted_output += evolution_weights[:, i:i+1, None] * output
        
        return weighted_output, evolution_weights
​

3.3.5 实际应用与性能分析

多任务学习中的注意力机制组合

class MultiTaskAttention(nn.Module):
    def __init__(self, input_dim, task_dims, shared_hidden_dim=256):
        super().__init__()
        self.input_dim = input_dim
        self.task_dims = task_dims
        self.num_tasks = len(task_dims)
        
        # 共享的底层注意力
        self.shared_attention = nn.MultiheadAttention(input_dim, 8)
        
        # 任务特定的注意力
        self.task_attentions = nn.ModuleList([
            nn.MultiheadAttention(shared_hidden_dim, 8)
            for _ in range(self.num_tasks)
        ])
        
        # 任务特定的输出层
        self.task_outputs = nn.ModuleList([
            nn.Linear(shared_hidden_dim, dim)
            for dim in task_dims
        ])
        
        # 任务权重网络
        self.task_weights = nn.Sequential(
            nn.Linear(input_dim, shared_hidden_dim),
            nn.ReLU(),
            nn.Linear(shared_hidden_dim, self.num_tasks),
            nn.Softmax(dim=-1)
        )
        
    def forward(self, x):
        """
        Args:
            x: input tensor [batch_size, seq_len, input_dim]
        Returns:
            task_outputs: dict of task outputs
        """
        # 共享底层处理
        shared_output, _ = self.shared_attention(x, x, x)
        
        # 计算任务权重
        task_weights = self.task_weights(shared_output)  # [batch, seq_len, num_tasks]
        
        # 任务特定的处理
        task_outputs = {}
        for i, (task_attention, task_output) in enumerate(
            zip(self.task_attentions, self.task_outputs)
        ):
            # 加权输入
            weighted_input = shared_output * task_weights[:, :, i:i+1]
            
            # 任务特定注意力
            task_feature, _ = task_attention(
                weighted_input, weighted_input, weighted_input
            )
            
            # 任务输出
            task_outputs[f'task_{i}'] = task_output(task_feature)
        
        return task_outputs
​

性能基准测试

def benchmark_attention_variants(input_configurations, attention_models):
    """
    基准测试不同注意力变体的性能
    
    Args:
        input_configurations: list of input shapes
        attention_models: dict of model instances
    
    Returns:
        benchmark_results: performance metrics
    """
    results = {}
    
    for input_shape in input_configurations:
        batch_size, seq_len, input_dim = input_shape
        
        for model_name, model in attention_models.items():
            model.eval()
            model.to('cuda')
            
            # 生成测试数据
            test_input = torch.randn(batch_size, seq_len, input_dim).cuda()
            
            # 预热
            with torch.no_grad():
                for _ in range(10):
                    _ = model(test_input)
            
            # 性能测试
            torch.cuda.synchronize()
            start_time = time.time()
            
            with torch.no_grad():
                for _ in range(100):
                    _ = model(test_input)
            
            torch.cuda.synchronize()
            end_time = time.time()
            
            # 计算指标
            avg_time = (end_time - start_time) / 100
            throughput = (batch_size * seq_len * input_dim) / (avg_time / 1000)
            memory_usage = torch.cuda.memory_allocated() / 1024 / 1024
            
            if model_name not in results:
                results[model_name] = {}
            
            results[model_name][f'input_{batch_size}_{seq_len}_{input_dim}'] = {
                'avg_time_ms': avg_time * 1000,
                'throughput': throughput,
                'memory_usage_mb': memory_usage,
                'flops': batch_size * seq_len * seq_len * input_dim * 2 / (avg_time / 1000)  # 简化计算
            }
    
    return results

# 使用示例
input_configs = [
    (32, 128, 512),
    (32, 256, 512),
    (32, 512, 512),
    (32, 1024, 512),
    (32, 2048, 512)
]

attention_models = {
    'self_attention': nn.MultiheadAttention(512, 8),
    'cross_attention': nn.MultiheadAttention(512, 8),
    'local_attention': LocalAttention(512, 64),
    'cross_modal_attention': CrossModalAttention(512, 512, 512),
    'hierarchical_attention': HierarchicalAttention(512, [256, 512], [4, 8]),
    'dynamic_attention': DynamicContentAttention(512, 256)
}

benchmark_results = benchmark_attention_variants(input_configs, attention_models)
​

3.3.6 总结与展望

本节深入探讨了注意力机制的高级变体：

1. 跨模态注意力: 实现不同模态间的信息融合，适用于多模态任务
2. 层次化注意力: 在多个粒度上应用注意力，捕获层次化信息
3. 动态注意力: 根据输入内容动态调整注意力模式
4. 组合演化: 多种注意力的组合和演化策略

这些高级变体为解决复杂任务提供了强大的工具。随着深度学习的发展，注意力机制将继续演化，出现更多创新的变体和应用。

本节详细介绍了注意力机制的高级变体，从跨模态到层次化，从动态到组合演化，为构建复杂的注意力系统提供了丰富的选择。下一节将深入探讨FlashAttention的具体实现和优化技术。