为何Adam优化器在近似模型上应用梯度耗时差异悬殊？

我有两个结构近乎一致的模型，仅最后一层存在差异：一个为Dense(1, "linear")，另一个为Dense(4, "softmax")。相关代码如下：

gradient_agent = tape.gradient(loss_agent, agent.trainable_weights)
gradient_value = tape.gradient(loss_value, value.trainable_weights)

# Do something to the gradients

optimizer_agent.apply_gradients(zip(gradient_agent, agent.trainable_weights))
optimizer_value.apply_gradients(zip(gradient_value, value.trainable_weights))

性能分析结果显示：

% Time       | Line
    29.3     |  optimizer_agent.apply_gradients(zip(gradient_agent, agent.trainable_weights))
    1.6      |  optimizer_value.apply_gradients(zip(gradient_value, value.trainable_weights))

为何第一条梯度应用语句耗时如此之久？
看起来这是某种TensorFlow优化导致的，因为交换两行代码顺序后，原本耗时短的那条会变慢。但这段代码并未使用@tf.function，这是什么情况？

这种现象和TensorFlowEager模式下的内核预热机制直接相关，哪怕没显式用@tf.function，底层依然存在即时编译开销：

• 首次执行的apply_gradients会触发对应权重更新操作的CUDA/CPU内核编译，这个编译过程的耗时远大于实际执行权重更新的时间，所有预热开销都被统计到了第一条执行的语句上。
• 第二次执行时，内核已经编译完成并被缓存，直接复用即可，所以耗时骤降。
• 交换代码顺序后，原本的第二条操作变成首次执行，自然会触发新的内核编译，耗时变长；原来的第一条操作复用已缓存的内核，耗时随之变短。

这个现象和模型最后一层的差异无关，本质是TensorFlow在Eager模式下，首次执行任何新的计算操作都会进行隐式的内核编译优化，只是不像@tf.function那样提前构建完整计算图，但首次执行的预热开销依然存在。

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