On Expressive Power of Looped Transformers: Theoretical Analysis and Enhancement via Timestep Encoding (2410.01405v3)

Abstract: Looped Transformers offer advantages in parameter efficiency and Turing completeness. However, their expressive power for function approximation and approximation rate remains underexplored. In this paper, we establish approximation rates of Looped Transformers by defining the concept of the modulus of continuity for sequence-to-sequence functions. This reveals a limitation specific to the looped architecture. That is, the analysis prompts us to incorporate scaling parameters for each loop, conditioned on timestep encoding. Experimental results demonstrate that increasing the number of loops enhances performance, with further gains achieved through the timestep encoding architecture.

• The paper provides a rigorous theoretical framework that quantifies the approximation rate of Looped Transformers using novel continuity metrics.
• It demonstrates that increasing loop iterations and incorporating timestep encoding significantly enhance performance in continuous function approximation.
• Empirical experiments on edit distance tasks validate the practical benefits and parameter efficiency of the enhanced Looped Transformer architecture.

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

The paper "On Expressive Power of Looped Transformers: Theoretical Analysis and Enhancement via Timestep Encoding" by Kevin Xu and Issei Sato from The University of Tokyo, conducts an in-depth theoretical analysis on the expressive power and the functional approximation capabilities of Looped Transformers. Looped Transformers are a variant of the Transformer architecture where the output is iteratively fed back into the input, leading to significant parameter efficiency and potential Turing completeness. This paper aims to bridge the gap in understanding the approximation capabilities of such architectures, especially for continuous sequence-to-sequence functions.

Background and Previous Work

Looped Transformers were first introduced to blend the strengths of Transformers and recurrent neural networks (RNNs) by \citet{dehghani2018universal}. They showed that Looped Transformers could match the performance of standard Transformers while using significantly fewer parameters by reusing computations across multiple iterations. The connection to weight-tying Transformers such as ALBERT~\citep{Lan2020ALBERT:} is also noted. Recent empirical studies by \citet{yang2024looped} and \citet{Giannou2023LoopedTA} have incremented the number of loop iterations, demonstrating improved performance on more complex tasks and theoretically proving the ability to simulate Turing machines, respectively.

The exploration of standard Transformers’ expressive power has revealed their capability as universal approximators for continuous permutation-equivariant functions~\citep{Yun2020Are}. However, the analogous properties for Looped Transformers had largely remained unexplored until now. The paper by \citet{zhang23} represents a significant stride in this direction, establishing approximation rates for simple looped ReLU networks.

Key Contributions and Theoretical Insights

In this paper, Xu and Sato introduce the notion of the modulus of continuity specifically adapted for sequence-to-sequence functions. By defining three types of continuity—sequence continuity, contextual continuity, and token continuity—they rigorously derive the approximation rate of Looped Transformers. These continuity parameters provide a quantitative measure of how input perturbations affect the output across different granularity levels within the sequence.

Main Theorem on Approximation Rate

The main theorem proves that a sufficiently deep Looped Transformer can approximate any continuous function within a defined error bound. The approximation rate heavily depends on the newly introduced moduli of continuity. Specifically, they show that the approximation error of a Looped Transformer is limited by a combination of sequence continuity, contextual continuity, and token continuity, yielding:

$\|\bm{\mathcal{L}_2\circ \mathrm{TF}^{\circ r}\circ \bm{\mathcal{L}_1-f\big\|_{L^p([0,1]^{d \times N})} \leq \omega^{tok}_f(\delta\sqrt{d}) + \omega^{cont}_f(\delta\sqrt{Nd}) + \omega_f(\delta \sqrt{Nd}) + \mathcal{O}(\delta^{d})$

for $\delta = \big((r-N)/2\big)^{-(N+1)d-1}$. Here, $\bm{\mathcal{L}_1}$ and $\bm{\mathcal{L}_2}$ represent token-wise applications of an affine transformation, and $\delta$ scales inversely with the number of loops $r$.

Implication of Timestep Encoding

The theoretical results indicate that increments in loops generally improve the approximation capability. However, the identified dependence on contextual and token continuity reveals a limitation inherent to the looped architecture. To mitigate this, the paper suggests integrating timestep-encoded scaling parameters for further refinement. By conditioning the model on timestep encodings, derived from sinusoidal positional encodings and processed via MLP to generate vector parameters, the architecture can adapt more dynamically to the iterative process, thereby effectively decoupling the dependency from continuity constraints.

Experimental Validation

The authors substantiate their theories with empirical studies on an edit distance problem. They compare Looped Transformers with varying loop counts (10, 50, 100) and demonstrate that increased loop iterations enhance performance. Additionally, models augmented with timestep encoding exhibit superior accuracy, validating the theoretical improvements. This enhanced model achieves better expressive power, emphasizing the potential practical benefits of the proposed timestep-encoded architecture.

Conclusion and Future Directions

This paper lays a robust theoretical foundation for understanding and enhancing the expressive power of Looped Transformers. The introduction of timesteps scaling parameters presents a novel way to address the architectural limitations, thus paving the way for more efficient and powerful Transformer models. Future research could explore identifying function classes with specific continuity properties to optimize architectures further. Additionally, extending these concepts to other variants of Transformers or exploring the potential of these theories in real-world applications such as NLP and beyond would be promising avenues for continued exploration.