GEMMA 4 ARCHITECTURE

Source: Hugging Face blog. Instruction-tuned variants. E4B included from blog (architecture page covers E2B, 26B-A4B, 31B in detail).

The 26B-A4B MoE achieves 98% of the 31B's Arena score with only 4B active parameters -- a 7.5x compute reduction. The E2B achieves competitive audio scores despite being a 2.3B-effective model.

Computed from: Q=dim×q_heads×head_dim, K/V=dim×kv_heads×head_dim, O=q_heads×head_dim×dim. GeGLU FFN=3×dim×hidden (gate+up+down). V proj=0 when K=V. MoE: per_expert=3×dim×expert_hidden, router=dim×num_experts.

Estimates: weight memory + KV cache (FP16). KV formula: 2 × kv_heads × head_dim × 2 bytes per token per layer, summed across sliding and full layers. E2B benefits from KV sharing (20 of 35 layers reuse cache).

• Dual-config attention (different head_dim, kv_heads per layer type) Sliding window layers use head_dim=256 with more KV heads, while full attention layers use head_dim=512 with fewer KV heads. This allows efficient local attention with fine-grained heads and powerful global attention with large heads, all within the same model.
• K=V weight sharing (eliminate V projection) In full attention layers, the V projection is removed (v_proj=None). The key tensor is cloned as the value before normalization, then K and V diverge: K gets k_norm (learned scale) + RoPE, V gets v_norm (no scale, no RoPE). So K and V start identical but develop different representations through their separate norm paths. Present in 31B and 26B-A4B.
• Proportional RoPE with partial rotation (25%) Full attention layers apply RoPE to only 25% of head dimensions (partial=0.25) with theta=1M, rather than the linear scaling used in Gemma 3. This provides better long-context extrapolation by leaving 75% of dimensions position-independent.
• Value normalization (without learned scale) All Gemma 4 models apply RMSNorm to value vectors without a learned scale parameter. This stabilizes attention outputs and prevents value magnitude drift across layers -- a feature absent in Gemma 3.
• Parallel dense FFN + MoE (26B-A4B) The 26B-A4B MoE model runs a dense GeGLU FFN (hidden=2,112) in parallel with 128-expert routed MoE (each expert hidden=704, top-8 routing). The outputs are summed and scaled by 1/sqrt(2), combining always-on capacity with sparse expert specialization.
• Per-layer input embeddings (E2B) The E2B model adds a second embedding table mapping tokens to (layers * 256) dimensions. Each layer receives a unique 256-dim slice, gated and added as a residual to the main hidden state. This accounts for the ~5B total vs ~2B effective parameter count.
• KV cache sharing across layers (E2B) In E2B, 20 of 35 layers reuse KV caches computed by earlier layers of the same type (sliding or full), dramatically reducing KV cache memory. Layers with shared KV compensate with 2x wider MLP (hidden=12,288 vs 6,144).
• Conformer audio encoder (E2B) E2B includes a 12-layer Conformer encoder (USM architecture) with chunked local attention, causal Conv1d, 2-stage Conv2d subsampling (4x temporal reduction), and projection to the text embedding space. This enables native audio understanding.
• Logit soft-capping All Gemma 4 models apply tanh(x/30) * 30 to final logits, bounding them to [-30, 30]. This prevents extreme logit values during generation without hard clipping, improving training stability. Gemma 3 does not use this technique.

Verified against transformers/models/gemma4/modeling_gemma4.py. Code snippets are exact quotes.

self.scaling = 1.0

# v_proj is None when attention_k_eq_v=True
value_states = self.v_proj(hidden_states).view(hidden_shape) \
    if self.v_proj is not None else key_states

# K path: scaled norm + RoPE
key_states = self.k_norm(key_states)          # with_scale=True
key_states = apply_rotary_pos_emb(key_states, cos, sin)

# V path: unscaled norm, NO RoPE
value_states = self.v_norm(value_states)      # with_scale=False

