Learning In-context $pmb{n}$-grams with Transformers: Sub-$pmb{n}$-grams Are Near-stationary Points

2025-08-18

AI Summary

This paper investigates the phenomenon of staged progress and prolonged plateaus observed during Transformer training for context-n-gram language modeling, focusing on the geometric structure of the loss landscape under cross-entropy loss. Method: We introduce the concept of “sub-n-gram estimators” and rigorously prove that, as sequence length tends to infinity, any k-gram parameter configuration (with k ≤ n) constitutes an approximate stationary point of the population cross-entropy loss—its gradient vanishes asymptotically, enabling discrete transitions. Our analysis integrates a simplified Transformer architecture, population-loss modeling, asymptotic gradient analysis, and numerical experiments. Contribution/Results: We reveal that learning proceeds via sequential discrete jumps along a chain of such approximate stationary points. This work provides the first unified, loss-landscape-based theoretical explanation for staged learning and phase transitions in Transformers, establishing a formal link between n-gram estimation capacity and optimization dynamics.

Technology Category

• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search
• Machine Learning: Structured Learning

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Motivated by empirical observations of prolonged plateaus and stage-wise progression during training, we investigate the loss landscape of transformer models trained on in-context next-token prediction tasks. In particular, we focus on learning in-context $n$-gram language models under cross-entropy loss, and establish a sufficient condition for parameter configurations to be stationary points. We then construct a set of parameter configurations for a simplified transformer model that represent $k$-gram estimators (for $k leq n$), and show that the gradient of the population loss at these solutions vanishes in the limit of infinite sequence length and parameter norm. This reveals a key property of the loss landscape: {sub-$n$-grams are near-stationary points of the population cross-entropy loss}, offering theoretical insight into widely observed phenomena such as stage-wise learning dynamics and emergent phase transitions. These insights are further supported by numerical experiments that illustrate the learning dynamics of $n$-grams, characterized by discrete transitions between near-stationary solutions.

Problem

• Analyzing loss landscape of transformer models
• Studying stationary points in n-gram learning
• Exploring stage-wise learning dynamics in transformers

Innovation

• Investigates loss landscape of transformer models
• Constructs parameter configurations for n-gram estimators
• Reveals sub-n-grams as near-stationary points