CausalLM is not optimal for in-context learning

Thank you for your questions. Please see our answers to the questions.

1. In a one-layer LSA, it can be proved (using the method in [1]) that the parameter configuration that we used in the paper is optimal for both prefixLM and causalLM in linear regression ICL. Although the proof does not necessarily generalize to the multi-layer case, we strongly suspect that, due to its architectural constraints, there is no parameter configuration that would allow causalLM to reach optimal in the linear setting given a finite number of examples. Specifically, the causalLM architecture causes the model to process the data in an online learning manner. An online learning may asymptotically converge to the optimum given infinitely many examples, when the step size of the learning algorithm is decreasing (which may be possible with a different LSA parameter configuration). However, even if an online learning system asymptotically converges to the optimal solution, it can be proved that at least O(1/\eps) examples are needed in order to reach the \eps-ball of the solution, which is still suboptimal compared to a batch optimizer when the number of examples is limited.

[1] Ruiqi Zhang, et al. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023.

2. In Appendix E.5, we study replacing the attention masks of a pretrained causal decoder in the PaLI-X model by a prefix decoder during finetuning. In our experiment, we found that the prefix-decoder improved the test accuracy over the original causal-decoder in both 4-shot and 8-shot cases. However, since such a swap introduces an architecture gap between pretraining and finetuning, we cannot guarantee that is always beneficial. On the other hand, we find that as far as the pretraining mixture contains both prefixLM and causalLM objectives, the prefixLM always outperforms causalLM on ICL.

Thank you for your questions. Please see our answers to your questions.

1. Although more and more recent models (especially the multimodal ones) are starting to apply the prefixLM (or encoder-decoder) architecture, there are still prominent models which do not, e.g. Flamingo, GPT-family, etc. Therefore, we believe that the comparison between causalLM and prefixLM remains important. Furthermore, while prefixLM is being adopted by many models recently, that line of work has not provided rigorous proofs that explain why prefixLM should be better than causalLM. To our best knowledge, our paper is the first to rigorously argue that from a theoretical perspective (in particular for multi-layer transformer ICL).

2. Following the reviewer’s concerns about the reproducibility of our results, we conduct a new set of large-scale experiments using the open source t5 models, which we describe in detail below. At the same time, please note that our first two experiments in Sec 5.1, 5.2 compare causalLM and prefixLM in a small experimental setup with synthetic data, and try to directly verify our theorems. The work done in these sections is readily reproducible for people with access to the jax/flax library. We will release the code to reproduce these main synthetic results of our paper after the anonymous period.

For the large scale experiments with LM or VLM, a pretraining procedure is needed to get reasonable results. To set up the experiments for comparing prefixLM and causalLM in a fair manner, two options are available: (A) two pretrained models that only differ in its attention masks (causalLM vs prefixLM) or (B) one pretrained model which was exposed to both causalLM and prefixLM during pretraining.

At the time of submission, there were no type-A models available, and the PaLM2 model was the only type-B model that we had access to. Unfortunately, it was not released and cannot be used for reproducibility.

In order to make sure that our main empirical results hold and are reproducible, we conducted another experiment based on the t5 models and c4 corpus (both publicly available). Note that the existing public t5 checkpoints are all based on EncDec models, which are similar to prefixLM. Thus, it would be unfair and unnatural to compare with causalLM by simply replacing the bidirectional attention of the encoder to the causal attention during the finetuning stage.

To make a more reasonable comparison, we reran the pretraining stages of t5 on the c4 corpus but with DecoderOnly models. Moreover, for each size (from base, large and xl) of the models, we pretrained two checkpoints, one with prefixLM and the other with causalLM, each for 1M steps using the same t5 pretraining recipe. In this way, we construct a type-A fair comparison between prefixLM and causalLM.

After pretraining, we use the FLAN recipe to finetune each checkpoint (for 20k steps) with its pretrained attention mask and compare them on two benchmarks: MMLU [1] and BBH [2], under the few-shot ICL setting in [3]. Their averaged accuracies are reported in the following tables:

As can be seen, except for the base-size MMLU case, where causalLM marginally wins over prefixLM, in all other five cases prefixLM wins by a significant margin. Furthermore, as the model size gets larger, the advantage of prefixLM tends to grow larger. We want to specifically thank the reviewer for pushing us in this direction, as we consider that these experiments strengthen the claims of our paper. We plan to add this set of new results into our main body of the paper, and will move the PaLM2 results into the appendix. Furthermore, we will release the corresponding checkpoints to make sure the above results are reproducible.

[1] Dan Hendrycks, et al. Measuring massive multitask language understanding. ICLR, 2020.

[2] Mirac Suzgun, et al. Challenging BIG-Bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09261, 2022.

[3] Hyung Won Chung, et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416, 2022.

Minor: We agree that the theoretical setup has its limitations, as the exact mathematical formalism is hard to generalize beyond LSA. Nevertheless, we use this simplified setup as a probe, to reveal an important consequence of the architecture differences, which could generalize to more general transformers. Specifically, the attention mask of prefixLM allows it to do batch learning, while the one of causalLM restricts it to online learning, which is suboptimal for few-shot learning. Our simplified setup allows us to conduct more rigorous analysis through the stationary points, so that such comparison is more obvious.

Thank you for your questions and the constructive suggestions on the two ablation tests. Please check our following experimental results to see whether they address your concerns.

1. We follow your suggestion by using the models that were trained with 40 in-context examples, but testing them on fewer (16, 24, 32) in-context examples. Note that 16 is the minimum number of examples to solve our 16-dim synthetic regression problems. The errors of prefixLM and causalLM are provided in the following table, where regression tasks (LR, N-LR) report mean squared errors; and the MC task reports the classification error. We see that in general prefixLM still consistently outperforms causalLM even when testing with a smaller number of in-context examples than the training time, a setting that favors causalLM.

2. To address your concern, we conduct a new set of experiments in which each sequence of data permutes the in-context examples in random order during the training time. The results of this experiment show that this style of causalLM training indeed improves over the fixed order training setting. However, prefixLM still outperforms causalLM in general. We will add these informative experiments to the appendix and discuss how for simple tasks, permutation of in-context examples might mitigate some of the problems with causalLM.

Thank you for your question about the relevance of the problem.

The interest of the NLP/ML community in the comparison between prefixLM and causalLM can be traced back as early as [1], where they showed prefixLM outperforms causalLM in varieties of NL tasks. Later, UL-2 [2] proposed to mix prefixLM and span corruption objectives, and found it to be more efficient than the causalLM objective alone. It was also shown in [3], that U-PaLM (a UL2-finetuning PaLM) outperforms PaLM (causalLM only) in various ICL tasks. Indeed, for the reasons above, some of the latest models have included prefixLM objectives in the pretraining mix (for example PaLM-2 [4]). On the other hand, prominent models such as Flamingo as well as the ones in the GPT-family are still based on the causalLM structure, so the comparison between prefixLM and causalLM remains important and relevant. Furthermore, all previous studies were done in an empirical manner, whereas we set out to explain their differences from a theoretical perspective and back the theory with empirical evidence. While we are not the first to follow this path, our work is the first to provide a theoretical justification for the advantage of prefixLM over causalLM in a multi-layer transformer ICL setting by analyzing their theoretical convergence properties. We will edit the paper to clarify the above points.

[1] Colin Raffel, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485–5551, 2020.

[2] Yi Tay, et al. Unifying language learning paradigms. arXiv preprint arXiv:2205.05131, 2022.

[3] Hyung Won Chung, et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416, 2022.

[4] Google, et al. Palm 2 technical report, 2023.