Generalization properties of the token perceptron: the role of embedding and vocabulary

Dario Bocchi; Matteo Negri; Flavio Nicoletti; Francesco D'Amico

cond-mat.dis-nnOct 23, 2026

Generalization properties of the token perceptron: the role of embedding and vocabulary

Authors: Dario Bocchi, Matteo Negri, Flavio Nicoletti, Francesco D'Amico

Organizations: Physics Department, University of Rome Sapienza, Rome, Italy · Institute of Nanotechnology, CNR-NANOTEC, Rome, Italy · Laboratoire de Physique Théorique et Modélisation, CY Cergy Paris Université · Department of Computer Science and Engineering Chalmers University of Technology and University of Gothenburg, Sweden

Abstract

Large language models are trained to predict tokens from a finite vocabulary, each represented by a vector in a $d$ -dimensional embedding space. Therefore, token prediction is a classification problem whose classes have a geometry determined by the embedding. We study how this geometry affects learning in a minimal solvable model: a linear predictor that maps a sequence of $N$ embedded tokens to a $d$ -dimensional output through one $d\times d$ matrix per position, trained in the teacher--student setting. We call this the token perceptron. We compute its generalization error with the replica method, taking $N\to\infty$ at fixed number of examples per position $α=P/N$ and keeping $d$ and the vocabulary size $D$ finite. As a first reference case, we consider continuous embeddings on the sphere, which isolate the role of the embedding dimension. The natural load is then $α/d$ , and for $d>1$ a normalized label fixes the direction of the teacher output, so the problem becomes a noiseless regression: the teacher is recovered exactly above $α_{\rm rec}=d^2/(d-1)$ , and spherical cross-entropy reaches this threshold at large inverse temperature. In the same setting the storage capacity of random associations scales as $d^2$ for margins $O(1/\sqrt{d})$ . We then turn to vocabulary embeddings, where a label only identifies the Voronoi cell of the teacher output, so that each example imposes inequalities rather than equations, as in the classical perceptron. We compute the Bayes-optimal error for a given embedding matrix through an approximate spectral closure and show that the optimal learner first resolves the directions in which the vocabulary has the largest variance, even when the teacher is isotropic.

Figures & tables

Figure 1: Schematic representation of the learning problem with continuous (top) and vocabulary (bottom) token embeddings. In both settings, the token perceptron maps the input sequence $\bm{\xi}=(\bm{\xi}_{1},\ldots,\bm{\xi}_{N})$ to the field $\bm{u}_{W}(\bm{\xi})=(Nd)^{-1/2}\sum_{i=1}^{N}\mathbf{W}_{i}\bm{\xi}_{i}$ . (Top) continuous embeddings, with input vectors on the unit sphere $\mathbb{S}^{d-1}$ and prediction $\hat{\bm{y}}(\bm{\xi})=\bm{u}_{W}(\bm{\xi})/\|\bm{u}_{W}(\bm{\xi})\|$ . (Bottom) vocabulary embeddings, with each input vector drawn from a fixed vocabulary $\{\bm{z}_{1},\ldots,\bm{z}_{D}\}\subset\mathbb{S}^{d-1}$ . The prediction is the vocabulary vector with maximal overlap with $\bm{u}_{W}(\bm{\xi})$ .

Figure 2: Using the rescaled load, the generalization error

\epsilon_{g}

curves converges for large embeding dimension. We compare analytical predictions (solid lines) to numerical results (dots) for the generalization error

\epsilon_{g}

of the spherical cross-entropy loss with

\tilde{\lambda}=\lambda/d=1

versus the rescaled load

\tilde{\alpha}

for different values of the embedding dimensions

d

Figure 3: Spherical cross-entropy learning curves at infinite

\widetilde{\lambda}

. Generalization error

\epsilon_{g}

versus

\widetilde{\alpha}=P/(Nd)

for different embedding dimensions, obtained by minimizing cross-entropy at finite

\lambda

and then taking

\widetilde{\lambda}\to\infty

. Vertical dashed lines mark the exact-recovery thresholds

\widetilde{\alpha}_{\rm rec}=d/(d-1)

for

d>1

, Eq. ( 30 ). The black dashed curve shows the

d\to\infty

limit, with threshold

\widetilde{\alpha}_{\rm rec}=1

Figure 4: For

d>1

, generalization error decreases as

\tilde{\lambda}

increases and eventually reaches zero for

\tilde{\lambda}\to\infty

