Predictor-Corrector Enhanced Transformers with Exponential Moving Average Coefficient Learning

Authors: bei  li
tong  zheng
rui  wang
jiahao  liu
qingyan  guo
junliang  guo
xu  tan
tong  xiao
jingbo  zhu
jingang  wang
xunliang  cai
microsoft  research
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 04:41:32

Abstract: Residual networks, as discrete approximations of Ordinary Differential Equations (ODEs), have inspired significant advancements in neural network design, including multistep methods, high-order methods, and multi-particle dynamical systems. The precision of the solution to ODEs significantly affects parameter optimization, thereby impacting model performance. In this work, we present a series of advanced explorations of Transformer architecture design to minimize the error compared to the true ``solution.'' First, we introduce a predictor-corrector learning framework to minimize truncation errors, which consists of a high-order predictor and a multistep corrector. Second, we propose an exponential moving average-based coefficient learning method to strengthen our higher-order predictor. Extensive experiments on large-scale machine translation, abstractive summarization, language modeling, and natural language understanding benchmarks demonstrate the superiority of our approach. On the WMT'14 English-German and English-French tasks, our model achieved BLEU scores of 30.95 and 44.27, respectively. Furthermore, on the OPUS multilingual machine translation task, our model surpasses a robust 3.8B DeepNet by an average of 2.9 SacreBLEU, using only 1/3 parameters. Notably, it also beats LLama models by 5.7 accuracy points on the LM Harness Evaluation.

