Keras MultiHeadAttention输出与《Python深度学习》自注意力实现结果不符问题

• 无参数与带参数实现的本质差异
  你手写的self_attention是无参数的基础自注意力实现，直接用输入向量同时作为查询(Q)、键(K)、值(V)进行计算。而Keras的MultiHeadAttention默认带有四组可训练的投影矩阵：分别对输入做Q、K、V的线性变换，以及对注意力输出做最终线性变换。你测试时仅设置了use_bias=False，并没有固定投影矩阵的取值，随机初始化的参数会完全改变Q、K、V的数值，最终得到的注意力权重和输出自然和手写版本不一致。
• 维度处理逻辑的默认差异
  你手写的实现针对单头、输入维度等于key_dim的场景设计，没有额外的维度拆分和投影逻辑。而MultiHeadAttention内部默认会先把输入投影到key_dim * num_heads的维度再按头拆分，只要投影矩阵不是单位矩阵，计算结果就会偏离手写实现。

你只需要在初始化MultiHeadAttention后，手动将所有投影权重设置为单位矩阵，即可和手写的无参数自注意力结果对齐，修改后的测试代码如下：

import numpy as np
from scipy.special import softmax
from tensorflow.keras.layers import MultiHeadAttention

def self_attention(input_sequence):
    output = np.zeros(shape=input_sequence.shape)
    for i, pivot_vector in enumerate(input_sequence):
        scores = np.zeros(shape=(len(input_sequence),))
        for j, vector in enumerate(input_sequence):
            scores[j] = np.dot(pivot_vector, vector.T)
        scores /= np.sqrt(input_sequence.shape[1])
        scores = softmax(scores)
        new_pivot_representation = np.zeros(shape=pivot_vector.shape)
        for j, vector in enumerate(input_sequence):
            new_pivot_representation += vector * scores[j]
        output[i] = new_pivot_representation
    return output

test_input_sequence = np.array([[[1.0, 0.0, 0.0, 1.0],
                                 [0.0, 1.0, 0.0, 0.0],
                                 [0.0, 1.0, 1.0, 1.0]]])

# 手写自注意力结果
print("手写自注意力输出：")
print(self_attention(test_input_sequence[0]))

att_layer = MultiHeadAttention(num_heads=1,
                               key_dim=4,
                               use_bias=False,
                               attention_axes=(1,))
# 手动设置所有投影矩阵为单位矩阵
att_layer.build(input_shape=(None, 3, 4))
att_layer._query_kernel.assign(np.eye(4, 4))
att_layer._key_kernel.assign(np.eye(4, 4))
att_layer._value_kernel.assign(np.eye(4, 4))
att_layer._output_kernel.assign(np.eye(4, 4))

# 调用MultiHeadAttention
output, scores = att_layer(test_input_sequence,
                          test_input_sequence,
                          test_input_sequence,
                          return_attention_scores=True)
print("Keras MultiHeadAttention输出：")
print(output[0].numpy())
print("Keras注意力权重：")
print(scores[0, 0].numpy())

运行上述代码后，两者的注意力权重和输出会完全匹配，仅存在浮点数精度导致的微小误差。

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