Gemma¶

Gemma is Google DeepMind's family of open-weight language models, released beginning in 2024. Built on the same research and technology that powers the Gemini models, Gemma brings Google's approach to efficient transformer design to the open-source community. The family spans three generations -- Gemma, Gemma 2, and Gemma 3 -- each refining the architecture with innovations like soft capping, alternating attention patterns, and GeGLU activations¹.

1. Architecture Overview¶

Gemma is a decoder-only autoregressive transformer. While it shares many design elements with LLaMA (RMSNorm, RoPE, GQA), it introduces distinctive features -- most notably soft capping of logits and alternating attention window patterns -- that represent Google's particular engineering philosophy.

2. Key Innovations¶

2.1 GeGLU Activation¶

Gemma uses the GeGLU (GELU-Gated Linear Unit) activation in its feed-forward layers. GeGLU is a gated variant of GELU, similar to SwiGLU but using GELU as the gating function:

2.2 Soft Capping¶

Gemma 2 introduced logit soft capping, a technique to prevent attention scores and final logits from growing unboundedly large:

Soft capping is applied at two points:

1. Attention logit capping: Before softmax, with a cap around 50
2. Final logit capping: Before the output softmax, with a cap around 30

2.3 Grouped-Query Attention (GQA)¶

Gemma uses GQA with varying numbers of KV heads:

This reduces the KV cache size and memory bandwidth requirements during inference while maintaining model quality.

2.4 Alternating Attention Patterns (Gemma 2+)¶

Gemma 2 alternates between local sliding window attention and global full attention across layers:

3. Architecture Diagram¶

flowchart TD
    INPUT["Token IDs"] --> EMB["Token Embedding\n(* sqrt(d_model) scaling)"]
    EMB --> BLOCKS

    subgraph BLOCKS["N x GemmaDecoderLayer"]
        direction TB
        RMS1["RMSNorm"] --> ATTN
        subgraph ATTN["Attention (alternating local/global)"]
            QKV["Q, K, V Projections (GQA)"]
            QKV --> ROPE["RoPE"]
            ROPE --> SCORES["Q K^T / sqrt(d_h)"]
            SCORES --> CAP["Soft Capping\n(c * tanh(z/c))"]
            CAP --> MASK["Causal Mask\n(local or global)"]
            MASK --> SM["Softmax"]
            SM --> VOUT["@ V -> Output Proj"]
        end
        ATTN --> RES1["Residual Add"]
        RES1 --> RMS2["RMSNorm"]
        RMS2 --> FFN["FFN (GeGLU)"]
        FFN --> RES2["Residual Add"]
    end

    BLOCKS --> FINAL_RMS["Final RMSNorm"]
    FINAL_RMS --> LM_HEAD["LM Head\n(with final logit soft capping)"]

4. Configuration Parameters¶

5. Mathematical Formulation¶

5.1 Embedding with Scaling¶

Unlike most models, Gemma scales the embedding output by \( \sqrt{d} \):

5.2 Soft-Capped Attention¶

5.3 GeGLU Feed-Forward¶

where \( W_{\text{gate}}, W_{\text{up}} \in \mathbb{R}^{d \times d_{\text{ff}}} \) and \( W_{\text{down}} \in \mathbb{R}^{d_{\text{ff}} \times d} \).

5.4 Final Logit Capping¶

6. Zig Implementation¶

6.1 GemmaConfig¶

pub const GemmaConfig = struct {
    n_layers: u32,
    d_model: u32,
    n_heads: u32,
    n_kv_heads: u32,
    d_ff: u32,
    vocab_size: u32 = 256000,
    max_seq_len: u32 = 8192,
    rope_theta: f32 = 10000.0,
    norm_eps: f32 = 1e-6,
    activation: ActivationType = .geglu,

    // Soft capping (Gemma 2+)
    attn_logit_cap: ?f32 = null,    // null for Gemma 1
    final_logit_cap: ?f32 = null,

    // Alternating attention (Gemma 2+)
    sliding_window: ?u32 = null,

    pub fn headDim(self: GemmaConfig) u32 {
        return self.d_model / self.n_heads;
    }

    pub fn isLocalLayer(self: GemmaConfig, layer_idx: u32) bool {
        return self.sliding_window != null and (layer_idx % 2 == 0);
    }
};

6.2 Soft Capping Function¶

pub fn softCap(value: f32, cap: f32) f32 {
    return cap * std.math.tanh(value / cap);
}

pub fn applySoftCapping(
    logits: []f32,
    cap: ?f32,
) void {
    const c = cap orelse return;  // no-op if capping disabled
    for (logits) |*val| {
        val.* = softCap(val.*, c);
    }
}

6.3 Gemma Attention with Soft Capping¶

pub const GemmaAttention = struct {
    wq: Linear,
    wk: Linear,
    wv: Linear,
    wo: Linear,
    config: GemmaConfig,
    layer_idx: u32,

    pub fn forward(self: *GemmaAttention, x: Tensor(f32), pos: u32) !Tensor(f32) {
        var q = self.wq.forward(x);
        var k = self.wk.forward(x);
        const v = self.wv.forward(x);

        // Apply RoPE
        applyRoPE(&q, &k, pos, self.config.headDim(), self.config.rope_theta);

        // Compute attention scores
        var scores = scaledDotProduct(q, k, self.config.headDim());

        // Apply soft capping before masking (Gemma 2+)
        if (self.config.attn_logit_cap) |cap| {
            applySoftCapping(scores.data, cap);
        }

        // Apply mask (local or global depending on layer)
        if (self.config.isLocalLayer(self.layer_idx)) {
            applySlidingWindowMask(scores, pos, self.config.sliding_window.?);
        } else {
            applyCausalMask(scores, pos);
        }

        const attn_weights = softmax(scores);
        const context = matmul(attn_weights, v);
        return self.wo.forward(context);
    }
};

7. Variants¶

8. Educational Value¶

9. References¶