Keras SpectralNormalization and stop_gradient

[dryglicki]

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.

[hertschuh]

@dryglicki ,

I believe this omission came from the fact that SpectralNormalization was ported to Keras 3 before ops.stop_gradient was available.

I've made the change myself by implementing ops.stop_gradient at both commented-out spots, and expected (stable) functionality returned

Great! Are you able to open a PR with this fix?

[dryglicki]

@hertschuh Apologies for the delay. Yes, I can try to get a PR together.

[dryglicki]

PR here: #20353

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[hertschuh]

@dryglicki

Thanks for the PR!