🐛 fix: numerical stability of tanh for GELU

CHANGELOG.md

@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 
 ## [unreleased]
 
+🐛 **fix**: numerical stability of tanh for GELU ([#83](https://github.com/owkin/GrAIdient/pull/83))\
 🔨 **layer_seq**: SelectSeq ([#82](https://github.com/owkin/GrAIdient/pull/82))\
 🚀 **examples**: AutoEncoder models ([#79](https://github.com/owkin/GrAIdient/pull/79))\
 🪜 **layer_2d**: VQ2D ([#81](https://github.com/owkin/GrAIdient/pull/81))\

Sources/GrAIdient/Core/Function/Activation.swift

@@ -832,8 +832,17 @@ public class GELU: ActivationFunction
     public override func apply(_ x: Double) -> Double
     {
         let cst = sqrt(2.0 / Double.pi)
-        let tmp = tanh(cst * (x + 0.044715 * pow(x, 3)))
-        return 0.5 * x * (1 + tmp)
+        let tmp1 = cst * (x + 0.044715 * pow(x, 3))
+        let tmp2: Double
+        if tmp1 >= 0
+        {
+            tmp2 = (1.0 - exp(-2.0 * tmp1)) / (1.0 + exp(-2.0 * tmp1))
+        }
+        else
+        {
+            tmp2 = (exp(2.0 * tmp1) - 1.0) / (exp(2.0 * tmp1) + 1.0)
+        }
+        return 0.5 * x * (1 + tmp2)
     }
 
     ///
@@ -845,9 +854,18 @@ public class GELU: ActivationFunction
     public override func derivate(_ x: Double) -> Double
     {
         let cst = sqrt(2.0 / Double.pi)
-        let tmp1 = tanh(cst * (x + 0.044715 * pow(x, 3)))
-        let tmp2 = cst * (1 + 3 * 0.044715 * x * x) * (1 - tmp1 * tmp1)
-        let derivative = 0.5 * (1 + tmp1 + x * tmp2)
+        let tmp1 = cst * (x + 0.044715 * pow(x, 3))
+        let tmp2: Double
+        if tmp1 >= 0
+        {
+            tmp2 = (1.0 - exp(-2.0 * tmp1)) / (1.0 + exp(-2.0 * tmp1))
+        }
+        else
+        {
+            tmp2 = (exp(2.0 * tmp1) - 1.0) / (exp(2.0 * tmp1) + 1.0)
+        }
+        let tmp3 = cst * (1 + 3 * 0.044715 * x * x) * (1 - tmp2 * tmp2)
+        let derivative = 0.5 * (1 + tmp2 + x * tmp3)
         return derivative
     }
 }

Sources/GrAIdient/Metal/Kernel/Activation.metal

@@ -255,9 +255,8 @@ kernel void forwardGELU(
     {
         tmp2 = (exp(2.0 * tmp1) - 1.0) / (exp(2.0 * tmp1) + 1.0);
     }
-    float tmp = tanh(cst * (x + 0.044715 * pow(x, 3)));
     tmps[id] = x;
-    outs[id] = 0.5 * x * (1 + tmp);
+    outs[id] = 0.5 * x * (1 + tmp2);
 }
 
 kernel void backwardGELU(