Layer-wise attention budgeting
Layer-wise attention budgeting means varying the attention cost by layer instead of giving every transformer layer the same full attention budget. One recent example is Laguna XS.2, the first open-weight model by Poolside, a Europe-based company focused on training LLMs for coding applications. Several of my former colleagues joined Poolside in recent years, and they have a great team with lots of talent. It’s just nice to see more companies also releasing some of their models as open-weight variants.
Anyways, the Laguna XS.2 architecture depicted below looks very standard at first glance. However, one detail that I didn’t show (/try to cram into there) is a concept we can refer to as “layer-wise attention budgeting”.
Where the budgeting happens
Laguna has 40 layers, but they don’t all use the same attention setup. Thirty are sliding-window layers, where each token can look back over a local window of 512 tokens. The other ten use full attention and can access the whole context. These global layers are more expensive, whereas the sliding-window layers keep the KV cache and attention computation cheaper.
This mixed sliding-window and full-attention pattern isn’t unique to Laguna XS.2. Gemma 4, for example, uses both types as well.
But what’s new is the use of per-layer query-head counts. The Hugging Face config.json includes a num_attention_heads_per_layer setting, so the number of query heads can change while the KV-cache shape stays compatible.
The head counts
The full-attention layers use 48 query heads. Sliding-window layers get 64. Both keep the KV-head count fixed at 8, which works out to 6 query heads per KV head for full attention and 8 for sliding-window attention.
That’s the layer-wise head budgeting encoded in Laguna’s config.
Why vary the budget?
The broader idea of varying model capacity by layer goes back to (at least) Apple’s 2024 OpenELM. And again, what’s the point of such a design?
Similar to KV sharing, the point is to spend attention capacity where it is most useful instead of giving every layer the same budget. Full-attention layers look across the whole context, so Laguna gives them fewer query heads than the cheaper sliding-window modules.
(Besides, another smaller implementation detail is that Laguna also applies per-head attention-output gating; this is somewhat similar to Qwen3-Next and others, which I omit here since I covered it in earlier articles.)
Sources