Main paper: 
 introduction: 
 Residual networks [16] , formally y t+1 = y t + F(y t , θ t ), represent a cornerstone in the development of deep neural networks [59, 10] , primarily due to their capacity to facilitate the flow of information across multiple layers. Beyond their pivotal role in convolutional networks, residual connections have become an essential element in the architecture of more complex models, including the Transformer [59] and its various derivatives. This concept can be likened to the discretization process in the Euler method [63, 37, 15, 7, 51, 29] , which serves as a first-order solver for ordinary differential equations (ODEs), where dy(t) dt = F(y(t), θ(t)). In both cases, the new state (be it the next layer's output in ResNets or the solution at the next time step in the Euler method) is computed by taking the current state and adding an adjustment term.Given this analogy with ODEs, there has been a surge of interest in improving residual network architectures by using more powerful numerical methods for ODEs. For instance, the linear multistep method [62, 37, 71, 24] has been employed to bolster the optimization of deep models. Other efforts have included redesigning the Transformer architecture from a multi-particle dynamical system perspective [38, 12] and improving parameter learning efficiency through high-order methods [29] . Additionally, ODEs have been extensively studied for their potential to accelerate diffusion processes, with multistep and high-order solvers offering more accurate predicted noise among each denoising process, generating comparable images but consuming much fewer NFEs [31, 35, 36] .In this work, we mainly focus on advancing the architecture design, specifically by minimizing the truncation error across each timestep. Building upon the ODE Transformers [29] , which replace the first-order Euler method with a high-order method for more precise numerical solutions, our focus extends to addressing two key limitations. First, high-order solutions, such as those from the Runge-Kutta method or multistep methods, are found not to lead to significant improvements when we scale up training data and/or model size. Second, the gated fusion coefficient learning method, which is widely used in previous work, is not well suited for higher-order solutions.Our work draws inspiration from the predictor-corrector method [11] , a well-established approach in numerical analysis known for its accuracy in solving differential equations. This method involves a two-step process: a prediction step that estimates the solution based on known conditions, followed by a correction step that refines the prediction for a more accurate result. We introduce a novel family of PCformers that embrace this predictor-corrector paradigm. Our approach integrates the final solution using an exponential moving average (EMA) method, capitalizing on the insight that higher-order intermediate approximations tend to be more accurate. This assertion is supported by the truncation error analysis presented in Section 3.1.2. Our method is not only readily extensible to arbitrary higher orders but also consistently outperforms the gated fusion method.Our contributions are summarized below:• We extend explicit ODE solutions to implicit ODE solutions via a predictor-corrector learning paradigm. This kind of iterative refinement can attain more accurate solutions than previous studies both theoretically and empirically. In particular, we choose the high-order method as the predictor and the multistep method as the corrector. • To further strengthen the learning ability and training stability for high-order methods, we propose an exponential moving average coefficient learning method to replace the constant coefficients. This leads to a much stronger predictor. • Our extensive experimental evaluation on several benchmarks, including WMT'14 English-German, WMT'14 English-French, WMT'16 Romanian-English, and the OPUS multilingual machine translation benchmark, demonstrates the superior effectiveness of our PCformer models. Notably, our model surpasses the 3.8B DeepNet by an average BLEU score of 2.9 with only 1/3 of the parameters. Furthermore, our model can be extended to other domains. Results on abstractive summarization, language modeling, and language understanding tasks demonstrate its generality. 
background: 
 We build our method upon Transformer [59] as it is one of the most popular models in NLP. The encoder is a stack of identical layers. Each layer consists of a self-attention block and a feedforward network (FFN) block. Both of them are equipped with a residual connection [16] and a layer normalization unit [25] . The output of a block can be defined asEQUATIONwhere F(•) is either the self-attention or FFN block. This equation illustrates that the layer output y t+1 is determined by the layer input y t and a learnable derivative estimated by the current function F. This approach aligns with the benefits of the Euler method, which we will discuss in the following sections. In this work, we use F t to denote the t-th layer representation and Fi to denote the i-th order intermediate approximations.The Euler Method The Euler method is the most basic solution to solve ODEs given the initial value, involving a function y(t) of a variable t and its derivatives. The Euler method defines the first-order derivative of y(t) dy(t) dt = f (y(t), t)where f (y(t), t) defines a time-dependent vector field if we know its value at all points of y and all instants of time t. Eq. 2 illustrates that the change of a variable is determined by its current value γ γ * (1 -γ) γ * (1 -γ) 2 γ * (1 -γ) 3 1 2 1 21Figure 1 : Illustration of several advanced numerical methods and our proposed predictor-corrector paradigm. The right part plots a 4-order method as the predictor to obtain P t+1 ; F t+1 is then estimated via a function F(•); A 4-step method as the corrector to obtain the y t+1 . and a time variable t. In deep learning, we can use a trainable function F(•) to estimate f (y(t), t).In a nutshell, residual networks could be regarded as a 1st-order discretization of the Euler method [63, 29] . The advantage is obvious since residual networks deliver consistent performance gains in the artificial intelligence, but the precision of y t+1 is limited. Fortunately, we can move forward along the numerical analysis perspective as more advanced numerical methods can alleviate this issue.The Linear Multistep Method Compared with the Euler method, the linear multistep method uses previously obtained "solutions" to estimate the current one, leading to more accurate results. Formally, a multistep method could be defined as: y t+1 = y t + t i=1 α i F i , where F t = F(y t , θ t ). The High-Order ODE Method Another family of numerical methods is high-order ODE solvers by repeatedly refining the solutions within a single step. Previous work [29] employed the Runge-Kutta methods [50, 23, 6, 2] for a higher-order solution to ODEs, where Runge-Kutta is a classic family of iterative methods with different orders of precision. More formally, the explicit Runge-Kutta methods of an n-order solution is defined to be:y t+1 = y t + n i=1 γ i Fi , where F1 = F(y t , θ t ), Fi = F(y t + i-1 j=1 β ij Fj , θ t ).Note that Fi is the intermediate approximation to the solution at an inner step. β and γ are coefficients to model the scale of the input and the output of Fi . This kind of method can be adapted to Transformer blocks by reusing F(•) within a block.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Please refer to Appendix B Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
predictor-corrector transformer: 
 In this section, we first show the core design of Predictor-Corrector paradigm to more accurately solve ODEs. Then we propose an alternative coefficient learning strategy that could be applied to arbitrary orders using the merit of the exponential moving average. At last, we show some additional training techniques for stable and well-performance training.A genuine problem-solving process involves the repeated use of available information to initiate exploration, which discloses, in turn, more information until a way to attain the solution is finally discovered. -Newell et al. [43] The Predictor-Corrector framework leverages an iterative process of using available information to refine approximations continuously. Initially, the Predictor generates a rough estimate, which is subsequently refined by the Corrector using newly available data. This cyclical process mirrors the problem-solving strategy described earlier, where each iteration uncovers additional information that enhances the final solution's accuracy.A predictor-corrector method typically uses an explicit method for the predictor and an implicit method for the corrector. Here we take the 4-step Adams-Bashforth-Moulton [1] methods as an instance, where Adams-Bashforth is the predictor and Adams-Moulton is the corrector. Adams-Bashforth is a 4-step method, which defined asEQUATIONwhere F t = F(y t , θ t ). F(•) denotes the t-th function, and θ t is corresponding parameters. Obviously, the Adams-Bashfroth methods are explicit since y t+1 only depends on the "observed" statistics (F i≤t ). Similarly, Adams-Moulton is also a 4-step method as belowEQUATIONFormally, both Eq. 3 and Eq. 4 reused F t , F t-1 , F t-2 to improve the accuracy, but the corrector necessitates an approximate current "solution" F t+1 to substitute F t-3 , which is an implicit method. This is because that y t+1 is the value to be solved, thus we cannot compute F t+1 . To solve this, the Adams-Bashforth-Moulton methods utilize Eq. 3 to obtain the approximate value (P t+1 ) for y t+1 .Then F t+1 could be approximated following F t+1 = F(P t+1 , θ t ). Concretely, the predictor provides a rough approximation, which is the combination of the preceding four layer representations. And the corrector then improves the approximation, offering a more precise sample derived from the data.However, applying the Adams-Bashforth-Moulton method directly to Transformer architecture design leads to unstable training and limited benefits due to the difficulty in optimizing constant coefficients. Similar issues have been observed in training a Runge-Kutta (RK4) network with numerically suggested coefficients [29] . To address these challenges, Wang et al. [62] proposed a Dynamic Linear Combination of Layers (DLCL) method. This approach utilizes learnable coefficients and adjusts steps based on layer depth, effectively transforming it into a variable multistep method. Additionally, the Adams-Bashforth method is not the only choice for the predictor; other numerical methods, such as high-order methods, are also considered strong alternatives as they often provide more accurate solutions [29, 72] .The aforementioned discussion motivates us to design a more powerful and stable architecture based on the above principles. Our preliminary experiments indicate that more accurate predictors indeed improve the performance, thus we choose a high-order method to serve as the predictor. A 2-order method could be defined asEQUATIONwhere F1 = F(y t , θ t ), and F2 = F(y t + F1 , θ t ). Rather than utilizing the previously obtained representations in multistep methods, high-order methods iteratively estimate the approximations upon the last timestep. Similarly, a 4-order method is:EQUATIONwhere F1 = F(y t , θ t ), F2 = F(y t + 1 2 F1 , θ t ), F3 = F(y t + 1 2 F2 , θ t ), and F4 = F(y t + F3 , θ t ). To break the limit of constant coefficients, Li et al. [29] employed a gated network to dynamically compute the coefficients of F1 and F2 , however, this method cannot applied to higher-order methods, e.g., RK4. To facilitate higher-order optimization, we design a more flexible coefficient learning method via an exponential moving average strategy.Predictor with Exponential Moving Average Coefficient Learning The Exponential Moving Average (EMA) method [20] is widely used for estimating time-series data by assigning variable weights to past observations, giving more importance to recent data compared to simple weighted averaging methods. We hypothesize that high-order approximations at each step should have a larger impact on the final output, as they provide a more accurate initial state than previous ones. To support this claim, we replaced Eq. 6 by y t+1 = y t + Fi , where i ∈ [1, ..., 4] . We used a single-layer decoder to compute the perplexity (PPL) on the validation set to simulate truncation errors. Our expectation is that the fewer truncation errors, the larger the coefficient it should own. V a n il la R K 4 , γ • (1 -γ), γ • (1 -γ) 2 , and γ • (1 -γ) 3 from F4 to F1, where γ is learnable and the initialization is 0.5 empirically. In this way, our n-order predictor approximates P t+1 as follows:EQUATION)Corrector with Parameterization Leveraging a robust predictor, our corrector is designed to be computationally lightweight, striking an optimal balance between performance and efficiency.Utilizing the Adams-Moulton method, we parameterize the coefficients of previous states with learnable parameters. These coefficients are initialized using an EMA value, where the newly estimated F t+1 is assigned a larger weight (α = 0.5), and the weights of previous states decrease in a descending order. In this way, we rewrite the Eq. 4 byEQUATION)where P t+1 is obtained with Eq. 7. Empirically, we found that when the dataset is limited, an Backward Euler method [5] as the corrector is enough to provide precise correction, where y t+1 = y t + F(P t+1 , θ t ). We will discuss this in the analysis for more insights.Algorithm 1 Predictor-Corrector Paradigm Step Normalization (RK-Norm) We built our PCformer following pre-norm architecture [62, 58] , by rewritten Eq. 1 to y t+1 = y t + F(LN(y t ), θ t ), where LN(•) denotes the normalization. This ensures that the representation is normalized before computing the derivative F i . To achieve this, we normalize the obtained intermediate approximations Fi at each inner step and then compute the offset, e.g., y t + 1 Sublayer Dropping Additionally, we observe that our models benefit from the rich information brought by high-order predictor and subsequent implicit multistep corrector. To prevent from overfitting (settling into sub-optimal solutions) as the learning ability is quite strong, we borrowed the sublayer dropping technique [28, 34] . The drop rate is empirically set as 0.1 which delivers robust results in previous studies.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification:Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experimental results: 
 We mainly evaluated the proposed method on machine translation, abstractive summarization, language modeling, and language understanding benchmarks. The details of datasets, and corresponding hyper-parameters please refer to Appendix C. For a clear comprehension, note that RK2-block (gated) is 2-order method with learnable coefficients in [29] 's work. And RK2-block (EMA) denotes our EMA strategy. by 1.00 and 1.05 BLEU points with 2-order and 4-order configurations. This demonstrates that the predictor-corrector paradigm is a more parameter-efficient option than pure high-order methods. Results of OPUS Table 3 provides the comparison of PCformer against existing state-of-theart models [69, 61] on the OPUS-100 testset. The findings here are three aspects: 1) Across all configurations, PCformer delivers significant BLEU gains over vanilla baselines. 2) Our 12-layer EMA Pre-Cor model attains an average SacreBLEU score of 32.6, which not only outperforms the 3.8B DeepNet but also beats its further optimized variant, BranchNorm, with only 1/8 model parameters.3) PCformer can benefit from the enlarging width and depth. Notably, our 1.2B model shows an average SacreBLEU of 35.0, thereby setting a new state-of-the-art on the OPUS-100 testset.Abstractive Summarization Table 4 presents the results of the abstractive summarization task. As shown, our PCformer consistently improves upon pure the high-order method [29] , in terms of three rouge scores. Notably, PCformer (2-order) even beats RK4-block which consumes less computation cost. Additionally, PCformer (4-order) sets a new state-of-the-art on the summarization task which excludes models based on pre-trained models. This result strongly supports our hypothesis that an appropriate coefficient learning schedule is essential for the effectiveness of higher-order methods, thereby enhancing the performance of our PCformer model.Language Modeling Table 5 presents a comparative analysis of our PCformer against vanilla Transformers in Adaptive Input Representation [3] and Shortformer [48] settings. Our 2nd-order configuration achieves significant reductions in perplexity (PPL), outperforming Adaptive and Shortformer by 1.79 and 1.23 PPL, respectively, even within identical model capacity constraints. Remarkably, PCformer surpasses the high-order method (RK2-block) by a substantial margin in both settings on both validation and test sets, demonstrating the superiority of PCformer. In response to the increasing significance of attention mechanisms in LLMs, we conducted a comprehensive evaluation of PCformer using established benchmarks, LM Evaluation Harness [13] , focusing on a diverse range of downstream tasks including common-sense reasoning and question-answering. Language Understanding We also validate our method on the widely used natural language understanding benchmarks, namely GLUE, which consists of 8 sub downstream tasks. The evaluation metrics are as follows: The result for STS-B is the Pearson correlation; Matthew's correlation is used for CoLA; Other tasks are measured by Accuracy. The results are presented in Table 7 . We can see that PCformer achieves 2.1 points (on average) improvement over the BERT-large, which demonstrates the effectiveness of PCformer.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We have summarized all details including the datasets, the model configurations and training details in the appendix. Please refer to Section C of the Appendix. And we also make a clear algorithm to illustrate the detailed computation flow of the proposed PCformer. And we will release the code after re-organizing the codebase, to provide detailed scripts for reproduce our results.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
analysis: 
 In this section, we explore several significant issues, comprising the visualization of truncation errors, COMET results, and a set of essential ablation studies. We primarily conducted ablation studies on machine translation tasks, but the conclusions are generalizable. Quantization of the Truncation Error Following the suggestion in [29] 's work, we use the perplexity between the single-layer Transformer decoder output and the ground truth to approximate the "truncation error". The results of Table 8 were conducted on the Penn Treebank dataset. We see that the proposed EMA method achieves a lower perplexity than the learnable coefficient (gated) learning method, similar observation in 4-order (EMA v.s. Rk4-block). Additionally, the Predictor-Corrector paradigm can further reduce the truncation error, which demonstrates the effectiveness of our method. Ablation Study on Predictor-Corrector Framework The predictor-corrector framework is crucial in our work, with the choice of ODE solutions for each component significantly impacting performance. We experimented with various combinations of predictors and correctors, including high-order methods, linear multistep methods, and the Backward Euler method. Table 10 summarizes these results. We chose RK2-block with EMA as the default for high-order solutions due to its performance and efficiency. Key insights include: 1) The predictor must be highly accurate, as it sets the performance lower bound. High-order predictors outperform the multistep method (DLCL) and the Euler method. 2) A complex corrector isn't always optimal; a Backward Euler method suffices for small and medium datasets (e.g., En-Ro and En-De), while more complex methods may cause overfitting. 3) Combining a multi-step method predictor with a high-order corrector performed worse than other combinations, highlighting the importance of predictor choice.Ablation Study on Core Design Technique Table 11 presents the performance of our PCformer with various initial coefficients for the EMA method and stable training techniques. Our default, with γ set to 0.5, performs best because smaller γ values disrupt the numerical bound, and larger values overly focus on recent approximations. As detailed in Section 3.2, RK-Norm is essential for training stability, as shown by the BLEU score drop without it. Testing different layer-wise coefficients and replacing the learnable scalar γ with a learnable matrix vector showed no significant performance difference, so these were excluded from our default settings. Inference Speed and Memory Consumption Table 12 presents a detailed comparison of inference speed and memory consumption across various large model configurations, revealing that the proposed PCformer models achieve satisfactory inference performance. This is primarily because the computational overhead is concentrated on the decoder side rather than the encoder, as demonstrated in our experiments. Additionally, PCformer is memory-efficient, as shown by the memory usage comparison between the baseline and the ODE Transformer. Despite the fact that our PCformer models are more than twice as slow as the vanilla baseline in encoder-only and decoder-only configurations, they deliver significantly superior performance, making the trade-off worthwhile. These performance gains are clearly illustrated in Table 6, where the substantial improvement in model effectiveness justifies the increased inference time.We acknowledge that further optimization to accelerate PCformer's inference speed is a promising direction for future research.More Analyses Due to the limited space in the main content, we summarized more detailed analyses in Appendix D, including the parameter efficiency (Figure 3 ), illustration of training and validation curves (Figure 4 ), and visualization of coefficients during the learning procedure (Figure 5 ). We anticipate that these analyses will offer a deeper and more comprehensive understanding of our method.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes]Justification: We will open-source the code as soon as we collect our scripts for an easy way to reproduce our results.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/ guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. Guidelines:• The answer NA means that the paper does not include experiments. • The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
