google/gemma-4-12B @1.0.0

gemma4 · 48 layers · hidden 3840 · 11.91 B parameters · text-only language-model path; vision and audio towers are out of scope

Trust

Compute graph

Domains

Theorems (Lean kernel-checked)

Versions

Advisories

No advisories.

Lean 4 definitions

TheFormalEra/Ops/Attention.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: softmax attention
-/

namespace TheFormalEra.Ops

open Finset

/-- Softmax over the reals, normalized only over visible positions -/
noncomputable def softmax {L : ℕ} (s : Fin L → ℝ) (visible : Fin L → Prop)
    [DecidablePred visible] : Fin L → ℝ :=
  fun l => if visible l then
    Real.exp (s l) / ∑ m ∈ Finset.univ.filter visible, Real.exp (s m) else 0

/-- Causal GQA: each KV head serves G query heads; pos s is the position of query s -/
noncomputable def causalGQA {S L Hkv G D : ℕ} (scale : ℝ) (pos : Fin S → ℕ)
    (q : Fin S → Fin Hkv → Fin G → Fin D → ℝ)
    (k v : Fin L → Fin Hkv → Fin D → ℝ) : Fin S → Fin Hkv → Fin G → Fin D → ℝ :=
  fun s c g d =>
    let logit : Fin L → ℝ := fun l => scale * ∑ e, q s c g e * k l c e
    ∑ l, softmax logit (fun l => l.val ≤ pos s) l * v l c d

/-- Sliding-window causal GQA: query at position `p` sees keys `l` with `p - w < l ≤ p`. -/
noncomputable def slidingGQA {S L Hkv G D : ℕ} (scale : ℝ) (window : ℕ) (pos : Fin S → ℕ)
    (q : Fin S → Fin Hkv → Fin G → Fin D → ℝ)
    (k v : Fin L → Fin Hkv → Fin D → ℝ) : Fin S → Fin Hkv → Fin G → Fin D → ℝ :=
  fun s c g d =>
    let logit : Fin L → ℝ := fun l => scale * ∑ e, q s c g e * k l c e
    ∑ l, softmax logit (fun l => l.val ≤ pos s ∧ pos s < l.val + window) l * v l c d

/-- Decode attention for one sequence of `L` cached tokens: the output of the query at the last
position `L - 1` (every cached key is visible).  Query head `c·G + g` reads KV head `c`. -/
noncomputable def decodeGQA {L Hkv G D : ℕ} (scale : ℝ) (q : Fin Hkv → Fin G → Fin D → ℝ)
    (k v : Fin L → Fin Hkv → Fin D → ℝ) : Fin Hkv → Fin G → Fin D → ℝ :=
  fun c g d => causalGQA (S := 1) scale (fun _ => L - 1) (fun _ => q) k v 0 c g d

/-- Sigmoid output gate of Gated Attention: `o ⊙ σ(gate)`. -/
noncomputable def sigmoidGate {ι : Type*} {N : ℕ} (o gate : ι → Fin N → ℝ) : ι → Fin N → ℝ :=
  fun t i => o t i * sigmoid (gate t i)

end TheFormalEra.Ops

TheFormalEra/Ops/Linear.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: linear maps, feed-forward activations, residuals, embeddings
-/

namespace TheFormalEra.Ops

open Finset

variable {ι : Type*}

/-- `y = x Wᵀ` (bias-free `nn.Linear(K, N)`; `W` is stored `[N, K]`). -/
noncomputable def linear {K N : ℕ} (W : Fin N → Fin K → ℝ) (x : ι → Fin K → ℝ) : ι → Fin N → ℝ :=
  fun t n => ∑ k, x t k * W n k

/-- The fused gate/up projection: `[x W_gateᵀ ‖ x W_upᵀ]`, gate first. -/
noncomputable def gateUp {K I : ℕ} (Wg Wu : Fin I → Fin K → ℝ) (x : ι → Fin K → ℝ) :
    ι → Fin (2 * I) → ℝ :=
  fun t n => if h : n.val < I then linear Wg x t ⟨n.val, h⟩
    else linear Wu x t ⟨n.val - I, by omega⟩

/-- SwiGLU activation of a fused `[gate ‖ up]` row: `silu(gate) ⊙ up`. -/
noncomputable def siluMul {I : ℕ} (gu : ι → Fin (2 * I) → ℝ) : ι → Fin I → ℝ :=
  fun t i => silu (gu t ⟨i.val, by omega⟩) * gu t ⟨I + i.val, by omega⟩

/-- GeGLU (tanh) activation of a fused `[gate ‖ up]` row. -/
noncomputable def geluMul {I : ℕ} (gu : ι → Fin (2 * I) → ℝ) : ι → Fin I → ℝ :=
  fun t i => geluTanh (gu t ⟨i.val, by omega⟩) * gu t ⟨I + i.val, by omega⟩

