Delayed Attention Training Improves Length Generalization in Transformer--RNN Hybrids
This work addresses length generalization for sequence modeling tasks, offering a solution to improve hybrid architectures, though it is incremental in nature.
The paper tackled the problem of length generalization in sequence models by addressing the complementary strengths of recurrent networks and Transformers, proposing a delayed attention training strategy that enabled hybrid models to achieve near-perfect accuracy (>90%) on sequences three times longer than training data.
We study length generalization in sequence models on a composite problem involving both state tracking and associative recall. Prior work finds that recurrent networks handle state tracking well but struggle with recall, whereas Transformers excel at recall yet fail to extend state-tracking capabilities to longer sequences. Motivated by the complementary strengths of these architectures, we construct hybrid models integrating recurrent and attention-based components, and train them on the combined task to evaluate whether both capabilities can be preserved. Our results reveal that, in such hybrids, the Transformer component tends to exploit shortcut solutions, leading to poor length generalization. We identify this shortcut reliance as a key obstacle and propose a simple yet effective training strategy -- delaying the training of the attention layers -- that mitigates this effect and significantly improves length generalization performance. Our experiments show that this approach enables hybrid models to achieve near-perfect accuracy ($>90\%$) on hybrid sequences three times longer than those used during training.