Two context windows per layer — and how they grow during training
Sliding-window attention limits each token to a fixed context window. The speedrun uses this for 9 of 11 layers, but asymmetrically: a short window for most layers and a longer window for layers 3 and 10 (wider context bridges). Layer 6 uses full attention.
The window sizes grow during training on a curriculum schedule — ws_schedule=(3,7,11). This is why YaRN is needed: when the window grows, the model encounters position indices further apart than anything it trained on. YaRN rescales RoPE angular frequencies smoothly.
short_bm = ws_short * args.block_size
long_bm = ws_long * args.block_size
bm_sizes = [short_bm, short_bm, short_bm, long_bm, # L0, L1, L2, L3↑
short_bm, short_bm, None, short_bm, # L4, L5, L6∅ (full), L7
short_bm, short_bm, long_bm] # L8, L9, L10↑
# YaRN applied when ws_long changes:
if new_ws_long != ws_long:
self.model.yarn.apply(ws_long, new_ws_long)
Most predictions are highly local — the last few hundred tokens contain the needed signal. Short-window layers handle this efficiently. Layers 3 and 10 act as wider context bridges. L6's full attention is the single fully-expressive layer in the stack middle.
When the window grows, the model encounters token pairs at distances it was never trained on. YaRN rescales RoPE angular frequencies: long-period components are left unchanged; short-period ones are scaled down. Context length extension without catastrophic forgetting.