如何提取MultiHeadAttention注意力？训练完成后如何获取分数？

以下是两种实用方案，针对你已训练好模型的场景：

方案1：训练前修改模型结构，直接输出注意力分数 #

如果还未完成最终训练，或可以重新训练，建议在构建模型时把注意力分数作为额外输出分支，后续提取更方便：

1.1 调整Transformer Encoder函数 #

确保你的transformer_encoder函数返回特征输出和注意力分数：

def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
    # 归一化与多头注意力
    x = layers.LayerNormalization(epsilon=1e-6)(inputs)
    x, attn_scores = layers.MultiHeadAttention(
        key_dim=head_size, num_heads=num_heads, dropout=dropout
    )(x, x, return_attention_scores=True)
    x = layers.Dropout(dropout)(x)
    res = x + inputs

    # 前馈网络部分
    x = layers.LayerNormalization(epsilon=1e-6)(res)
    x = layers.Conv1D(filters=ff_dim, kernel_size=1, activation="relu")(x)
    x = layers.Dropout(dropout)(x)
    x = layers.Conv1D(filters=inputs.shape[-1], kernel_size=1)(x)
    return x + res, attn_scores  # 返回特征+注意力分数

1.2 构建带注意力输出的模型 #

在主模型中收集每个Transformer块的注意力分数，作为额外输出：

def build_model(input_shape, head_size, num_heads, ff_dim, num_transformer_blocks, mlp_units, num_classes, dropout=0, mlp_dropout=0):
    inputs = keras.Input(shape=input_shape)
    x = inputs

    attn_scores_list = []
    for _ in range(num_transformer_blocks):
        x, attn_scores = transformer_encoder(x, head_size, num_heads, ff_dim, dropout)
        attn_scores_list.append(attn_scores)

    # 分类头部分
    x = layers.GlobalAveragePooling1D(data_format="channels_last")(x)
    for dim in mlp_units:
        x = layers.Dense(dim, activation="relu")(x)
        x = layers.Dropout(mlp_dropout)(x)
    outputs = layers.Dense(num_classes, activation="softmax")(x)

    # 返回包含分类结果和所有注意力分数的模型
    return keras.Model(inputs=inputs, outputs=[outputs] + attn_scores_list)

1.3 编译与预测 #

编译时只需关注分类损失，其他输出不计算损失：

model = build_model(...)
model.compile(
    loss=["sparse_categorical_crossentropy"] + [None]*num_transformer_blocks,
    optimizer=keras.optimizers.Adam(learning_rate=1e-4),
    metrics=["accuracy"]
)

# 训练模型（步骤略）

# 预测时提取注意力分数
sample = X_test[0:1]
predictions, *all_attn_scores = model.predict(sample)

# 第一个Transformer块的注意力分数，形状：(batch_size, num_heads, seq_len, seq_len)
first_block_attn = all_attn_scores[0]
# 取第一个样本、第一个注意力头的注意力矩阵
attn_matrix = first_block_attn[0, 0, :, :]

方案2：基于已训练模型提取注意力分数 #

如果已经训练好模型且不想重新训练，可以通过创建新模型来提取注意力分数：

2.1 定位Transformer Encoder层 #

先查看模型结构，找到Transformer Encoder层的名称或索引：

trained_model = keras.models.load_model("your_trained_model_path.h5")
trained_model.summary()  # 从输出中找到Transformer Encoder层的名字，比如"transformer_encoder"

2.2 创建注意力提取模型 #

基于原模型的输入和Transformer层的输出构建新模型：

# 获取目标层（替换成你的层名）
encoder_layer = trained_model.get_layer(name="transformer_encoder")

# 创建新模型，输出为Transformer层的特征和注意力分数
attn_extractor = keras.Model(inputs=trained_model.input, outputs=encoder_layer.output)

# 提取注意力分数
sample = X_test[0:1]
_, attn_scores = attn_extractor.predict(sample)

注意力矩阵可视化（可选） #

拿到注意力矩阵后，可以用matplotlib可视化，直观观察序列位置间的注意力关系：

import matplotlib.pyplot as plt

# 取第一个样本的第一个注意力头矩阵
attn_mat = attn_scores[0, 0, :, :]

plt.figure(figsize=(10, 10))
plt.imshow(attn_mat, cmap="viridis")
plt.xlabel("Sequence Position")
plt.ylabel("Sequence Position")
plt.title("Attention Matrix (Head 1)")
plt.colorbar()
plt.show()

关键说明 #

注意力分数的形状为(batch_size, num_heads, sequence_length, sequence_length)，其中：

• num_heads：多头注意力的头数
• sequence_length：时间序列的长度
• 矩阵中[i,j]的值代表序列第i个位置对第j个位置的关注程度

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