related work: 
 Ordinary Differential Equations The connection between ResNet and ODEs was first proposed by Weinan [63] , while Neural ODENet [8] introduced a new perspective on neural architecture design.Several architectures [71, 24, 37, 17, 72, 38, 53] can be interpreted from the ODE perspective. Recent studies leverage ODE benefits for Transformers. Lu et al. [38] proposed MacaronNet using the Strang-Marchuk Splitting Scheme, and Zhang et al. [70] introduced continuous self-attention models. Dutta et al. [12] redesigned Transformer architecture for efficiency from a multi-particle dynamic system view. Li et al. [29] showed that first-order ODE blocks could cause error accumulation, and high-order methods were suggested as solutions. In this work, we advance the Transformer design with a more accurate Predictor-Corrector paradigm and a general coefficient learning strategy inspired by the exponential moving average, showing significant performance improvements on NLP benchmarks.ODEs and numerical methods are also popular in diffusion models, reducing prediction errors in denoising processes. Text-to-image generation typically uses a two-stage model, including a text-to-image diffusion model and super-resolution models. The standard diffusion model, DDPM [19] , requires up to 1000 iterations to recover images from Gaussian noise. Subsequent work accelerated DDPMs using denoising equations [57] or scheduled variance [44] , though often at the cost of performance. Liu et al. [31] proposed treating DDPMs as solving differential equations on manifolds, introducing a pseudo linear multi-step method for efficiency and performance. Further diffusion acceleration efforts [35, 36] were motivated by ODE benefits. More recently, Xu et al. [67] presented a deep generative model solving the Poisson equation, opening new possibilities in text-to-image generation. ODEs also show promise in discrete diffusion models, as demonstrated by Lezama et al. [26] , who used a predictor-corrector paradigm to enhance the accuracy. 
