Repository: google/jax

Issue:
Use exact GELU
jax.nn.gelu uses the approximate form of the GELU, but [tensorflow](https://github.com/tensorflow/tensorflow/blob/938cc7bf9c4f354361e18e3ba485af53e602d341/tensorflow/python/keras/activations.py#L311) and [pytorch](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html?highlight=gelu#torch.nn.GELU) use the exact version.

I believe the exact form is more numerically stable and similarly fast. Figure 16 of the [Performer paper](https://arxiv.org/pdf/2009.14794.pdf) (@xingyousong) showed the GELU running into NaN issues, and I suspect this is because jax uses the approximate version.


Discussion:
Interesting! https://github.com/google/jax/pull/1556 switched to the approximate version. @trevorcai @jekbradbury 
Exact GeLU is significantly slower on TPUs (easily noticeable even in end-to-end step time). We’d be happy to take a PR adding the exact implementation as an option, but keeping the approximate one as the default?
> Exact GeLU is significantly slower on TPUs

Interesting. For PyTorch this was not the case: https://github.com/pytorch/pytorch/issues/39853#issuecomment-658806898 The exact version was slightly faster.

> We’d be happy to take a PR

Sadly I don't know what optimizations would make them similarly fast like in PyTorch.
In JAX under a JIT, both versions are quite well optimized (fused, etc.). But TPUs are not very fast at certain kinds of vector math, and the exact GeLU happens to hit some of those cases (I think).
https://github.com/google/jax/pull/1556#issuecomment-545130783 says
> Confirming that my benchmarks are showing jax.grad(jax.jarrett(gelu)) as slower than jax.grad(gelu) on GPU as well.

So do you think the issue is with both TPUs and GPUs yet PyTorch on GPUs doesn't have a problem? Or do you think it's more of a TPU-specific issue?
I just tried some new JAX timings on TPUv2 (two generations old) and V100 (one generation old), mostly because I had easy access to them via Colab.

I found that on V100 when compiled with `jax.jit`, the approximate formulation is 1.12x faster on the forward pass (which seems relatively insignificant), but on TPUv2 the approximate formulation is 1.75x faster. The difference on the backward pass was much smaller. I would guess this is related to `erf` being much more expensive to compute than `tanh` on TPU. It's possible we could optimize our `erf` implementation, which would presumably improve the relative performance of the exact formulation on both platforms.

Since the performance differences are so large, at the moment it does seem like we might do best to let users choose which they want.

 

Resolving patch:
diff --git a/jax/nn/functions.py b/jax/nn/functions.py
--- a/jax/nn/functions.py
+++ b/jax/nn/functions.py
@@ -171,22 +171,33 @@ def selu(x):
   scale = 1.0507009873554804934193349852946
   return scale * elu(x, alpha)
 
-def gelu(x):
+def gelu(x, approximate: bool = True):
   r"""Gaussian error linear unit activation function.
 
-  Computes the element-wise function:
+  If ``approximate=False``, computes the element-wise function:
+
+  .. math::
+    \mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{erf} \left(
+      \frac{x}{\sqrt{2}} \right) \right)
+
+  If ``approximate=True``, uses the approximate formulation of GELU:
 
   .. math::
     \mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{tanh} \left(
       \sqrt{\frac{2}{\pi}} \left(x + 0.044715 x^3 \right) \right) \right)
 
-  We explicitly use the approximation rather than the exact formulation for
-  speed. For more information, see `Gaussian Error Linear Units (GELUs)
+  For more information, see `Gaussian Error Linear Units (GELUs)
   <https://arxiv.org/abs/1606.08415>`_, section 2.
+
+  Args:
+    approximate: whether to use the approximate or exact formulation.
   """
-  sqrt_2_over_pi = np.sqrt(2 / np.pi).astype(x.dtype)
-  cdf = 0.5 * (1.0 + jnp.tanh(sqrt_2_over_pi * (x + 0.044715 * (x ** 3))))
-  return x * cdf
+  if approximate:
+    sqrt_2_over_pi = np.sqrt(2 / np.pi).astype(x.dtype)
+    cdf = 0.5 * (1.0 + jnp.tanh(sqrt_2_over_pi * (x + 0.044715 * (x ** 3))))
+    return x * cdf
+  else:
+    return jnp.array(x * (lax.erf(x / np.sqrt(2)) + 1) / 2, dtype=x.dtype)
 
 def glu(x, axis=-1):
   """Gated linear unit activation function."""


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -15,12 +15,14 @@
 """Tests for nn module."""
 
 import collections
+from functools import partial
 import itertools
 
 from absl.testing import absltest
 from absl.testing import parameterized
 
 import numpy as np
+import scipy.stats
 
 from jax import core
 from jax import test_util as jtu
@@ -93,9 +95,21 @@ def testGluValue(self):
     val = nn.glu(jnp.array([1.0, 0.0]))
     self.assertAllClose(val, jnp.array([0.5]))
 
+  @parameterized.parameters(False, True)
+  def testGelu(self, approximate):
+    def gelu_reference(x):
+      return x * scipy.stats.norm.cdf(x)
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng((4, 5, 6), jnp.float32)]
+    self._CheckAgainstNumpy(
+      gelu_reference, partial(nn.gelu, approximate=approximate), args_maker,
+      check_dtypes=False, tol=1e-3 if approximate else None)
+
   @parameterized.parameters(*itertools.product(
       (jnp.float32, jnp.bfloat16, jnp.float16),
-      (nn.gelu, nn.relu, nn.softplus, nn.sigmoid)))
+      (partial(nn.gelu, approximate=False),
+       partial(nn.gelu, approximate=True),
+       nn.relu, nn.softplus, nn.sigmoid)))
   def testDtypeMatchesInput(self, dtype, fn):
     if dtype is jnp.float16 and jtu.device_under_test() == "tpu":
       self.skipTest("float16 not supported on TPU")