如何在PyTorch的MultiHeadAttention层中获取值投影

你在实现Graphormer架构时，需要在自注意力的注意力权重中加入边偏置，再与值投影后的张量做矩阵乘法，但PyTorch原生的nn.MultiheadAttention没有直接暴露值投影模块，这里提供几种可行的实现方式：

方法1：直接拆分原生MultiheadAttention的投影参数 #

PyTorch的nn.MultiheadAttention的in_proj_weight是query、key、value三个投影层的权重拼接而成（形状为[3*C, C]），in_proj_bias同理（形状为[3*C]）。我们可以拆分出value对应的权重和偏置，手动完成值投影：

import torch as th
from torch import nn

# 变量初始化
B, T, C, H = 2, 3, 4, 2
self_attn = nn.MultiheadAttention(C, H, batch_first=True)

# 张量定义
x = th.randn(B, T, C)
attn_bias = th.ones((B, T, T))

# 1. 拆分q、k、v投影的权重和偏置
query_weight = self_attn.in_proj_weight[:C, :]
key_weight = self_attn.in_proj_weight[C:2*C, :]
value_weight = self_attn.in_proj_weight[-C:, :]

query_bias = self_attn.in_proj_bias[:C] if self_attn.in_proj_bias is not None else None
key_bias = self_attn.in_proj_bias[C:2*C] if self_attn.in_proj_bias is not None else None
value_bias = self_attn.in_proj_bias[-C:] if self_attn.in_proj_bias is not None else None

# 2. 手动完成q、k、v投影
q = th.matmul(x, query_weight.T)
k = th.matmul(x, key_weight.T)
v = th.matmul(x, value_weight.T)

if query_bias is not None:
    q = q + query_bias
if key_bias is not None:
    k = k + key_bias
if value_bias is not None:
    v = v + value_bias

# 3. 计算未归一化的注意力分数，添加边偏置后做softmax
attn_scores = th.matmul(q, k.transpose(-2, -1)) / th.sqrt(C / H)
if attn_bias is not None:
    attn_scores = attn_scores + attn_bias
attn_wei = th.softmax(attn_scores, dim=-1)

# 4. 注意力权重乘投影后的value得到最终输出
x = attn_wei @ v

print(x)

方法2：完全手动实现query/key/value投影 #

这种方式完全脱离原生nn.MultiheadAttention的内部逻辑，自己定义投影层，灵活性更高：

import torch as th
from torch import nn

# 变量初始化
B, T, C, H = 2, 3, 4, 2
head_dim = C // H

# 自定义投影层
q_proj = nn.Linear(C, C)
k_proj = nn.Linear(C, C)
v_proj = nn.Linear(C, C)
out_proj = nn.Linear(C, C)  # 对应原生MultiheadAttention的输出投影

# 张量定义
x = th.randn(B, T, C)
attn_bias = th.ones((B, T, T))

# 1. 完成q、k、v投影并适配多头维度
q = q_proj(x).view(B, T, H, head_dim).transpose(1, 2)  # [B, H, T, head_dim]
k = k_proj(x).view(B, T, H, head_dim).transpose(1, 2)
v = v_proj(x).view(B, T, H, head_dim).transpose(1, 2)

# 2. 计算注意力分数并添加边偏置
attn_scores = th.matmul(q, k.transpose(-2, -1)) / th.sqrt(head_dim)
# 扩展边偏置到多头维度
attn_scores = attn_scores + attn_bias.unsqueeze(1)
attn_wei = th.softmax(attn_scores, dim=-1)

# 3. 计算注意力输出、合并多头并做输出投影
attn_output = th.matmul(attn_wei, v).transpose(1, 2).contiguous().view(B, T, C)
x = out_proj(attn_output)

print(x)

方法3：自定义MultiheadAttention类，暴露投影层 #

如果需要复用类似逻辑，可以自定义一个继承自nn.Module的MultiheadAttention类，把q/k/v投影层作为类属性暴露：

import torch as th
from torch import nn

class CustomMultiheadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, batch_first=True):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.batch_first = batch_first
        
        # 暴露投影层
        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, query, key, value, attn_bias=None):
        if self.batch_first:
            query, key, value = query.transpose(0,1), key.transpose(0,1), value.transpose(0,1)
        
        # 投影
        q = self.q_proj(query).view(-1, self.num_heads, self.head_dim).transpose(0,1)
        k = self.k_proj(key).view(-1, self.num_heads, self.head_dim).transpose(0,1)
        v = self.v_proj(value).view(-1, self.num_heads, self.head_dim).transpose(0,1)
        
        # 计算注意力分数并添加偏置
        attn_scores = th.matmul(q, k.transpose(-2, -1)) / th.sqrt(self.head_dim)
        if attn_bias is not None:
            attn_scores = attn_scores + attn_bias.transpose(0,1).unsqueeze(0) if self.batch_first else attn_bias.unsqueeze(0)
        
        attn_wei = th.softmax(attn_scores, dim=-1)
        attn_output = th.matmul(attn_wei, v)
        
        # 合并多头并投影输出
        attn_output = attn_output.transpose(0,1).contiguous().view(-1, self.embed_dim)
        if self.batch_first:
            attn_output = attn_output.transpose(0,1)
        
        return self.out_proj(attn_output), attn_wei, v

# 使用示例
B, T, C, H = 2, 3, 4, 2
custom_attn = CustomMultiheadAttention(C, H, batch_first=True)
x = th.randn(B, T, C)
attn_bias = th.ones((B, T, T))

# 直接获取输出、注意力权重、投影后的value
output, attn_wei, v = custom_attn(x, x, x, attn_bias=attn_bias)
print(output)

关键注意点 #

• 原代码直接对softmax后的attn_wei加偏置不合理，应该在未做softmax的注意力分数上添加边偏置，再进行softmax归一化，否则会破坏注意力权重的概率分布。
• 多头注意力中，边偏置需要扩展到多头维度，保证每个头都能应用对应的偏置。

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