Adds full gelu without approximation

Project.toml

@@ -18,6 +18,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [extensions]
@@ -27,6 +28,7 @@ NNlibCUDAExt = "CUDA"
 NNlibEnzymeCoreExt = "EnzymeCore"
 NNlibFFTWExt = "FFTW"
 NNlibForwardDiffExt = "ForwardDiff"
+NNlibSpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
 AMDGPU = "1"

docs/src/reference.md

@@ -10,6 +10,8 @@ Non-linearities that go between layers of your model. Note that, unless otherwis
 celu
 elu
 gelu
+gelu_tanh
+gelu_erf
 hardsigmoid
 sigmoid_fast
 hardtanh

ext/NNlibSpecialFunctionsExt.jl

@@ -0,0 +1,15 @@
+module NNlibSpecialFunctionsExt
+
+using NNlib: NNlib, oftf
+using SpecialFunctions: erf
+
+# Full gelu (gelu_erf)
+NNlib.gelu_erf(x) = x/2*(1 + erf(x/sqrt(oftf(x,2))))
+
+function NNlib.deriv_gelu_erf(x)
+    SQRT2 = sqrt(oftf(x,2))
+    Φ = (1 + erf(x/SQRT2))/2
+    Φ + x/SQRT2*exp(-(x^2)/2)/sqrt(oftf(x,π))
+end
+
+end

src/NNlib.jl

@@ -26,7 +26,7 @@ include("activations.jl")
 for f in ACTIVATIONS
     @eval export $(f)
 end
-export sigmoid, hardsigmoid, logsigmoid, thresholdrelu # Aliases
+export sigmoid, hardsigmoid, logsigmoid, thresholdrelu, gelu # Aliases
 
 include("attention.jl")
 export dot_product_attention, dot_product_attention_scores, make_causal_mask

src/activations.jl

