JAX instability when factoring out matmul

[max-vladymyrov]

I'm working with matrix multiplication in JAX and encountering a precision issue. Specifically, I want to factor out a specific portion of the multiplication.

Given two matrices:

P with dimensions (m x n), where a is the last row of P.
Q with dimensions (n x p), where b is the last column of Q.
I expect the following to be mathematically equivalent:

P @ Q = A @ B + a @ b

However, in JAX, these calculations are not exactly equal. The difference is small (around 1e-7), but it becomes significant in my workflow due to the repeated use of this operation.

I've tried enabling 64-bit precision (jax_enable_x64), which improves the situation, but the difference is still not zero. Changing jax.lax.Precision=Highest doesn't seem to help at all.

The problem seem to be worse for smaller matrices (up to 32x32):

Is it possible achieve exact equivalence (or a much smaller difference without using float64) between these two calculations in JAX?

Exact code to reproduce:

jax.config.update('jax_enable_x64', False)

def mult(x,y):
  return jnp.einsum('...qk,...kd->...qd', x, y, precision=jax.lax.Precision.HIGHEST)

rng = jax.random.PRNGKey(0)
rngs = jax.random.split(rng,4)

n_dim = 30
batch_size = 100

x = jax.random.normal(rngs[0], shape=(batch_size, n_dim, n_dim))
y = jax.random.normal(rngs[1], shape=(batch_size, n_dim, 1))

a = jax.random.normal(rngs[2], shape=(batch_size, n_dim, n_dim))
b = jax.random.normal(rngs[3], shape=(batch_size, 1, n_dim))

xy = jnp.concatenate([x, y], axis=2)
ab = jnp.concatenate([a, b], axis=1)

res1 = mult(xy, ab)
res2 = mult(x, a) + mult(y, b)
print(np.abs(res1 - res2).mean())

[jakevdp]

This looks like it's behaving as expected: you're executing your code in float32 precision. float32 has 23 bits of mantissa, so different ways of computing the "same" floating-point result would be expected to differ by 1 part in $2^{23}$, which is roughly 1E-7.

Note that this is not a JAX issue, but an issue with any library that does floating point computations. For example, here's NumPy doing float32 vector multiplication two ways: front-to-back and back-to-front:

In [1]: import numpy as np
In [2]: np.random.seed(0)
In [3]: x = np.random.randn(1000).astype('float32')
In [4]: result1 = x @ x
In [5]: result2 = x[::-1] @ x[::-1]
In [6]: np.testing.assert_array_equal(result1, result2)
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<...>
Arrays are not equal

Mismatched elements: 1 / 1 (100%)
Max absolute difference: 0.00048828
Max relative difference: 5.001431e-07
 x: array(976.28265, dtype=float32)
 y: array(976.28314, dtype=float32)

You see that NumPy shows similar differences if you compute the "same" result two different ways.

There is no fix to this: it's fundamental to floating point computations. The real solution here is to write your code such that it's robust to this kind of expected floating point rounding error.