conclusions: 
 This paper advances the design of parameter-efficient neural network backbones through a numerical analysis perspective. Previous work has utilized high-order ODE solutions for more accurate approximations at each block, yielding promising results on various sequence generation tasks. However, challenges remain, such as the scalability of learnable coefficients to RK4-blocks and the lack of exploration into implicit ODE methods. To address these issues, we introduce a predictor-corrector framework to improve estimation precision. Additionally, we proposed an EMA coefficients learning strategy to promote coefficients learning for high-order methods with high flexibility.  
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: This paper only uses publicly available open-source code and datasets, and the authors have properly credited the original creators and respected the licenses and terms of use.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: We will provide detailed documentation in our open-source code repository after publication.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.• According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification:Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A Broader ImpactWe do not anticipate any specific negative impacts from our work. However, as with any machine learning method, we recommend exercising caution. Our primary contribution is advancing neural model design from a numerical perspective, aiming to facilitate models' parameter learning and enhance the performance, which we believe to be environmentally friendly. This approach encourages the community to enhance neural models by integrating knowledge from various domains, such as physiology, physics, and mathematics, thereby further benefiting AI research.The proposed PCformer has demonstrated significant performance gains across a variety of tasks, including sequence generation, natural language understanding, and language modeling. While the encoder-decoder PCformer is computationally efficient, as it only employs the predictor-corrector paradigm on the encoder side, the primary computational overhead during inference remains in the decoder. However, when our method is applied to encoder-only models (for NLU) or decoder-only models (for LM/LLM), the additional computational complexity becomes non-negligible. Addressing how to further accelerate inference in these scenarios remains critical, and we aim to explore this in future work.We plan to address this issue from two primary aspects:• High-order Computation in Latent Space: When adopting high-order methods, the iterative computations among inner steps cannot be skipped. Therefore, we aim to investigate whether high-order computations can be performed in a latent space with a reduced dimensionality, such as 64 or 128 hidden dimensions, compared to the original 1024 dimensions. This approach presents a significant theoretical challenge: maintaining the stability of the ODE while computing higher-order intermediate approximations in the latent space.• High-order Training and Inference in a First-order Manner: Another alternative is to achieve high-order training and inference using a first-order (Euler-like) approach. This might involve employing distillation techniques or treating high-order computations as a form of training regularization.We will continue to advance the design of PCformers to balance performance gains with inference efficiency, particularly in the context of large language models. This research could provide valuable insights to the community on whether such modeling methods can further enhance the performance of LLMs.Datasets We present experimental results across three WMT benchmarks, including WMT'14 English-German (En-De), WMT'14 English-French (En-Fr), and WMT'16 English-Romanian (En-Ro), as well as on a large-scale, challenging multilingual machine translation benchmark (OPUS).• For the En-De task, the training data consisted of approximately 4.5M tokenized sentence pairs, as in [59] . All sentences were segmented into sequences of sub-word units [54] with 32K merge operations using a shared vocabulary. We selected newstest2013 as the validation data and newstest2014 as the test data.• For the En-Fr task, we used the dataset provided within Fairseq, i.e., 36M training sentence pairs from WMT'14. newstest2012+newstest2013 was the validation data and newstest2014 was the test data.• For the En-Ro task, we replicated the setup of [39] , which used 600K/2K/2K sentence pairs for training, evaluation and inference, respectively.• Apart from the aforementioned three bilingual translation benchmarks, OPUS is a multilingual benchmark that contains more challenges for the model to serve. We choose OPUS [69] , an English-centric corpus covering 100 languages, which is randomly sampled from the OPUS collection. For a fair comparison with prior work, Evalutation We measured performance in terms of BLEU. Both tokenized BLEU and SacreBLEU5 scores were reported on the En-De and En-Fr tasks. For the OPUS task, we average the SacreBLEU scores among 94 languages for both X→En and En→X. The beam size and length penalty of each task are summarized in Table 13 . In order to attain results that are more compelling, it is imperative to acknowledge that BLEU may not be a suitable metric for evaluating a model's performance. To supplement our evaluation, we have included a summary of the COMET scores [49] of the top-performing model, utilizing a reference-based model 6 .Training Details In accordance with the recommendations provided by [27] , we have incorporated relative positional representation [55] , namely RPR for short, into our model architecture to establish stronger baselines.To ensure stable learning during training, we have also borrowed the merit of dense connections among layers [62] for stable optimization within FP16 training. Our models were trained on 8 GPUs with 4, 096 tokens per GPU. For the En-De, En-Fr and OPUS tasks, we have observed that larger batching schemas often result in better convergence [45] . Therefore, we accumulated gradients every 2 and 8 steps for En-De and En-Fr/OPUS, respectively. Adam optimizer [22] with (0.9, 0.997) for β1 and β2 is adopted. The hyperparameters including the learning rate, the warmup step and the total training steps of three tasks could be found in Table 13 . Note that we trained Base/Deep and Big models for 50K and 100K steps on the En-De task. We regarded merging SAN and FFN as the default ODE block. In addition, main results were the average of three times running with different random seeds (1, 42 and 2024), and we averaged the last several checkpoints towards the robustness. The detail of Base/Deep/Wide configurations is as follows:• Base/Deep Model. The hidden size of self-attention was 512, and the dimension of the inner-layer in FFN was 2, 048. We used 8 heads for attention. For training, we set all dropout to 0.1 as default, including residual dropout, attention dropout, ReLU dropout. Label smoothing ϵ ls = 0.1 was applied to enhance the generation ability of the model. For deep models, we only enlarged the encoder depth considering the inference speed.• Wide (or Big) Model. We used the same architecture as Transformer-Base but with a larger hidden layer size 1, 024, more attention heads (16) , and a larger feed forward inner-layer (4, 096 dimensions). The residual dropout was set to 0.3 for the En-De task and 0.1 for the En-Fr task.We also tested the models' ability to process long sequences on the CNN-DailyMail summarization task [42, 18] . The preprocessed method was the same as in [46] . We used a shared BPE with 30K operations, resulting in a vocabulary of 32, 580 entries. The evaluation metric was F1-Rouge [30] (Rouge-1, Rouge-2 and Rouge-L). Datasets As mentioned above, the truncation error analysis is conducted on the Penn Treebank [41] , which is a widely-used language model dataset. It contains 88K, 3, 370 and 3, 761 sentences for training, validation and test. The vocabulary size is 10K. We set the layer depth of the language model to 1 or 2 to make a fair comparison. Assuming the layer depth is 1, then the loss between the model output and the true label could be regarded as the truncation error. In this way, we alleviate the influence of the error accumulation across different layers. Apart from PTB, we also evaluate our approach on another widely acknowledged language modeling dataset, Wikitext-103 dataset [40] , which is the largest available word-level language modeling benchmark with long-term dependency. WikiText-103 consists of 103M training tokens from 28K articles on Wikipedia, and the average length of tokens per article is about 3.6K. The data is can be easily obtained and preprocessed following Baevski et al. [3] 's work.Setups For the PTB dataset, we used the transformer_base configuration, whose hidden size is 512, and the filter size of the FFN is 2, 048. All the dropout rates are 0.1, including the residual dropout, attention dropout and ReLU dropout. Each model was trained up to 20 epochs, while most models arrived at convergence on the validation set when the epoch is 10. Then the validation PPL began to increase, though the training PPL is still declining. This is due to the small model capacity. The warmup step was 2, 000 and the batch size was 4, 096. The max learning rate was set to 0.0007. For the Wikitext-103 dataset, all models were based on the original open-source Fairseq toolkit, and the corresponding configuration is transformer_lm_baevski_wiki103. It inherits the configuration of transformer_big with 1, 024 hidden size and 4, 096 filter size. Besides these common hyper-parameters, it adopts the adaptive input representation to reduce the embedding matrix by decreasing the embedding size of low-frequency words or sub-words. This is also the most common choice when building large-scale language models.Our proposed PCformer by continue training the BERT model using the provided pretrained model, BERT-Large-cased. And we evaluated on the General Language Understanding Evaluation (GLUE) [60] . Concretely, the GLUE benchmark has 8 different text classification or regression tasks including MNLI, MRPC, QNLI, QQP, RTE, SST-2, SST-B, and CoLA.We also evaluate our PCformer on a wide range of downstream tasks covering common-sense reasoning and question-answering, including LAMBADA (LMB. [47] ), PIQA [4] , SciQ [64] , Winograde (Wino. [52] ), HellaSwag (Hella. [68] ), and ARC-challenge [9] . We mainly evaluate language models via zero-shot evaluation. We randomly sampled 100B tokens from Slimpajama dataset for training.Parameter Efficiency Figure 3 summaries the results of several efficient Transformer variants, including Lite Transformer [66] , DeLight [39] , a light version of the Evolved Transformer [56] , and our PCformer. It is clear to see that PCformer is significantly more parameter efficient than others. We make detailed comparisons of PCformer within different hyper-parameters, comprising of hidden size and model depth. Concretely, RK2-block with gated fusion coefficient learning strategy delivers stronger performance than DeLight within the same model parameters. And it is on par with DeLight in terms of BLEU, but having 9M fewer parameters. Moreover, as expected, our newly proposed EMA coefficient learning method slightly improves the performance of gated fusion in almost all scenarios. A notable bonus it brought is higher-order solutions, e.g., RK4-block is superior to RK2-block with the help of EMA. It may offer a new choice for deploying NMT systems on edge devices.Training and Validation Perplexity Apart from the BLEU scores, we also compare our methods with baselines regarding perplexity on both training and validation sets. Figure 4 plots the training and validation PPL curves of the PCformer (RK2-block (EMA + Pre-Cor)) and the baseline on two representative translation tasks. All models were in big configurations for more convincing conclusions. The proposed RK2-block with EMA coefficient learning and Predictor-Corrector framework delivers much lower training and validation PPLs within the same configurations. More specifically, our method still benefits from increasing model depth, as the 12-layer model outperforms the 6-layer one. For both the En-De and En-Fr tasks, we observed that our 6-layer method even shows lower PPLs than a 12-layer Transformer. This phenomenon is more evident in the WMT En-Fr task, which again demonstrates the high parameter efficiency of our PCformer.We also collect the learning process of learnable coefficients γ during training. As we discussed above that the inspiration of EMA coefficient learning method comes from the prior that the most current approximation is more precise. This assumption has already been clarified by the comparisons of truncation errors. Here, we want to figure out how coefficients learned if removing the constraint. To achieve this goal, we set all coefficients to be independently initialized by the mean of 1, e.g., a 2-order block with γ1 = 0.5, γ2 = 0.5. Figure 5 plots the learning curves of RK2-block and RK4-block within these two learning strategies. As we can see that, for the independent initialization scenarios, coefficients vary dramatically within the first several epochs, and then show convergence in a small range. Both two figures show that the contribution of the most current coefficient is larger than others.To our surprise, γ2 in RK2-block (Independent) converges to large than 1 and oppositely, γ1 is even goes to a negative value. Meanwhile, γ3 in the RK4-block (Independent) shows a negative impact to the final solution, but the other 3 coefficients vary within our expectation. We notice the order relation among these 4 coefficients are consistent with the numerical suggested coefficients in the Adams-Bashforlth method (-9/24, 37/24, -59/24, (a) RK4-block (Independent), BLEU=29.00 55/24). This indicates the underlying relationship between the numerical analysis and the neural network optimization. While, after employing our EMA method, coefficients are optimized along our expected direction, and these models are empirically better than those without constraints.Results on Time-Series Forecasting We also followed the reviewer's suggestion to evaluate PCformer on time-series forecasting tasks. We selected 10 multivariate datasets from UEA Time Series Classification Archive following the setting and the codebase provided by Flowformer [65] . Thus we choose the Flowformer as the baseline, which is also a strong model on these testsets. For the details, we build the PCformer upon Flowformer and report the 2-order predictor and Euler corrector as the training data is very small. Also, we use RK-Norm to avoid the model suffering from the overfitting problem as the authors of Flowformer trained their models upon to 100 epochs (or even 400 epoch on some tasks.). The results are evaluated by the best accuracy. We can see that PCformer can beat the Flowformer by 2 average score on 10 testsets, which demonstrates the effectiveness on time-series forecasting tasks. Justification: We have highlighted our contributions in Section 1. We proposed a predictor-corrector paradigm for Transformer architecture design, and also an EMA coefficient learning algorithm to augment the learning of high-order predictor. Experimental results have also shown the superiority of PCformer.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.LN( F1 ) to obtain the F2 for the next timestep. Meanwhile, the Fi in Eq. 7 is rewritten by LN( Fi ).If not, this oversight can cause instability when computing the final ODE solution, where we will make ablations in the analysis section. The algorithm (right part) presents a more detailed computation flow of a single layer in our PCformer, where H stores the previously obtained F t+1 .we followed the setting in Gu and Dao[14]'s work.https://object.pouta.csc.fi/OPUS-100/v1.0/opus-100-corpus-v1.0.tar.gz.BLEU+case.mixed+numrefs.1+smooth.exp+tok.13a+version.2.3.1The default setting which uses a reference-based regression model built on top of XLM-R (Large). 


