Towards Auto-Regressive Next-Token Prediction: In-context Learning Emerges from Generalization

Q2: Although the theoretical analysis establish the bounds of ICL loss, but the results apply to general prompt, rather than in-context examples only. It also implies that the ICL loss should be lower and performance should be better with more training data (in terms of topic and sequence). But only well-structured in-context samples work empirically. I am uncertain if the presented theory is sufficiently tailored to address the problem of ICL.

A: Thanks for your thoughtful question! We agree with you that well-structured in-context samples can signicantly enhance performance. Although our theoretical results are derived for general prompts (with concatenated in-context examples, without assuming any specific order of concatenation), we emphasize that the effects of sample order are already implicitly captured within our framework and results. To clarify this point, we would like to make detailed discussions mainly from the following three aspects, as well as offer some inspiring guidance for future research on ICL.

(1) The characterization of conditional probability. In the setting of a more realistic AR-NTP, individual tokens are generated based on the preceding tokens. Compared to most ICL research that simply considering $(x_i, y_i)$ pairs in each sequence/prompt, we capture the dependency between tokens by ultilizing conditional probability, i.e., $\mathbb{P}(x_t \mid x_1, x_2, \cdots, x_{t-1})$, which serves as our optimization objective. When considering the expectation over sequence (generalization of sequence), we specially divide it into two expections: one is the expection over token conditioned on its prefix sequence, and the other is the expectation over prefix sequence itself. Clearly, the order of samples affects the distribution of the prefix sequence, which in turn influences the optimization process of the conditional probability. Thus, the representation of token-dependency provides a theoretical foundation for capturing the influence of sample order.

(2) The impacts of sample order on generalization can be induced from our theorems. As introduced above, in ICL, the prediction not only depends on the current input but is also influenced by the preceding context. When the model processes the in-context examples, if the sample order causes significant fluctuations (e.g., abruptly switching from a sample highly related to the query to one with very low relevance), the model may need to continuously adjust its parameters to adapt to the new input. This could lead to increased gradient fluctuations and training instability. In our theory, based on Assumption 4.5, we use $\sigma$ to describe the upper bound of the expected gradient. If $\sigma$ is large, it indicates that the magnitude of gradient changes is uncontrolled at each update, reflecting the instability of the optimization process.

Combining the above analysis and the optimization insights based on $\sigma$, it demonstrates that, a good sample order can help the model converge more stably, which means a smaller $\sigma$, a smaller generalization upper bound, and better generalization performance. In total, the impacts of sample order on generalization can be induced from our theorems.

(3) Our theory provides guiding principles for in-context examples. Our generalization theoretical analysis essentially provides an upper bound for the population loss when considering two-level expectation over topics and sequences. While this upper bound applies to general prompts, it can be used to guide the construction of better in-context examples. For instance, reducing irrelevant samples in the context (optimizing the sample order) can better satisfy Assumption 4.5, thus indirectly improving generalization performance.

Inspiring guidance for future research on ICL. We strongly agree with you that a more detailed consideration of sample order is of significant research value for a deeper understanding of ICL. In our initial exploration of both considering a more realistic AR-NTP paradigm and establishing a systematic generalization framework, we hope to provide new insights for this field!

Here, we leave a more detailed analysis of sample order for future work, but based on our current framework, we offer a few potential directions for further exploration:

My point is not about the effect of K, N, T, but controlling the train data size when you use different K/N/T. When you use a larger K, the train size is larger than the one with smaller K, all else equal. The size of training data also impact the performance which cannot be attributed to the effect of large K.

(2) Local Attention Mechanism.

• Building on the above discussion regarding optimizing sample order, when samples more closely related to the query are grouped together (a well sample-order prompt sequence), the dependencies between adjacent tokens become stronger. It enables the design of a local attention mechanism, which focuses only on a fixed-length prefix sequence of the current token during computation. This approach reduces the memory and efficiency overhead associated with long-range token dependencies.

• Although our theoretical analysis of token dependency assumes reasoning over the entire sequence of preceding tokens, the existing framework can be extended to handle scenarios where the prefix sequence $E_t$ in $\mathbb{P}(x_{t+1} \mid E_t, w_k)$ is constrained to a certain length, akin to an n-gram model. Furthermore, the theoretical techniques for addressing token-dependency remain applicable. It provides great inspiration for future work that we can indeed extend the current theory to computation-limited scenarios!

Thanks again for your suggestion! We have complement this analysis to Appendix D in the revised version! Hope that could address your concerns!

Q: Why particularly choose PAC-Bayesian framework for ICL? There are numerous works study ICL using statistical frameworks like [1] Why not based on uniform convergence or algorithmic stability, could help situate the results within a broader landscape of generalization theory?

