[efficient] Make ReLU numeric stable.

PiperOrigin-RevId: 448695305

[efficient] Make ReLU numeric stable.
PiperOrigin-RevId: 448695305
7fdde945 · Frederick Liu · A. Unique TensorFlower · fa268452 · 7fdde945
Commit 7fdde945 authored May 14, 2022 by Frederick Liu Committed by A. Unique TensorFlower May 14, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

official/nlp/modeling/layers/kernel_attention.py official/nlp/modeling/layers/kernel_attention.py +4 -1

No files found.
--- a/official/nlp/modeling/layers/kernel_attention.py
+++ b/official/nlp/modeling/layers/kernel_attention.py
@@ -117,7 +117,10 @@ _TRANSFORM_MAP = {
            h=lambda x: 1),
    "relu":
        functools.partial(
-            _generalized_kernel, f=tf.keras.activations.relu, h=lambda x: 1),
+            _generalized_kernel,
+            # Improve numerical stability and avoid NaNs in some cases by adding
+            # a tiny epsilon.
+            f=lambda x: tf.keras.activations.relu(x) + 1e-3, h=lambda x: 1),
    "square":
        functools.partial(
            _generalized_kernel, f=tf.math.square, h=lambda x: 1),