Decision: This paper explores transformer architecture design to minimize the error compared to the solutions from ODE formulation. It introduces a predictor-corrector learning framework to minimize the truncation errors. It also proposes an exponential moving average-based coefficient learning method to further strengthen the higher-order predictor. Extensive experiments show the superiority of the approach on different benchmarks. All the reviewers acknowledge that this paper provides novel contribution and has valuable contribution to transformer architecture innovation from ODE perspective. However, the reviewers have also pointed out several weakness of the paper, such as restrictions of the application to natural language domain, relatively small model size in experiments, and potential complexity of the method. During the rebuttal, the authors have successfully addressed most of the major concerns by providing additional results on vision domain and time series domain. It also provides additional experiments on the scaling behavior of the method, which shows its potential effectiveness when the FLOPs go up. One suggestion from my side to present the scaling law results is to also show the results in the form of “performance vs FLOPs”, which is another standard visualization in scaling law study.

Reviews: review: paper_summary: paper_summary: The paper presents advancements in Transformer architecture to minimize errors in approximating solutions to Ordinary Differential Equations (ODEs). 

The contributions are:
- Introducing a learning paradigm with a high-order predictor and multistep corrector to reduce truncation errors.
- Proposing an exponential moving average-based method to enhance the predictor's learning ability and stability by replacing constant coefficients with dynamic ones.
- Demonstrating superior performance across various benchmarks, including machine translation, abstractive summarization, language modeling, and natural language understanding, achieving notable improvements in BLEU scores and parameter efficiency.