A: Thanks for your question! Uniform convergence and algorithmic stability are indeed great techniques to analysis generalization. We specially choose PAC-Bayesian technique mainly considering the following reasons.

1. Which characterization is relatively more accurate for modeling language learning, as well as much easier to theoretically deal with token-dependency under the AR-NTP setting?

• The characterization with conditional probability in PAC-Bayesian. With language tasks like text generation or question answering, the next-token prediction paradigm is generally adapted in current LLMs. This method carefully considers the dependency in tokens for better context understanding. Thus, it's necessary to characterize this AR-NTP paradigm, and conditional probability in statistic is indeed an useful technique, i.e., $\mathbb{P}(x_t \mid x_1, x_2, \cdots, x_{t-1})$ serves as our optimization objective. As introduced in "Summary of Challenges" in Section 4, constructing ghost sequences provides a good chance that enables taking expectation over each token within all sequences, under the case of token-dependency.
• Uniform Convergence. From the perspective of problem modeling, uniform convergence theory typically relies on assumptions about the hypothesis class, especially requiring that the hypothesis class exhibits some "controllable" complexity. For large models, defining the hypothesis class or achieving good generalization performance even with high-complexity hypothesis classes often limits the practical guidance of conclusions derived from this method. Additionally, and more critically, uniform convergence assumes the training data is i.i.d. When dealing with sequential data, modeling these dependencies and deriving corresponding generalization bounds is theoretically challenging and often infeasible based on current knowledge on uniform convergence technique.
• Algorithmic Stability. The algorithmic stability technique has been ultilized by Li et al.[1] which focuses on ensuring that small changes in training data do not lead to significant fluctuations in predictions. From this perspective, compared to uniform convergence, algorithmic stability at least considers data aspect. However, it seems to bypass the modeling of token dependencies by characterizing parameter changes caused by perturbing a single token. The traditional approach can be relatively easily extended to ICL scenarios. While this method is feasible for providing generalization bounds in AR-NTP settings, it appears to skip detailed modeling of token-dependency, making the characterization of NTP less precise and somewhat limited.

Q2, Q5: The experimental setting is not fully discussed, including the detailed information about GINC datasets [...] How do you define the topic examples in your experiments? Can you give some examples of the topic in linear dynamic systems and GINC?

A: Thanks for your great consideration! In the revised paper, we have added detailed experimental settings and moved GINC experiments to the main text. Meanwhile, we have supplemented more insightful experiments on real-world language datasets, observing the optimization process and the potential effects of prior model initialization (guided by our therorems). The revised paper has been updated and we strongly encourage to read the "Experiments" section for more details! In the following, we would like to answer your questions one by one!

More experimental details about GINC.

• GINC Dataset. GINC is a small-scale language dataset generated from a uniform mixture of Hidden Markov Models (HMMs) over a family of topics/concepts. The generation steps are as follows: (1) Prepare transition matrix for HMMs: The topic/concept determines the state transition matrix in HMM. For simulation, the transition matrix is randomly generated for each topic (each HMM), respectively; (2) Prepare vocabulary: The vocabulary is generated as combinations of letters starting from a' to z', aa' to az', and so on. We can obtain vocabularies with different sizes; (3) Prepare memory matrix: A unique matrix is created that records the mapping of vocabulary and state; (4) Generate sequences: Given a fixed topic and an initial state, generate the next state based on the transition matrix, and then obtain the observed token using the memory matrix. In total, each sequence is sampled from a random HMM in the family.

• Model and Hyperparameter Settings. Our transformer model is based on the GPT-2 architecture with 4 layers, 12 attention heads, and 768-dimensional embeddings. Limited by computation resources and training data, we use "AutoModelForCausalLM.from_pretrained("gpt2")" to load a pretrained version from the HuggingFace model hub, not training from scratch, like most ICL research does. We train the model for 5 epochs using the AdamW optimizer with a batch size of 8 and a linear learning rate schedule. The schedule includes a warmup phase of 1000 steps, up to the learning rate of 8e-4. All experiments on GINC are conducted using a single 24GB NVIDIA GeForce RTX 3090.

• We empirically explore the separate effects of the number of topics ($K$), number of sequences per topic ($N$), sequence length ($T$) and prompt length ($T_p$). We detail $K \in \\{10,20,30\\}$, $N \in \\{20,40,60,80,100\\}$, $T \in \\{1280, 2560, 5120, 10240\\}$, $T_p \in \\{8, 16, 32, 64\\}$, where ranging $T$ with directly masking the token that exceeds the specified length and do not taking special consideration. In totoal, we arrange groups of comparative experiments to verify that increasing $K, N, T, T_p$ individually improves the model's generalization performance as demonstrated in our Theorems. Additionally, we discuss the effect of vocabulary size and provide an interesting case involving with a failed ICL.

