On Expressive Power of Looped Transformers: Theoretical Analysis and Enhancement via Timestep Encoding

2024-10-02

arXiv.org

AI Summary

This work investigates the expressive power and structural bottlenecks of Looped Transformers in sequence-to-sequence function approximation. We identify a fundamental theoretical limitation arising from their fixed-loop architecture and, for the first time, establish an approximation-rate analysis framework—proving that expressive capacity strictly increases with loop count. To overcome this architectural ceiling, we propose a learnable loop-scaling mechanism grounded in timestep-aware positional encoding. Our method integrates modulus-of-continuity theory for functional approximation analysis, parameterized timestep encoding, and empirical validation. Experiments across multiple reasoning benchmarks demonstrate significant performance gains over strong baselines. Collectively, this work advances both the theoretical understanding and practical optimization of recurrent-style Transformers, establishing a new paradigm for analyzing and enhancing loop-based sequence modeling architectures.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Looped Transformers provide advantages in parameter efficiency, computational capabilities, and generalization for reasoning tasks. However, their expressive power regarding function approximation remains underexplored. In this paper, we establish the approximation rate of Looped Transformers by defining the modulus of continuity for sequence-to-sequence functions. This reveals a limitation specific to the looped architecture. That is, the analysis prompts the incorporation of scaling parameters for each loop, conditioned on timestep encoding. Experiments validate the theoretical results, showing that increasing the number of loops enhances performance, with further gains achieved through the timestep encoding.

Problem

• Transformer
• Sequence-to-Sequence
• Performance Enhancement

Innovation

• Cyclic Transformers
• Parameter Resizing
• Time Step Encoding