The work shows improvements in translation tasks and highlights the general applicability of the proposed methods across different natural language processing domains.
strength_weakness: strength_weakness: The complexity of implementing these methods might be a barrier for practical adoption. The paper could benefit from providing more detailed guidelines or code to facilitate easier implementation and replication of the results.

The experiments primarily focus on natural language processing tasks. While the results are impressive, it remains unclear how well the proposed methods generalize to other domains, such as time series forecasting. Including preliminary results or discussions on the potential applicability to these other areas could strengthen the paper.
questions: questions: Can you provide more details on the computational cost of the proposed predictor-corrector framework and EMA coefficient learning method? Specifically, how do these methods impact training time and resource utilization compared to traditional Transformer models?

How does the proposed method scale with increasing model size and dataset complexity? Have you encountered any challenges in scaling up your approach, and if so, how did you address them?
limitations: limitations: Provide a detailed analysis of the computational complexity and resource requirements of the proposed methods. This should include a comparison with standard Transformer models and an explanation of any trade-offs between performance improvements and computational costs.

Discuss the scalability of the proposed methods in more detail. Explain any challenges encountered when scaling up to larger datasets or models and how these were addressed. Provide insights into potential limitations in terms of scalability and how future work could overcome these challenges.
paper_summary: paper_summary: This paper presents an approach to improve the performance of Transformer models for conditional natural language generation (machine translation and summarization).   The authors introduce a predictor-corrector framework, inspired by numerical methods for
solving ordinary differential equations (ODEs), to enhance the accuracy and
stability of the  models. The proposed model is evaluated against standard Transformer models on multiple benchmark datasets. The presented results show improvements in prediction accuracy and efficiency when compared to standard models.
strength_weakness: strength_weakness: The paper does not compare its results with state-of-the-art  models, which could have provided a more comprehensive evaluation of its effectiveness.  For instance for WMT14 En to FR, we can reach 43 in BLEU, instead of 41 for Attention is all you need.  Moreover, the difference with ODE Transformer is very limited. 


The theoretical justification for the predictor-corrector framework could be more detailed, particularly in explaining how it relates to and improves upon existing methods. 

