Predictor-Corrector Enhanced Transformers with Exponential Moving Average Coefficient Learning

Following your suggestion, we have scaled up our large language model (LLM) experiments, focusing on how our PCformer performs in the LLM setting. We aim to demonstrate the capabilities of PCformer from both model capacity and data volume perspectives. The updated results are as follows:

• In our initial submission, we reported results on a 340M parameter configuration trained with 6B Slimpajama data. Here we included additional perplexity metrics on Wikitext and Lambada to provide more comprehensive results. Furthermore, for more convincing results, we scaled the model size from 340M to 1.3B parameters, which is the maximum size feasible within the limited time. Scaling to larger models requires additional engineering efforts, such as implementing tensor or pipeline parallelism for robust and efficient training (using Megatron codebase). We plan to report results on 7B or larger models in our next version and we believe PCformer can indeed perform strongly on those settings.

• For the data, we scaled from 6B to 16B and 100B tokens. Our findings show that with more training data, the average performance improves significantly. This demonstrates that our model benefits from increased data size, without experiencing diminishing returns. This indicates that model bias persists and plays a crucial role in our setting (Larger settings still worth to be explored in future work). As the model size increases (from 340M to 1.3B), PCformer shows substantial performance gains, both in terms of accuracy on the LM harness evaluation and lower perplexity on Wikitext and Lambada tests.

• We trained the 1.3B model on a cluster of 256 A100 GPUs. Note that 1.3B models consist of 24 layers, where the hidden size is 2048 and the FFN size is 5432 （8/3 * hidden size, 16 attention heads and SiLU activation functions. The baseline (1.3B + 100B tokens) was trained up to 20 hours and nearly 40 hours for our PCformer. Thus PCformer (1.3B + 50B tokens) was trained within similar time.

• Additionally, given that PCformer consumes more FLOPS per forward pass, we compared models with similar total FLOPS consumption. Our PCformer trained on less than 50B data outperformed Transformer++ trained on 100B data by a significant margin. Moreover, PCformer (340M & 16B) achieved results comparable to Transformer (1.3B & 16B) with nearly 1/4 of the parameters. These results demonstrate that PCformer remains competitive in settings with larger model capacity and data volumes, highlighting the potential for further research in model designs that fully utilize available data.

• These results demonstrate that PCformer remains competitive in settings with larger model capacity and data volumes, highlighting the potential for further research in model designs that fully utilize available data. While the increased inference and training costs in the decoder-only paradigm are notable, the substantial performance gains justify continued exploration of PCformer, including parameter-efficient training algorithms.

• We believe these updated results address concerns about the relevance and sufficiency of our experiments. Thank you for your valuable feedback. Please let us know if you have any further concerns.

Q1: 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.

This question is a common issue among reviewers, we hope you can see the global response W1 for more details!

Due to the limited page in global response, we provide a comparison of the total inference time between PCformer and the baseline in the 1.3B setting (trained on 100B tokens). Using 8 A100 GPUs with CUDA 12.2, we measured the inference time across all benchmarks listed in the table. The PCformer took 1428 seconds, while the Transformer++ took 1037 seconds. Although PCformer's inference time is longer, the performance gains make the additional time investment worthwhile.

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.

We thanks for your recognition for our motivation and the organization of our method. We also appreciate for your constructive feedback towards the shortcomes of the current manuscript, and we think all the concerns would be well addressed in our improved version. Here, we would like to show the details of your main concerns:

W1: Computational overhead: e.g., training and inference cost.

We hope you can refer to the details of the general response W1. We report the training and inference cost of machine translation tasks and give a detailed analysis of large language models.

W2: Parameter efficiency: performance with a smaller model size on the rest of the datasets.

• Besides OPUS, actually, we have already done simialr experiments in Figure 3 (Appendix), where a 28M PCformer can beat other models larger than it.
• In our general response W2 for our newly proposed results on LLMs, we found that Pcformer can beat the vanilla Transformer by only training on half of the training data! (1.3B models and 50B vs 100B tokens); More excited, our 340M PCformer (16B tokens) can achieve comparable results with 1.3B models with only 1/4 parameters!
• For summarization tasks, a baseline with 12/18 encoder-layer configuration achieves RG-1, RG-2, and RG-L scores of 41.18, 18.23, 38.03 and 41.33, 18.31, 38.20, respectively. In comparison, our 2-order method yields scores of 41.96, 18.99, 38.74, outperforming the baseline while utilizing only 1/3 to 1/2 of the layers.

We'll add more comprehensive comparisons next.

W3: Results on larger models such as Mistral 7B.