@@ -5,7 +5,7 @@
 
 ACTIVATIONS = [
     :σ, :hardσ, :hardtanh, :relu,
-    :leakyrelu, :relu6, :rrelu, :elu, :gelu, :swish, :hardswish, :selu,
+    :leakyrelu, :relu6, :rrelu, :elu, :gelu_tanh, :gelu_erf, :swish, :hardswish, :selu,
     :celu, :softplus, :softsign, :logσ, :logcosh,
     :mish, :tanhshrink, :softshrink, :trelu, :lisht,
     :tanh_fast, :sigmoid_fast,
@@ -301,14 +301,14 @@
 deriv_elu(Ω, α=1) = ifelse(Ω ≥ 0, one(Ω), Ω + oftype(Ω, α))
 
 """
-    gelu(x) = 0.5x * (1 + tanh(√(2/π) * (x + 0.044715x^3)))
+    gelu_tanh(x) = 0.5x * (1 + tanh(√(2/π) * (x + 0.044715x^3)))
 
-Activation function from ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415).
+Activation function from ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415) using tanh approximation.
 
 ```julia-repl
-julia> lineplot(gelu, -2, 2, height=7)
+julia> lineplot(gelu_tanh, -2, 2, height=7)
            ┌────────────────────────────────────────┐        
-         2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊│ gelu(x)
+         2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊│ gelu_tanh(x)
            │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⠀⠀│        
            │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀│        
    f(x)    │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⣀⡠⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│        
@@ -319,11 +319,11 @@
            ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀        
            ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀        
 
-julia> lineplot(gelu, -5, 0, height=7);
+julia> lineplot(gelu_tanh, -5, 0, height=7);
 
 julia> lineplot!(ans, swish)
              ┌────────────────────────────────────────┐         
-           0 │⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠒⠒⠤⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸│ gelu(x) 
+           0 │⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠒⠒⠤⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸│ gelu_tanh(x) 
              │⠑⠒⠢⠤⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇│ swish(x)
              │⠀⠀⠀⠀⠀⠈⠉⠒⠤⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⢆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣸⠁│         
    f(x)      │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠒⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⢄⠀⠀⠀⠀⠀⠀⠀⠀⢠⡇⠀│         
@@ -335,7 +335,7 @@
              ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀         
 ```
 """
-function gelu(x)
+function gelu_tanh(x)
     α = oftf(x, 0.044715)
     # λ = oftf(x, gelu_λ)
     # x/2 * (1 + tanh(λ * (x + α * x^3)))  # Standard implementation, for reference
@@ -346,7 +346,7 @@
 const gelu_λ = √(2 / π)
 const gelu_2λ = √(8 / π)
 
-function deriv_gelu(x)
+function deriv_gelu_tanh(x)
     α = oftf(x, 0.044715)
     α2 = oftf(x, 0.08943)
     λλ = oftf(x, gelu_2λ)
@@ -357,6 +357,24 @@
     muladd(dσ * λλ * muladd(x2, α2, t), x, Ω)
 end
 
+"""
+    gelu_erf(x) = xΦ(x) = 0.5x * (1 + erf(x/√2))
+
+Activation function from ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415) without approximation.
+The SpecialFunctions.jl package needs to be loaded to use this function.
+"""
+function gelu_erf end
+function deriv_gelu_erf end
+
+"""
+    gelu(x) = gelu_tanh(x)
+
+Activation function from ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415). 
+See [`gelu_tanh`](@ref).
+"""
+const gelu = gelu_tanh
+const deriv_gelu = deriv_gelu_tanh
+
 """
     swish(x) = x * σ(x)
 
@@ -874,7 +892,8 @@
     (:relu6,        :((Ω>0) & (Ω<6))),
     # rrelu is random, can't write a rule.
     (:elu,          :(deriv_elu(Ω))),
-    (:gelu,         :(deriv_gelu(x))),
+    (:gelu_tanh,    :(deriv_gelu_tanh(x))),
+    (:gelu_erf,     :(deriv_gelu_erf(x))),
     (:swish,        :(Ω + sigmoid_fast(x) * (1 - Ω))),
     (:hardswish,    :(deriv_hardswish(x))),
     # lisht

test/Project.toml

@@ -22,3 +22,4 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"

test/activations.jl

@@ -12,6 +12,8 @@ BINARY_ACTIVATIONS = filter(f -> hasmethod(f, Tuple{Float64, Float64}), ACTIVATI
 @test rrelu(0.0) == 0.0
 @test elu(0.0) == 0.0
 @test gelu(0.0) == 0.0
+@test gelu_tanh(0.0) == 0.0
+@test gelu_erf(0.0) == 0.0
 @test swish(0.0) == 0.0
 @test hardswish(0.0) == 0.0
 @test lisht(0.0) == 0.0
@@ -36,6 +38,8 @@ BINARY_ACTIVATIONS = filter(f -> hasmethod(f, Tuple{Float64, Float64}), ACTIVATI
 @test rrelu(1.0) == 1.0
 @test elu(1.0) == 1.0
 @test gelu(1.0) == 0.8411919906082768
+@test gelu_tanh(1.0) == 0.8411919906082768
+@test gelu_erf(1.0) == 0.8413447460685429
 @test swish(1.0) == sigmoid(1.0)
 @test hardswish(1.0) == hardsigmoid(1.0)
 @test lisht(1.0) ≈ 1.0 * tanh(1.0)
@@ -58,6 +62,8 @@ BINARY_ACTIVATIONS = filter(f -> hasmethod(f, Tuple{Float64, Float64}), ACTIVATI
 @test -1/3.0 <= rrelu(-1.0) <= -1/8.0
 @test elu(-1.0) == exp(-1.0) - 1.0
 @test gelu(-1.0) ≈ -0.15880800939172324
+@test gelu_tanh(-1.0) ≈ -0.15880800939172324
+@test gelu_erf(-1.0) == -0.15865525393145707
 @test swish(-1.0) == -sigmoid(-1.0)
 @test hardswish(-1.0) == -hardsigmoid(-1.0)
 @test lisht(-1.0) ≈ -1.0 * tanh(-1.0)
@@ -114,7 +120,7 @@ end
         a == softsign && continue
         @test !isnan(a(Inf32))
 
-        a in [gelu, swish, hardswish, logcosh, mish] && continue
+        a in [gelu, gelu_tanh, gelu_erf, swish, hardswish, logcosh, mish] && continue
         @test !isnan(a(-Inf32))
     end
 end

test/runtests.jl

@@ -16,6 +16,7 @@ using KernelAbstractions
 using FFTW
 import ReverseDiff as RD        # used in `pooling.jl`
 import Pkg
+using SpecialFunctions
 
 const Test_Enzyme = VERSION <= v"1.10-"