• More details are provided in Section 5 and Appendix C.

Topics in linear dynamic systems and GINC.

• For linear dynamic systems, we provide a detailed discussion in Appendix C.1. We adopt numerical experiments to simulate the gold state equation $x_{t+1} = W x_t + \zeta_t$, where different $W \in \mathbb{R}^{d\times d}$ represent different tasks and different dimensions represent the varying difficulities of tasks. Thus, the topics indeed correspond to the randomly generated weight matrices $W$.
• For the sythetic dataset GINC, given the above introduction of its generating process, we would prepare transition matrix for each HMM. For simulation, the transition matrix is randomly generated for each topic (each HMM), respectively. Since the topic/concept uniquely determines the state transition matrix in HMM, we would use the transition matrix to represent the concept, rather than explicitly defining the concept.

Dear Reviewer ARXg,

Thanks for your further questions! We are more than willing to address all of your concerns! Regarding the assumption of bounded gradient, we find that there exists a mismatch between our views. For the connection of "generalization" and "ICL emergence", we further dissect the complete process from pre-training to ICL, including the role played by generalization performance metrics. Finally, we summarize the impact of our paper from research question, theoretical framework and techniques, and practical applications supported by theoretical and experimental evidence.

Further clarification for Q4: Is the Bounded Gradient a strong assumption for the Transformer and recent large language model? After multiple layers of nonlinear superposition in the Transformer, the network output may be highly sensitive to input changes.

A: Thanks for your question! We have reconsidered your questions that there exists a mismatch between your concerns about "the network output being highly sensitive to input changes" and our assumption of bounded gradient. These correspond to the gradient over input data ($\nabla_x L(x; \theta)$) and the gradient over model parameters ($\nabla_\theta L(x; \theta)$), respectively. Below, we will provide a more detailed explanation!

1. The gradient over input data reflects model robustness.

The gradient over input data ($\nabla_x L(x; \theta)$) reflect the sensitivity of model outputs to input changes. As you mentioned, transformers are built on multiple layers of complex nonlinear transformations, causing large changes in the model's output. In practice, there are many methods to improve a model's stability to input perturbations, such as data augmentation and adversarial training. Theoretically, if the gradient over input data is bounded, the model would be robust to input noise, allowing it to generalize better to unseen data.

However, we have not analyzed model generalization from the perspective of robustness. Instead, we consider optimization-dependent generalization bounds, which are closely linked to the optimization process. Our assumption, the boundedness of gradients over model parameters reflects the stability of the training process.

2. The bounded gradient over model parameters is reasonable.

Our bounded gradient assumption focuses on the gradient over model parameters, $\nabla_\theta \mathcal{L}(x;\theta)$, as we answered before, we mainly explain the reasonableness of the bounded gradient over model parameters. From both practical and theoretical perspectives:

• Practically, gradient explosion is a common issue in deep learning optimization, leading to unstable convergence or failure to converge. To address this, many training strategies (e.g., gradient clipping and regularization) explicitly or implicitly enforce constraints on gradient magnitudes. This highlights that bounding parameter gradients is a widely adopted engineering practice. For LLMs which involve numerous parameters, controlling gradient magnitudes is especially critical. Bounded gradients on parameters ensure training stability and avoid numerical issues.

• Theoretically, in the optimization process, especially when using gradient descent, assuming that the gradient is bounded helps ensure the convergence of the algorithm. If the gradient is too large, it can lead to excessively large parameter updates, causing instability in the training process and preventing convergence. When the gradient is bounded, it ensures that each update step is within a reasonable range, allowing the algorithm to gradually approach the optimal solution and ultimately achieve convergence.

In summary, the boundedness of gradients over input data (as you considered) is more directly related to model robustness. However, our bounded gradient assumption pertains specifically to gradient over parameters and is justified by its role in ensuring optimization stability, when considering optimization-dependent generalization bounds.

We hope the above effectively addresses the mismatch between our understandings and resolves your concerns!

(2) Token-dependency setting brings great proof difficulties. As "Summary of challenges" introduced in Section 4, under the setting of AR-NTP, we make great efforts to address the dependency between the current token and its preceding tokens by constructing ghost sequences (see the detailed construction in Appendix G.2.1, where we summarize the proof sketch), thereby enabling the possibility of taking expectation over each token within all possible sequences.

