tf.keras.layers.MultiHeadAttention的key_dim参数匹配问题咨询

为什么支持传入不符合论文规则的参数 #

首先要明确：key_dim = embed_dim / num_heads 是原论文为了平衡计算量与效果选择的实验配置，不是多头注意力机制的数学强制约束。原论文这么设，本质是将q/k/v投影的总维度设置为和模型嵌入维度相等，以此保证多头注意力的总计算量和单头同维度注意力基本持平，属于效率优先的工程选择。
Keras的实现是工业级通用设计，没有把输入嵌入维度和注意力头的维度、投影总维度做强制绑定，只要传入的num_heads、key_dim是正整数，层内权重都能正常初始化、完成前后向传播，自然不会报错。
很多人误以为必须满足整除关系，是因为不少教学场景下的简易手写多头注意力为了省代码，省略了独立投影步骤，直接把原始输入嵌入按维度切分成多个头，这种简易实现才要求输入维度能被头数整除，和Keras的正式实现逻辑不一样。

传入不匹配参数时的内部处理流程 #

不管输入的特征维度是多少，层内部会按固定流程完成计算，全程不会依赖“输入维度与头数/头维度整除”的条件，具体步骤如下：

• 第一步：独立初始化投影权重
  层会单独创建query、key、value三个线性投影层的权重，投影输出的总维度固定为num_heads * key_dim（如果不手动指定value_dim参数，value的投影维度默认和key_dim保持一致）。
  举个测试场景的例子：输入最后一维是10，设置num_heads=20, key_dim=9时，q/k/v三个投影矩阵的形状都是(10, 20*9) = (10, 180)，直接把10维的输入特征映射到180维的投影空间，和原始输入维度能不能被头数整除完全无关。
• 第二步：拆分注意力头计算
  投影得到的q/k/v张量形状为(batch_size, 序列长度, num_heads*key_dim)，内部会先reshape为(batch_size, 序列长度, num_heads, key_dim)，再转置为(batch_size, num_heads, 序列长度, key_dim)，每个头独立完成缩放点积注意力计算，这一步每个头的q/k维度都是你指定的key_dim，维度完全匹配，不会有计算错误。
• 第三步：拼接输出并做最终投影
  所有头计算完的注意力结果形状为(batch_size, num_heads, 序列长度, value_dim)，转置回(batch_size, 序列长度, num_heads, value_dim)后，会被reshape为(batch_size, 序列长度, num_heads*value_dim)，之后再过一个输出线性投影层：如果没有手动指定输出维度，这个投影层会把拼接后的特征映射回和query输入一致的最后一维（也就是测试场景里的10维），保证输出形状和输入注意力层的形状对齐，后续接其他层自然不会报错。

补充说明 #

测试里设置的num_heads=123, key_dim=17这类参数，本质是把投影总维度设为了123*17=2091，远大于输入的10维，相当于升维做注意力计算，参数量和计算量会比常规配置高很多，训练可以正常跑，但实际效果和性价比通常不如按经验配置的参数。

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