A Law of Next-Token Prediction in Large Language Models (2408.13442v1)

Abstract: LLMs have been widely employed across various application domains, yet their black-box nature poses significant challenges to understanding how these models process input data internally to make predictions. In this paper, we introduce a precise and quantitative law that governs the learning of contextualized token embeddings through intermediate layers in pre-trained LLMs for next-token prediction. Our findings reveal that each layer contributes equally to enhancing prediction accuracy, from the lowest to the highest layer -- a universal phenomenon observed across a diverse array of open-source LLMs, built on architectures such as Transformer, RWKV, and Mamba. We demonstrate that this law offers new perspectives and insights to inform and guide practices in LLM development and applications, including model scaling, pre-training tasks, and information flow. Overall, our law enables more fine-grained approaches to the design, training, and interpretation of LLMs through scrutinizing their internal data processing mechanisms.

A Law of Next-Token Prediction in LLMs

The paper "A Law of Next-Token Prediction in LLMs" by Hangfeng He and Weijie J. Su presents a rigorous and quantitative analysis of the learning process of contextualized token embeddings across the layers of pre-trained LLMs used for next-token prediction. This paper is motivated by the significant opacity associated with the internal workings of LLMs, despite their widespread applications in diverse fields such as medical diagnosis, genomic research, and education.

The central contribution of the paper is the identification and characterization of a universal law termed as the "law of equi-learning". This law governs how LLMs progressively enhance their predictive accuracy across intermediate layers. The authors demonstrate through extensive experiments that each layer in these models, from the lowest to the highest, contributes equally to the process of next-token prediction. This finding is consistent across a variety of open-source LLMs, including those based on Transformer, RWKV, and Mamba architectures.

Main Findings

1. Exponential Improvement Across Layers: The authors found that LLMs exhibit an exponential improvement in the ability to predict the next token across layers. This phenomenon is quantified by the prediction residual (PR), a metric that measures the variance of the actual next token unexplained by the layer’s token embeddings. The relationship is expressed as:

$\text{PR}_l \approx \rho^{l-1} \times \text{PR}_1$

where $0 < \rho < 1$. The consistent and roughly constant reduction in the logarithm of the PR value across each layer led to the term "equi-learning".

1. Universality Across Models and Data: This law holds for an array of open-source LLMs including older models like GPT-1 and GPT-2, as well as more recent ones like Llama-3, Mistral, phi series, and RWKV. The paper also extended to different probing datasets (e.g., BookCorpus, C4, OpenWebText) and various domains including medicine, law, and politics, reinforcing the universality of the equi-learning law.
2. Impact of Training Dynamics: The emergence of the equi-learning law is investigated across various training stages. Analyses revealed that the law becomes apparent after a certain number of training steps, epochs, and data repetitions, highlighting the influence of training dynamics on the manifestation of this geometric pattern.

Implications and Future Directions

Practical Implications:

Understanding the equi-learning law has significant practical implications:

• Model Scaling: The law provides nuanced insights into how increasing the model size affects the feature learning process, suggesting that larger models, although demonstrating superior overall predictive capabilities, may distribute learning more uniformly across layers.
• Training Task Choices: The choice of pre-training tasks is critical. The paper suggests that tasks such as next-token prediction (NTP) may be more conducive to the emergence of the equi-learning law compared to masked LLMing (MLM) or span corruption (SC).
• Information Flow: The findings offer a deeper understanding of internal information flow in LLMs. As models process layers, they tend to forget prior information while increasingly focusing on refining predictions of future tokens.

Theoretical Implications:

Theoretically, the equi-learning law challenges existing views on feature learning across layers in LLMs:

• Layer Contribution: The law contradicts the assumption that certain layers might disproportionately contribute to prediction capabilities. Instead, it emphasizes equal importance across all layers for characterizing features from embeddings.
• Training Objectives: Analyzing how the decay ratio $\rho$ is influenced by model depth, pre-training data, and architecture could lead to optimizing the training process by minimizing the overall decay ratio ($\rho^{L-1}$).

Future Research Directions:

Several promising avenues for future research emerge from these findings:

• Layer-wise Learning Rates: Investigating varying learning rates across layers to achieve faster convergence and equilibrium as described by the equi-learning law.
• Model Pruning and Fine-tuning: Preserving the equi-learning law during model pruning and fine-tuning may enhance model robustness and efficiency.
• Transfer Learning: Leveraging the law in transfer learning contexts, where lower layers may be frozen, and upper layers re-trained to adapt to new domains.
• Robust Evaluation Metrics: Further exploring the role of high-quality probing data in capturing the true capabilities of high-quality pre-trained LLMs could refine evaluation metrics.

In summary, the equi-learning law presents a fundamental and insightful perspective on the internal dynamics of LLMs. The paper's identification of this quantitative pattern across diverse models and datasets offers valuable theoretical and practical benefits, paving the way for more transparent and efficient development of LLMs in the future.