Learning Representations through Token Prediction: Geometry, Approximation, and Downstream Guarantees
Abstract
Token prediction is a central pre-training objective for modern language models. Despite its empirical success, why token prediction learns broadly useful representations remains incompletely understood. We develop a statistical framework connecting token prediction with representation geometry, encoder approximation, and downstream performance. Under a softmax prediction head, we show that accurate token prediction organizes token embeddings according to similarities between the distributions of contexts in which different token types appear, as measured by Hellinger distance, with explicit errors governed by prediction accuracy and token frequency. Meanwhile, the contextual representation provides a low-dimensional coordinate for the conditional distribution of the target token relative to these embeddings. We further introduce a self-consistency principle showing that repeated applications of a shared representation block can progressively refine the contextual representation without introducing additional block parameters. Among representations with the same prediction accuracy, this recurrent construction favors those that can be stably reconstructed from their contexts. Finally, we establish downstream guarantees for token generation, token community recovery, and classification by a linear probe, showing how prediction accuracy and recovered geometry translate into performance beyond the pre-training objective. Together, these results explain how the simple objective of predicting tokens can recover semantic geometry and produce broadly useful representations. A controlled simulation illustrates the theoretical mechanisms.
Related
- Improved Approximation Guarantees for Periodic-Review Stochastic Inventory Control
- Online Line Aggregation with Deadlines: Randomized Guarantees and Learning-Augmented Tradeoffs
- Learning in Inverse Games: Tractable Training with Probabilistic Guarantees
- Provable Guarantees and Efficient Learning of Structural Equation Models with Latent Confounders
- Learning to Solve Two-Stage Stochastic Unit Commitment Problems with Quality Guarantees
BibTeX
@misc{wang2026learning,
title = {{Learning Representations through Token Prediction: Geometry, Approximation, and Downstream Guarantees}},
author = {Shulei Wang},
howpublished = {arXiv},
year = {2026},
month = aug,
doi = {10.48550/arxiv.2608.30072},
eprint = {2608.30072},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2608.30072},
}
RIS
TY - UNPB
TI - Learning Representations through Token Prediction: Geometry, Approximation, and Downstream Guarantees
AU - Shulei Wang
PY - 2026
DA - 2026/08/30
PB - arXiv
DO - 10.48550/arxiv.2608.30072
UR - https://arxiv.org/abs/2608.30072
AB - Token prediction is a central pre-training objective for modern language models. Despite its empirical success, why token prediction learns broadly useful representations remains incompletely understood. We develop a statistical framework connecting token prediction with representation geometry, encoder approximation, and downstream performance. Under a softmax prediction head, we show that accurate token prediction organizes token embeddings according to similarities between the distributions of contexts in which different token types appear, as measured by Hellinger distance, with explicit errors governed by prediction accuracy and token frequency. Meanwhile, the contextual representation provides a low-dimensional coordinate for the conditional distribution of the target token relative to these embeddings. We further introduce a self-consistency principle showing that repeated applications of a shared representation block can progressively refine the contextual representation without introducing additional block parameters. Among representations with the same prediction accuracy, this recurrent construction favors those that can be stably reconstructed from their contexts. Finally, we establish downstream guarantees for token generation, token community recovery, and classification by a linear probe, showing how prediction accuracy and recovered geometry translate into performance beyond the pre-training objective. Together, these results explain how the simple objective of predicting tokens can recover semantic geometry and produce broadly useful representations. A controlled simulation illustrates the theoretical mechanisms.
ER -