The computational overhead induced by the proposed method is only lightly discussed in the Appendix and would require more in depth discussion.
questions: questions: - How does the predictor-corrector framework perform in comparison to
state-of-the-art pretrained Transformer models?
- Do you think that all the 72 references are necessary? 
- The quotation of Newell 59 is maybe a bit too much. Don't you think?
limitations: limitations: The authors should discuss the potential computational overhead introduced by the predictor-corrector framework.
paper_summary: paper_summary: This paper takes inspiration from established high-order approaches in numerical analysis for solving differential equations to improve the architectural design of Transformers. Prior work has shown that residual networks can be seen as discrete approximations of Ordinary Differential Equations (ODE) and explored methods to improve the quality of the solution. Specifically, this paper introduces a predictor-corrector learning framework to minimize approximation errors and an exponential moving average-based coefficient learning method. These advancements are used within Transformer architectures that are evaluated on several tasks such as translation, summarization, language modeling and language understand tasks, improving over standard and previous ODE-based Transformers.
strength_weakness: strength_weakness: - W1. There is lack of emphasis in the experiments on the computational overhead introduced by the high-order predictor and multi-step methods. Achieving better quality is not sufficient for adoption in practical settings. I'd suggest quantifying the training and inference cost compared to standard and ODE Transformers.
- W2. In terms of parameter efficiency, the proposed model achieves better performance with 1/3 of parameters only on one of the examined datasets. It's not clear if this is by chance or if the result holds on other settings. It would be useful to report performance with a smaller model size on the rest of the datasets to better establish the underlying claim. 
- W3. The impact to research communities beyond the ones focusing on ODE transformers is somewhat limited because the results are with relatively small model sizes. Showcasing that the results hold on larger architectures such as 7B Mistral would make the results more convincing. 
- W4. The improvement of the proposed multi-step high-order method compared alternative predictor-corrector methods is small (Table 10) and the benefit compared to simpler first-order methods from the predictor-only paradigm in a controlled setting is not provided (similar to the ablation in Table 10).
questions: questions: - Q1. Where does the parameter efficiency of the model stems from and how a fair comparison with other models was ensured? In Appendix D, there is a section that discusses this but the comparison does not seem to be apples-to-apples (i.e. same config and #params) and whether the result holds in all the tasks in the experiment section.
- Q2. The method description wasn't entirely clear to me and I have the following questions: 
   - The parameterization of $\mathcal{F}(P_{t+1}, \theta_t)$ is not clear from the textual description. What kind of transformation is used here? 
   - How are the coefficients are parameterized exactly? Do you parameterize a single coefficient for each order for an $n$-order predictor or it is the same coefficient shared?  
    - The text mentions that a larger weight is assigned ($a=0.5$) to the estimated $F_{t+1}$, is this value fixed or you have experimented with different values? It would be useful to clarify which coefficients among $\alpha$ and $\gamma$ are hyper-parameters and which ones are learned.
- Q3. Is the observed improvement worth the the additional computational cost for obtaining a better approximation? It would be useful to show what is the difference compared to a standard or a simpler ODE Transformer. 
- Q4. Do you mean "ROUGE results" instead of "Rough results" in Table 4? In the same table, what is the number of parameters used by each model? 
- Q5. What is the size of the models in Table 8 and is it the same for all variants? 
- Q6. In the ablation experiment and other results, there is lack of comparison with a simple 1-order single/multi-step method or high-order 1-step method. Is the proposed method substantially better than them?
- Q7: Did you perform continued training of BERT directly on language understanding tasks or on some pre-training data first in Table 7? Please also report the number of parameters for both models and the exact configuration for PCformer.
limitations: limitations: The limitation section covers the shortcomings when the proposed methods are applied to different model designs but lacks discussion about the experiment shortcomings related to model size, computational cost, and comparison with simpler methods in a controlled setting.
paper_summary: paper_summary: This work presents a transformer architecture inspired by predictor-corrector methods for solving ODE problems. It adopts the Adams-Bashforth-Moulton method and utilizes an exponential moving average (EMA) method to compose the predictor, discussing two correctors (the EMA-based and the simple backward Euler method). Without the background of ODE, the proposed method could be seen as a special case of a transformer with cross-layer connections. Consider that the residual connection only links two layers, while the proposed method links multiple layers in a high-order way.
strength_weakness: strength_weakness: The experiments are insufficient and out-of-date. As the main contribution is the new transformer architecture, scaling law style experiments and testing on modern LLMs' benchmarks are recommended. I saw that Appendix D has some results, but it's better to have a formal experiment.
questions: questions: What's the efficiency of the proposed method, such as inference time and GPU memory consumption? I saw that inference efficiency has been mentioned in the limitations; it's better to have quantitative results and also memory states.
limitations: limitations: Agree with the limitation section.

Comments: title: Looking Forward to follow-up discussion
Further Results on a 3B PCformer
Any Further Concerns About Our Work
Thanks for the reconsideration
Response to authors
Any Other Concerns on PCformer
Further response to the remained questions
comment: Dear Reviewer 1pKC,

Thank you once again for your suggestions and valuable feedback. We apologize for reaching out again as the deadline approaches. To address your concerns, we have conducted experiments on both the time-series forecasting task and the image classification task. Additionally, we have provided data on the practical costs of training and inference for both the encoder-decoder and decoder-only paradigms. We have now uploaded the results of a 3B PCformer using 16B and 50B tokens, respectively. The results show that there are no major challenges for its scaling up. If our results and explanations help address your concerns, we would be grateful if you could acknowledge our rebuttal and consider adjusting your score accordingly.

Best wishes

The Authors.
| Model(Params & Tokens)     | Wiki.(ppl) | LMB.(ppl) | LMB.     | PIQA     | Hella.   | SCIQ     | ARC-c    | Winograde| Avg.     |
| -------------------------- | ---------- | --------- | -------- | -------- | -------- | -------- | -------- | -------- | -------- |
| Transformer++ (340M & 6B)  | 38.5       | 96.1      | 21.4     | 60.3     | 29.1     | 69.2     | 21.5     | 50.4     | 41.9     |
| PCformer (340M & 6B)       | 35.3       | 78.8      | 23.6     | 61.6     | 30.1     | 71.6     | 22.9     | 51.8     | 43.6     |
| Transformer++ (340M & 16B) | 28.3       | 65.3      | 29.8     | 63.2     | 33.9     | 73.2     | 23.1     | 51.4     | 45.8     |
| PCformer (340M & 16B)      | 25.6       | 39.7      | 34.5     | 65.2     | 36.9     | 79.6     | 23.2     | 52.2     | 48.6     |
| Transformer++ (1.3B & 16B) | 23.8       | 26.2      | 37.3     | 65.7     | 37.6     | 78.6     | 23.7     | 51.5     | 49.0     |
| PCformer (1.3B & 16B)      | 20.9       | 23.2      | 42.5     | 68.3     | 43.4     | 81.5     | 25.1     | 52.4     | 52.2     |
| Transformer++(1.3B & 100B) | 16.3       | 11.8      | 51.6     | 71.0     | 51.7     | 86.7     | 28.1     | 54.6     | 57.2     |
| PCformer (1.3B & 50B)      | 16.2       | 9.38      | 55.1     | 71.9     | 54.8     | 88.6     | 29.6     | 57.2     | 59.5     |
| PCformer (1.3B & 100B)     | **14.0**   | **7.46**  | **59.6** | **73.8** | **60.0** | **90.7** | **31.7** | **61.7** | **62.9** |
| PCformer (3B & 16B)        | 17.8       | 13.6      | 48.0     | 69.7     | 48.4     | 83.1     | 26.4     | 55.2     | 55.1     |
| PCformer (3B & 50B)        | **13.6**   | **6.5**   | **62.1** | **74.4** | **61.9** | **90.6** | **32.4** | **61.9** | **63.9** |



Since the discussion deadline is approaching, we would like to once again thank all the reviewers for their efforts in helping to further improve our work. Given the limited time, we have focused on reporting the results of PCformer on a 1.3B model using 100B tokens in the rebuttal phase. Additionally, we have attempted to further scale the model to over 3B parameters, with a hidden size of 3200, an intermediate size of 8640, and 26 layers, following the LLama3B configuration. We trained this model using 128 A100 GPUs and, up to now, it has been trained on nearly 50B tokens. Below are the results for the 3B model at different stages of training (16B tokens and 50B tokens):

- PCformer (3B & 16B) outperforms its 1.3B & 16B counterpart by an average score of 2.9, and PCformer (3B & 50B) surpasses PCformer (1.3B & 50B) by an average score of 4.4. This indicates that with increased data consumption, PCformer continues to benefit from the data and demonstrates superior results.

- PCformer (3B & 50B) achieves an average score 8.8 points higher than PCformer (3B & 16B), and PCformer (1.3B & 50B) exceeds PCformer (1.3B & 16B) by an average score of 7.3. This suggests that the performance gap widens with increasing data volumes, which is intuitive as larger models naturally require more data for effective training. 

The above two phenomena demonstrate the scalability of PCformer, and we are enthusiastic about the opportunity to demonstrate its potential to the community. Thank you once again to all the reviewers, ACs, and SACs for your invaluable feedback and support.
Thank you for all your comments and discussions. I acknowledge I have read them, but I stand by my original review.
Dear Reviewer 1pKC,

We apologize for reaching out as the discussion deadline approaches. We are grateful for the comprehensive and useful feedback you provided, and we have responded in detail to your comments during the rebuttal phase, e.g., new results on time-series forecasting, the practical cost of training and inference, more strong results on larger LLMs using PCformer and so on.

We are eager to know if our proposed results and clarifications have adequately addressed your concerns. If so, we would appreciate it if you could reconsider your score.

Thank you once again for your time and effort.

Best regards,
The Authors
Thank you for reconsidering our work. We greatly appreciate the valuable suggestions you have provided, which we believe will further enhance our research. We plan to include these results and make thorough revisions in our next version. Thank you once again for your efforts!
Thank you for the detailed responses to my questions and additional results! My main concerns have been addressed and I decided to increase my score:

- The scaling experiments with LLMs are quite promising and increased my confidence that the results hold on larger model sizes.
- In terms of efficiency, the results show that there is no significant overhead introduced which addressed my main concern. It was also great to see results that justify the parameter efficiency (e.g. performance of PCFormer 340B vs Transformer 1.3B). 
- The new ablation results clarify the differences compared to simpler baselines. It appears that the proposed methods have a fairly good improvement for translation that is greater than 0.5-1 point.   

The rest of the answers provided helpful clarifications, it would be great if they are reflected in the final version.
We extend our heartfelt gratitude to the four reviewers for their insightful feedback and to the Area Chairs for facilitating the rebuttal discussion. We appreciate the time and effort you have invested in evaluating our submission.

Since uploading our rebuttal, we have not yet received any further feedback. We eagerly look forward to engaging in discussions with the reviewers and believe that we have effectively addressed the major concerns raised. We also welcome any other further concerns!

We are aware that the ACL conference is approaching, and we hope you all have wonderful days
> Q2.1 The parameterization is not clear from the textual description. What kind of transformation is used here?

As we emphasized that $\mathcal{F}_(P_{t+1},\theta_t)$ is the function to computed the derivate upon the input. Thus, either a self-attention network or an FFN, or even the whole encoder block could be regarded as a F function, also they could be seen fuctions in different granularities. In this work, for a fair comparison with ODE Transformer, we choose the whole block as the function.

> Q2.2 Details of $\alpha$

Just a single coefficient. the code is as follows:
```
self.alpha = torch.nn.Parameter(torch.Tensor(1))
self.alpha.data.fill_(0.5)
```
Take a 2-order EMA as an instance, the final layer ouput is computed as follows:
```
x = residual + self.alpha*(1-self.alpha) * runge_kutta_list[0] + self.alpha*runge_kutta_list[1]
```
where runge_kutta_list is a list to stores the previous obtained intermediate approximations ($\hat{F}_i$).

We have uploaded our code to a anonymous github, more details could found in our codebase. 


> Q2.3 Whether $\alpha$ is fixed

Similar with the previous question, we use 0.5 as an initail value for $\alpha$, while it is learnable. Thus $\alpha$ would be changed according to the graident decsent during the training phase. Sorry for the missing details for $\alpha$ and $\gamma$, and these two are both learnable tensors with an initial value 0.5. We will include this in our next version of the paper.


> Q3. Is the observed improvement worth the the additional computational cost for obtaining a better approximation?  

As for the sequence generation models which follow a encoder-decoder paradigm,  our models are quite competive and the observed improvement is indeed worth the additional computational cost. Actually, we have already compared our PCformer with a simpler ODE Transformer in sequence generation tasks. We can see that, PCformer (RK2) can beat ODE Transformer (RK4) with  stronger performance and less computation cost. For example, in Table 1, PCformer (2-order) beats RK4-block (EMA) in both En-De and En-Fr tasks, with one less computation (The former 2 times predictor and 1 times corrector, while the later consumes 4 times forward computation for the high-order solution). SSimilar phenomena could be observed in abstractive summarization tasks (Table 2).

> Q4: typos and the parameters of the models displayed in Table 4.

Yes, we apologize for the typo; we mean "ROUGE results" instead of "Rough results" and will correct it in the next version. Regarding the number of parameters, all models listed in Table 4 have a similar parameter number, approximately 63M. This corresponds to a Transformer-base configuration, which includes 6 encoder layers and 6 decoder layers, each with a hidden size of 512 and 8 attention heads.


> Q5: What is the size of the models in Table 8 and is it the same for all variants?

The configurations of the models used to approximate the truncation errors are detailed in Appendix C.3. Specifically, we used a Transformer language model (decoder-only) with a hidden size of 512, tested in both 1-layer and 2-layer settings. Note that all all variants are the same. By comparing the results of the 1-layer and 2-layer models, we can see that both the ODE Transformer and our PCformer significantly reduce truncation errors, as measured by PPL. For instance, the PCformer (2nd order but with a 1-layer parameter) outperforms the Residual-Block (2-layer) and even surpasses the RK4-block (EMA) with fewer forward computations. This also provides context and a partial answer to your Q3 regarding computational efficiency.

> Q6: Comparison with a simple 1-order single/multi-step method or high-order 1-step method.

Yes, our PCformer significantly outperforms single/multi-step methods and high-order 1-step methods. Perhaps there was some oversight, but it is important to note that the single 1-step 1-order method corresponds to the vanilla Transformer (denoted as "Residual-block" in most tables). Transformer-DLCL represents a specific case of a multi-step method. We compared these models with our PCformer in Table 1. Additionally, the high-order 1-step method is represented by the ODE Transformer. We have already included comparisons with these methods in our MT experiments. We will ensure that these results are included in subsequent versions to provide a more comprehensive evaluation.

> Q7: Details of the BERT training.

Both the PCformer and the BERT models in Table 7 share the same parameter count, with approximately 335M parameters each. Specifically, both models have a hidden size of 1024, an FFN filter size of 4096, and 16 attention heads. We have aligned all settings with those specified in the original BERT paper to ensure a fair comparison. For PCformer, we pre-trained the model from scratch using a combination of WikiText and BookCorpus datasets. After the pre-training phase, we fine-tuned PCformer on the GLUE downstream tasks.