In the following, the token-dependency introduces a crucial connection of negative logarithm likelihood, KL divergence and TV distance (which is important for the definition and upper bounds of population loss). Specifically, we begin by examining the primary optimization objective: negative logarithm likelihood. Naturally, this leads to a connection with KL divergence, thereby formalizing the expression of population loss. Furthermore, in addressing the aforementioned token-dependency, we establish connections between TV distance and the expectation over a single token when given its predecessors. Therefore, it's necessary to establish connections between the two key distribution metrics: TV distance and KL divergence (see in Lemma G.7), to obtain our final generalization error bounds. The AR-NTP setup necessitates the establishment of the above series of connections, which are not considered in the previous ICL work.

(2) Using Small Model Training as Warm-up for Large Models: The prior model initialization strategy discussed above considers training the model once in advance with the same architecture as the formal training. This approach can be further extended to provide insights for training large models.

Specifically, prior knowledge can be acquired by first training a relatively smaller model with a different architecture. It enables effective initial parameters at a lower computational cost, providing a solid foundation for larger models and avoiding the instability and non-convergence issues that may arise from random initialization. The detailed analysis of $D_{KL}(\mu \parallel \nu)$ presented earlier serves as the theoretical understanding for the "small model warm-up" strategy. Furthermore, this approach has been successfully applied in engineering practices, including AutoML [2], Neural Architecture Search (NAS) [3] and current LLMs training.

(3) Gradual Expansion of Training Data: The strategy of expanding the training data involves beginning training with a small subset and gradually increasing the dataset size. In this process, the model's initial learning can be seen as based on a "data-dependent prior", and each expansion of the training data can be understood as the injection of a new model prior. Based on the analysis of $D_{KL}(\mu \parallel \nu)$ above, gradual expansion of training data similarly leads to improved generalization, faster convergence, and better handling of complex features. This guiding principle is also reflected in Curriculum Learning [4] and Progressive Networks [5].

2. Practical Guidance for Training Data Selection and Deduplication. It is well-known that the vast amount of data obtained from the internet serves as input for compressing world knowledge into LLMs [6]. In the redundant data, high-quality data determines the upper limit of the performance of LLMs. Therefore, considering a data-dependent pre-training and ICL generalization framework has immense potential for guiding data. In our theory, to explicitly show the impact of data, we adopt a data-dependent and topic-dependent prior $\nu_J$ and further detail $D_{KL}(\mu \parallel \nu)$ with optimization analysis. We have discussed this in detail before: in Practical Guidance for Model Training part, we emphasize the advantages of prior model initialization over random initialization in model training. Here, we aim to further explore its guidance for training data from the perspective of compression.

Specifically, we can select a subset of size $K^\prime$ from the $K$ pre-training topics to estimate a prior distribution. If a smaller $K^\prime$ can estimate a prior that is very close to the posterior distribution, it indicates that the information from the $K$ topics can actually be compressed into a smaller subset of $K^\prime$ topics. This reflects the compressibility of the data, and can thus backward guide pre-training data to further undergo data selection and deduplication, such as through topic clustering, data diversity, or information gain metrics (e.g., $D_{KL}(\nu(D) \parallel \nu(D_i))$, if this value is small, the data block $D_i$ is considered redundant and can be reduced in weight or removed to decrease the model's reliance on redundant information.) The reprocessed pretraining data may exclude some noise interference, further improving model performance, saving computational resources, and facilitating training for new models.

In summary: Based on the analysis above, the bounds derived by Li et al.[7], based on algorithmic stability techiniques, mainly reveal the impact of multi-task learning and sequence length. Zhang et al.[8] approach the problem by considering model distribution, however their results involve uniform priors, which are similar to random initialization, relatively having limitations in practical guidance. In contrast, our Bayesian approach characterizes model performance from a statistical perspective. In addition to the statistical insights regarding the number of pre-training topics, sequences and sequence length, the data-dependent prior also provides practical guidance on model training, data selection and deduplication, which is a unique contribution.

The above analysis of practical impacts has been supplemented in Appendix D. The revised paper has been updated accordingly!

Reference:

[1] Shannon, C. E. (1948)A mathematical theory of communication.

[2] He, X. (2021)AutoML:A survey of the state-of-the-art.

[3] Elsken, T. (2019)Neural architecture search: A survey.

[4] Bengio, Y.(2009)Curriculum learning.

[5] Rusu, A. A. (2016)Progressive Neural Networks.

[6] Delétang, G.(2023)Language modeling is compression.

[7] Li, Y. (2023)Transformers as algorithms: Generalization and stability in in-context learning.

[8] Zhang, Y. (2023)What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization.