Global vs Local Layer Routing

Gemma 3

Easy

Route layers to local or global attention

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Gemma 3 is a decoder-only stack, but the layers do not share one attention span. Most layers run sliding-window (local) attention; a minority run full-context (global) attention. The report’s reason is KV-cache memory at 128K context: only the global layers store a cache that grows with the full sequence. Local layers keep a short span (1024 tokens on the 4B–27B models). The chosen mix is five local layers for every global layer, and the first layer is local.

Gemma 2 already interleaved local and global attention, but at a 1:1 ratio and with a 4096-token local window. Gemma 3 tilts the ratio toward local and shortens the window. Section 5.2 reports that raising the local:global ratio to 5:1, and even 7:1 in an ablation, barely moves validation perplexity, while the KV-cache footprint drops sharply versus a global-only stack.

Routing is a static per-layer flag, not a learned router. Each layer is constructed as either AttentionType.LOCAL_SLIDING or AttentionType.GLOBAL and keeps that type for the life of the model.

How it works

The official pattern is a six-slot tile that starts with local:

P = (L,L,L,L,L,G).

make_attention_layers_types in gemma/gm/nn/_config.py repeats (P) to cover num_layers. If the depth is not a multiple of 6, it appends a prefix of (P):

\mathrm{type}(\ell) = P[\ell \bmod |P|]\quad\text{for }\ell = 0,\ldots,N-1,

with the remainder implemented as pattern[: N % 6] rather than a wrap that would insert a stray global. Released depths:

So the paper’s “5:1, starting local” is the design target; the code’s tiling is exact only when (N) is a multiple of 6. The leftover layers are local, never an extra global.

That flag then selects two other hyperparameters in Transformer.setup (gemma/gm/nn/_transformer.py):

• Local layers: RoPE base (10^4), scale 1.
• Global layers: RoPE base (10^6), and on 4B/12B/27B a RoPE scale of 8 (the positional-interpolation factor used when stretching pre-training from 32K to 128K).

The 1B and 270M configs do not set global_scale_factor, so it stays at the default 1. The 1B model is also the one the report caps at 32K context rather than 128K.

The paper specifies the 5:1 ratio, the “first layer is local” rule, and the two RoPE bases. The exact tiling function, the leftover-prefix rule, per-size depths, and the global RoPE scale of 8 are from the official configs.

Official code

• gemma/gm/nn/_gemma.py — GEMMA3_ATTENTION_PATTERN and each Gemma3_* attention_types=... call.
• gemma/gm/nn/_config.py — make_attention_layers_types.
• gemma/gm/nn/_transformer.py — maps each attn_type onto RoPE base and scale when constructing Block.
• gemma/gm/nn/_modules.py — AttentionType enum (GLOBAL, LOCAL_SLIDING).

The same pattern is re-exported as gm.nn.config.GEMMA3_ATTENTION_PATTERN for custom stacks (docs/research.md).

Watch-outs

• Index from zero and start with local. A 1-based “every 6th layer is global” rule that treats layer 1 as global inverts the stack.
• Do not force an exact 5:1 on 1B/4B/27B. Trailing locals from the prefix rule are part of the released models.
• Routing also switches RoPE. Giving every layer (\theta=10^4) (or every layer (\theta=10^6)) ignores the global-context extension the report describes.
• Gemma 2’s 1:1 pattern and 4096-token window are the wrong defaults if you copy an older config.

Sources

• Paper: Gemma 3 Technical Report, §2 (5:1 interleaving) and §5.2 (ratio ablation)
• Code: google-deepmind/gemma gemma/gm/nn/_gemma.py, gemma/gm/nn/_config.py, gemma/gm/nn/_transformer.py