Transformer Attention: Wy does Gemma4 let most layers see only a small window?

Published

This is a series where I answer the questions my tutoring students really ask – the ones textbooks tend to skip. I teach machine learning, deep learning, statistics, and math to learners of all levels.

The question. After the causal masking post, a student stared at the masked attention matrix and asked: “If most of the matrix is already zeros, why do we compute full attention at all? And I read that Gemma uses a sliding window – doesn’t throwing away the distant past hurt the model?”

Great instinct. The full causal matrix costs memory that grows with every token you feed the model. Gemma’s answer is hybrid attention: most layers look only at a small local window, and every few layers one global layer looks at everything. Here is how it works, on our usual six tokens.

The local mask: causal plus a window

A sliding-window layer adds one more condition to the causal mask from post #4: token i may only attend to token j if j <= i (the causal rule) and i - j < W (the window rule). Everything else becomes minus infinity before the softmax:

scores = inputs @ inputs.T   # the same six-token example as always

W = 3
i = torch.arange(6).unsqueeze(1)  # query positions
j = torch.arange(6).unsqueeze(0)  # key positions
local_mask = (j &lt;= i) & (i - j &lt; W)

local_weights = torch.softmax(scores.masked_fill(~local_mask, float("-inf")), dim=-1)
print(local_weights[5])  # the row for "step"
# tensor([0.000, 0.000, 0.000, 0.330, 0.230, 0.440])

Compare with the plain causal row for “step”, which was [0.138, 0.218, 0.213, 0.142, 0.099, 0.190]. The windowed row keeps only the last three entries – “with”, “one”, and “step” itself – and renormalizes them to sum to one. “Your”, “journey”, and “starts” are not down-weighted; they are gone from this layer’s view entirely.

Gemma-style hybrid attention. Most layers apply a local sliding-window mask (a band around the diagonal), and every few layers one global layer applies the full causal mask. The journey row is outlined in both masks.

The obvious objection: what about long-range information?

If every layer were local, information could only travel W - 1 tokens per layer: a token’s context vector mixes its window, the next layer mixes windows of windows, and so on. After L local layers with window W, a token has heard about at most 1 + L * (W - 1) positions. With W = 1024 that sounds like plenty, but it is still a hard ceiling, and long documents would exceed it quickly.

This is why the design is a hybrid. Gemma interleaves the two masks: several cheap local layers, then one global layer with the plain causal mask that re-syncs every token with the entire past. Gemma 2 alternated one local, one global, with a 4096 window. Gemma 3 pushed the ratio to five local layers for every one global layer and shrank the window to 1024 – and its technical report shows the perplexity barely moves. Most language understanding is local; the occasional global layer is enough to keep the whole document within reach.

There is one subtle supporting detail: local and global layers use different RoPE base frequencies, 10k for local and 1M for global. The local layers only ever see short distances, so they keep the frequency tuned for short-range positions, while the global layers get a frequency that stretches across the full 128K context.

The arithmetic: why the cache stops growing

Here is the payoff. A local layer never needs keys or values older than its window, so its KV cache stops growing once the context passes W. Only the global layers keep a full-length cache. Take a 60-layer model with Gemma’s 5:1 split – 50 local layers, 10 global – and GQA with 8 key-value groups of width 128:

ctx = 128 * 1024          # 128K context
d_kv = 2 * 8 * 128        # K + V, 8 KV groups, 128 dims each
nbytes = 2                # bf16

global_only = 60 * ctx * d_kv * nbytes / 1e9
hybrid = (10 * ctx + 50 * min(ctx, 1024)) * d_kv * nbytes / 1e9
print(f"global-only: {global_only:.1f} GB   hybrid: {hybrid:.1f} GB")
# global-only: 32.2 GB   hybrid: 5.6 GB

The local layers contribute a fixed 1024-token slab no matter how long the document gets, so the hybrid cache is almost six times smaller at 128K, and the gap widens with every extra token. This is exactly the memory curve from the Gemma 3 report, where KV-cache overhead drops from about 60% of total inference memory to under 15%.

KV cache in GB versus context length for a global-only 60-layer model and a 5:1 hybrid with a 1024 local window. The hybrid curve nearly flattens because local layers stop growing after the window fills.

One thing to try

In the toy code, set W = 2 and rerun. Now “step” sees only “one” and itself. Then stack two windowed layers by hand – apply the mask, mix, feed the result through one more round – and check whether “step” can reach “starts” now. That small experiment is the whole argument for the hybrid design: locality is cheap, but reach must be bought back, either with depth or with an occasional global layer.

The code in this article is adapted from the teaching notebooks I use with my students, and the architecture details follow the Gemma 3 technical report, which I recommend for the full ablations on window size and layer ratios.

References

  1. Gemma Team, Google DeepMind (2025). “Gemma 3 Technical Report.” arXiv:2503.19786.
  2. Beltagy, I., Peters, M. E., & Cohan, A. (2020). “Longformer: The Long-Document Transformer.” arXiv:2004.05150.
  3. Vaswani, A., et al. (2017). “Attention Is All You Need.” Advances in Neural Information Processing Systems 30 (NeurIPS 2017).
  4. Raschka, S. (2024). Build a Large Language Model (From Scratch). Manning.

I tutor machine learning, statistics, and math. If attention – or anything else – is not clicking for you, you can find me on Wyzant.