AlgoFormer: An Efficient Transformer Framework with Algorithmic Structures

TL;DR

AlgoFormer introduces a structured transformer framework designed to learn and execute algorithms by embedding prior algorithmic knowledge into its architecture. It decomposes into pre-, looped-, and post-transformers to handle preprocessing, iterative solving, and postprocessing, enabling efficient in-context algorithm learning. Theoretical results demonstrate that AlgoFormer can emulate gradient-descent style updates, autoregressive learning, and chain-of-thought processes, with proofs provided in appendices, and empirical results show advantages over standard and vanilla looped transformers on synthetic tasks and language benchmarks. This work highlights the potential of task-informed architectural priors to yield more efficient and capable transformers for scientific computing and natural language tasks, while acknowledging design- and scalability-related limitations and offering avenues for future automation and scaling analyses.

Abstract

Besides natural language processing, transformers exhibit extraordinary performance in solving broader applications, including scientific computing and computer vision. Previous works try to explain this from the expressive power and capability perspectives that standard transformers are capable of performing some algorithms. To empower transformers with algorithmic capabilities and motivated by the recently proposed looped transformer, we design a novel transformer framework, dubbed Algorithm Transformer (abbreviated as AlgoFormer). We provide an insight that efficient transformer architectures can be designed by leveraging prior knowledge of tasks and the underlying structure of potential algorithms. Compared with the standard transformer and vanilla looped transformer, the proposed AlgoFormer can perform efficiently in algorithm representation in some specific tasks. In particular, inspired by the structure of human-designed learning algorithms, our transformer framework consists of a pre-transformer that is responsible for task preprocessing, a looped transformer for iterative optimization algorithms, and a post-transformer for producing the desired results after post-processing. We provide theoretical evidence of the expressive power of the AlgoFormer in solving some challenging problems, mirroring human-designed algorithms. Furthermore, some theoretical and empirical results are presented to show that the designed transformer has the potential to perform algorithm representation and learning. Experimental results demonstrate the empirical superiority of the proposed transformer in that it outperforms the standard transformer and vanilla looped transformer in some specific tasks. An extensive experiment on real language tasks (e.g., neural machine translation of German and English, and text classification) further validates the expressiveness and effectiveness of AlgoFormer.

Figures (5)

• Figure 1: Algorithmic structure of the AlgoFormer. Here, $\text{TF}_{\text{pre}}$, $\text{TF}_{\text{loop}}$, and $\text{TF}_{\text{post}}$ are multi-layer transformers; "statements" represent some fundamental operations in classical algorithms.
• Figure 2: The validation error of trained models (the standard transformer, the vanilla looped transformer, and the AlgoFormer), assessed on regression with representation, AR(q) with representation, and CoT with MLPs tasks. By choosing suitable hyperparameters (i.e., we set $(T,\Delta T)=(20,15)$), the AlgoFormer has significantly better performance than the standard transformer and the vanilla looped transformer on those tasks.
• Figure 3: The validation error of trained models, evaluated on regression with representation task, with varying hyperparameters $T$ and $\Delta T$. The AlgoFormers are trained for $T$ loops, defined in Equation \ref{['empirical_loss']}, and the evaluation focuses on square loss at longer iterations, where the number of loop iterations far exceeds $T$.
• Figure 4: The validation error of trained models, evaluated on regression with representation task, with varying numbers of layers (denoted as $L$) and heads (denoted as $h$). In the context of AlgoFormer, the number of layers $L$ corresponds to the layers in the pre-, looped, and post-transformers, all of which are $L$-layer transformers. The AlgoFormers are trained with $(T, \Delta T)=(20,15)$, defined in Equation \ref{['empirical_loss']}.
• Figure 5: The validation error of trained AlgoFormer models and the linear regression models optimized by gradient descent and Newton's method. The AlgoFormers are trained with $(T, \Delta T)=(20,15)$ and $(T, \Delta T)=(10,10)$, defined in Equation \ref{['empirical_loss']}.

Theorems & Definitions (14)

• Theorem 3.1
• Theorem 3.2
• Theorem 3.3
• Theorem 4.1
• Theorem 4.2
• Lemma A.1: Quasi-orthogonal vectors
• Lemma A.2
• proof
• Lemma A.3
• proof