Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds full gelu without approximation #629

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

Adds full gelu without approximation #629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
cc4b9eb
Add SpecialFunctions as dependency
se-schmitt Feb 6, 2025
86256b2
add full `gelu`, changes old `gelu` -> `gelu_fast`
se-schmitt Feb 6, 2025
8275b18
Add tests and docs
se-schmitt Feb 6, 2025
bdfa0f0
Change names: `gelu_fast` -> `gelu`, `gelu` -> `gelu_full`
se-schmitt Feb 9, 2025
44ee4f5
Rename gelus and use OpenLibm_jll instead of SpecialFunctions.jl for erf
se-schmitt Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ 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"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

Expand Down Expand Up @@ -40,6 +41,7 @@ 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"
Statistics = "1"
cuDNN = "1"
Expand Down
2 changes: 2 additions & 0 deletions docs/src/reference.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/NNlib.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ using LinearAlgebra: AdjOrTransAbsMat, Adjoint, BlasFloat, Transpose
using Random
using Statistics
using Statistics: mean
using OpenLibm_jll

const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}

Expand Down
54 changes: 44 additions & 10 deletions src/activations.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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).
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) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⣀⡠⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
Expand All @@ -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) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠒⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⢄⠀⠀⠀⠀⠀⠀⠀⠀⢠⡇⠀│
Expand All @@ -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
Expand All @@ -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λ)
Expand All @@ -357,6 +357,39 @@ function deriv_gelu(x)
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.
"""
gelu_erf(x) = x/2*(1 + _erf(x/sqrt(oftf(x,2))))

function deriv_gelu_erf(x)
SQRT2 = sqrt(oftf(x,2))
Φ = (1 + _erf(x/SQRT2))/2
Φ + x/SQRT2*exp(-(x^2)/2)/sqrt(oftf(x,π))
end

_erf(x::Number) = _erf(float(x))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This intends to catch integers but isn't a safe pattern, as there are other weird numbers out there:

StackOverflowError:
  Stacktrace:
   [1] _erf(x::ForwardDiff.Dual{ForwardDiff.Tag{Zygote.var"#141#142"{var"#170#178"{typeof(gelu_erf)}}, Float64}, Float64, 1}) (repeats 79984 times)
     @ NNlib ~/work/NNlib.jl/NNlib.jl/src/activations.jl:373

This particular case could be allowed via NNlibForwardDiffExt

_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)

Expand Down Expand Up @@ -874,7 +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_tanh, :(deriv_gelu_tanh(x))),
(:gelu_erf, :(deriv_gelu_erf(x))),
(:swish, :(Ω + sigmoid_fast(x) * (1 - Ω))),
(:hardswish, :(deriv_hardswish(x))),
# lisht
Expand Down
8 changes: 7 additions & 1 deletion test/activations.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down
Loading