[Depth] クロスアテンションの仕組み

クロス アテンション メカニズムはクロス アテンションとしても知られ、あるシーケンス内の特定の位置が別のシーケンス内のすべての位置に対してアテンション計算を実行するアテンション メカニズムを指します。

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

class CrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim):
        super(CrossAttention, self).__init__()
        self.query_dim = query_dim
        self.context_dim = context_dim

        self.linear_q = nn.Linear(query_dim, query_dim)
        self.linear_c = nn.Linear(context_dim, query_dim)

    def forward(self, query, context):
        # Query和Context的维度分别为 [batch_size, query_len, query_dim] 和 [batch_size, context_len, context_dim]
        # 首先将Query和Context分别通过线性变换
        query_proj = self.linear_q(query)  # [batch_size, query_len, query_dim]
        context_proj = self.linear_c(context)  # [batch_size, context_len, query_dim]

        # 计算注意力权重
        attention_weights = torch.bmm(query_proj, context_proj.transpose(1, 2))  # [batch_size, query_len, context_len]
        attention_weights = F.softmax(attention_weights, dim=-1)

        # 对Context序列进行加权求和
        attended_context = torch.bmm(attention_weights, context)  # [batch_size, query_len, context_dim]

        return attended_context, attention_weights