Understanding the Generalization of In-Context Learning in Transformers: An Empirical Study

W1: (also see Q4): The paper claims that “transformers lack inter-problem generalization with ICL” (l. 19), yet Fig. 2 left or Fig. 3 left show that transformers generalize to y-axis shifted sinusoids (to a certain extent). This seems contradictory.

R1: Thank you for your suggestion. Our primary focus is on the generalization ability of transformers under the three shift patterns we investigate, as outlined in Section 2.2. These patterns, namely inter-problem, intra-problem, and inter-task, represent a decreasing order of generalization difficulty. Our findings indicate that transformers can achieve intra-task and intra-problem generalization while fail in inter-problem generalization.

Regarding Figures 2 and 3, we have clarified the research questions addressed by the different parts of the figures. As stated in Lines 230-233, "The performance of the Baseline model on base function classes answers Question 2.3 on intra-task generalization, while its performance on combinations of base function classes with T+ answers Question 2.1 on inter-problem generalization...". Furthermore, the left-hand side portions of Figures 2 and 3 demonstrate the generalization capability of the transformer, which is one of our key conclusions. In these instances, the model encounters the same types of functions during both training and testing, illustrating the strong intra-task generalization ability of transformers.

This observation aligns with our conclusion that "Transformers fail to show inter-problem generalization yet succeed in achieving strong intra-problem generalization during ICL" (Lines 64-66). The left-hand side plots in Figures 2 and 3 represent intra-task generalization within our experimental setup, where the model is tested on the same types of functions it encountered during training. This confirms the generalization capability of transformers in this context, consistent with our findings.

Regarding the y-axis shifts, we provide a detailed description of our data generation mechanism in Appendix B.1, Lines 759-769. To avoid inaccuracies in analyzing the model's generalization ability caused by mutual approximation between base functions after weighting, we ensure that the expected function values of different base functions have significant differences. In the function fitting experiments (Sections 3.1, 3.2, and 3.3), we added a bias to each base function. Specifically, the whole equation for sine and cosine functions is:

$f(x) = \alpha sin(\beta*x) + (-1)^\beta * \beta / 2$.

W2: Other popular research areas (e.g., model robustness) also found that data is a key factor. That this is also the case for ICL generalization is expected.

R2: It's important to note that we are not simply stating that more or more diverse data generally improves model generalization. Instead, we provide a nuanced analysis of the generalization behavior of transformers in ICL. We categorize different generalization scenarios and investigate specific areas where transformers exhibit limitations and how data design can be used to address these limitations.

Specifically, we find that increasing the diversity of tasks within a problem does not lead to significant improvements in generalization. However, exposing the model to a wider range of problems can broaden the boundaries of its capabilities across different tasks. This detailed analysis is presented in Line 64-78.

W3: The paper ... I think it is clear that the baseline transformer won’t generalize to convex combinations with basis functions with different frequencies (e.g., sin(x)+sin(2x)) when trained only on sinusoidals with the same frequency, as these combinations are not well-supported by the training data.

R3: We believe there might be a misunderstanding. Please note that our experiments demonstrate the ability of the transformer to generalize to combinations of functions with different frequencies, even when trained only on functions with the same frequency. For instance, as shown in Appendix D.1, Lines 1372-1377, after being exposed to functions like sin(x) + cos(x), the model can generalize to sin(x) + sin(2x).

As you pointed out, this might seem counterintuitive. However, this is precisely what we aim to demonstrate. In our definition (Section 2.3, Lines 155-180), fitting convex combinations of base functions constitutes a single problem, regardless of the frequencies of the individual functions. Different frequencies simply represent different base functions. Therefore, our conclusion is that once the model has been trained on one type of convex combination, it can generalize to other convex combinations, irrespective of the frequencies involved.

W4: It seems that the combined ICL problems in Eq. 3 will have function values outside the range of those in models trained only on the basis function from Eq. 2. However, this may also be a misunderstanding, as the values for \phi in the experiments are not provided (see Q2).