for

\tilde{\alpha}>\frac{d}{d-1}=\tilde{\alpha}_{rec}(d)

. Generalization error

\epsilon_{g}

versus

\widetilde{\lambda}=\lambda/d

\widetilde{\alpha}=0.5

(left) and

\widetilde{\alpha}=5

(right). At low load, all curves retain a positive error. At high load, cross-entropy recovers the teacher for every

d>1

, differently from the

d=1

case where error stays finite.

Figure 5: Memorization capacity of random associations. (Left) Critical load per degree of freedom

\alpha_{c}/d^{2}

for standard alignment, Eq. ( 35 ). The curve shows the Gardner prediction in Eq. ( 35 ). Markers show simulations at

d=1,2,3,4

with

N=1000

. (Right) Critical load per degree of freedom for vector alignment, Eq. ( 38 ). Theoretical curves for

d=2,3,4,8,16,64

(continuous lines) and the

d\rightarrow\infty

limit of eq. ( 39 ) (dashed curve) are shown together with numerical simulations for

d=2,3,4

of finite systems of size

N=1000

Figure 6: Learning curves of generalization error for vocabulary classification task. We compare the approximate Bayes-optimal analytical learning curves derived with the spectral anisotropic replica computation with real students trained minimizing cross-entropy loss without constraints on weights norm. The two minimization procedures are different in principle, but the fact that they are close imply that cross-entropy is able to find a solution with

\epsilon_{g}

close to the optimal, similarly to [ 2 , 13 ] . Both results are explicitly dependent on the extracted vocabulary

Z

, and therefore we generate 16 different random spherical vocabularies for each

d,D

couple and we obtain prediction and numerical curves by averaging the error over the same 16 vocabularies. A numerical simulation with increasing number of variables

N

would have the Bayes-optimal predictions increasingly becoming lower bounds for

\epsilon_{g}

. We note that when

d<D

numerical results have higher

\epsilon_{g}

then predictions as expected, while instead when

d\geq D

some of the numerical points have lower

\epsilon_{g}

. This is compatible with the fact that our computation is an approximation of the full

d\times d

problem, that we expect to work best in the regime

d<D

Figure 7: Learning curves scale correctly by normalizing amount of data via

\widetilde{\alpha}

. (Left) D=16 case; we note that here

d=32>D

and in that case

\mathrm{r}_{Z}

saturates and the curves do not shift anymore when plotted against

\alpha

, as occurs in Fig. 6 . (Right) D=128; in this case all curves have

d<D

and therefore the problem remains full rank. We observe that, even if the curves are rescaled correctly, there is still a dependence on

d,D

due to the term

H_{Z}

that depends explicitly on the vocabulary.

Panel of Figure 8 — Figure 8: The optimal strategy is to learn high variance vocabulary embeddings direction first, then weaker directions as more data is available. (top left) case $d=8,D=128$ , (top right) case $d=16$ , $D=128$ , (bottom left) case $d=8$ , $D=16$ , (bottom right) case $d=16$ , $D=16$ . In all cases we observe that the normalized overlap $q_{m}^{*}$ increase monotonically with the corresponding $\lambda_{m}$ . Therefore, even if the task is isotropic (learning an isotropic spherical teacher), the anisotropy of the vocabulary affects the geometry of the inference. Saddle-point from replica computation defines implicitly the relations $q^{*}_{m}(\lambda_{m})$ that control how learning happens in each spectral direction $m$ .

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Schematic representation of the minimization of the potential . We show how the square distance to points $(z,\bar{t})$ is minimised under the different constraints of SA and VA storage. Left panel : the minimisation for SA storage. For different points (full dots) in the $z,\bar{t}$ half-plane, we show the projections (crosses) on the satisfability half-plane $z>\sqrt{d}\kappa$ , indicated in green, corresponding to the points $h_{*},t_{*}$ that minimise the potential on the constraint set $\mathcal{C}=\{h\in\mathbb{R},t\geq 0:h\geq\sqrt{d}\kappa\}$ . Right panel : the minimisation for VA storage. Again, points in the $z,\bar{t}$ plane (full symbols) are associated to the corresponding points $h_{*},t_{*}$ that minimise the potential on the constraint set $\mathcal{C}=\{h\in\mathbb{R},t\geq 0:h\geq\kappa_{\rm eff}t\}$ (crosses). In the blue region $\mathcal{C}=\{h\in\mathbb{R},t\geq 0:-\frac{t}{\kappa_{\rm eff}}\leq h\leq\kappa_{\rm eff}t\}$ , to minimise the potential each point is projected onto the line $z=\kappa_{\rm eff}t$ . In the orange region ( $z<-\frac{\bar{t}}{\kappa_{\rm eff}}$ ), all points are projected to the origin.

