Learning What's Missing: Attention Dispersion and EMA Stabilization in Length Generalization
This addresses the problem of transformers failing to generalize to longer sequences, which is crucial for tasks like board-game reasoning, though the improvements are incremental.
The paper tackled length generalization in transformers for tasks like set complement and board-game reasoning, proving theoretical bounds and showing that balanced logit displacement at short lengths leads to generalization with reduced precision, while validating that dropout and EMA improve performance in experiments.
We study length generalization in transformers through the set complement task, where a model must predict a uniform distribution over tokens absent from an input sequence -- an ability central to board-game style reasoning. Our main theoretical result establishes two statements. First, we prove tight bounds on embedding and value dimensions for single-layer attention-only transformers. Second, we show that if such a model achieves balanced logit displacement at lengths 1 and 2, then it must generalize to longer sequences, though with reduced precision. A mechanistic reading of the proof explains this limitation: as more tokens are attended to, softmax compresses logit displacements, eroding separation between valid and invalid outputs. Training dynamics also suggest a second obstacle: when many next tokens are possible, updates become noisy. We hypothesize that dropout can counteract the first effect and Exponential Moving Average (EMA) the second. We validate these hypotheses through random hyperparameter search on the set complement task, which confirms both mechanisms. We then test OthelloGPT, a GPT-1 style model trained on random Othello moves, and find that EMA again improves length generalization in this more complex setting.