R4: During training, if the sampled combination function yields values exceeding the range of the base functions, we discard that sample. This effectively prevents any OOD issues arising from mismatched value domains.

W5: The experiments on convex combinations (Sec. 3.1) and product combinations (Sec. 3.2) are unrelated to the real-word experiments in Sec. 4. It is unclear whether the findings on the synthetic data also hold in the real-world setting for these combinations.

R5: The motivation and purpose behind the design of Our synthetic and real-world experiments are identical. The research questions we raise in Sections 2.2 and 2.3, along with our definitions of base functions and function combinations, apply to both the synthetic data in Section 3 and the real-world settings in Section 4.

Trigonometric functions are simply an example we use in our function fitting experiments. Our goal is to analyze the generalization behavior of in-context learning in transformers broadly, and provide guidance for real-world applications. For instance, in Section 4.1, we state: "Mirroring the definitions of generalization problem and task, we define single-API calls as the basic generalization problem $\mathcal T_{base}$, and the combination of APIs as the composite generalization problem $\mathcal T_{com}$" (Line 423 - Line 431 ). Here, single-API calls correspond to the base functions in Section 3. Similarly, in Section 4.2, translation tasks between single languages, such as En2De, also correspond to the base function fitting in Section 3.

W6: Code is not provided and certain details are missing.

R6: Thank you for your suggestion. To ensure reproducibility, we have open-sourced our code at https://anonymous.4open.science/r/Generalization-of-Transformers-8227/.

Q1: What is the meaning of the superscripts in Eq. 2 and other equations? I.e., what does (0) mean? It appears that it is always the same throughout the paper and seems unnecessary (e.g., superscripts are absent in Eq. 3).

R1: The superscript 0 referring to the trigonometric family of functions is simply one instance of the function set F. Our experiments and considerations extend far beyond trigonometric functions, as described in Lines 143-145.

Our research is not limited to trigonometric functions. For example, in Section 4.1, Lines 426-429, we state: "Mirroring the definitions of generalization problem and task, we define single-API calls as the basic generalization problem $\mathcal T_{base}$, and the combination of APIs as the composite generalization problem $\mathcal T_{com}$." This demonstrates the broader applicability of our framework.

Furthermore, we also consider fitting other families of functions. As shown in Appendix D5, Lines 1508-1654, we explore fitting more complex functions, further demonstrating the generalizability of our findings.

Q2: What is the range of ϕ in the experiments?

R2: For the weight function $\phi$, we initially sample from a standard normal distribution $N(0, 1)$ and clip it to the range $[-1, 1]$. We simply discard combination function samples with function values exceeding the range of the base functions.

Q3: Are the GPT-2 models trained with permutation invariance?

R3: Yes, we train GPT-2 with permutation invariance because the order of samples presented to the model does not affect the underlying function being learned.

Q4: How are test functions generated.

R4: Please refer to the response to W1.

Q5: l. 255-258: should it be Baseline instead of ComFuncLearner? If not, could you please clarify what is meant in that sentence?

R5: We believe "ComFuncLearner" is indeed the correct term here, and there is no error. This passage provides an analysis of the experimental results. We observe that ComFuncLearner, after being trained on only the base functions and one type of convex combination of functions, can generalize to other unseen convex combinations.

For example, in our experiment, ComFuncLearner is trained on the following function types: {sin(x), cos(x), sin(2x), cos(2x), α*sin(x)+(1-α)*sin(2x)}, where α ∈ (0,1). However, it successfully generalizes to unseen convex combinations like sin(2x)+cos(2x), sin(x)+cos(2x), etc.

Conversely, the baseline model only demonstrates fitting ability on the base functions. This aligns with our findings: the transformer model exhibits intra-problem generalization but fails to achieve inter-problem generalization.

Q9: How do the authors account for synonyms/paraphrases in the translation experiments (Sec. 4.2)?

