Rename gelus and use OpenLibm_jll instead of SpecialFunctions.jl for erf

Author: se-schmitt

Project.toml

@@ -9,8 +9,8 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+OpenLibm_jll = "05823500-19ac-5b8b-9628-191a04bc5112"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
@@ -41,8 +41,8 @@ ForwardDiff = "0.10.36"
 GPUArraysCore = "0.1, 0.2"
 KernelAbstractions = "0.9.2"
 LinearAlgebra = "<0.0.1, 1"
+OpenLibm_jll = "0.8.1"
 Random = "<0.0.1, 1"
-SpecialFunctions = "2.5.0"
 Statistics = "1"
 cuDNN = "1"
 julia = "1.9"

docs/src/reference.md

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

src/NNlib.jl

@@ -15,7 +15,7 @@ using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
 using Random
 using Statistics
 using Statistics: mean
-using SpecialFunctions
+using OpenLibm_jll
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 

src/activations.jl

@@ -5,7 +5,7 @@
 
 ACTIVATIONS = [
     :σ, :hardσ, :hardtanh, :relu,
-    :leakyrelu, :relu6, :rrelu, :elu, :gelu, :gelu_full, :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 @@ elu(x, α=1) = ifelse(x ≥ 0, float(x), @fastmath oftf(x, α) * (exp(x) - 1))
 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) (see also [`gelu_full`](@ref)).
+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 @@ julia> lineplot(gelu, -2, 2, height=7)
            ⠀-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 @@ julia> lineplot!(ans, swish)
              ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀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 @@ end
 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λ)
@@ -358,18 +358,38 @@ function deriv_gelu(x)
 end
 
 """
-    gelu_full(x) = xΦ(x) = 0.5x * (1 + erf(x/√2))
+    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.
 """
-gelu_full(x) = x/2*(1 + erf(x/sqrt(oftf(x,2))))
+gelu_erf(x) = x/2*(1 + _erf(x/sqrt(oftf(x,2))))
 
-function deriv_gelu_full(x)
+function deriv_gelu_erf(x)
     SQRT2 = sqrt(oftf(x,2))
-    Φ = (1 + erf(x/SQRT2))/2
+    Φ = (1 + _erf(x/SQRT2))/2
     Φ + x/SQRT2*exp(-(x^2)/2)/sqrt(oftf(x,π))
 end
 
+_erf(x::Number) = _erf(float(x))
+_erf(x::Float64) = ccall((:erf, libopenlibm), Float64, (Float64,), x)
+_erf(x::Float32) = ccall((:erff, libopenlibm), Float32, (Float32,), x)
+_erf(x::Float16) = Float16(_erf(Float32(x)))
+_erf(x::BigFloat) = begin
+    z = BigFloat(x)
+    ccall((:mpfr_erf, :libmpfr), Int32, (Ref{BigFloat}, Ref{BigFloat}, Int32), z, x, Base.MPFR.ROUNDING_MODE[])
+    return z
+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
+export gelu
+const deriv_gelu = deriv_gelu_tanh
+
 """
     swish(x) = x * σ(x)
 
@@ -887,8 +907,8 @@ UNARY_ACTS = [ # f, dfdx
     (:relu6,        :((Ω>0) & (Ω<6))),
     # rrelu is random, can't write a rule.
     (:elu,          :(deriv_elu(Ω))),
-    (:gelu,         :(deriv_gelu(x))),
-    (:gelu_full,    :(deriv_gelu_full(x))),
+    (:gelu_tanh,    :(deriv_gelu_tanh(x))),
+    (:gelu_erf,     :(deriv_gelu_erf(x))),
     (:swish,        :(Ω + sigmoid_fast(x) * (1 - Ω))),
     (:hardswish,    :(deriv_hardswish(x))),
     # lisht

test/activations.jl

@@ -12,7 +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_full(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
@@ -37,7 +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_full(1.0) == 0.8413447460685429
+@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)
@@ -60,7 +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_full(-1.0) == -0.15865525393145707
+@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)
@@ -117,7 +120,7 @@ end
         a == softsign && continue
         @test !isnan(a(Inf32))
 
-        a in [gelu, gelu_full, swish, hardswish, logcosh, mish] && continue
+        a in [gelu, gelu_tanh, gelu_erf, swish, hardswish, logcosh, mish] && continue
         @test !isnan(a(-Inf32))
     end
 end