Looped Transformers for Length Generalization

2024/09/24 by Ying Fan, Yilun Du, Fan, Ying +5 · 14 citations
Computer Science · #Neural Networks and Applications

paper · pdf · doi:10.48550/arxiv.2409.15647

Abstract

Recent work has shown that Transformers trained from scratch can successfully solve various arithmetic and algorithmic tasks, such as adding numbers and computing parity. While these Transformers generalize well on unseen inputs of the same length, they struggle with length generalization, i.e., handling inputs of unseen lengths. In this work, we demonstrate that looped Transformers with an adaptive number of steps significantly improve length generalization. We focus on tasks with a known iterative solution, involving multiple iterations of a RASP-L operation - a length-generalizable operation that can be expressed by a finite-sized Transformer. We train looped Transformers using our proposed learning algorithm and observe that they learn highly length-generalizable solutions for various tasks.

Cited by

• VersatileFFN: Achieving Parameter Efficiency in LLMs via Adaptive Wide-and-Deep Reuse
• Teaching Pretrained Language Models to Think Deeper with Retrofitted Recurrence
• Transformers as Intrinsic Optimizers: Forward Inference through the Energy Principle
• Quantitative Bounds for Length Generalization in Transformers
• Parallel Loop Transformer for Efficient Test-Time Computation Scaling
• To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models
• Bridging the Gap Between Latent and Explicit Reasoning with Looped Transformers
• MeSH: Memory-as-State-Highways for Recursive Transformers
• Encode, Think, Decode: Scaling test-time reasoning with recursive latent thoughts
• Coevolutionary Continuous Discrete Diffusion: Make Your Diffusion Language Model a Latent Reasoner
• Two-Scale Latent Dynamics for Recurrent-Depth Transformers
• A Formal Comparison Between Chain-of-Thought and Latent Thought
• Neural Scaling Universality: If Exponents Are Fixed, Time to Understand Coefficients
• Channel-Wise MLPs Improve the Generalization of Recurrent Convolutional Networks

Related

• What Algorithms can Transformers Learn? A Study in Length Generalization
• Quantitative Bounds for Length Generalization in Transformers
• Transformers Can Achieve Length Generalization But Not Robustly
• Improving Length-Generalization in Transformers via Task Hinting
• Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count