R9: Thank you for your suggestion. We adopted additional metrics beyond BLEU, including BertScore and GPT-4o scoring, to assess the accuracy of the model's output. The detailed experimental setup is described in Appendix B.5 Line 1054 - Line 1133. The experimental results are shown in the table below. Our conclusions remain valid.

Bert score:

gpt4o eval:

S1: Most experimental Figs. & Tabs: ** I suggest writing out the functions.

R1: Thank you for your suggestion. We have specified the exact forms of the composite functions in the relevant section, as detailed in Figure 2, 3, 4, 5 in the paper.

S2: The paragraph in l. 263-265 appears to address two separate things that could be separated for clarity.

R2: Thank you for pointing this out. We have revised the statement in Line 263 - Line 265 as follows:

"For a detailed experimental setup, please refer to Appendix B.1. Except for convex combinations, we also observe remarkably similar conclusions regarding the model's ability to learn multiplicative combinations of functions as shown in the following section. "

S3: It'd be good to add the description of the ComFuncLearner model variants in l. 364-366 also in Sec. 3.2.

R3: We apologize, but we did not fully understand your point. Could you please elaborate further? Thank you for your clarification.

S4: The meaning of the red numbers in Table 4 should be clarified, especially since the number for Baseline+ICL_s under Average_s is not directly computable from the table (likely due to rounding).

R4: We have indicated in the table caption that "The difference highlighted refers to the improvement of ICL." The improvement of 1.8 for Baseline + ICL_s over Baseline that you mentioned is indeed due to rounding. Specifically, the red number means improvement over its corresponding baseline. Baseline + ICL -> Baseline, CFL + ICL -> CFL. The sum of previous items is 89.6, and 89.6 / 3 = 29.866. And we discarded the fractional part instead of rounding it to 29.9.

Q6: Do the findings from Sec. 3.1 also hold when we train the ComFuncLearner on the convex combination F1(0)+F2(0) instead of F1(0)+F3(0)?

R6: Yes, the conclusion still holds. In the original manuscript, we presented the results of ComFuncLearner trained on F2(0)+F4(0) in Appendix D.1, Lines 1372-Line 1376. Following your suggestion, we conducted additional experiments with F1(0)+F2(0), and the results are shown in the table below. We have also updated these experimental results in the appendix, as detailed in Appendix D.1, Lines 1376-1377.

---

Q7: Why was Llama-3 chosen in Sec. 3.4? The experiments would be more consistent if they also used Llama 2, as in Sec. 4 (as they must for fair comparison with ToolLlama).

R7: We chose LLaMA-3 because it represented the state-of-the-art base model at the time of our study. Our intention was to verify that our conclusions hold even for advanced models pre-trained on large-scale datasets.

To further validate our findings on LLaMA-2 and maintain consistency in our experimental setup, we conducted additional experiments using LLaMA-2. The results are presented in the table below. We have also updated the paper with these results in Appendix D.2, Lines 1387-1397.

Q8: Can you provide results for the mixed translation task (Sec. 4.2)?

R8: Thank you for your suggestion. The complete results are shown in the table below. We have also updated the paper with these results in Appendix B.5, Lines 1026-1042.

---

W1: More in-depth analysis or discussion on this would enhance the work's generalizability and provide greater insights into the workings of intra-problem generalization.

R1: Thank you for your suggestion. We have explored the practical implications of our findings and how they can be applied to improve ICL generalization in Section 4.1 and Section 4.2. We believe that identifying base tasks within different domains and exposing the model to diverse combinations of these base tasks is crucial for improving generalization in real-world applications. For instance, in API call learning (as shown in Section 4.1), training on a small set of API combinations can significantly enhance the model's ability to generalize to unseen combinations. Here, we consider individual API calls as analogous to base functions, and API combinations as operators.

However, we acknowledge the difficulty in establishing a universal definition of base tasks and operators applicable to all downstream tasks. This requires task-specific analysis. Our work primarily provides intuitive guidance and suggestions on this foundational level. Further investigation into the interplay between base tasks and operators within specific downstream tasks is an important direction for future work.