Language Model Generalization

See all 29 papers

Aug 30, 2026

Learning Representations through Token Prediction: Geometry, Approximation, and Downstream Guarantees

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.

Shulei Wang

Department of Statistics, University of Illinois at Urbana-Champaign, 605 E. Springfield Ave., Champaign, IL 61820

Language Model Pretraining Next-Token Prediction Predictive Representation Learning Representation Geometry in Language Models Semantic Representation Learning

Ask PDF

Oct 2, 2026

Do Language Models Need a Trainable Input Embedding Table? Fixed Minimal Token Codes at 1.7B-Class Scale

A trainable input embedding table assigns each vocabulary item an independently adjustable vector. We investigate whether this token-specific parameterization is required for substantial language-modeling capability, or whether a shared Transformer can learn from fixed token identities. We compare three decoder-only language models trained from scratch with the same tokenizer, contextual backbone, untied output-head architecture, and training recipe, with a target budget of 100 billion prediction tokens per model. Their input interfaces are a learned table, canonical 16-bit token-ID codes, and one fixed invertible recoding over GF(2). The fixed codes are repeated to model width without an additional trainable input projection. Both fixed-code models acquire substantial capabilities: canonical codes achieve 52.40% HellaSwag normalized accuracy, 70.51% PIQA accuracy, and 42.75% LAMBADA accuracy. The learned-input control performs better on several evaluations, including HellaSwag and LAMBADA, so these results establish viability rather than performance parity. The fixed interfaces remove 100.7 million trainable parameters, yielding 1.711B-parameter models, but parameter reduction is not the central result. These single-run experiments distinguish architectural necessity from empirical utility: independently trainable token-specific input vectors are not required for the observed capabilities. A fixed identity interface also provides a controlled setting for studying representation learning downstream of an immutable input, without establishing where particular capabilities are localized.

A. Bochkov

LLM Training Word Embeddings Language Model Pretraining Decoder-Only Language Models

Ask PDF

Jun 1, 2026

Representational Capacity: Geometric Limits on Feature Representation in Transformer Language Models

Model dimension (

d_{model}

) is a fundamental hyperparameter in transformer language models, yet its role in setting the geometric limits of feature representation remains under-explored. Grounded in the Linear Representation and Superposition Hypotheses - which propose that models encode features as near-orthogonal directions in latent space - we develop a framework for estimating how many such directions a model can support. We first establish the embedding matrix as a measurable proxy for near-orthogonality constraints across the latent space: the boundary between meaningful token relationships and incidental similarity in the pairwise cosine similarity distribution gives a concrete estimate of the model's accepted deviation

\varepsilon

from perfect orthogonality. Applying this metric across dozens of open-source models reveals two classes: models with high

\varepsilon

whose embeddings lack near-orthogonal structure, and models with low

\varepsilon

that maintain it. We then show that the standard Johnson-Lindenstrauss lemma greatly underestimates the packing efficiency of trained representations, and derive an adjusted capacity formula in which the number of near-orthogonal directions depends on the ratio of vectors to dimensions (

k/d

) rather than the raw count - a single modification that cuts prediction error by two orders of magnitude with no extra parameters. Combining these results, we define representational capacity as an upper bound on the number of distinguishable directions available for features and embeddings in a model's latent space. Capacity is exponentially sensitive to

\varepsilon

, and larger models favor tighter orthogonality constraints over maximizing raw capacity - a pattern compatible with several explanations (a stability-capacity trade-off, a ceiling on usable concepts, or confounds with model scale) that we leave to future work.

Alexander Guha

Arizona State University

Feature Superposition Transformer Expressivity Neural Representation Geometry Linear Representation Hypothesis Representation Disentanglement Representation Geometry in Language Models

Ask PDF

Abstract

Language Model Generalization

Explore similar work

Learning Representations through Token Prediction: Geometry, Approximation, and Downstream Guarantees

Do Language Models Need a Trainable Input Embedding Table? Fixed Minimal Token Codes at 1.7B-Class Scale

Representational Capacity: Geometric Limits on Feature Representation in Transformer Language Models