Paper: CausalLM is not optimal for in-context learning (ICLR-2024)

Reviewer comment (Theory):
To the reviewer's understanding, this paper hasn't excluded the possibility that there exists a parameter configuration that can allow CausalLM to reach optimal solution in the linear setting.

Author response:
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.

Reviewer comment (Theory):
(See weakness) If the reviewer's understanding is correct, then a question is whether proving such a statement is possible.

Author response:
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.

Reviewer comment (Novelty):
My understanding is that this work aims at demonstrating that prefixLM is superior to causalLM for in context learning. While I believe they do a good job at it, I am under the impression that most projects already used prefixLM when possible. For example, InstructBLIP and Llama2 use prefixLM as far as I can understand. If this is the case that most influential language model or VLM already use prefixLM then I am unclear about the intended impact of this work.

Author response:
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).

Reviewer comment (Theory):
The considered theoretical setup seems contrived to me. Since I don't expect the author to come up with a new theoretical setup, this point is not reflected in my evaluation of this paper. However, I believe that this is a weakness of this line of works and I would appreciate to be convinced why any conclusion made on this setup (LSA, linear regression and hand-crafted parameters) should necessary translate to more realistic setups.

Author response:
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.

Reviewer comment (Novelty):
Does prior works generally use causalLM instead of prefixLM?

Author response:
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).

Reviewer comment (Experiments):
The experimental setting is limited to a given number of in-context examples which seems to naturally favor prefix LMs. Casual LMs would train the model with different numbers of in-context samples simultaneously while prefix LM using all possible in-context and query partitions with the same in-context length. Testing with fewer in-context examples could be beneficial to provide more comprehensive results.

Author response:
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. 

|16 test ICEs | LR (MSE) | N-LR (MSE) | MC (classification error) |
|---|---|---|---|
|PrefixLM-SL |1.01 | 2.1e-2 | 42.8 |
|CausalLM-SL | 1.76 | 2.7e-2 | 43.3 |
|PrefixLM-UL | 0.97 | 1.9e-2 | 42.9 |
|CausalLM-UL | 1.12 | 3.2e-2 | 46.6 |

|24 test ICEs | LR (MSE) | N-LR (MSE) | MC (classification error) |
|---|---|---|---|
|PrefixLM-SL |1.4e-1 | 2.0e-3 | 33.4 |
|CausalLM-SL | 7.0e-1 | 1.0e-2 | 35.9 |
|PrefixLM-UL | 1.0e-1 | 1.7e-3 | 37.1 |
|CausalLM-UL | 1.3e-1 | 1.0e-2 | 41.2 |

|32 test ICEs | LR (MSE) | N-LR (MSE) | MC (classification error) |
|---|---|---|---|
|PrefixLM-SL |2.4e-2 | 4.7e-4 | 32.4 |
|CausalLM-SL | 3.1e-1 | 5.0e-3 | 34.6 |
|PrefixLM-UL | 9.5e-3 | 3.4e-4 | 36.2 |
|CausalLM-UL | 4.0e-2 | 5.7e-3 | 37.3 |

Reviewer comment (Experiments):
Another unfairness in the experimental setting is that the nature of casual LMs would basically reduce the number of examples with the same in-context samples compared to prefix LM (i.e. for one input sequence, prefix LMs would enumerate all possible in-context and query partitions but casual LM would only have one for a specific context length). See Q1.

Author response:
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. 
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. 

| Permuted ICEs | LR (MSE) | N-LR (MSE) | MC (classification error) |
|---|---|---|---|
|PrefixLM-SL |9.0e-3 | 1.5e-4 | 24.1 |
|CausalLM-SL | 1.9e-1 | 2.5e-3 | 26.9 |
|PrefixLM-UL | 2.6e-3 | 9.5e-5 | 26.1 |
|CausalLM-UL | 1.1e-2 | 1.8e-3 | 26.2 |

Reviewer comment (Experiments):
A simple augmentation could be used to manually feed those other possibilities of partitions to casual LM, so prefix LM and casual LM would be trained with the same number of in-context and query possibilities for a given context length. Under this setting, would prefix LM still perform better than casual LM?

Author response:
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. 
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. 

| Permuted ICEs | LR (MSE) | N-LR (MSE) | MC (classification error) |
|---|---|---|---|
|PrefixLM-SL |9.0e-3 | 1.5e-4 | 24.1 |
|CausalLM-SL | 1.9e-1 | 2.5e-3 | 26.9 |
|PrefixLM-UL | 2.6e-3 | 9.5e-5 | 26.1 |
|CausalLM-UL | 1.1e-2 | 1.8e-3 | 26.2 |

Reviewer comment (Novelty):
The reviewer is unsure of the relevance of the problem and suggests that the paper's impact could be improved if the motivations were clarified better.

Author response:
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.

Reviewer comment (Novelty):
The reviewer requests clarification on what models in current literature actually use a PrefixLM architecture.

Author response:
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.

Reviewer comment (Experiments):
The reviewer suggests demonstrating or citing papers that show models using a PrefixLM architecture have a different qualitative behavior compared to a standard causalLM architecture.

Author response:
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.