OlmPool: How small architectural choices compound to undermine long context extension

Ask most people why Llama 3 stretches so gracefully to long context windows and they'll shrug and say "good data." AI2's new OlmPool study says that's mostly wrong. The researchers built 26 near-identical 7B models, all trained on the same 140 billion tokens with the same extension recipe, and only varied four architectural knobs. The gaps that opened up were not small.

The four culprits all touch attention, the part of a transformer that decides what to look at. QK normalization stabilizes training but, combined with other choices, quietly hurts retrieval over long inputs. Grouped-query attention saves memory by sharing key-value heads across attention heads, at the cost of flexibility. Sliding window attention limits most layers to a local neighborhood of tokens instead of the whole sequence. And pretraining at a short context length, then bolting on length later, leaves the model with less intuition for long-range patterns. Individually, each costs a few points on benchmarks like HELMET and RULER. Stack three or four together and scores fall by up to 47%.

The more unsettling finding is that none of the usual warning signs show up in advance. Training loss looks fine. Validation perplexity looks fine. Sixteen different short-context benchmarks look fine. Models that are statistically indistinguishable on every short-context metric AI2 tried can differ by 26 points once extended to 32K tokens. The only reliable early signal the team found was literally running a cheap context-extension test partway through pretraining, long before anyone would normally check.

AI2 also poked at why QK norm in particular seems to backfire for long context. Models without it develop stronger "attention sinks" — a handful of early tokens that soak up disproportionate attention regardless of relevance. Normally researchers treat sinks as an annoyance, something that gets in the way of compressing models. Here, stronger sinks correlated with better long-range retrieval, as if the model without QK norm invents its own crude mechanism for managing attention across a long input, and that crude mechanism works better than the polished one.

Throwing more data at the problem doesn't fix it either. AI2 scaled context extension from 1 billion to 50 billion tokens and the architectural gap barely moved. The weakest architecture, even after 50 billion tokens of extension, still lagged behind what the Llama-style setup achieved with just 1 billion. Same story across pretraining runs up to 2 trillion tokens — architecture rankings locked in early and stayed put.

AI2 is releasing all 26 models with 38 checkpoints apiece, spanning pretraining through extension, which is a genuinely useful gift to anyone trying to build or debug long-context models without guessing which Llama-specific trick will or won't transfer to their own setup.