在回归任务深度神经网络中添加MultiHeadAttention遇IndexError问题

问题根源 #

1. 输入维度不匹配：tf.keras.layers.MultiHeadAttention要求输入是三维张量，格式为(batch_size, sequence_length, feature_dim)，但你的x1是二维张量(None, 58)，缺少序列长度维度，导致内部Softmax计算时维度索引越界。
2. 错误设置output_shape参数：该参数无需手动指定，模型会根据输入特征数、头数等自动推导输出形状，强行传入x1.shape反而会引发维度冲突。

修正方案 #

1. 给输入张量增加序列维度，将二维的(None, 58)转换为三维的(None, 1, 58)（这里序列长度设为1，如果你原本有序列结构，可根据实际情况调整序列长度值）。
2. 移除output_shape参数，让模型自动处理输出形状。
3. 若后续层需要二维输入，可通过Reshape或Flatten将注意力输出的三维张量转回二维。

修正后的代码 #

x1 = Dense(58, activation='relu')(x1)
x1 = Dropout(0.1)(x1)
print(x1.shape)  # (None, 58)

# 增加序列维度
x1 = tf.expand_dims(x1, axis=1)  # 现在形状是(None, 1, 58)

# 移除output_shape参数，传入三维的x1作为query、key、value
attention = tf.keras.layers.MultiHeadAttention(
    num_heads=2, 
    key_dim=29,  # key_dim建议设为feature_dim/num_heads，这里58/2=29，也可自定义
    dropout=0.1
)(x1, x1)

# 将三维输出转回二维
attention = tf.keras.layers.Reshape((-1,))(attention)  # 形状回到(None, 58)

x1 = Dropout(0.2)(attention)
x1 = Dense(59, activation='relu')(x1)
output = Dense(1, activation='linear')(x1)
model = tf.keras.models.Model(inputs=input1, outputs=output)

补充说明 #

• key_dim的取值建议与头数匹配，比如总特征数58，2个头的话每个头对应29维，这样计算更合理，也能避免维度浪费。
• 如果你的输入原本就应该包含序列（比如文本、时序数据），那应该在更早的环节保留序列维度，而不是用expand_dims强行添加。

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