class Gemma4RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6, with_scale=True):
        super().__init__()
        self.eps = eps
        self.with_scale = with_scale
        if self.with_scale:
            self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, hidden_states):
        normed = self._norm(hidden_states.float())
        if self.with_scale:
            normed = normed * self.weight.float()
        return normed.type_as(hidden_states)

class Gemma4TextScaledWordEmbedding(nn.Embedding):
    def __init__(self, num_embeddings, embedding_dim,
                 padding_idx, embed_scale=1.0):
        super().__init__(num_embeddings, embedding_dim, padding_idx)
        self.register_buffer(
            "embed_scale", torch.tensor(embed_scale))

    def forward(self, input_ids):
        return super().forward(input_ids) * self.embed_scale

# Main embedding: scale = sqrt(5376) ≈ 73.3
# Per-layer embedding: scale = sqrt(256) = 16.0

# After attention + FFN residuals:
if self.hidden_size_per_layer_input:
    residual = hidden_states
    hidden_states = self.per_layer_input_gate(hidden_states)  # d → 256
    hidden_states = self.act_fn(hidden_states)
    hidden_states = hidden_states * per_layer_input  # gated element-wise
    hidden_states = self.per_layer_projection(hidden_states)  # 256 → d
    hidden_states = self.post_per_layer_input_norm(hidden_states)
    hidden_states = residual + hidden_states

# Dense path (standard MLP)
hidden_states = self.pre_feedforward_layernorm(residual)
hidden_states = self.mlp(hidden_states)

if self.enable_moe_block:
    hidden_states_1 = self.post_feedforward_layernorm_1(
        hidden_states)

    # MoE path (parallel, from pre-MLP residual)
    hidden_states_flat = residual.reshape(-1, residual.shape[-1])
    _, top_k_weights, top_k_index = self.router(
        hidden_states_flat)
    hidden_states_2 = self.pre_feedforward_layernorm_2(
        hidden_states_flat)
    hidden_states_2 = self.experts(
        hidden_states_2, top_k_index, top_k_weights)
    hidden_states_2 = self.post_feedforward_layernorm_2(
        hidden_states_2)

    # Sum both paths, then final post-norm
    hidden_states = hidden_states_1 + hidden_states_2
    hidden_states = self.post_feedforward_layernorm(
        hidden_states)

hidden_states = residual + hidden_states

def forward(self, hidden_states):
    hidden_states = self.norm(hidden_states)        # RMSNorm, no scale
    hidden_states = hidden_states * self.scale * self.scalar_root_size

    expert_scores = self.proj(hidden_states)        # Linear → num_experts
    router_probs = softmax(expert_scores, dim=-1)

    top_k_weights, top_k_index = torch.topk(
        router_probs, k=self.config.top_k_experts)

    top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True)
    top_k_weights = top_k_weights * self.per_expert_scale[top_k_index]

    return router_probs, top_k_weights, top_k_index

hidden_states *= self.layer_scalar  # buffer, not nn.Parameter

if self.config.final_logit_softcapping is not None:
    logits = logits / self.config.final_logit_softcapping
    logits = torch.tanh(logits)
    logits = logits * self.config.final_logit_softcapping

# From Gemma4TextConfig:
rope_parameters = {
    "sliding_attention": {
        "rope_type": "default",
        "rope_theta": 10_000.0
    },
    "full_attention": {
        "rope_type": "proportional",
        "partial_rotary_factor": 0.25,
        "rope_theta": 1_000_000.0
    }
}

# Last layer forced to full_attention:
if self.layer_types[-1] != "full_attention":
    self.layer_types[-1] = "full_attention"

# In shared layer forward:
if self.is_kv_shared_layer and past_key_values is not None:
    key_states, value_states = \
        past_key_values.shared_layers[self.kv_shared_layer_index]

# In non-shared layer: store KV for later reuse
if self.store_full_length_kv:
    past_key_values.shared_layers[self.layer_idx] = \
        key_states, value_states