Keras SpectralNormalization and stop_gradient #20309
Description
Keras Version: 3.5.0
Hello.
A while ago, I had some issues with PyTorch and SpectralNormalization in an RNN layer: #19527
It looks like some things have changed with SpecNorm, but there's something in the code where I'm seeing degraded performance. So degraded, that I'm getting exploding gradients. The difference comes down to some commented-out code in the SpectralNormalization wrapper:
"""Generate spectral normalized weights.
This method returns the updated value for `self.kernel` with the
spectral normalized value, so that the layer is ready for `call()`.
"""
weights = ops.reshape(self.kernel, [-1, self.kernel_shape[-1]])
vector_u = self.vector_u.value
for _ in range(self.power_iterations):
vector_v = normalize(
ops.matmul(vector_u, ops.transpose(weights)), axis=None
)
vector_u = normalize(ops.matmul(vector_v, weights), axis=None)
-> # vector_u = tf.stop_gradient(vector_u)
-> # vector_v = tf.stop_gradient(vector_v)
sigma = ops.matmul(
ops.matmul(vector_v, weights), ops.transpose(vector_u)
)
kernel = ops.reshape(ops.divide(self.kernel, sigma), self.kernel_shape)
return ops.cast(vector_u, self.vector_u.dtype), ops.cast(
kernel, self.kernel.dtype
)
The stop_gradient() call is used in both the TF Addons implementation (https://github.com/tensorflow/addons/blob/v0.20.0/tensorflow_addons/layers/spectral_normalization.py#L120-L121) and the official PyTorch implementation (https://pytorch.org/docs/stable/_modules/torch/nn/utils/spectral_norm.html)
I've made the change myself by implementing ops.stop_gradient at both commented-out spots, and expected (stable) functionality returned.