Gemma 4 31B-it Text Architecture

Figure 1

The Forward Pass at a Glance

The language model consumes a sequence of token IDs with shape [batch, time]. After lookup, every token is represented as a 5376-dimensional vector. The decoder preserves that width all the way through the stack, then maps each final hidden vector back to 262,144 vocabulary logits.

Model constants

The Numbers That Shape the Model

Most of the architecture can be understood by tracking a few dimensions: the model width D = 5376, the feed-forward width D_ff = 21504, 32 query heads, and two different attention geometries.

What the checkpoint is

Scope: Dense Text Decoder Inside a Multimodal Model

The public model card describes Gemma 4 31B as the dense 31B member of the Gemma 4 family, with 60 layers, a 1024-token sliding window on local layers, 256K context length, a 262K vocabulary, and text+image input support. It also lists a roughly 550M-parameter vision encoder and no audio encoder for the 31B dense model.

Interactive tools

Memory and Compute Estimators

These calculators are architecture-level estimates for the Gemma 4 31B-it text path. They account for the tied text weights and the K/V cache implied by the layer pattern. They do not include framework overhead, allocator fragmentation, vision-model memory, temporary activations, or hardware throughput limits.

Figure 2

The 5-Sliding / 1-Full Layer Rhythm

Gemma 4 31B-it uses 50 sliding-window attention layers and 10 full-attention layers. If layers are counted from 1, the full layers are 6, 12, 18, 24, 30, 36, 42, 48, 54, and 60. This creates a practical long-context pattern: most computation is local, but information periodically passes through a global layer.

Interactive figure

Layer Explorer

Move through the 60 layers to see how the attention geometry changes. The repeated pattern is easy to miss in code because the same module class covers both layer types; this explorer makes the active dimensions visible.

More diagrams

How the Pieces Work Together

The architecture is easiest to understand as a set of conserved shapes and repeated transformations. These diagrams emphasize the flow of one token, how grouped-query attention shares KV heads, how the residual stream is updated, and why long-context cache memory is uneven across layer types.

Decoder layer

One Layer: Norm, Attention, Norm, Residual, MLP

A Gemma 4 text decoder layer has two residual sublayers. The attention output is normalized before it is added back to the residual stream; the MLP output is also normalized before its residual add. That gives four residual-stream RMSNorms per layer.

Attention anatomy

Two Attention Geometries, One Grouped-Query Idea

Every layer has 32 query heads. The key/value head count is smaller, so several query heads share each key/value head. This is grouped-query attention: it keeps the query side expressive while reducing the cache and projection cost for keys and values.

Position information

RoPE: Local Heads Rotate Fully, Global Heads Rotate Partly

Rotary position embeddings are applied to Q and K after their per-head RMSNorms. Sliding layers use the ordinary RoPE setup with rope_theta = 10000 over 256-dimensional heads. Full layers use proportional RoPE with rope_theta = 1000000 and partial_rotary_factor = 0.25.

The full-layer proportional design gives the global layers larger head vectors while limiting how much of each vector is explicitly position rotated. In implementation terms, the non-rotated dimensions receive zero inverse frequency, so cos(0) = 1 and sin(0) = 0.

Feed-forward network

The Gated MLP Carries Most of the Parameters

Each layer has a bias-free gated MLP with three matrices: gate, up, and down. The gate and up streams both expand from 5376 to 21504, then the down projection returns to the residual width. The activation is PyTorch's tanh-approximate GELU.

Parameter budget

Where the Text-Path Weights Live

The checkpoint is called 31B because the large dense matrices in the text path add up to roughly that scale. This figure is approximate and ignores small norm vectors, but it gives the right intuition: the dense gated MLPs dominate.

Normalization

RMSNorm Is the Stabilizer Everywhere

RMSNorm divides a vector by its root mean square and, when enabled, applies a learned scale. Gemma 4 uses the scale directly. It is not the older Gemma-style 1 + weight parameterization.

Output head

Final RMSNorm, Tied Embeddings, and Softcapped Logits

After layer 60, Gemma 4 applies a final RMSNorm. The LM head projects each hidden vector to the vocabulary. Because word embeddings are tied, the LM head uses the same weight table as the input token embedding matrix.

Inference behavior

Prefill, Cache, and Long Context

During prefill, the model computes Q, K, and V for all prompt positions. During generation, the new query positions attend to cached K/V tensors from previous tokens. The same masks still apply: sliding layers are limited to the latest 1024 causal positions, while full layers can see all previous positions.

What is not active here

Useful Non-Features in This Checkpoint

The implementation supports optional components that are not active in the 31B-it text configuration. Calling these out matters because otherwise the code can look more complicated than the checkpoint's actual forward pass.

Compact form

The Whole Text Forward Pass

# Text-only input IDs path. If inputs_embeds are supplied, start with those instead.
H = embedding_table[input_ids] * sqrt(5376)
positions = past_seen_tokens + arange(T_current)

full_mask = causal_mask                 # keys s <= query t
sliding_mask = causal_mask & last_1024  # keys t - 1024 < s <= t
# Supplied padding or packed-sequence masks further restrict these masks.

for layer_index in 0..59:
    kind = "full" if (layer_index + 1) % 6 == 0 else "sliding"
    head_dim = 512 if kind == "full" else 256
    kv_heads = 4 if kind == "full" else 16
    group_size = 32 // kv_heads
    mask = full_mask if kind == "full" else sliding_mask

    U = input_rmsnorm[layer_index](H)

    Q = reshape_heads(q_proj(U), heads=32, dim=head_dim)
    K_raw = reshape_heads(k_proj(U), heads=kv_heads, dim=head_dim)
    V_raw = reshape_heads(v_proj(U), heads=kv_heads, dim=head_dim) if kind == "sliding" else K_raw

    Q = to_attention_layout(rope_kind(q_norm(Q), positions))
    K = k_norm(K_raw)
    K = to_attention_layout(rope_kind(K, positions))
    V = to_attention_layout(v_norm_without_scale(V_raw))

    # In cached generation, append the new layer K/V before attention.
    if use_cache:
        K, V = cache[layer_index].append_and_return(K, V)

    # Shown explicitly; the eager backend performs this repeat internally.
    K = repeat_kv_to_32_query_heads(K, repeats=group_size)
    V = repeat_kv_to_32_query_heads(V, repeats=group_size)

    scores = (Q @ K.transpose(-1, -2)) * 1.0 + mask
    P = softmax(scores, dtype=float32)
    A = o_proj(concat_heads(P @ V))

    Y = H + post_attention_rmsnorm[layer_index](A)

    Z = pre_ffn_rmsnorm[layer_index](Y)
    F = down_proj(gelu_tanh(gate_proj(Z)) * up_proj(Z))
    H = Y + post_ffn_rmsnorm[layer_index](F)

H = final_rmsnorm(H)
logits = H @ embedding_table.T  # tied LM head
logits = 30 * tanh(logits / 30)
next_token_probs = softmax(logits[:, -1, :])

if labels are supplied:
    loss = cross_entropy(logits[:, :-1, :], labels[:, 1:], ignore_masked_labels=True)

Primary references

Sources Used

The page is based on public model metadata and the Hugging Face reference implementation. Links are intentionally direct so the exact config and code paths can be inspected.