Vision Transformer中的Cross Attention机制：从原理到实践指南

1次阅读

没有评论

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

在传统的计算机视觉任务中，Self-Attention 机制在 Vision Transformer（ViT）中表现出色，但当我们需要处理多模态数据（如视觉 - 文本联合任务）时，它就显得力不从心了。

模态差异问题 ：图像和文本数据具有完全不同的特征分布和结构，直接使用 Self-Attention 难以有效对齐不同模态的特征。
计算效率低下 ：传统的多模态处理方法通常需要分别处理不同模态的数据，然后再进行融合，这种分步操作增加了计算开销。
信息交互不足 ：Self-Attention 只能在同模态内进行特征交互，缺乏跨模态的深层特征融合能力。

Cross Attention（简称 cat）是 Self-Attention 的扩展，专门用于处理不同模态间的特征交互。

架构差异
Self-Attention：Q、K、V 都来自同一个输入序列
Cross Attention：Q 来自一个模态，K、V 来自另一个模态
计算复杂度
Self-Attention：O(n²)（n 是序列长度）
Cross Attention：O(nm)（n 是 Q 序列长度，m 是 K / V 序列长度）
内存占用
Cross Attention 通常比 Self-Attention 更节省内存，特别是当两个模态的序列长度差异较大时

下面我们实现一个完整的 Cross Attention 层，包含多头注意力和必要的正则化操作。

import torch
import torch.nn as nn
import torch.nn.functional as F

class CrossAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # 定义 Q、K、V 的投影矩阵
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)

        # 输出投影
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, context, mask=None):
        """
        参数:
            x: 查询序列 [batch_size, seq_len_q, embed_dim]
            context: 上下文序列 [batch_size, seq_len_kv, embed_dim]
            mask: 可选注意力掩码 [batch_size, seq_len_q, seq_len_kv]
        """
        batch_size = x.size(0)

        # 1. 投影计算
        Q = self.q_proj(x)  # [B, L_q, D]
        K = self.k_proj(context)  # [B, L_kv, D]
        V = self.v_proj(context)  # [B, L_kv, D]

        # 2. 多头拆分
        Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # [B, H, L_q, D/H]
        K = K.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # [B, H, L_kv, D/H]
        V = V.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # [B, H, L_kv, D/H]

        # 3. 缩放点积注意力
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)  # [B, H, L_q, L_kv]

        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_output = torch.matmul(attn_weights, V)  # [B, H, L_q, D/H]

        # 4. 多头拼接
        attn_output = attn_output.transpose(1, 2).contiguous()  # [B, L_q, H, D/H]
        attn_output = attn_output.view(batch_size, -1, self.embed_dim)  # [B, L_q, D]

        # 5. 输出投影
        output = self.out_proj(attn_output)

        return output