Wrong gradients in simple NNs

[SNMS95]

I was creating a new package with Jax (and its ecosystem) and was writing tests to check that the gradients are correct.
I created a simple fully connected neural network using Haiku and then checked the gradients using the following code:

import jax
jax.config.update('jax_enable_x64', True)
import jax.numpy as np
import haiku as hk
import numpy as onp


class DenseNetParameterization(hk.Module):

    def __init__(self, dtype = np.float32):
        super().__init__()
        self.layers = self._build_layers()
        self.dtype = dtype

    def _build_layers(self):
        """Create list of network components."""

        hidden_layer_activation = jax.nn.relu
        output_activation = jax.nn.sigmoid

        layers = []
        # Iterate over 5 hidden layers
        for index in range(5):
            layers.append(
                hk.Linear(output_size=20, name=f'linear_{index}'))
            layers.append(
                hk.BatchNorm(create_scale=True, create_offset=True,
                             decay_rate=0.999, name=f'batch_norm_{index}'))
            layers.append(hidden_layer_activation)

        # Output layer
        layers.append(
            hk.Linear(output_size=1, name='linear_output'))
        layers.append(output_activation)
        return layers

    def __call__(self, model_input,
                is_training = True):
        x = model_input#.astype(self.dtype)

        for layer in self.layers:
            if isinstance(layer, hk.BatchNorm):
                x = layer(x, is_training=is_training)
            else:
                x = layer(x)
        return x.reshape(-1, 1)

def model(x):
    net = DenseNetParameterization(dtype=np.float64)
    return net(x)

# Create a random input
model_inp = jax.random.uniform(key=jax.random.PRNGKey(42), shape=(100, 2))
# Use haiku to transform to pure_functions
forward_pass_pure = hk.without_apply_rng(
                            hk.transform_with_state(model))
# Initializes the shapes of the parameters and state
init_params, init_state = forward_pass_pure.init(x=model_inp,
                                                 rng=jax.random.PRNGKey(42))
nn_forward = jax.jit(forward_pass_pure.apply)

# Test gradients
def dummy_loss(params):
    y, _ = nn_forward(params, init_state, model_inp)
    return y.sum()

from jax.test_util import check_grads
check_grads(dummy_loss, (init_params,), order=2)

However, this results in an AssertionError stating that the gardients are wrong!

Mismatched elements: 1 / 1 (100%)
Max absolute difference: 0.99549156
Max relative difference: 0.03364104
 x: array(-30.587076)
 y: array(-29.591584)

Can anyone tell me what is going wrong with this?

[jakevdp]

check_grads works by comparing the analytical gradient to the numerical gradient, by evaluating the function at x and x + eps. The default eps is 1E-3, which is appropriate for float32 computations. You're using float64 computations, so a smaller eps would more appropriate and lead to more accurate numerical gradients. I'd try this:

check_grads(dummy_loss, (init_params,), order=2, eps=1e-6)

[SNMS95]

Hey Jake, the thing is I did try with float32 (with default values) and still got an assertion error

[jakevdp]

If I run your example in 64-bit precision with eps=1E-6, I find that the gradients match. In float32 precision, if I use a smaller eps then the numerical result diverges.

This makes me think that your function has a fast-varying second derivative, which makes the numerical gradient inaccurate, but nevertheless the analytic gradient is probably producing the correct value.

[SNMS95]

Thanks Jake. I did not know that this can happen!