替换Crossformer原生注意力为PyTorch scaled_dot_product_attention遇错求解

问题背景 #

将Crossformer中的原生注意力前向传播代码替换为PyTorch的scaled_dot_product_attention以借助Flash Attention加速时，出现维度不匹配错误。

原生注意力代码 #

def forward(self, queries, keys, values):
    B, L, H, E = queries.shape
    _, S, _, D = values.shape
    scale = self.scale or 1./sqrt(E)

    scores = torch.einsum("blhe,bshe->bhls", queries, keys)
    A = self.dropout(torch.softmax(scale * scores, dim=-1))
    V = torch.einsum("bhls,bshd->blhd", A, values)
    
    return V.contiguous()

改写后的错误代码 #

def forward(self, queries, keys, values):
    # B 表示批量大小
    # L 是查询的序列长度（或目标序列长度）
    # H 是注意力头数量
    # E 是每个查询/键注意力头的深度（维度）
    # S 是键/值的序列长度（或源序列长度）
    # D 是每个值注意力头的深度（维度）
    B, L, H, E = queries.shape
    _, S, _, D = values.shape

    y = torch.nn.functional.scaled_dot_product_attention(
        queries, keys, values, dropout_p=self.dropout_p if self.training else None)
    y = y.contiguous()
    return y

错误信息与张量尺寸 #

• 错误提示：

RuntimeError: The size of tensor a (10) must match the size of tensor b (4) at non-singleton dimension 1

• 调试得到的张量尺寸：
  • keys: (2048, 4, 16, 32)
  • queries: (2048, 10, 16, 32)
  • values: (2048, 4, 16, 32)

问题根源 #

PyTorch的scaled_dot_product_attention对4维输入的维度顺序有严格要求：需遵循(batch_size, num_heads, seq_len, embed_dim_per_head)格式（注意力头维度H位于序列长度维度L/S之前）。而Crossformer原生代码的输入维度顺序是(batch_size, seq_len, num_heads, embed_dim_per_head)（序列长度在前，注意力头在后），两者维度顺序不匹配，导致函数误将序列长度维度识别为注意力头维度，引发维度不匹配错误。

解决方法 #

在调用scaled_dot_product_attention前，先将输入张量的维度顺序调整为函数要求的格式；得到结果后，再将维度顺序转回原生代码的输出格式，确保前后逻辑一致。

修改后的正确代码 #

def forward(self, queries, keys, values):
    B, L, H, E = queries.shape
    _, S, _, D = values.shape

    # 调整维度顺序为 (B, H, L, E)，适配scaled_dot_product_attention的要求
    queries = queries.transpose(1, 2)  # 从 (B, L, H, E) 转为 (B, H, L, E)
    keys = keys.transpose(1, 2)        # 从 (B, S, H, E) 转为 (B, H, S, E)
    values = values.transpose(1, 2)    # 从 (B, S, H, D) 转为 (B, H, S, D)

    # 调用scaled_dot_product_attention，若原代码有自定义scale需传入scale=self.scale
    y = torch.nn.functional.scaled_dot_product_attention(
        queries, keys, values, 
        dropout_p=self.dropout_p if self.training else None,
        scale=self.scale if self.scale else None
    )
    
    # 将结果维度转回原生的 (B, L, H, D) 格式
    y = y.transpose(1, 2).contiguous()
    return y

补充说明 #

• 若原代码中self.scale未使用默认的1/sqrt(E)，需在scaled_dot_product_attention中显式传入scale=self.scale参数；若使用默认值则可省略。
• 转置操作transpose(1,2)仅交换维度顺序，不会改变张量的存储内容，性能开销可忽略。

内容的提问来源于stack exchange，提问作者Serge Rogatch