更新 LSTM 网络的自定义输出层

Question

我有一个文本生成任务，学习使用具有多个输出层的 LSTM 网络来预测下一个单词。 句子生成完成后，我计算整个句子的奖励，并尝试更新参与生成的输出层（贡献层获得计算的奖励值，其他层获得0）。 我的问题是，即使我只更新选定的输出层，似乎其他层的权重也被更新了。

我有一个带有虚拟数据的最小化示例来呈现问题：

import random

import numpy as np
import tensorflow as tf

from keras.layers import Input, LSTM, Dense, Embedding
from keras.utils import pad_sequences
from tensorflow.keras.models import Model


def policy_gradient_loss(y_true, y_pred):
    return tf.reduce_mean(tf.math.log(y_pred) * float(y_true))

# Define the model with 3 output layers (named 'a', 'b' and 'c').
input_layer = Input(shape=(4,))
embedding_layer = Embedding(input_dim=10, output_dim=4)(input_layer)
lstm_layer = LSTM(4)(embedding_layer)
output_layers = [Dense(3, activation='softmax', name=name)(lstm_layer) for name in ['a', 'b', 'c']]
model = Model(inputs=input_layer, outputs=output_layers)
model.compile(loss=[policy_gradient_loss] * 3, optimizer='adam', run_eagerly=True)

# Dummy input data.
input_data = np.array([[2, 3, 4, 5]])

# Create target data to reward only the 'b' output layer.
target_data = [np.array([0]) for _ in range(len(model._output_layers))]
target_data[1] = np.array([10]) 

# Save initial weights.
initial_weights = model.get_weights()

model.train_on_batch(input_data, y=target_data)

# Save weights after the learning.
updated_weights = model.get_weights()

# Compare the before-after weights.
for layer_idx, (layer_name, initial_w, updated_w) in enumerate(zip([layer.name for layer in model.layers], initial_weights, updated_weights)):
    if not tf.math.reduce_all(tf.equal(initial_w, updated_w)):
        print(f'The weights in layer {layer_idx} ({layer_name}). has changed.')

结果：

The weights in layer 0 (input_1). has changed.
The weights in layer 1 (embedding). has changed.
The weights in layer 2 (lstm). has changed.
The weights in layer 3 (a). has changed.

我的期望是更新第 4 层（输出层“b”）而不是层“a”（或至少在“a”旁边）。

我错过了什么？我的期望或我的实施是错误的吗？ （或者两者都...？）