AlgoFormer: An Efficient Transformer Framework with Algorithmic Structures
Authors: Yihang Gao
Chuanyang Zheng
Enze Xie
Han Shi
Tianyang Hu
Yu Li
Michael Ng
Zhenguo Li
Zhaoqiang Liu
Venue: Accepted by TMLR ()

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 (Yang et al., 2024; Giannou et al., 2023), 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.

Published in Transactions on Machine Learning Research (01/2025)
AlgoFormer: An Efficient Transformer Framework with Al- gorithmic Structures
Yihang Gao
Chuanyang Zheng?* Enze Xie Zhaoqiang Liušt
Han Shi
Tianyang Hu'
Yu Li²
Michael Ng4
Zhenguo Li
(gaoyh, @nus. edu. sg {shi ha, zhenguo JChuawei. com
cuhk. edu..edu...ee edu. hk (xieenzz))connemen hku.hk
michael-ngohkbu. edu.hk zqliu120gmail. com
NNational University of Singapore "Chinese University of Hong Kong 'Huawei Noah's Ark Lab
*Hong
Kong Baptist University Uniivesiity of Electronic Science and Technology of China
Reviewed on OpenReview:
Abstract
Besides natural language processing, transformers exhibit extraordinary performance in solving broader applications, including scientific computing and computer vision. Previ- ous works try to explain this from the expressive power and capability perspectives that standard transformers are capable of performing some algorithms. To empower transform- with algorithmic capabilities and motivated by the recently proposed looped transformer (Yang et al. 2024 Giannou et al. 2023 we design a novel transformer framework, dubbed Algorithm Translormer (abbreviated as AlgoFormer). We provide an insight that efficient transformer architectures can be designed by leveraging prior knowledge of tasks and the underlying structure potential algorithms. Compared with the standard transformer and vanilla looped transformer, the proposed AlgoFormer can perform efficiently in algorithm representation in specific tasks. In particular, inspired by the structure of human- designed learning algorithms, our transformer framework consists of pre-transformer that is responsible for task preprocessing, a looped transformer for iterative optimization al- gorithms, and post-transformer for producing the desired results after post-processing. We provide thcoretical evidence of the expressive power of the AlgoFormor in solving some challenging problems, mirroring human-designed algorithuss. Furthermore, some theoretical and empirical results are presented to show that the designed transformer has the poten- tial 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 offectiveness of AlgoFormer.
1 Introduction
The emergence of the transformer architecture (Vaswani et al. 2017 marks the onset of a new era in natural language processing. Transformer-basee large language models (LLMs), such as BERT (Devlin et al. 2019
and GPT-3 Brown et al.) 2020), revolutionized impactful language-centin centric applications, including language
translation Vaswani et al. 2017 Raffel al.] 2220)), text completion/ /generatinn (Radford et al.] 2019
Brown al. 2020), sentiment analysis (Devlin al. 2201)) and mathematical reasoning (Imani et al.
2023
Yu et al. 2024). Beyond the initial surge in LLMs, these transformer-based models have found extensive
applications in diverse domains such as computer vision (Dosovitskiy et al. 2021 time series (Li et al.
2019
"Equal contribution.
i "Corresponding author.

---

Published in Transactions on Machine Learning Research (01/2025)
While many studies have concentrated on employing transformer-based models to tackle challenging real-world tasks, yielding superior performances compared to earlier models, the mathematical understanding of transformers remains incomplete.
bioinformatics (Zhang et al. 2023b), and addressing various physical problems CCoo 2021
Garg et al. (2022) empirically investigate the performance of transformers in-context learning, where the input tokens are input-label pairs generated from classical machine learning models, e.g.; (sparse) linear regression and decision tree. They find that transformers can perform comparably standard human- designed machine learning algorithms. Some subsequent works try to explain the phenomenon. [Akyürek 2023 characterize decoder-based transformer employing stochastic gradient descent for lincar ct al. regression. Bai et al. (2023) demonstrate that transformers can address statistical learning problems and employ algorithm selection, such as ridge regression, Lasso, and classification problems. Zhang et al. (023a) and Huang al. (2023) simplify the transformer model with reduced active parameters, yet reveal that the simplified transformer retains sufficient expressiveness for in-context linear regression problems. In Inh al.,
(2023), transformers are extended to implement preconditioned gradient. deescntt The looped transformer is proposed in Giannou et al. (2023), and is shown to have the potential to perform basic operations (e.g..
addition and multiplication), well as implicitly learn iterative algorithms Yang et al. (2024). More related
and interesting studies can be found in Huang et al. 2023 Von Oswald et al. (2023
Mahankali et al.
(0244)
In this paper, inspired by the recently proposed looped transformer (Yang et al. 2024 Giannou et al.
2023
we propose a novel transformer framework, which we refer to as AlgoFormer, and strictly enforce it as an algorithm learner by the structure regularization on its architecture. The transformer framework consists of three modules (sub-traanformeer)) i.c., the pro-, looped, and post-transformers, designed to perform distinct roles. The pre-transformer is responsible for preprocessing the input data, and formulating it into some mathematical problems. The looped transformer acts as an iterative algorithm in solving the hidden problems. Finally, the post-transformer handles suitable postprocessing to produce the desired results. In contrast to standard transformers, the AlgoFormer is more likely to implement algorithms, due to its algorithmic structures shown in Figure
2 Motivation and Proposed Method
In this section, mainly discuss the construction and intuition of the AlgoFormer. Its advantages over the standard transformer is then conveyed. Before going into details, we first elaborate the mathematical definition of transformer layers.
2.1 Preliminaries
A lne-r transformer mathematically formulated as:
Σ w X softmax x (w (TT w x).
Attn (X) X +
(1)
TF (X) Attn (X) W2ReLU (W, (X) + b1) b2,
where X D the input tokens; h is the number of heads; (W()) ((()) ((())) denote value, key and query matrices al i-th head, respectively; (W2, W1, by, b1} are parameters of the shallow ferd-forward ReLU neural network. The attention layer with softmax activation function mostly exchanges information between different tokens by the attention mechanism. Subsequently, the feed-forward ReLU neural network applies nonlinear transformations to each token vector and extracts more complicated and versatile representations.
2.2 Algorithmic Structures of Transformers
As discussed in the introduction, rather than simply interpreting it as an implicit function approximator, the transformer may in-context execute some implicit algorithms learned from training data. However, it is still unverified whether the standard multi-layer transformer is exactly performing algorithms.
AlgoFrrmer. As shown in the green part of Figure ! vanilla looped transformers Yaagg et al. 2024 Giannou et al. 2023 admit the same structure as iterative algorithms. However, real applications usually
2

---

Published in Transactions on Machine Learning Research (01/2025)
Inpul prompt
Fiir
Pre-transfemmer
Algorithm! 1
statements 1 for ē to T do statements 2 end for statements 3
Tresup
Loopal tranaformer
TV
Post-trassfemmrr
oulput
Figure 1: Algorithmic structure of the AlgoFormer. Here, TF pre, TF Toupp and TF Topt multi-layer trans- formers; "stalements" represent some fundamental operalions in classical algorithms.
much more complicated. For example, given a task with data pairs, a well-trained researcher may first pre- process the data under some prior knowledge, and then formulate a mathematical (optimization) problem. Following that, some designed solvers, usually iterative algorithms, are performed. Finally, the dcsired results are obtained after further proc-procsssing. The designed AlgoFormer (Algorithm Transformer). visualized in Figure ! enjoys the same structure wide classes of algorithms. Specifically, separate the transformer framework into three parts, i.e., pre-transformer TF pre: looped transformer ToFoo and post-transformer TF post. Here, those three sub-transformers are standard multi-layer transformers in Equation (1] Given the input token vectors X and the number of iteration steps T, the output admits:
TF post (TF (Toooop loop (... TF loop
Toooo
(TF pre (p)e(X)))
(2)
iterations
The looped transformer layers (TFlo share the same set of weights, they perform identical computa- tions during each iteration. In contrast, the pre-transformer (TF pre) and post-transformer (TF utilize distinct weights and architectures to handle their specific computational roles. The hyperparameters for each transformer module (e.g,, number of heads, layers, and hidden dimensions) are configured based on prior knowledge the computational complexity of the target algorithms. Although we present AlgoFormer as Equation 2 which consists of three modules, the key insight here is that the AlgoFormer can be designed flexibly based on the prior knowledge of the algorithm structure for the given task. Importantly, we do not restrict AlgoFormer to the specific form outlined in Equation (2). Compared with standard trans- formers, the AlgoFormer acts more as the algorithm learner, by strictly regularizing the loop structure. In contrast to Giannou et al. (2023) and Yang et al. (2024), which primarily focus on tasks solvable by itcrative algorithms, our approach introduces additional transformer modules (e.g., pre- and post-transformers)) for processing. These components are crucial for addressing the processing needs of real-world applications. This design makes the AlgoFormer capable of representing more complex algorithms and solving challenging tasks more efficiently. Additionally, one of the core insights of our work is that transformer architectures be designed more efficiently, flexibly, and diversely by leveraging prior knowledge and the pre-defined structure of potential algorithms. This approach enables AlgoFormer to generalize across broader range of applications while maintaining high elficiency and adaptability, such as representing algorithms involving nested loops, multiple loops, or multi- mroti-processing.
2.3 Training Strategy
Our training strategy builds upon the methodology introduced in Yang et al. (2024). Let P [x1, fm1)).... (oon))....,,,.11) 22:11,((;;11) æ; represents the input prompt for < i S N. We denote the AlgoFormer as TF A (:; θ), where f(-) a task-specific function that varies across different sequences and f indicates the number of loops (iterations) in Equation 2 and represent the transformer parameters. Instead of evaluating the loss solely on TFT Algo (-+) with T' iterations, we minimize the expected loss over averaged iteration mumbers:
T
min Ep
TF Algo Algo (P f(x:))1
(3)
T To
t=T
3

---

Published in Transactions on Machine Learning Research (01/2025)
where To max AT, 0} is the initial step for evaluating performance, T is the maximal loop iterations during training, and AT is the number of loop iterations included in the loss. To stabilize the training, we adopt the moving average over loop iterations from To to T in the loss. Incorporating the loss from To to T cnsures that the transformer module faithfully applies iterative algorithms and generalizes beyond the number of Training loops. This approach helps enhance the transformer's 's capability for lasks requiring a varying number of iterations. Experimental results in Sections 5.2 and 5.3 further demonstrate that the transformer not only applies certain iterative algorithms but also generalizes well to longer loop iterations beyond those encountered during training.
3 Expressive Power
In this section, we theoretically show by construction that AlgoFormer is capable of solving some challenging tasks, akin to human-designed algorithms. The core idea is as follows. Initially, the pre-transformer under- takes the crucial task of preprocessing the input data, such as representation transformation. The looped transformer is responsible for iterative algorithms in optimization problems. Finally, it is ready to output the desired result by the post-transformer. Through the analysis of AlgoFormer's expressive power in addressing these tasks, we expect its potential to make contributions to the communities of scientific computing and machine learning. Throughout the section, we assume that the maximal number of data samples is N and all sample observations (e.g., T2 and y1) are bounded.
3.1 Regression with Representation
We consider regression problems with representation, where the output behaves a linear function of the input with fixed representation function. Here, adopt the L-layer MLPs with (leaky) ReLU activation function as the representation function (-(-). Specifically, we generate each in-context sample by first sampling the linear weight A from the prior PA, and then generating the input-label pair ((x, y:)} with a; E [d ~ Pas y; = A& (x; + E; and E; ~ ~ (0,o21)) We aim to find the test label ytest := A (Drest)) given the in-context samples and lest data {ar "1...., NN, yN, Trast). Here, the weight matrix varies across different sequences of in-context samples but remains constant within a single sequence, and is learned in-context. The representation function Φ", on the other hand, is fixed across all samples of sequences and is learned during training. A reliable solver is expected first to identify the representation function and transform the input data æ to its representation (x))). Then it reduces to regression problem, and some optimization algorithms are performed to find the weight matrix from in-context samples. Finally, it outputs the desired result Yeest by applying transformations on the test data. We prove by construction that there exists an AlgoFormer that solves the lask, akin to the human-designed reliable solver.
Theorem 3.1. There erists a designed AlyoFormer with TFPm (an (L + 1)-lyye two-head trunsformer), Foooo (a one-layer two-head transformer), and yossst (a one-layer one-head transformer), that outputs Ad" (Itett) from the input-label pairs {I1, ",.""",,"",,.. UN, Diest) by fitting the representation. function and applying gradient descent for malli-variale regression. The emeulation of each slep is not eracl, as there is some error introduced in each step. However, the error can be made arbitrarily close to zero by increasing the temperature of the softmax and adjusting another free parameter, neither of which affects the size of the network.
Remarks. The detailed proof available in Appendix A.1 Our construction of the transformer framework involves three distinct sub-trande nub-transformers, each assigned specific responsibilities. The pre-trausformer, char- acterized by identity attention, is dedicated to representation transformation through feed-forward neural networks. This stage reduces the task to multivariate regression problem. Subsequently, the looped trans- former operates in-context to determine the optimal weight, effectively acting as an iterative solver. Finally, the post-transformer is responsible for the post-processing and generate the desired result A$* (Irest). Here,
EN
0
Etest
the input prompt to the transformer formulated as
9
YN
0
where p.
P. and P denote positional embeddings and will be specified in the proof. Due to the differing dimensions of the input x and its corresponding label y, zero padding is incorporated to reshape them into vectors of the
pf PY
PVy
P
4

---

Published in Transactions on Machine Learning Research (01/2025)
same dimension. The structure of the prompt P aligns with similar formulations in previous works (Bai
jtt al. 2023 Akyürek et al. 2023; Garg et al. 2022). For different input prompts P, the hidden linear weights A are distinct but the representation function (.) is fixed. In comparison with the standard transformer adopted in Guo et al.] (2024), which investigates similar tasks, the designed AlgoFormer has a significantly lower parameter size, making closer to the envisioned human-designed algorithm. Notably, we construct the looped transformer to perform gradient descent for the multi-variate regression. However, the trans- former exhibits remarkable versatility, as it has the capability apply (ridge) regularized regression and more effective optimization algorithms beyond gradient descent. For more details, please refer to Section
3.2 AR(q) with Representation
We consider the autoregressive model with representation. The dynamical (time series) system is generated by E++1 = A ([ 2z1) + et, where t5 (-) is a fixed representation function (e.g., we take the L- layer MLPs), and the weight A varies from different sequences but remains constant within a single sequence. Our goal is to learn 4 (.) during training and to perform in-context learning of the weight matrix A. In standard AR(q) (multivariate autoregressive) models, the representation function t* (.) is identity. Here, we investigate a more challenging situation in which the representation function is fixed but unknown. A well- behaved solver should first. find the representation function and then translate it into a modified autoregressive model. With standard Gaussian priors on the white noise Et, the Bayesian estimator of the AR(q) model parameters admits arg max A I ( (2-11.".., = arg minA Σ || - A ((ze-g... ... 21)12) where f(x, "4- , 3
is the conditional density function of Tt, given previous g observations. A practical solver initially identifies the representation function and transforms the input time series into its reprosentation, denoted as (xx)). Then the problem is reduced to an autoregressive form. Similar to the previous subsection, we prove by construction that there exists a AlgoFormer, akin to numan-designa algorithms, capable of effectively solving the given task.
Theorem 3.2. There erists a designed AlgoFormer with TFpre (a one-layer q-head transformer with an (laye)-layer one-head transformer), Taioop (a one-layer two-head transformer), and post (a one-layer one- head transformer), that predicts ØN+1 from the data sequence {x1, .NN by copying, transformation of the representation function and opplying gradient descent for multi-variate regression. The emauaation of each step is not eract, as there is some error introduced in each step. However, the error can be made arbitrarily close to zero by increasing the temperature of the softmax and adjusting another free porameter, neither of which affects the size of the network.
Remarks. The detailed proof can be found in Appendix A.2 The technical details are similar to Theorem
3.1
The input prompt to the transformer is formulated as P
x1
C2
xN
where p denote
pf
p2
PK
positional embeddings and will be specified in the proof. Additionall,, the pre-transformer copies the feature from the previous q tokens, utilizing q heads for parallel processing.
3.3 Chain-of- Thouuhh with MLPs
Chain-of-Thoughh (CoT) demonstrates exceptional performances in mathematical reasoning and text genera- tion (Wei et al. 2022). The success of CoT has been theoretically explored, shedding light on its effectiveness in toy cases (Li et al. 2023) and on its computational complexity (Feng et al.) 2023). In this subsection, we revisit the intriguing toy examples of CoT generated by leaky ReLU MLPs, denoted as CoT with MLPs, as discussed in Li et al. (2023). We bogin by constructing an L-layer MLP with Icaky ReLU activation. For an initial data point ac ~ P.,, the CoT point of represents the output of the l-th layer of the MLP. Consequently, the COT sequence {xr, s² ..... s } is exactly generated as the output of each (hidden) layer of the MLP. The implicit L-layer MLP remains the same within a single sequence but varies across different sequences, and it is learned in-context. The target of CoT with MLPs problem is to find the next state gf41 based on the CoT samples (æ1, 8|. ... sf. .22.... ,N,, s,, 8% Btest, CoT prompting of Prest- We establish by construction in Theorem 3.3 that the AlgoFormer adeptly solves
.8)). where (8).... 8²} denotes the
the CoT with MLPs problem, exhibiting a capability akin to human-designed algorithms.
5

---

Published in Transactions on Machine Learning Research (01/2025)
Theorem 3.3. There erists a designed AlgoFormer with TFp pre (a seven-layer two-head transformer). TFtoop (a one-layer bwo-hend transformer), and TF post {a one-layer one-head transformer), that finds 8 from samples {01, s1, sf. x2.... ,,,, ,',, ".....,, ,8 Drest, st...., . by filtering and applying gradient descent
for malti-variate regression. The emulation of each step is not exact, as there is some error introduced in each slep. However, the ervor can be made urbilrarily close zero by increusing the temperabure of the softmaz and adjusting another free porameter, neither of which affects the size of the network.
Remarks. We put the proof in Appendix A.3 The prentransfrrmer first identifies the positional number e, and subsequently filters the input sequence into sf SSy s² ..... sy sy .s... This filtering transfor- mation reduces the problem to multi-variate regression problem. Compared with Li et al. (2023), where assumption is made, elaborate on the role of looped transformers in implementing gradient descent. While the CoT with MLPs may not be explicilly equivalent to CoT tasks in real applications, Theorem 3.3 somewhat implies the potential of the AlgoFormer in solving CoT--rllted problems.
4 Extensions and Further Analysis
In this section, we provide complementary insights to the results discussed in Section 3 Firstly, as discussed in the remark following Theorem 3.1 we construct the looped transformer that employs gradient deseent to solve (regularized) multi-variate regression problems. However, in practical scenarios, the adoption of more efficient optimization algorithms often preferred. Investigating the expressive power of transformers beyond gradicnt descent is both intriguing and appealing. As stated in Thcorem 4.1 we demonstrate that the AlgoFormer can proficiently implement Newton's method for solving linear regression problems. Secondly, the definition in Equation [ implics the encoder-based transformer. In practical applications, decoder- based transformer with causal attention, as seen in models like GPT-2 (Radford et al. 2019), may also be favored. For completeness, it is also compelling to examine the behavior of decoder- based transformers in algorithmic learning. Our findings, presented in Theorem 4.2, reveal that the decoder-based AlgoFormer can also implement gradient descent in linear regression problems. The primary distinction lies the fact that the becoder-based transformer utilizes previously observed data to evaluate the gradient, while the encoder-baat based transformer calculates the gradient based on the full data samples.
4.1 Beyond the Gradient Descent
Newton's (second-order) methods enjoy superlinear convergence under some mild conditions, outperforming gradient descent with lincar convergence. This raiscs a natural question:
Can the transformer implement algorillmss beyond gradient. descent, including higher-orler oplimization algorithms?
In this section, we address this question by demonstrating that the designed AlgoFormer can also realize Newton's 's method in regression problems. Consider the linear regression problem given by:
N
arg min w0
(w xz - y1)² -y))²
(4)
2N
Denote X = [211 22 EN) E Nxd y [y1, U2,
_N] E Nxl and S XTX. A typical Newton's
method for linear regression problems follows the update scheme:
My = oS, where a E 00
2 |SST
Mx+1 = 2M, - M, uo
Newton M X y.
(5)
As described in Söderström & Stewart 1974]) : Pan & Schreiber 1991] the above update scheme (Newton's method) enjoys superlinear convergence, in contrast. to the linear convergence of gradient. descent. The following theorem states that Newton's method in Equation 5 can be realized by the AlgoFormer.
Thcorem 4.1.
There erists a designed AlgoFormer with TF Fre (a one-lager two-head transformer). TFooop (a one-layer two-head transformer). and Topt (a two-layer two-head transformer). that implements Newton's method described by Equation in solving regression problems. The emulation of each step is not exact,

---

Published in Transactions on Machine Learning Research (01/2025)
there is some error introduced in each step. However, the ervor can be made arbitrarily close to zero by increasing the lemperature. of the sofbmax and adjusting onother free parameter, neilher of which affects the size of the network.
Remarks. The proof can be found in Appendix A.4 The pre-transformer performs preparative tasks, such copying from neighboring tokens. The looped-transformor is responsible for updating and calculating Myæ; for each token at every step k. The post-transformer compute the final estimated weight wi Newton and outputs the desired the results w.Nwwoon Ttest, where 7' is the iteration number in Equation 2 and Equation 5 In a related study by Fu et al. (2023). similar topics are explored, indicating that transformers exactly perform higher-order optimization algorithms. However, our transformer architectures differ, and technical details are distinct.
4.2
Decoder-based Transformer
In the preceding analysis, the encoder-based AlgoFormer (with full attention) demonstrates its capability solve problems by performing algorithms. Previous studies (Giannou et al. 2023 Bai et al.] 2023: Zhang et al. 2023a Huang et al. 2023 Ahn et al 2023) also focus the mncoder-baset models. We opted for encoder-based translormer because full-batch data is available for estimating gradient and Hessian information. However, in practical applications, decoder-based models, like GPTT2, are sometimes more provalent. In this subsection, we delve into the performance of the docoder-based model when exccuting iterative optimization algorithms, such as gradient descent, to solve regression problems.
We consider the linear regression problem in Equation Due the limitations of the decoder-based transformer, which can only access previous tokens, implementing iterative algorithms based on the entire batch data is not feasible. However, it is important to note that the current token in Secoder-based transformer can access data from all previous tokens. To predict the label based on the input prompt p 1i " tn æi], the empirical loss for the linear weight at æ; is given by
w² E arg min (w: Pt) = 2(i E w
(w x - y3)
(6)
In essence, the linear weight is estimated using accessible data from the previous tokens, reflecting the restricted information available in the decoder-based transformer.
Theorem 4.2. There erists a designed AlgoFormer with TFpre (a one-layer two-head transformer), tooo (a one-layer two-head transformer), and Toss (a two-layer two-head transformer), that oulpuls wif w for each input data æi, where w comes from arg min L (w; P') ofter T steps of gradient descent. The emulation of each step is not exact, as there is some error introduced in each step. However, the error can be made arbitrarily close to zero by increasing the temperature of the softmax and adjusting another free parameler, neilher of which affects the size of the nelwork.
Remarks. The detailed proof is available in Appendix
The technical details closcly resemble those in Theorem 3.1 with the key instinntion being that the decoder-based transformer can solely leverage data from previous tokens to determine the corresponding weight w'. Our findings align with those in Guo et al. (2024). although our transformer architectures and attention mechanisms differ. In a related study by Akyürek et al. (2023), similar topics are explored, demonstrating that the decoder-basse transformer performs single-sample stochastic gradient descent, while our results exhibit greater strength with w arg min L (wr; Pt). The construction of a decoder-based transformer for representing Newton's method is more challenging. We left it potential topie for future investigation.
5 Experiments
In this section, we conduct. a comprehensive empirical evaluation of the performance of AlgoFormer in tackling challenging tasks, specifically addressing regression with representation, AR(q) with representation, and CoT with MLPs, as outlined Section 2 Additionally, we also implement AlgoFormer on the neural machine translation of German and English and AG News classification, demonstrating its expressiveness and iveness in real-world language tasks.

---

Published in Transactions on Machine Learning Research (01/2025)
10a
Standard
Suazzrdd
Stansard
10-1
square error
10-1
square error
10-1
square
10-2
10-2
10
10a²
10
15
20
25
10 15
in-context samples
in-context samples
in-context samples
(a) Regression with reprosentation
(b) AR(()) with representation
(c) CoT with MLPs
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, AT) = (20, 15)), the AlgoFormer has significantly better performance than the standard transformer and the vanilla looped transformer on those tasks.
Figure 2:
5.1 Experimental Settings and Hyperparameters
Experimental settings. In all experiments, we adopt the fecoder-based AlgoFormer, standard transformer (GPT-2), and vanilla looped transformer Yang et al. (2024). For synthetic tasks, we utilize N = 40 in in-context samples as input prompts and d = 20 dimensional vectors with D = 256 dimensional positional embeddings for all experiments. To ensure fairness in comparisons, all models are trained using the Adam optimizer, with learning rate 7 le 4 and total of 500K iterations to ensure convergence. The standard transformer is designed to have L = 12 layers while pre- looped and post-trausformers are all implemented in one- layer. The default setting for the AlgoFormer, as well as the vanilla looped transformer, involves setting (T) = (20, 15). Here, the prompt formulation and the training loss in Equation (3) may slightly differ for different tasks. For example, in the AR(q) task, the input prompt is reformulated as P = [x1,''. 1; æ;,] and 0i+1- = [(0
2x:)) is the target for prediction. However, the training strategy can be easily transmitted to other tasks. Here, both the iteration numbers To and T are hyperparameters, which will be analyzed in the next subsection.
Regression with representation. In this task, we instantiate 3-layer leaky ReLU MLPs, denoted Φ- (-), which remains fixed across all tasks. The data generation process involves sampling weight matrix A IR 1x20, Subsequently, input-label pairs {< ]/ are generated, where æ; N (0, 22)), ~ (0, a²) and y; = Aab"(x;)+>> In Figure 2a we specifically set o = O. Additionally, we explore the impact of different noise levels by considering a = 0.1 and a = 1. Here, our larget to learn the representation function (-(() during training and to perform in-context learning of the weight matrix A.
AR(q) with representation. For this task, we q 3 and employ 3-layer leaky ReLU MLP denoted as t" (*), consistent across all instances. The representation function accepts 60-dimensional vector as input and produces 20-dimensional feature vectors. The time series sequence { is gener- ated by initially sampling A ~ (O, I20x20). Then the sequence is auto-regressively determined, with 0++1 = A ([x(+1-q:'-. xx)) + Es, where E, ~ (0, I20). Here, our target is to learn the representation function (-) during training and to perform in-context learning of the weight matrix A.
CoT with MLPs. In this example, we generate a 6-layer leaky ReLU MLP to serve as CoT sequence generator, determining the length of CoT steps for each sample to be six. The CoT sequence, denoted {a, .²) is generated by first sampling æ ~ N (0, I20), where sf E 2o20 represents the intermediate state output from the l-th layer of the MLP. Here, our target is to perform in-context learning of the generator function (i.e., the 6-layer leaky ReLU MLPP)
5.2 AlgoFormer Exhibits High Expressiveness
In this subsoction, we conduct a comparative analysis of the AlgoFormer against the standard and vanilla looped transformers across challenging tasks, as outlined in Section Figure@illustrater the validation error trends, showcasing a decrease with increasing number of in-context samples, aligning with intuition.
8

---

Published in Transactions on Machine Learning Research (01/2025)
Crucially, the AlgoFormer consistently outperforms both the standard and the vanilla looped transformer across all tasks, highlighting its superior expressiveness in algorithm learning. Particularly in the CoT with MLPs task, both the AlgoFormer and the standard transformer significantly surpass the vanilla looped transformer. This further underscores the significance preprocessing and postprocessing steps handling complex real-world applications. The carefully designed algorithmic structare of the AlgoPormer emerges as an effective means of structural regularization, contributing to enhanced algorithm learning capabilitics.
5.3
Impact of Hyperparameters
In this subsection, we conduct the empirical analysis of the impact of the hyperparameters on the AlgoFormer.
1.2
1.0
T-12
1.0
T-20
T-20
square error
0.8
T-522
T=50
square error
Tu50
T:30
0.6
0.6
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
25
100 125
200
150 275 200
75
100 150 175 200
loop
loop iteration
loop
(a)
AT 5
(b) AT 10
(c) AT 15
Figure The validation error of trained models, evaluated on regression with representation task, with varying hyperparameters T and AT. The AlgoFormers are trained for loops, defined Equation 3 and the evaluation focuses on square loss at longer iterations, where the number of loop iterations far exceeds T.
Stadtr:: Ln3
Standar L-12
Square error
10
Sundarc 1-14
Standars
square error
10-1
10-2
10
Twed Iwo
10-3
10
h-
10
30
35
15
40
context samples
cantext samples
(a) Varying numbers of layers
(b) Varying numbers of heads
Figure The validation error of trained models, evaluated on regression with representation task, with varying numbers of layers (denoted L) and heads (denoted 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, (20, 15), defined in Equation
10a
104
Nowion Ratws
square errоr
100
Nostan hawdd
Newto Bar=20
square afror
Neatin h5e-15
109
Newnos Bee-16
Weaton Ite
Neetos ter 20
CD lure
square
@ Roww
10
lar-25
100
lam15
10-2
Agatomer T-20
Agstomer T-20
10-2
10-2
10a
in-context samples 25
25
25
In-context samples
In-context samples
(a) Linear regression,
(b) Linear regression, 0.1
(c) Linear regression, 1.0
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, AT) = (20, 15) and (T, (T) = (10, 10), defined in Equation 3
Loop iterations. We conduct comprehensive experiments on the AlgoFormer with varying loop numbers, solving the regression with representation task. The results highlight the crucial role of both 7 and AT

---

Published in Transactions on Machine Learning Research (01/2025)
in the performance of the AlgoFormer. It is observed that a larger AT contributes to the stable inference of transformers. Comparing Figure 3a with Figures 3b and 3c it is evident that larger AT enhances stable long-term inference. The number of loop iterations T determines the model capacity in expressing algorithms. However, it is important to note that there exists a trade-off between the iteration numbers (T) and computational costs. Larger (T, AT) certainly increases model capacity but also leads to higher computational costs and challenges in model training, as reflected in Figure B
Number of heads and layers. In our experiments the AlgoFormer, we vary the numbers of heads and layers in Toop while addressing the regression with representation task. The results reveal a consistent trend that an incrcase in both the number of heads and layers leads to lower errors. This aligns with our intuitive understanding, as transformers with greater numbers of heads and layers exhibit enhanced expressiveness. However, a noteworthy observation is that 4-layer and 16-head transformers may exhibit suboptimal performance, possibly due increased optimization challenges during model training. This finding underscores the importance of carcfully selecting the model size, as a larger model, while offering higher expressiveness, may present additional training difficulties. The visualized results are shown in Figure Moreover, compared with the standard transformer, even with the same number of layers, the AlgoFormer cxhibits better performance in those tasks, mainly due to the introduced algorithmic structure. This finding highlights the role of the structure regularization of the model. Therefore, we have reasons to believe that the good performance of the AlgoFormer not only comes from the higher expressiveness with deeper layers but also from the structure regularization of model architecture, which facilitates easier training and good generalization.
In Figure 8 we also provide insights into the behavior of AlgoFormer during inference. Although the Al- goFormer is trained with T-20 loop iterations, the inference is conducted with siguificantly longer loop iterations (e.g., T=200), As T increases beyond the training length, AlgoFormer demonstrates improved performance and eventually stabilizes (i.e., converges). This behavior suggests that the inner looped trans- formers are effoctively performing itcrative computations that converge to a stable solution.
5.4 AlgoFormer and Hluman-Designed Algorithms
In this subsection, we compare the AlgoFormer with Newton's method and gradient descent in solving linear regression problems. We adopt the same default hyperparameters, with their selection grounded in a comprehensive grid search.
As illustrated in Figure 5 we observe that in the noiseless case, the AlgoPormer outperforms both Newton's method and gradient descent in the beginning stages. However, Newton's 's method suddenly achieves nearly zero loss ( machine precision) later on, benefiting from its superlinear convergence. In contrast, our method maintains an error level around le- 3. With increasing noise levels, both Newton's mothod and gradient descent converge slowly, while our method exhibits betier performance.
Several aspects contribute to this phenomenon. Firstly, in the noiseless case, Newton's method can precisely recover the weights through the linear regression objective in Equation 6 capitalizing its superlinear convergence. On the other hand, the AlgoFormer operates as a black-box, trained from finite data. While we demonstrate good model expressiveness, the final generalization error of the trained transformer results from the model's expressiveness, the finite number of training samples, and the optimization error. Despite exhibiting high expressiveness, the trained AlgoFormer cannot eliminate the last two errors entirely. This observation resonates with similar findings solving partial differential equations (Raissi et al. 2019 and
large linear systems (Gu& & Ng. 2023) using deep learning models.
Secondly, with larger noise levels, Newton's method shows suboptimal results. This is partly due to the inclusion of noise, which slows down the convergence rate, and Newton's method experiences convergence challenges when moving away from the local solution. In terms of global convergence, the AlgoFormer demonstrates superior performance compared to Newton's method.
Human-designed algorithms, backed by problem priors and precise computation, achieve irreplaccable per- formance. It's important to note that deep learning models, including transformers, are specifically designed for solving black-box tasks where there is limited prior knowledge but sufficient. observation samples. We
10

---

Published in Transactions on Machine Learning Research (01/2025)
expect that transformers, with their substantial expressiveness, hold the potential to contribute to designing implicit alms in solving scientific computing tasks.
5.5 Applications to Language Tasks
In this subsection, extend the evaluation of the proposed AlgoFormer to real-world language tasks, com- plementing its performance in real applications. Specifically, we focus on Neural Machine Translation using the IWSLT 2015 German-English dataset. The experimental setup includes standard Transformer with 12 layers, 8 attention heads, feature dimension of 256, and a learning rate of 5e-5. The pre-, looped, and post-transformers are all implemented as single layers, with T set to 10 and AT to 10 (see the hyperpa- rameters in Equation 3 for training). The input German text treated as prefix, and the output English text is generated autoregressively using decoder-based transformers for all three models. The translation performance evaluated using cross-entropy loss, where lower values indicate better results. Additionally, BLEU (Bilingual Evaluation Understudy) is widcly used metric for mcasuring the quality of translation tasks, with higher scores indicating better performance. The results are presented in the table below:
Model
Standard
Looped
AlgoFormer
Cross entropy
4.99
4.73
4.61
BLUE
9.30
10.56
14.72
Table 1:
Quality of Transformer models machine translation.
We also implement the proposed AlgoFormer on the text classification task using various datasets (AG News, IMDB, DBPedia, Yelp Review, and Yahoo News) to evaluate its performance on a real-world language application. The experimental setup includes a standard Transformer with 12 layers, 1 attention head, a feature dimension of 32, and a learning rate of 1e-3. The pre-, looped, and post-transtormers are all implemented as single layers with one attention head, with (T, = (10, 10). The input news text data is processed using encoder-based transformers, and the output classification generated. Classification accuracy, where higher accuracy indicates better performance, is used the evaluation metric. The results
are summarized in the table below:
Model AG News
Standard
Looped
A AlgoFormer
92.07
91.04
97.92
IMDB
75.00
77.50
77.50
DBPedia
99.21
99.11
98.21
Yelp Review
66.25
53.57
57.50
Yahoo News
73.96
69.79
81.25
Table 2: Accuracy (%) of Transformer models text classification across different datasets.
As shown in the results, AlgoFormer outperforms and performs comparably with the standard Transformer and the vanilla looped Transformer the Neural Machine Translation and Text Classification lasks, sug- gesting that conventional models may have inherent redundancies. This aligns with recent findings on the rodundancy in large language models (Chen et al. 2024 Frantar & Alistarh 2023 Xia. et al. 2024). These preliminary results indicate the potential of AlgoFormer, which is designed as algorithric conductor, in real-world language tasks.
5.6 Computation Comparison
The additional computational costs and complexity introduced by AlgoFormer over the standard transformer and the vanilla looped transformer depend on the specific structure of the AlgoFormer framework. In our experimental settings, these costs are negligible. Compared to the standard transformer, the total computation overhead is primarily determined by the loop iterations T and the number of loops AT' included in the training loss. Since T our experiments is comparable transformer, the dominant additional computation comes from AT. This overhead can be mitigated by using smaller batch sizes, ensuring that the proposed method does not introduce significant computational
the number of layers in the standard
11

---

Published in Transactions on Machine Learning Research (01/2025)
overhead compared to the standard transformer. The main complexity lies in the training strategy outlined in Equation 3 which ensures stable training.
When compared to the vanilla looped transformer, the computational overhead comes from the pre- and cost-transformer modules. However, if these modules are similar size to the looped transformer, the additional computational cost negligible, especially when the loop iteration T much greater than the number of layers in these modules. For instance, in our experiments, we set the pre, looped-, and post- transformer layers to 1, with T 20 during training. In this scenario, the additional computation introduced by the pre- and testtraanformer modules constitutes approximately 1/10 of the total computation. During inference, where T is typically much larger than during training (e.g., T-200 loop iterations in Figure 3 the additional computational overhead is further reduced to about 1/100.
6 Conclusion and Discussion
In this paper, we introduce novel AlgoFormer, an algorithm learner designed from the looped trans- former, distinguished by its algorithmic structures. We provide an insight that the efficient transformer framework can be designed by considering the prior knowledge of the task and the structure of potential algorithms. Comprising three sub-transformers, each playing distinct role in algorithm learning, the Al- goFormer demonstrates expressivenes and efficiency while maintaining low parameter size. Theoretical analysis establishes that the AlgoFormer can tackle challenging in-context lcarning tasks, mirroring human- designed algorithms. Our experiments further validate our claim, showing that the proposed transformer outperforms both the standard transformer and the vanilla looped transformer in specific algorithm learning and real-world language tasks.
While the proposed AlgoFormer framework offers flexibility and efficiency, several potential limitations and challenges warrant discussion.
First, the dependence on prior knowledge for designing concrete architecture. The design of the AlgoFormer framework relies heavily prior knowledge of tasks and the structure of potential algorithms. In scientific computing tasks, where problems are often well-formulated and informed by domain knowledge, this reliance can be an advantage, enabling the design of task-specific, efficient architectures. However, for real-world language lasks, the abstract and unstructured nature of the problems often limits the availability of such prior knowledge. This can make it challenging to design an optimal AlgoFormer architecture for these tasks. While the insights of AlgoFormer can significantly benefit the scientific computing community, extending its design principles to language tasks may require more advanced techniques, such automated archilecture search or meta-learning, to determine suitable architectures without extensive prior knowledge.
Second, the training complexity of AlgoFormer with more complicated algorithmic structures. Although incorporating more complex algorithmic structures into AlgoFormer, such as nested looped transformers, enhances its expressiveness, training AlgoFormer can become more challenging. inevitably introduces additional computational costs and requires adjustments to the training strategy. These challenges could limit the usability of AlgoFormer in scenarios where algorithmic structures are complicated to model.
Third, the scalability of AlgoFormer to large-scale applications. The scaling behavior of AlgoFormer has not been thoroughly investigated in this paper. The tasks studied are primarily hand-generated scientific problems or toy cases, and the language task considered is of medium scale. Similarly, the model sizes used in this study are relatively small compared to state-of-theeare large language models (LLMs). To fully assess the potential of AlgoFormer, future research should explore its scalability to larger, more challenging tasks and datasets, as well as the associated computational and memory requirements when deploying AlgoFormerr at scale.
Acknowledgements
M. Ng's rescarch partly supported by the National Key Research and Development Program of China under Grant 2024YFE0202900; HKRGC GRF 17201020 and 17300021, HKRGC CRF C7004-21GF, and Joint NSFC and RGC N-HKUU69//11
12

---

Published in Transactions on Machine Learning Research (01/2025)
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement prc- conditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 2023. URL http:s///pppaarrevee....ee//o/aa wttp:e/.orrrmmdddaaaaaaa.aa....
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Repre- senlations, 2023. URL httppp//ppeaarreeeeeeeeerrrrrrro
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in- context learning with in-context algorithm selection. Advances in Neural Information Processing Systems, 2023. URL
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind
Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners.
Advances in Newral Information Processing Systems, 33:1877 1901, 2020. URL htttss://proceednggs
Shuhao Cao. Choose transformer: Fourier or galerkin. Advances in Neural Information Pro-
cessing Systems, 34:24924 24940, 2021. URL
html
Xiaodong Chen, Yuxuan Hu, and Jing Zhang. Compressing large language models by streamlining the unimportant layer. arXiv preprint arXiv:2403.0 19135, 2024.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidircctional transformers for language understanding. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, 1:4171-4188 2019. URL htttps:/
-14-13
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition scale. In International Conference Learning Representations, 2021. URL httppp///pppprrooee..ee/orraaaaa
Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Newral Information Processing Systems, 2023. URL htttpp///peeneeeesseeeeerrrrurtp http:///oorrrooo.dddggrrr//AAtt
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pp. 10323-10337. PMLR, 2023.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiw:2310.11000 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in- context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35: 30583-30598, 2022. URL http::///paaaeeeeee.eee/ortttttr net./forumPdd=f11222eee
Angeliki Giannou, Shashank Rajput, Yong Sohn, Kangwook Loc, Jason D. Lee, and Dimitris Papailiopou- los. Looped transformers as programmable computers. In International Conference on Machine Learning, volume 202, pp. 11398-11444 PMLR, 2023. URL htttttts//ppooeedinggs mlr. press/ v202 2w.00/2giannou2320 html
Yiqi Gu and Michael K Ng. Deep neural networks for solving large linear systems arising from high- dimensional problems. STAM Journal on Scientific Computing, A535): 2023.
13

---

Published in Transactions on Machine Learning Research (01/2025)
Tianyu Guo, Wei Hu, Song Mei, Huan Wang. Caiming Xiong, Silvio Savarese, and Yu Bai. How do transform- ers learn in-context beyond simple functions? case study on learning with representations. International Conference on Learning Representations, 2024. URL
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context. convergence of transformers. ar preprint arXiv:2810.05244 2023.
Shima Imani, Liang Du, and Harsh Shrivastava. MathPrompter: Mathematical reasoning using large lan- guage models. In Proceedings of the 61st Annual Meeting of the Association for Computational Lin- guistics (Volume 5: Industry Track), pp. 37-42. Association for Computational Linguisties, 2023. URL
attps://aclanthollogy org/20233 acl- industry.u-- .4
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy, The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 2023. URL
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. Advances in Newral Information Processing Systems, 32, 2019. URL. htttps://papers. neapperrr.euuipp.... www..eeuw....ppppeeeeeee wwtpecep_aeeeeeeeeeeeeee fiile// www.rabt.ttttta...
Yingeong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting thain-of-thooght:: Compositionality through in-context filtering and lcarning. In Advances in Neural
Information Processing Systems, 2023. URL, http:///ppppeeeeeeeeeee/rrrrd//e
Arvind Mahankali, Tatsunori B Hashimolo, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. International Conference on Learning Representations, 2024. URL
Santiago Ontanon, Joshua Ainslic, Zachary Fisher, and Vaclav Cvicok. Making transformers solve composi- tional tasks. In Proceedings of the 60th Annual Meeting of the Association for Compulational Linguistics (Volume 1: Long Papers), pp. 3591-3607, 2022. URL
wrg.:/22...200222..... acl-long -ong 251
Victor Pan and Robert Schreiber. An improved newton iteration for the generalized inverse of a matrix, with applications. STAM Journal on Scientific and Statistical Computing, 12(5): 1009):100-1130, 1991.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL https:
//openreview.o prrpenrrrieennnt/foouuiieeeeee pett/forrunidddBBspPpPCC
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
Colin RAffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-lo-text trans- former. Journal of Machine Learning Research, 21(1):5485 5551, 2020. URL https://jmlr: org/papers/
volume21/20-074/20 01100000/220077221000 pdf
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learn- ing framework for solving forward and inverse problems involving nonlinear partial dilferential equations. Journal of Computational Physics, 378:686- 707, 2019.
Torsten Söderström and GW Stewart. On the numerical properties of an iterative method for computing the peoro-e generalized inverse. STAM Journal Numerical Analysis, 11(1):61-74, 1974.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances Neural Information Pro- cessing Systems, 30. 2017. URL Abstract html 3t5
14

---

Published in Transactions on Machine Learning Research (01/2025)
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, An- drey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Interna- tional Conference on Machine Learning, pp. -51144. PMLR, 2023. URL https://proceedingg
mlr. vaaass//202/vonoossaaaa 201/20//vooooo0
html
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thoughl prompling elicits reasoning in large language models. Advances in Neural Information Processing Systems, 01224244244832 2022. URL https://openrevienn wtw.://wwee//ooooiii--e- Latpiidd__viiMBB______
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared LLaMA: Accelerating language model pro-training via structured pruming. In The Twelfth International Conference Learning Representations, 2024. URL htctps///ppeerreveeeeeeeeefrrrre
Liu Yang, Kangwook Lee, Robert D Nowak, and Dimitris Papailiopoulos. Looped transformers are better learning learning algorithms. In International Conference Learning Regresenlations, 2024. URL
https::/openreviie met:/oorumid
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. International Conference on Learning Representations, 2024. URL https:
//openreview.net/ hwt.:/reww....oommmm////nhhhhhaa
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprinl 2rx6: 09927, 2023a.
Shnang Zhang, Rui Fan, Yuti Liu, Shnang Chen, Qiao Liu, and Wanwen Zeng, Applications of transformer- based language models in bioinformatics: a survey. Bioinformatics Advances, 3(1)vbba000,, 2023b.
15

---

Published in Transactions on Machine Learning Research (01/2025)
A Technical Proofs
In this section, provide comprehensive proofs for all theorems stated in the main content.
Notation. We use boldface capital and lowercase letters to denote matrices and vectors respectively. Non- bold letters represent the elements of matrices vectors, or scalars. For example, Aij denotes the (i,5)-th element of the matrix A. We use || 2 to denote the 2-norm (or the maximal singular value) of a matrix.
A.1 Proof for Theorem 3.1
Positional embedding. The role of positional embedding is pivotal in the performance of transformers. Several studies have investigated its impact on natural language processing tasks, as referenced in (Kazem- nejad al., 2023 Ontanon al. 2022 Press et al., 2022). In our theoretical construction,
deviate from empirical settings by using quasi-orthooonal vectors positional embedding in each token vector. This choice, also employed by Li et al.
(2023) Giannou et al. (2023), is made for theoretical convenience.
Lemma A.1 (Quasi-orthogonal voctors). For any fixed e > O, there erists set of vectors {Pi, P2. .pr of dimension (log N) such that p, p; P; P, +e e for all i j.
Before going through the details, the following lemma is crucial for understanding transformers as algorithm learners.
Lemma A.2. A one-layer two-head transformer ezhibits the capability to implement single step of gradient descent. in multivariate regression.
Proof. Let us consider the input prompt with positional embedding follows:
æ1
Trest
0
0
y/
YN
0
0
N
Ag
1/1
A N
UN
A
Frest
0
:=
1
0
0
0
0
0
1
0
0
0
0
1
where the 0-1 indicators are used to identify leatures, labels, and the test data, respectively. We denote the loss function for the multi-variate regression given samples {x1, M1.". EN yN as
1
((A)
|Ax, - 19112
2N
j=1
then
N
1
(At)
(A yj) xJ.
0A
N
Now, let us define
WqP P
CET
0
C
0
Caetett
0
1
1
1
1
1
WK p=
N
01
0
N
0
0
C
and
Wy c
A
y1
A
UN
Arersst
0
16

---

Published in Transactions on Machine Learning Research (01/2025)
for some scalers C, 0. Here, denote
#0 #1
N
w Ptest
Z : W
N NEN
0
N
x.
NNNN
N a Testt
0
0
0
0
C
C
C
C
C
C
then
ec softmax (Z2i- ~ +
N ah z/j-
where the two sides of the "se" can be arbitrarily close if C 0 is sufficiently large and c 0 is sufficiently small. The constant here can be canceled by introducing another head. Therefore, the output of the attention layer is
E w P. softmax (PTWk) W P)
= cl
aC 8A
(A
8A (A
O
0C 8A
Trest
The transformer layer's output, afler passing through the leed-forward neural network, expressed as:
0
0
0
0
Y1
YN
0
A
Ag+10
3/1
0
A
YN
0
A Ag+1
Trest
0
0
0
1
1
0
0
O
0
0
0
0
O
0
1
This signifies the completion of one step of gradient descent with A A -JAA (Ax) and positive step size 7) O.
We start by showing that L-layer transformer represent L-layer MLPs. It is observed that the identity operation (i.e., Attn (X) X) be achieved by setting Wv = due to the residual connection in the attention layer. Each feed-forward neural network in transformer layer can represent. a one-layer MLP. Consequently, the representation function (*(-) can be realized by T--layer transformers. At the output layer of the L-th layer transformer, let Ag be an initial guess for the weight. The current output loken vectors are then given by:
Proof for Theorem 3.1
s (I1)
0
O" (xN)
+ (Irest)
0
0
71
UN
0
Ao
(x1)
O
Ao t-
(ØN)
0
A t- 1- (xrest)
0
P1
P1
PN
PN
PN+1
PN+1
=
0
N
N
0
1
1
0
0
0
0
0
0
1
Here, the set of quasi-orthogonal vectors (PI..... .N4+1} is generated, according to Lemma transformer layer is designed to facilitate the exchange of information between neighboring tokens. Let
A.
The next
WxP=Wap
PI
P Pn
PN PN+1 PN+1
and
WyP
y1
O UN
0
17

---

Published in Transactions on Machine Learning Research (01/2025)
then
WvP. softmax W
W 2
2n 301
JUN NUN 0 0
where the two sides of the "z" can be arbitrarily close the temperature of the softmax function is sufficiently large, due to the nearly orthogonality of positional embedding vectors. It's important note that the feed- forward neural network capable of approximating nonlinear functions, such as multiplication. Here, we construct a shallow neural network that caleulates the multiplication between the first d elements and the value N in each token. Passing through the wordad neural network together with the indicators, we obtain the final output of the first (T. + llyye transformer TF pre:
& (
Φ (ØN)
0
4' (Xrrst)
0
y/
YN
0
0
(211
(xN)
0
0
N
Ao
(x1)
y/1
Ao 2
(x:N)
UN
O
Ao dh (Drest)
0
P1
P1
PNV
PN
PN+1
PN+1
0
0
1
0
0
1
0
1
100
0
0
0
1
0
0
0
0
1
According to the construction outlined in Lemma A.2 there exists one-layer, two-head transformer Tooo independent of the exact values of imput data samples (tokens), that implement gradient descent for finding the optimal weight A in the context of multivariate regression. The optimization aims to minimize the following empirical loss:
N Σ AA (æ,) -vell
min A 2N
After k-steps of looped transformer TF loop; which corresponds to applying steps of gradient descent, the resulting token vectors follows
4
(x1)
Φ- (NN)
0
Φ- (xtest)
Y1
YNN
(xt)
0
6 (WN)
0
A,0
31
0
A
(EEN)
UN
0
AF (Deest)
0
P1
Pi
PN
PN
PN+1
PN+1 0
1
0
0
0
1
0
0
0
0
0
These token vectors then ready for processing by the output transformer layer TF post The post- transformer is designed to facilitate communication between the last two tokens and position the desired result A² (Bees) in the appropriate position. We can similarly set
WKP WoP
P1 Pt
PN
PN PN+1
PN+1 ].
WvP A (x))
91
A 4' (œN) UN
Ag (Xrest)
and pass it through the feed-forward neural network. This results in the final output:
Φ (xt)
0
4 (xN)
0
4- (Itest)
0
].
0
U1
0
UN
O
A (Irest
which completes the proof.
18

---

Published in Transactions on Machine Learning Research (01/2025)
A.2 Proof for Theorem 3.2
The following lemma highlights the intrinsic "copying' capability of translormers, pivotal feature for autoregressive models, especially in the context of time series analysis. Lemma A.3. A one-layer transformer with y heads possesses the ability to efeecively copy information, from the previous q tokens to the present token.
Proof. We construct the input prompt with positional embedding as follows:
Pu
P1
P2
PN
p
Po
P
PN-1
p-q
P1-g
P2-q
PN-9
0
2
N
The i-th head aims to connect and communicate the current token with the previous i-th token. Specifically, we let.
W
PN ].
Po
P1
P2
p-
Pi-i
P2-1
PN-i
and
1 P=
30
21
D2
NN
we have
W (1) P softmax P
W
w
P
21
22
NN ]-
Here, we to mask some unimportant Loken values. Therefore, the g-head attention layer outputs
3
32
NN
DD
NN 21
NN
To
IN
2
N
Here, the samples are only supported {x,} with 0 0st5 N. It is common to alternatively define æ, = 0 for t < O. Passing through the feed-forward neural network, together with the indicators at the last row, we can filter out the undefined elements, i.e.,
32
ØN
ED
21
ØN-1
0
0
20
EN
0
0
C NN
2
Proof for Theorem 3.2 According to Lemma and Theorem 3.1 is able to do copying and transformation of the representation function. The output after the preprocessing is given by
the (L+2)-layer pre-transformer TF pre
ED
I1
C2
ON
( (#1-g:0) g:
(x2-g:1)
4 (as (13--:2) 9:2)
$
(#N+1--:N)
where we denote aay as the concatenation [x,, 2:41. tion in Lemma A.2. a one-layer two-head transformer TF loop capable of implementing gradient descent on the multivariate regression. Finally, the post-transformer TF post moves the desired result to the output.
2zJ for notational simplicity. Similar the construc-
19

---

Published in Transactions on Machine Learning Research (01/2025)
A.3 Proof for Theorem 3.3
According to Lemma 5 in Li et al. 2023 seven-layer, two-head pre-transformer TF pre is introduced for preprocessing the input CoT sequence. This pre-transformer performs filtering, transforming the input sequence into the structured form given by Equation 7
0
0
82
0
0
8
0
0
8
0
8
0
0
0
1
0
1
0
0
0
1
(7)
0
0
1
0
1
0
0 1
0 0
0
0
0
Specifically, it identifies the positional index 8-1 of the last token 8 retains only sf- and 8f for <is N, and filters out all other irrelevant tokens. this context, the representation function Φ-(-) corresponds to L-layer leaky ReLU MLPs. Notably, the transformation si =o (W' oW'''' the weight matrix the Ł-th layer, and af.) represents the leaky ReLU activation function. Given the reversibility and piecewise linearity of the leaky ReLU activation, that of = s in Equation Consequently, the problem is reduced to multi-variate regression, and a one-layer two-head transformer TF is demonstrated to effectively implement gradient descent for determining the weight matrix W as shown in Lemma A.2 Subsequently, the post-transformer TF post produces the desired result a ( s(--)
expressed, where W denotes
assume, without loss of generality,
A.4 Proof for Theorem 4.1
In this section, we first show that the one-layer two-head transformer can implement single step of Newton's method in Equation 5 with the special form of input token vectors. Then, we introduce the pre-transformer, designed to convert general input tokens into the prescribed format conducive to the transformer's 's operation. Finally, the post-transformet facilitates the extraction of the desired results through additional computations, given that the output from the tropeeetranssormer corresponds to an intermediate product.
Lemma A.4. A transformer with one layer and two heads is capable of implementing step of Newton's method in the linear regression problem in Equation 4
Proof. Let us consider the input prompt with positional embedding follows:
31
ØN
Frest
0
y/1
0
UN
0
M
0
M.
0
0
P:
0
0
0
0
1
0
Lot
ME
O
MEØN
0
0
0 1
WoP
1
1
1
1
1
WkP-
CCC
0
N
0
Carrest
0
0
0
0
C'
and
Wyp= M
MENN
0
0
0
Similarly, denote Z = PT W WoP, we can establish that
é softmax (72i-1,23 z ca. MLng.
20

---

Published in Transactions on Machine Learning Research (01/2025)
To mullify the constant term, an additional attention head can be incorporated. Therefore, the output takes the form:
W. Psoftmax (P' W WoP) x cM,.MM..
cM,
*
Here, we use w0n mask some unimportant token values. Upon passing through the feed-forward neural and weight 1/c, the resulting output is
network with indicator
1
EN
0
3 Fest
y1
0
UN
0
M l:
M++1UN
EN
0
0
1
0
0
0
0
1
0
0
0
0
0
where M 2M, MM
Proof for Theorem 4.1
For TF pres adopt the following configurations:
WqP
0
EN
0
0
0
1
C
0
CEN
0
Corest
0
WRP
0
0
C
and
WvP C
EN
0
0 0
Denote = PT W WoP, we can show that
softmmax (22--125-1) ca. xy.
We may include another attention head to remove the constant. Therefore, the output is formulated as
Psoftmax (P" W WoP) =
eSxx
eSEN
After passing through the feed-forward neural network with indicators weight a/cc, the resulting output becomes:
10 0
and
31
EN
Ltest
91
UN
0
M
MOÆ
0
0
0
0
0
0
1
0
where Mo a.S.
As illustrated in Lemma A.4 output:
after T iterations of the looped transformer TF ('Flopp: it produces the following
EN
Dtest
0
0
y1
UN
0
M Mr
MT
0
0
1
1
0
0
0
0
0
1
0
0
0
0
0
0
1
0
0
0
0
0
1
21

---

Published in Transactions on Machine Learning Research (01/2025)
In the post-transformer, additional positional embeddings are introduced to address technical considerations. The input is structured as follows:
31
Utest
0
y/1
UN
0
0
MT
0
MrNN
0
0
0
Pi
P1
PN
PN
PN+1
PN+1
1
0
0
0
0
0
1
0
1
0
0
0
0
0
1
0
0
0
1
where the positional embedding vectors P1. PP++1 are designed be nearly orthogonal (Lemma A.1) To
initiate the weight wi wewtton propagate the target label y to adjacent tokens using the following attention mechanism:
and W = P1 P
PN
PN PN+1
PN+1
WyP = 2
21
0
UN 0
0].
This operation results in the attention layer producing the following output:
21
IN
0
Etest
0
111
5/1
UN
YN
0
0
My 21
0
MTØN
0
0
0
P
P
PN
PN
PN+1
PN41
1
0
0
0
0
0
1
0
0
0
1
0
0
0
1
In the next layer, analogous to the construction in Equation
S
define the following transformations:
WoP-
C
0
CEN
0
Ca
0
1
1
1
MT
0
MTØN
0
0
WkP=
0
0
and
WvP = e
1/1
y1
UN
UN
0
0].
for some C,c> 0. Defining the matrix
cæ MT
0
ca
MTÆN
0
cæ Mresest
0
Z : P-P'TWWOP-
ca M
CE
N
M.TE
car Ty My Prest
0
0
0
0
0 C
C
C
C
C
C
we can show that
softmmx (Zzi-122-)) R cx. Mræj,
where the closeness of the two sides of the approximation "z" can be achieved by selecting C > 0 sufficiently
large and c > 0 sufficiently small. The constant term can be removed by introducing another head. Therefore, the output of the attention layer expressed as
Σ W X softmax x
(i)
(2) X
2 c
W Nowton'
Dtest
1=1
Finally, the feed-forward neural network yields the desired prediction w
Newton Prest-
22

---

Published in Transactions on Machine Learning Research (01/2025)
A.5 Proof for Theorem 4.2
In this section, we extend the realization of gradient descent, demonstrated in Lemma A.2 for encoder- based transformers, to decoder- based transformers. Although the construction is similar, the key distinction lies in the docoder-based transformer's utilization of previously viewed data for regression, consistent with our intuitive understanding. The following lemma is enough to conchude the proof for Theorem 4.2
Lemma A.5. The one-lager two-head decoder-basee transformer con implement one step of gradient descent in linear regression problems in Equation 0
Proof. We consider the input prompt with positional embedding follows:
31
-;-1
02
V1
yi
æ
27;
1
w
W
w
2-2
We construct the atlention layer with
WoP
k
0
k
1--1
0
W
1
1
WRP'
CE1
0
CT
0
0
C
0
0
and
Wvp' e
/cc
21-1 x;
Here, we adopt causal attention, where the attention mechanism can only attend to previous tokens. The output
Σ T,
w
1-1T
Σ oo a,
For the second head, we similarly let
1 0
0
1
WoP
1 0
0
1
WKP*
0
cy1
0
cy:-1
0 0
C
0
and
WvPi ///cl
3;-1
Then, we have the output
**
EE
E
The atlention layer outputs
æc
Di-1
0
4;
0
y1
y/
O
E;
W
un
2-2
M)j
*
Σ ==
(w
T
O
0
0
2
0
1
O
1
0
0
0
0
0
23

---

Published in Transactions on Machine Learning Research (01/2025)
Since the feed-forward layer is capable of approximating nonlinear functions, e.g., multiplication, the trans- former layer outputs
æ1
x2
0
91
Yi
1
0
w
k++
0
W
O
0
1-1
1
0
0
I
0
0
0
1
0
0
where k/
w/
11 DC
(w
pi
Σ 1=
(w
2hh
yh
24


Reviews:
[{"review_id": "d3oFzgQErj", "reviewer": "Reviewer_E2zR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. Provides rigorous theoretical analysis with detailed proofs for the expressive power of AlgoFormer\n2. Demonstrates capability to implement both first-order (gradient descent) and second-order (Newton's method) optimization algorithms\n3. Novel three-part structure (pre-transformer, looped transformer, post-transformer) mimics human-designed algorithms\n\n**Weaknesses**\n\n1. While the German-English translation task provides real-world validation, additional practical applications would strengthen the paper.\n2. The three-part structure, while theoretically elegant, may introduce additional complexity in training\n3. Continuing the last point, how do the authors see AlgoFormer being used in real-world scenarios, and how would one decide on the architecture (how big is pre-processing/looped/post-processing transformer), and how many loops to execute?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces AlgoFormer, a novel transformer architecture designed to learn and implement algorithms through a structured approach combining pre-processing, iterative optimization, and post-processing steps. The authors demonstrate both theoretically and empirically that AlgoFormer can effectively learn and execute various algorithms, particularly for tasks involving optimization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I am proposing some suggestions here, with the understanding that this paper is primarily a theoretical work, but I think these would make it a better paper:\n\n1. Include more real-world applications beyond German-English machine translation\n2. Include analysis of training stability and computational cost\n3. Discuss potential challenges in practical applications\n\nWriting:\nIn general, the writing could use some polishing, some paragraphs could've been better.\n\n\"However, it is still unverified that the standard multi-layer transformer is exactly performing algorithms.\"\n\n\"A One-layer transformer is mathematically formulated as:\" -> \"A one-layer...\"", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "iI4SjLdtcl", "reviewer": "Reviewer_jrwC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n+ The intuition behind the proposed method is clearly explained and makes sense to me. In particular, the section linking the proposed model design to how humans solve these problems provides a clear intuitive basis for why the model works better than the standard looped transformers.\n+ The use of toy examples to empirically demonstrate the advantages of the proposed method (in addition to the theoretical analysis) demonstrates that this method can indeed be useful in real-world scenarios.\n\nWeaknesses:\n+ Key implementation details are missing, making it difficult to understand how everything works end-to-end, and the additional complexity introduced by AlgoFormer over the looped Transformer.\n+ Scalability of the proposed method to more difficult tasks is unclear", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "A new Transformer-based model, called AlgoFormer, is introduced for improving algorithmic capabilities of Transformers. AlgoFormer adds pre-and post-Transformer blocks to the previously proposed looped Transformer structure to enhance expressive power.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical for acceptance:\n+ Add specific details on how training and inference are performed. (a) Are the pre-, looped- and post-Transformer blocks jointly trained end-to-end? Or are they trained on different inputs and/or with different objective functions? (b) The dashed arrows in Figure 1 are very confusing. Do they indicate that new tokens are introduced after each block? For instance, after the pre-Transformer blocks are completed, are there new tokens that are added to represent the loop function before feeding as input to the looped Transformer block? (c) How exactly is the prompt constructed for a given input in the experiments? Does it follow the P matrix defined in section 3.1? If so, is there any intuition behind this? Is the same prompt generation strategy used in all baselines? \n+ Discuss the training and inference time overheads introduced by AlgoFormer over looped transformer. Do the pre- and post-Transformer blocks add substantial overheads? Or is the runtime dominated by the looped Transformer blocks?\n+ Discuss the scalability of the proposed method to more complex tasks. For instance, can AlgoFormer solve (or can be trivially modified to solve) problems that requires the use of nested loops? Or are there any assumptions made on the structure of tasks/algorithms when computing expressive power?\n\nWould strengthen the paper:\n+ A figure (using a simple toy example) showing the end-to-end flow from prompt generation to output production would be very helpful to understand the inner workings of the proposed method\n+ A discussion on the limitations of the proposed method and the types of tasks it can solve", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "sBcBgh31P0", "reviewer": "Reviewer_K1bS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Yang et al (2024) previous showed experimentally that “vanilla” looped transformers can learn to do linear regression in-context. This paper provides some theoretical constructions to back up these findings, and also extends them to more difficult variations of the problem.\n\n# Expressivity\n\nThe AlgoFormer differs from the \"vanilla\" looped transformer only by the addition of a layer at the bottom and the top. Additionally, since the number of iterations of the loop is fixed (T), the AlgoFormer really amounts to a standard transformer with T+2 layers, where the parameters of layers 2 through T+1 are tied.\n\nOn page 4, some differences with \"vanilla\" looped transformers are listed: AlgoFormer is claimed to be \"closer to the transformer in real applications,\" but this isn't really explained or justified as far as I could see; it's claimed that the vanilla model uses \"task-specific knowledge...like token order switching\" whereas AlgoFormer is \"task-independent,\" but again I didn't really understand this claim.\n\nAt the end of section 2, it's stated that AlgoFormer is not actually more expressive than the \"vanilla\" looped transformer, but it can use fewer parameters. I believe this is correct, but this is a statement about parameter efficiency, not expressivity.\n\n# Section 3.1/A.1\n\nThe position embedding uses \"quasi-orthogonal\" vectors with dimension O(log N) (Lemma A.1). Since N is the number of data points, doesn't this contradict the claim on page 17 that the transformer is independent of the data?\n\nLemma A.2 is similar to proofs by van Oswald et al, and Akyurek et al, but uses softmax instead of linear attention. As a result, it's only approximate, making use of the fact that exp x \\approx 1 + x. What value of C is needed to make the approximation in the second equation of page 16 work? Doesn't C have to depend on the x's, and wouldn't that contradict the claim on page 17 that the transformer is independent of the data?\n\n# Learning setup\n\nI was a little confused about what is fixed by hand, what is learned in training, and what is learned-in-context.\n\nI gathered that the representation function \\Phi^* is simulated by the transformer but *not* learned in-context. However, the statement of Theorem 3.1 does not make this clear at all. On the contrary, it says that the transformer “fit[s] the representation and appl[ies] gradient descent,” making the two sound like one single learning process. Unless I’ve deeply misunderstood, the statement of this and other theorems needs to be corrected.\n\nSimilarly, in the experiments, it was not clear which parameters are fixed and which are learned. Did you fix the parameters as in the construction of Theorem 3.1, while learning the parameters for the FFNs that fit \\Phi^*? Or did you learn all the parameters?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper describes AlgoFormer, a variant of the looped transformer that adds a \"pre-transformer\" layer and a \"post-transformer\" layer. It shows theoretically that several in-context learning problems can be (approximately) expressed by AlgoFormers, and carries out experiments to validate those results.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "# Critical\n\n1. Since the AlgoFormer is not really that different from a looped transformer or indeed a standard transformer, perhaps it should not be given a brand new name. And since its expressivity is essentially the same as the other two models, perhaps the title of the paper should not focus on expressivity.\n\n2. There are quite a few places where the constructions are approximate, not exact, and where dependencies on the input (length) come in. The statements of the theorems should be precise about the goodness of the approximation and about the dependencies on the input.\n\n3. In both the theoretical constructions and experiments, please clarify what is fixed by hand, what is learned in training, and what is learned-in-context.\n\n# Non-critical\n\n## Terminology\n- The word \"block\" is often used to refer to a single self-attention layer and FFNN; your use of the word \"block\" to refer to the whole network was confusing at first.\n- The word \"regularizing\" makes me think of regularization as in L2 regularization, so I'd suggest not using this word.\n\n## Organization\n- The definition of AlgoFormer is buried in the Motivation section.\n- Why are there new results in the Discussion section?\n- The training algorithm is deferred to the Experiments section, which is a common choice, but since this seems to be a key part of the method, I'd suggest presenting it earlier.\n\n## Experiments\n- \\Delta T is introduced without being defined. Is it just the maximum number of layers (iterations) that are included in the loss function?\n- Please explain how the looped transformer and AlgoFormer are used for MT. Is it an encoder-decoder, and if so, is the loop in the encoder, the decoder, or both?\n- It seems strange to me to use cross-entropy to evaluate machine translation.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviews unanimously agree that the paper is above acceptance criteria, albeit somewhat borderline.

Reviewer K1bS and jrwC raised many points which I agree with, mainly on the imprecise language of their architecture contributions:
  * What are the differentiating mechanics between the "pre", "loop", and "post" layers? Are these layers somehow trained differently (e.g. different loss functions) or have different architecture hyperparameters (e.g. different numbers of heads)? Or do the "loop" layers all share the same set of weights?
    * This absolutely needs to be made clear in the paper. At the moment, the authors' post-rebuttal explanations still haven't satisfactorily resolved this, and it can appear as if the paper is introducing additional novelty when none exists. If they are purely for naming purposes, then the paper needs to reflect this as well, by stating explicitly that they are purely distinguished by name.
  * "vanilla model uses "task-specific knowledge...like token order switching" whereas AlgoFormer is "task-independent,": Unless these descriptions can be rigorously justified or explained, then they should be removed. It's not clear what any of these mean at all.

I expect that after a proper re-write of Sections 1 to 2.2, it can be made much more concise, down from 4 pages to possibly even 2.
Decision: Accept (poster)