多头注意度计算

Question

我创建了一个多头关注层的模型，

import torch
import torch.nn as nn
query = torch.randn(2, 4)
key = torch.randn(2, 4)
value = torch.randn(2, 4)
model = nn.MultiheadAttention(4, 1, bias=False)

model(query, key, value)

我试着匹配得到的注意力输出，

softmax_output = torch.softmax(((query@model.in_proj_weight[:4])@((key@model.in_proj_weight[4:8]).t()))/2, dim=1)
intermediate_output = softmax_output@(value@model.in_proj_weight[8:12])
final_output = intermediate_output@model.out_proj.weight

但是final_output与注意力输出不匹配。

Answer 1

能够与输出相匹配，

q_w = query@model.in_proj_weight[:4].t()
k_w = key@model.in_proj_weight[4:8].t()
v_w = value@model.in_proj_weight[8:12].t()

softmax_output = torch.softmax((q_w@k_w.t())/2, dim=1)

attention = softmax_output@v_w

final_output = attention@model.out_proj.weight.t()

丢失了先前的转位