W2: Simple task formulation that may not apply to complex natural language tasks

R2: Thank you for your suggestion. To better align with real-world applications, such as common NLP tasks, we investigated the inter / intra-problem generalization in an NLP translation task, as detailed in Section 4.2 (Line 464-Line 513). Similar to the function fitting tasks, we trained models from random initialization to avoid potential overlap between pre-training corpora and the test data.

As described in Line 506-Line 514, our findings on this task indicate that: "We observe that on simple tasks, such as translating sentences composed of a single language into another single language with simple sentence structures and shorter lengths, the improvement brought by ICL is less significant, and the gap between ComFuncLearner and Baseline is not substantial. However, on hard tasks, such as those involving mixed-language input or complex sentence structures, ICL can bring more significant improvements to ComFuncLearner, while the Baseline shows considerably less improvement. This indicates that during the training phase, if the model has been exposed to a wider range of problems, it significantly raises the upper limit of improvement achievable through ICL, thereby realizing generalization capabilities on complex tasks."

This observation is consistent with our findings in the function-fitting tasks. Therefore, the experiments presented in this paper offer valuable guidance for real-world applications.

Q1: To ensure a fair comparison, did you control the training data size to have an equal number of samples related to the target simple tasks?

R1: Yes, to ensure a fair comparison, we trained both the baseline and ComFuncLearner using the same amount of training data, including the combined data, as described in Line 313 - Line 315.

While maintaining the same total amount of training data, we explored two different approaches for selecting the training functions for the baseline:

1. Increased samples from the same base functions: The baseline was trained with a larger number of samples drawn from the same 4 base functions used to train ComFuncLearner. Detailed results are shown in Line 1345-Line 1349 in Appendix D.1.

2. Additional base function: The baseline was trained with an additional base function beyond the 4 used for ComFuncLearner. Detailed results are shown in Line 234-Line 245 in Section 3.1.

---

Both approaches yielded results consistent with our main findings: the transformer model exhibits intra-problem generalization but fails to achieve inter-problem generalization.

Q2: Have you tried a reverse inter- and intra-problem generalization setup to explore how effectively the model can learn directly from combination data?

R2: Thank you for your suggestion. To investigate whether the observed inter- and intra-problem generalization behaviors hold when the roles of base functions and their combinations are reversed, we conducted the following supplementary experiment, and have updated our paper with these findings in Appendix D.3, Line 1399-Line 1492.

Experimental Setup:

In contrast to the experimental setup in the main text, we now explore whether a model exposed to diverse combinations of base functions can generalize to base function fitting tasks.

The Baseline model is trained solely on combinations of 4 functions: { $\mathcal{F}^{(0)}_1 + \mathcal{F}^{(0)}_2$, $\mathcal{F}^{(0)}_1 + \mathcal{F}^{(0)}_4$, $\mathcal{F}^{(0)}_2 + \mathcal{F}^{(0)}_3$, $\mathcal{F}^{(0)}_3 + \mathcal{F}^{(0)}_4$ }.

The ComFuncLearner is trained on the same 4 function combinations as the Baseline model, but is additionally exposed to one base function: { $\mathcal{F}^{(0)}_1$ }.

The results demonstrate that our findings hold even when we reverse the roles of composite functions and base functions, treating function combinations as T0 and the base functions (which can be approximated as combinations of multiple combinations of functions) as T1. This further reinforces our conclusion that transformers with ICL enable intra-problem generalization but not inter-problem generalization.

Q3: Regarding the ComFuncLearner setting, did you observe any impact of data point order, such as when using a curriculum learning schedule?

R3: Thank you for your suggestion. To investigate the influence of data point order and sample difficulty on the generalization ability of the transformer, we designed the following experiments:

