为何MATLAB selfAttentionLayer总键维度相同但参数数量不同？

问题描述 #

我在将多头自注意力层插入DigitDataset的CNN示例中，发现两种总键维度相同的配置，MATLAB报告的可学习参数数量不一致：

• 配置1：NumHeads = 4, NumKeyChannels = 784（总键维度 4×784=3136）
• 配置2：NumHeads = 8, NumKeyChannels = 392（总键维度 8×392=3136）

网络结构如下：

layers = [
    imageInputLayer([28 28 1])
    convolution2dLayer(3,32,'Padding','same')
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(2,'Stride',2)
    convolution2dLayer(3,64,'Padding','same')
    batchNormalizationLayer
    reluLayer
    maxPooling2dLayer(2,'Stride',2)
    flattenLayer
    selfAttentionLayer(NUM_HEADS, NUM_KEY_CHANNELS)
    fullyConnectedLayer(10)
    softmaxLayer
    classificationLayer];

根据多头注意力的论文逻辑，Q/K/V投影的总参数应仅与总键维度相关，与头的拆分方式无关，但实际结果却不同，想了解MATLAB该层的参数计算逻辑。

核心原因与计算逻辑 #

根据MATLAB Deep Learning Toolbox中selfAttentionLayer的官方实现逻辑，该层的可学习参数由Q/K/V投影矩阵、输出投影矩阵，以及对应的偏置项组成，具体计算方式如下：

假设输入到自注意力层的特征维度为D（在你的示例中，经过CNN和flatten后，D=7×7×64=3136），NumValueChannels默认等于NumKeyChannels，则：

1. Q/K/V投影参数：
   • 每个投影矩阵维度为 D × (NumHeads × NumKeyChannels)
   • 三个投影矩阵的总参数为 3 × D × (NumHeads × NumKeyChannels)
2. 输出投影参数：
   • 将多头拼接后的特征（维度NumHeads × NumKeyChannels）投影回输入维度D，矩阵维度为 (NumHeads × NumKeyChannels) × D
   • 参数数量为 (NumHeads × NumKeyChannels) × D
3. 偏置项：
   • Q/K/V投影的偏置各为NumHeads × NumKeyChannels个，输出投影偏置为D个
   • 总偏置数量为 3 × (NumHeads × NumKeyChannels) + D

两种配置的参数计算对比 #

代入你的配置参数：

• 配置1与配置2的NumHeads × NumKeyChannels均为3136，且输入维度D=3136
• Q/K/V投影参数：3×3136×3136=29,859,840（两种配置一致）
• 输出投影参数：3136×3136=9,828,608（两种配置一致）
• 偏置项：3×3136 +3136=12,544（两种配置一致）
• 总参数：29,859,840+9,828,608+12,544=39,700,992（两种配置理论上完全相同）

可能的差异来源 #

如果你的MATLAB报告参数数量不同，大概率是以下原因之一：

1. 网络层参数不一致：检查两次实验中，是否除了NumHeads和NumKeyChannels外，其他层（如卷积通道数、池化步长）被意外修改。
2. 参数统计范围错误：使用analyzeNetwork(layers)单独查看selfAttentionLayer的参数数量，排除其他层（如BatchNorm、全连接层）的干扰。
3. MATLAB版本差异：不同版本的selfAttentionLayer实现可能存在细节调整，确认你使用的是R2023a版本。
4. 自定义参数修改：是否手动修改了NumValueChannels参数（默认与NumKeyChannels一致），若该参数不同会导致参数数量变化。

总结 #

从官方实现逻辑来看，总键维度相同的情况下，selfAttentionLayer的可学习参数数量应完全一致。若出现差异，优先排查上述操作或配置问题，而非层的实现逻辑本身。

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