Documentation

LeanMlir.Proofs.Architectures.GeluForm

Either GELU #

The GELU comes in two forms: the tanh approximation (geluScalar, the default of jax.nn.gelu) and the exact x · Φ(x) (geluErfScalar, PyTorch's nn.GELU). GeluForm names the choice. A net that uses the GELU, and every theorem about it, takes a GeluForm and holds for both, so the renders that trained under the tanh form and the ones that train under the exact form are instances of one statement.

GeluForm.scalar, GeluForm.scalarDeriv: the scalar activation and its derivative (a deriv; the closed forms are geluScalarDeriv_eq and geluErfScalarDeriv_eq).
GeluForm.map, GeluForm.pdiv_map, GeluForm.hasVJP: the activation on a vector, its diagonal Jacobian, its VJP.
GeluForm.map_tanh, GeluForm.map_erf, GeluForm.hasVJP_tanh, GeluForm.hasVJP_erf: at each form these are gelu / geluHasVJP and geluErf / geluErfHasVJP, by rfl.

Nothing here cases on the form except scalar and its differentiability, so a definition built on GeluForm.map unfolds the same way at either form.

inductive Proofs.GeluForm :

Which GELU: the tanh approximation or the exact x · Φ(x).

tanh : GeluForm
0.5 · x · (1 + tanh(√(2/π) · (x + 0.044715 · x³))), geluScalar.
erf : GeluForm
x · Φ(x), geluErfScalar.

Instances For

@[instance_reducible]

instance Proofs.instDecidableEqGeluForm :

DecidableEq GeluForm

Equations

Proofs.instDecidableEqGeluForm x✝ y✝ = if h : x✝.ctorIdx = y✝.ctorIdx then isTrue ⋯ else isFalse ⋯

def Proofs.instReprGeluForm.repr :

GeluForm → ℕ → Std.Format

Equations

Proofs.instReprGeluForm.repr Proofs.GeluForm.tanh prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Proofs.GeluForm.tanh")).group prec✝
Proofs.instReprGeluForm.repr Proofs.GeluForm.erf prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "Proofs.GeluForm.erf")).group prec✝

Instances For

@[instance_reducible]

instance Proofs.instReprGeluForm :

Equations

Proofs.instReprGeluForm = { reprPrec := Proofs.instReprGeluForm.repr }

@[instance_reducible]

instance Proofs.instInhabitedGeluForm :

Inhabited GeluForm

Equations

Proofs.instInhabitedGeluForm = { default := Proofs.instInhabitedGeluForm.default }

noncomputable def Proofs.GeluForm.scalar :

GeluForm → ℝ → ℝ

The scalar activation of a form.

Equations

Instances For

@[simp]

theorem Proofs.GeluForm.scalar_tanh :

tanh.scalar = geluScalar

@[simp]

theorem Proofs.GeluForm.scalar_erf :

erf.scalar = geluErfScalar

theorem Proofs.GeluForm.scalar_differentiable (gf : GeluForm) :

Differentiable ℝ gf.scalar

Either form is differentiable.

noncomputable def Proofs.GeluForm.scalarDeriv (gf : GeluForm) (x : ℝ) :

Scalar derivative of a form — Mathlib's deriv, as geluScalarDeriv and geluErfScalarDeriv are.

Equations

gf.scalarDeriv x = deriv gf.scalar x

Instances For

theorem Proofs.GeluForm.scalarDeriv_tanh (x : ℝ) :

tanh.scalarDeriv x = geluScalarDeriv x

theorem Proofs.GeluForm.scalarDeriv_erf (x : ℝ) :

erf.scalarDeriv x = geluErfScalarDeriv x

noncomputable def Proofs.GeluForm.map (gf : GeluForm) (n : ℕ) (x : Vec n) :

The GELU of a form, applied componentwise to a vector.

Equations

gf.map n x i = gf.scalar (x i)

Instances For

theorem Proofs.GeluForm.map_tanh (n : ℕ) :

tanh.map n = gelu n

theorem Proofs.GeluForm.map_erf (n : ℕ) :

erf.map n = geluErf n

theorem Proofs.GeluForm.map_differentiable (gf : GeluForm) (D : ℕ) :

Differentiable ℝ (gf.map D)

Differentiability of gf.map D as a function on Vec D.

theorem Proofs.GeluForm.pdiv_map (gf : GeluForm) (n : ℕ) (x : Vec n) (i j : Fin n) :

pdiv (gf.map n) x i j = if i = j then gf.scalarDeriv (x i) else 0

Partial derivative of either GELU — diagonal, pdiv_elementwise at gf.scalar.

noncomputable def Proofs.GeluForm.hasVJP (gf : GeluForm) (n : ℕ) :

HasVJP (gf.map n)

VJP of either GELU: elementwise multiply by the scalar derivative,

back(x, dy)_i = dy_i * gf.scalarDeriv(x_i).

Equations

gf.hasVJP n = { backward := fun (x dy : Proofs.Vec n) (i : Fin n) => dy i * gf.scalarDeriv (x i), correct := ⋯ }

Instances For

theorem Proofs.GeluForm.hasVJP_tanh (n : ℕ) :

tanh.hasVJP n = geluHasVJP n

theorem Proofs.GeluForm.hasVJP_erf (n : ℕ) :

erf.hasVJP n = geluErfHasVJP n

theorem Proofs.GeluForm.hasVJP_correct (gf : GeluForm) (n : ℕ) (x dy : Vec n) (i : Fin n) :

(gf.hasVJP n).backward x dy i = ∑ j : Fin n, pdiv (gf.map n) x i j * dy j

Public correctness theorem for GeluForm.hasVJP: the backward (diagonal scaling by gf.scalarDeriv) equals the pdiv-contracted Jacobian.