• We appreciate your suggestion and completely agree. Our current results demonstrate that our model performs exceptionally well on smaller model sizes (less than 1B), indicating that PCformer has significant potential to scale up to larger model sizes. Given the interest from reviewers in seeing how PCformer performs in LLM evaluations, we have scaled up the model size from 340M to 1.3B parameters and increased the training data from 6B tokens to 100B tokens. These efforts have pushed the limits of our current computational resources. The summarized results are provided in the general response W2. We hope these results address your concerns.

• We recognize the importance of showcasing results on even larger architectures like the 7B Mistral. However, achieving this would require additional computational resources. At present, we plan to train a PCformer based on the Mistral checkpoints and restart the training process with significantly less data compared to what was originally consumed in their pretraining. We are committed to conducting these more challenging experiments as soon as we have access to the necessary computational resources. Thank you for your valuable feedback.

W4: More details of the ablation (Table 10)

Very good suggestions. Here we add the simpler first-order methods and ODE Transformer for each setting. The results are below:

We observe that all settings, except for the Multistep Predictor and High-order Corrector, show consistent BLEU improvements over the baseline (#1). This phenomenon is explained in our current manuscript: the predictor needs to be concise enough, otherwise, the initial value may cause the subsequent corrector computations to deviate. Additionally, both multistep and backward Eulere correctors show substantial improvements to ODE Transformer. Our proposed PCformer is a framework that you can choose a proper corrector according to the training date complexity, and RK2-block with EMA as the predictor always behaves well. We hope the results here can address your concern.

Q1: About fair comparison in experiments.

We mainly want to compare the PCformer with the vanilla Transformer in two settings.

• Firstly, as we only introduce quite small parameters, just a 1D tensor for EMA coefficient learning and several layernorm to achieve RK-Norm, we compare these two models in similar (near the same) parameters. The PCformer beats the vanilla Transformer by a quite large margin nearly in all scenarios.
• Secondly, we can compare them in the same FLOPs, for example, a 6-layer PCformer with much fewer parameters can beat a 12-layer vanilla Transformer, meanwhile fewer computations.
• Thirdly, for LLMs, we find that PCformer also beats Transformer even the former trained with 50B tokens and the latter trained with 100B tokens. As we all know, increasing the training data is the most effective way and also the simplest way to improve performance. But at this serve setting, PCformer still behaves strongly.

Due to the limited time and page, we will add more comparisons in all experiments for a much stronger claim.

Q2.1 The parameterization of 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 compute the derivate upon the input. Thus, either a self-attention network or an FFN, or even the whole encoder block could be regarded as an F function, also they could be seen functions in different granularities. In this work, for a fair comparison with ODE Transformer, we choose the whole block as the function.

Thanks for your constructive advice and we think all the conerns would be well addressed in our improved version. We would like to address your concerns as follows:

W1: The paper does not compare its results with state-of-the-art models, which could have provided a more comprehensive evaluation of its effectiveness.

Thank you for pointing out this issue. Due to page limitations, we focused our comparisons on the most closely related work, particularly those designed based on numerical methods, such as Transformer-DLCL, MacroNet, and ODE-Transformer. Notably, the ODE Transformer is a strong baseline on the WMT En-De and En-Fr benchmarks, used to achieving state-of-the-art results in these tasks. In addition, for the OPUS task, we compared PCformer with SoTA models like DeepNet and BranchNorm, which are leading models on this benchmark. To the best of our knowledge, PCformer indeed achieves SoTA results on WMT En-De (without pretrained models), En-Fr, En-Ro, and OPUS in these machine translation tasks. As for other tasks, such as summarization, PCformer shows competitive results compared to baselines learning from scratch. The current SoTA models in these tasks are primarily finetuned versions of BART and other advanced pretrained models. We will aim to make a fair comparison with these models in our next version. Thank you for your understanding and your valuable feedback.

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

• Compared to existing high-order methods such as Runge-Kutta (ODE Transformer), our PCformer leverages a predictor-corrector paramdigm to estimate intermediate approximations more accurately. Notably, we introduced an EMA-based coefficient learning method to enhance the robustness and stability of our high-order predictor, and the corrector then applies a multistep method to further reduce truncation errors.

• Theoretical Improvements: Truncation error is a crucial factor in numerical solutions. In Section 3.1.2, we provide a theoretical analysis demonstrating that higher-order intermediate approximations tend to be more accurate, leading to more reliable final results. And in Table 8, we show that PCformer can achieve lower ppl (analogous to truncation errors) than the 1-step and high-order methods.

• Practical Superiority: As demonstrated in our experimental results section, our framework significantly outperforms existing models, including the robust 3.8B DeepNet, on several large-scale datasets while using only one-third of the parameters.

W3: The computational overhead induced by the proposed method is only lightly discussed in the Appendix and would require more in depth discussion.

Thank you for your advice. We regret that we were unable to include this discussion in the main body of this paper due to time constraints. We would like to provide the following summary on the complexity and computational overhead induced by our proposed method, the details please refer to the general response W1!

In a nutshell, PCformer only introduces slight inference overhead as we mainly applied the predictor-corrector in the encoder side, which takes a very small amount of computation in the process of auto-regressive decoding.

Q1: How does the predictor-corrector framework perform in comparison to state-of-the-art pretrained Transformer models?

We have conducted extensive experiments across several benchmarks. Our proposed PCformer not only significantly outperforms the vanilla Transformer but also shows a substantial improvement over the strong ODE Transformer. We believe your question may be directed towards understanding how PCformer performs in comparison to other large pretrained language models (LLMs). If our interpretation is incorrect, please let us know, and we will address the specifics accordingly.

Due to the limited number of response characters, please refer to the general response W2 for the new proposed results on LM harness evaluation. We can see that PCformer beat the Transformer baseline across all settings, we will further scale the model to a much larger capactiy in the future, e.g. 7B/8B model to make a fair comparison with LLAMA models. While training even larger LMs from scratch is costly and time-consuming, it is a promising direction for further research.

Q2: Do you think that all the 72 references are necessary?

We have carefully reviewed all the references cited in our work and believe that almost all reference are necessary. The large number of references is due to the comprehensive nature of our study, which spans five distinct tasks: machine translation, abstractive summarization, language modeling, natural language understanding, and LM harness evaluation. For each task, we report results on several widely used benchmarks, such as WMT En-De, En-Fr, En-Ro, OPUS, nine test sets for BERT, and eight test sets for LM harness evaluation. To ensure a fair comparison and to validate our results, it was essential to cite the datasets and the results from related work. Consequently, the extensive referencing reflects the breadth and depth of our study rather than an excess. We will futher check the details to avoid unnecessary reference included!

Q3: The quotation of Newell 59 is maybe a bit too much. Don't you think?

Thank you for pointing this out. Our intention was to illustrate that the predictor-corrector paradigm aligns with the seminal ideas of Newell, as also mentioned in the Tree-of-Thought paper. Specifically, the predictor first provides an initial estimate. This estimate is then refined by the corrector for greater accuracy, mirroring the problem-solving process proposed by Newell. We appreciate your feedback and will carefully reconsider and reorganize this section in the next version of our manuscript to ensure it is succinct and relevant.

Thank you for all your comments and discussions. I acknowledge I have read them, but I stand by my original review.

Thanks for your recognition on our writting and the style to present our idea. We would like to answer your questions as follows:

W1: 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.

We open-source our core code at https://anonymous.4open.science/r/Neurips-PCformer/. This allows others to easily reproduce our results and apply PCformer to additional benchmarks. We will also release the complete codebase after finalizing the README.md to facilitate better reproducibility.

W2: 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.

We mainly evaluate our PCformer on NLP tasks, including both the natural language understanding and natural language generation. While, our method is quite general that it could be applied to other Transformer variants in other areas, e.g., Swin-Transformer in computer vision. For example, our PCformer can also improve Swin-Transformer in tiny and base configurations. The results were conducted on 224*224 image size, and 2-order predictor with a multistep corrector.

Beside this, we also follow your suggestion to evaluate PCformer on time-series forecasting tasks. We select 10 multivariate datasets from UEA Time Series Classification Archive following the setting and the codebase provided by Flowformer [1]. 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. We hope these results can address your concern, and we'd like to add them into the updated version of this paper.

[1] Flowformer: Linearizing Transformers with Conservation Flows, ICML 2022.

Q1: 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?

Our proposed PCformer is parameter-efficient as we share the parameters between the predictor and the corrector, especially for the F functions in the high-order predictor. Please refer to the general response W1 for more details.

Q2: 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?

Our PCformer could be easily scaled from both the model capacity and training tokens. We summarize the results of larger model capacity (from 340M to 1.3B) and training tokens (from 6B to 100B) in the general response W2, where PCformer can beat LLama-like models in all scenarios. Additionally, the results in Table 3 have already shown the superiority of PCformer in a 1.2B setting on OPUS multilingual translation tasks.

We think there is only a major challenge: the computational overhead of PCformer is somewhat higher than that of the vanilla Transformer (e.g., LLama models). This limitation has been discussed in our paper. We are actively working to address the training and inference overhead in our ongoing research.

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