/-- Residual connection with multiplier `r` (`r = 1` except for Granite). -/
noncomputable def residual {H : ℕ} (r : ℝ) (x h : ι → Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t i => x t i + r * h t i

/-- Token embedding with multiplier `c`. -/
noncomputable def embed {V H T : ℕ} (c : ℝ) (E : Fin V → Fin H → ℝ) (tokens : Fin T → Fin V) :
    Fin T → Fin H → ℝ :=
  fun t i => c * E (tokens t) i

/-- Language-model head with logit multiplier and optional soft-capping. -/
noncomputable def lmHead {V H : ℕ} (mul : ℝ) (cap : Option ℝ) (W : Fin V → Fin H → ℝ)
    (x : ι → Fin H → ℝ) : ι → Fin V → ℝ :=
  fun t v => match cap with
    | none => mul * linear W x t v
    | some c => softcap c (mul * linear W x t v)

end TheFormalEra.Ops

TheFormalEra/Ops/Norm.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: normalisation

Every norm normalises the last axis of a tensor whose rows are indexed by an arbitrary type
`ι` (tokens `Fin T`, or `(token, head)` pairs for per-head norms).
-/

namespace TheFormalEra.Ops

open Finset

variable {ι : Type*} {H : ℕ}

/-- `√(mean(x²) + ε)` of one row. -/
noncomputable def rms (eps : ℝ) (x : Fin H → ℝ) : ℝ := Real.sqrt ((∑ j, x j ^ 2) / (H : ℝ) + eps)

/-- RMSNorm with the weight applied as `w` (Qwen3, Granite, Gemma 4). -/
noncomputable def rmsnorm (eps : ℝ) (x : ι → Fin H → ℝ) (w : Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps) * w h

/-- RMSNorm in the Qwen3.5 family: the weight is applied as (1 + w) -/
noncomputable def rmsnormOffset (eps : ℝ) (x : ι → Fin H → ℝ) (w : Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps) * (1 + w h)

/-- RMSNorm without a weight (Gemma 4 value norm, Muse Glimmer q/k and embedding norms). -/
noncomputable def rmsnormUnit (eps : ℝ) (x : ι → Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps)

/-- Gated RMSNorm of the Gated DeltaNet output: `w ⊙ x̂ ⊙ silu(z)`. -/
noncomputable def rmsnormGated (eps : ℝ) (x z : ι → Fin H → ℝ) (w : Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => w h * (x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps)) * silu (z t h)

/-- L2 normalisation `x / √(Σ x² + ε)` (queries and keys of the delta rule). -/
noncomputable def l2norm (eps : ℝ) (x : ι → Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt (∑ j, x t j ^ 2 + eps)

end TheFormalEra.Ops

TheFormalEra/Ops/Rope.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: rotary position embeddings

Rotate-half RoPE on the first `rot` coordinates of every head (`rot = D` for full RoPE,
`rot = 64` of `256` for Qwen3.8's partial RoPE), with inverse frequencies
`ω_i = θ^(-2i / denom)` for `i < active` and `ω_i = 0` beyond (Gemma 4's proportional RoPE
rotates the whole head but keeps only the first `active` frequencies).
-/

namespace TheFormalEra.Ops

variable {ι : Type*} {D : ℕ}

/-- Inverse frequency of pair `i`. -/
noncomputable def invFreq (theta : ℝ) (denom active i : ℕ) : ℝ :=
  if i < active then theta ^ (-(2 * (i : ℝ) / denom)) else 0

/-- RoPE at position `pos t` of row `t`; pairs `(i, i + rot/2)` for `i < rot/2`. -/
noncomputable def rope (theta : ℝ) (rot denom active : ℕ) (pos : ι → ℕ) (x : ι → Fin D → ℝ) :
    ι → Fin D → ℝ :=
  fun t d =>
    let half := rot / 2
    if d.val < half then
      let a := (pos t : ℝ) * invFreq theta denom active d.val
      x t d * Real.cos a - (if h : d.val + half < D then x t ⟨d.val + half, h⟩ else 0) * Real.sin a
    else if d.val < 2 * half then
      let a := (pos t : ℝ) * invFreq theta denom active (d.val - half)
      x t d * Real.cos a + x t ⟨d.val - half, by omega⟩ * Real.sin a
    else x t d

/-- Standard (partial) RoPE over the first `rot` coordinates. -/
noncomputable def ropeStd (theta : ℝ) (rot : ℕ) (pos : ι → ℕ) (x : ι → Fin D → ℝ) : ι → Fin D → ℝ :=
  rope theta rot rot (rot / 2) pos x

end TheFormalEra.Ops