For the ComFuncLearner, we designed a two-stage training process. In the first 25,000 iterations, the model learns only the base functions. Subsequently, in the next 25,000 iterations, it is exposed to the combined functions. We then compare the performance of this model with a ComFuncLearner trained on a random sequence of samples. The results are presented in the following table. CFL1_CL refers to CFL1 trained with curriculum learning.

---

We find that training the model on base functions first, followed by combined functions, does not further improve its generalization ability. In fact, it results in worse generalization performance compared to a model trained with randomly sampled data. This could be attributed to the need for careful tuning of the ratio of iterations dedicated to learning base functions versus combined functions to achieve optimal performance.

Q4: Why choose $\mathcal F_5^{(0)}$ to be included in Baseline? It seems unrelated to the functions of interest and may hinder the performance.

R4: This is to ensure that the baseline and ComFuncLearner are exposed to the same diversity of data. While maintaining the same total amount of training data, we explored two different approaches for selecting the training functions for the baseline. For the comparison with Baseline trained on 4 base functions, please refer to the response to Q1 above.

Q1: Why did you not just sample more data from the 4 base function classes?

R1: Thank you for your constructive feedback. Our initial motivation for this experimental setup was to ensure that both models were exposed to the same number of functions, in order to demonstrate that cross-problem generalization cannot be achieved simply by increasing the diversity of tasks within a problem. To verify whether cross-problem generalization can be achieved solely by increasing the amount of data, we conducted additional experiments where the model was trained with more data from only 4 base functions and tested on various function fitting tasks. The results are shown in the following table. The conclusion drawn from these experiments is consistent with the main text: the transformer model only achieves intra-task and intra-problem generalization, but not inter-problem generalization.

Furthermore, we have included the supplementary results in Line 1345-Line 1349 in Appendix D.1.

---

Q2: Do you expect your results to hold for a different class of functions, s.a. polynomials?

R2: Thank you for your feedback. We have further investigated the generalization ability of the transformer model on other base functions. To ensure significant differences between the functions, we selected the following base functions:

• $y^{\text{(1)}}_{\text{L}} = \frac{\sqrt{30}}{50} \cdot \mathbf{x} \cdot |\mathbf{w}|$
• $y^{\text{(2)}}_{\text{L}} = \frac{\sqrt{2}}{50} \cdot \left( 3 \mathbf{x}^2 - 25 \right) \cdot \mathbf{w}$
• $y^{\text{(3)}}_{\text{L}} = \frac{\sqrt{70}}{500} \cdot \left( \mathbf{x}^3 - 15 \mathbf{x} \right) \cdot \mathbf{w}$
• $y^{\text{(4)}}_{\text{L}} = \frac{3 \sqrt{10}}{10000} \cdot \left( 7 \mathbf{x}^4 - 150 \mathbf{x}^2 + 375 \right) \cdot \mathbf{w}$

Using these base functions, we trained and tested both the baseline and ComFuncLearner using convex combinations as the function composition method. The results are shown in the following table. We found that using different base functions, as long as there are significant differences between them, leads to conclusions consistent with those presented in the paper. We include this experiment in Appendix D.5.

---

Q3: Do you expect similar results for non-transformer architectures?

R3：Thank you for your insightful suggestion. Our main focus was to investigate the generalization ability of transformer models in the context of ICL, and other models were not originally within the scope of our study. However, prompted by your valuable feedback, we conducted further experiments with other model architectures, including MLP and KAN. The results are presented in the tables below.

Results with MLP:

---

Results with KAN:

---

We are surprised to find that the generalization behavior observed in transformers persists in these other network architectures. This opens up exciting avenues for future research.

It is important to note that this should not be strictly termed "in-context learning" for MLP and KAN. Since these models require fixed input and output dimensions, our experimental setup necessitates feeding 39 (x, f(x)) pairs along with the 40th x as input, and the model predicts only the 40th f(x). This differs from the sequential prediction capability of transformers in ICL. However, we believe that our findings regarding generalization limitations may extend beyond transformers and ICL, potentially applying to other deep learning architectures and tasks.