PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.32788v1cond-mat.dis-nncond-mat.stat-mechcs.LGstat.ML

TL;DR

The paper studies how embedding geometry shapes token learning in a minimal sequence-prediction model. Using replica analysis, it finds exact recovery for continuous spherical embeddings above a critical load, while vocabulary learning first resolves high-variance embedding directions.

  • Problem

    How embedding geometry affects learning and generalization in token prediction matters because vocabulary embeddings define the geometry of the classification classes.

    Source: §I · p. 3, §IV · p. 14

  • Method

    The token perceptron assigns a d × d linear predictor to each sequence position, and its generalization error is analyzed with the replica method as N tends to infinity at fixed α = P/N.

    Source: §I · p. 2

  • Results

    α_rec = d^2/(d −1) is the exact-recovery threshold for d > 1 with continuous spherical embeddings; for vocabulary embeddings, the optimal learner first resolves directions with the largest variance.

    Source: §I · p. 2

  • Takeaways & Limitations

    Continuous labels yield noiseless regression, whereas vocabulary labels constrain weights through inequalities and make learning depend on the geometry of vocabulary cells.

    Source: §I · p. 2

    The model uses explicit positional dependence, does not jointly learn embeddings, and omits sequence temporal structure needed for next-token prediction.

    Source: §V · p. 20

Abstract

from arXiv · show

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 \textit{token perceptron}. We compute its generalization error with the replica method, taking $N\to\infty$ at fixed number of examples per position $\alpha=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 $\alpha/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 $\alpha_{\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.

I. INTRODUCTION AND RELATED WORK

Token prediction is classification over classes whose geometry comes from their embeddings, a feature absent from classical scalar-perceptron models. The paper studies this geometry in a sequence-based teacher–student token perceptron, separating continuous spherical embeddings from finite vocabularies.

I. INTRODUCTION AND RELATED WORK

  • The token perceptron extends the perceptron to embedded sequences with one d × d weight matrix per position, taking N → ∞ at fixed α = P/N while keeping d and D finite.This scaling targets the sequence length as the largest scale and uses the replica theory of the perceptron.Source: §I · p. 2
  • Continuous spherical embeddings isolate embedding dimension, while vocabulary embeddings make the learning problem depend on the geometry of a discrete set of token vectors.The paper studies how embedding dimension, vocabulary size, and vocabulary-vector arrangement affect learning and the amount of data needed.Source: §I · pp. 1–2

A. Data: sequences of embedded tokens

  • Input tokens are sampled independently across positions and examples from either the unit sphere or a fixed vocabulary, making the model supervised rather than a model of natural-sequence statistics.For the vocabulary case, the embedding matrix is fixed during training; the paper notes that this simplifies settings where embeddings are learned jointly.Source: A. Data: sequences of embedded tokens · pp. 3–4

B. Task and training procedure

  • A teacher assigns labels by projecting its field onto the closest allowed output, and a same-architecture student learns weights by minimizing a per-example loss.For vocabulary embeddings, the predicted and correct labels are the vocabulary vectors with maximal overlap with the student and teacher predictions, respectively.Source: B. Task and training procedure · pp. 4–5

C. Exponential losses and learning rules for embeddings

The paper uses a common conditional-distribution construction to define losses for continuous and discrete embeddings. For continuous spherical embeddings, a replica calculation reduces the analysis to saddle-point equations for student–teacher and replica overlaps.

C. Exponential losses and learning rules for embeddings

  • The learning rule uses a locally normalized conditional distribution, with embedding geometry entering through the measure in its normalization term.This construction is presented as a supervised generalization of cross-entropy to multidimensional labels.Source: C. Exponential losses and learning… · pp. 5–6
  • For continuous unit-norm outputs, the conditional distribution is von Mises–Fisher and yields spherical pseudolikelihood, also called spherical cross-entropy.When d = 1, the continuous and two-token vocabulary forms reduce to common binary cross-entropy (logistic loss).Source: C. Exponential losses and learning… · p. 6

A. Analytical overview for the continuous case

  • The continuous-embedding analysis uses the replica method and an isotropic replica-symmetric ansatz to characterize the thermodynamic-limit overlaps.The replicated free energy separates an entropic term for student-weight volume from an energetic term for training patterns; at zero temperature its saddle points characterize training-loss minimizers.Source: A. Analytical overview for the… · pp. 7–8

B. Embedding regression generalization error and comparison of losses

The section relates continuous-embedding generalization error to the student–teacher overlap and compares analytical predictions with numerical simulations. In the rescaled-load plot, the error curves converge for large embedding dimension.

B. Embedding regression generalization error and comparison of losses

  • For fixed d, the continuous-embedding generalization error is determined by student–teacher overlap r, yielding ϵg as a function of α and d through r(α, d).The error averages over unseen inputs; its expression is given explicitly in the paper.Source: B. Embedding regression generalization… · pp. 8–9
  • Figure 2 plots spherical cross-entropy generalization error against rescaled load, using solid lines for analytical predictions and dots for numerical results.The numerical student is obtained by optimizing the loss for a problem realization with L-BFGS-B.Source: B. Embedding regression generalization… · p. 9
  • The generalization-error curves converge for large embedding dimension when plotted against rescaled load.Figure 2 compares analytical predictions (solid lines) with numerical results (dots) for spherical cross-entropy at rescaled inverse temperature ˜λ = 1, across embedding dimensions d.Source: Figure 2 · p. 9

C. Zero generalization error at finite load

For continuous embeddings, normalized labels permit exact teacher recovery above a finite load, and spherical cross-entropy reaches recovery at large inverse temperature. Random-association storage capacity depends on the alignment criterion and margin.

C. Zero generalization error at finite load

  • α_rec = d²/(d−1) is the exact-recovery threshold: above it, the teacher is the only solution and generalization error is zero.The counting argument uses d−1 independent linear constraints per example and Nd²−1 student degrees of freedom.Source: C. Zero generalization error at finite… · p. 10
  • At loads above the recovery threshold, spherical cross-entropy reaches zero error as λ → ∞; recovery requires λ ≫ d.At lower loads, the error remains positive even as the rescaled inverse temperature increases.Source: C. Zero generalization error at finite… · p. 10
  • In the infinite-temperature limit, the d > 1 cross-entropy curves reach zero error at eα_rec = d/(d−1), while the d = 1 curve remains positive at finite load.At eα = 0.5 all dimensions retain positive error; at eα = 5, error tends to zero for d > 1 as ˜λ → ∞.Source: C. Zero generalization error at finite… · p. 10

D. Memorization of random inputs

  • For standard alignment, each pattern imposes one scalar constraint, and the zero-margin critical load is α_c = 2d².The capacity increase reflects the Nd² weight parameters.Source: D. Memorization of random inputs · p. 12
  • For vector alignment, capacity is O(d²), but becomes O(d) for κ_eff=O(1); at κ=1/2 (κ_eff=1), α_c^VA=2d, and as κ→1− it tends to d²/(d−1/2).Source: D. Memorization of random inputs · p. 13
  • Simulations confirm dimension-independent rescaled standard-alignment capacity, while vector-alignment capacity retains a dependence on d from orthogonal fluctuations.The reported simulations use N = 1000.Source: D. Memorization of random inputs · pp. 13–14

IV. VOCABULARY CLASSIFICATION

Vocabulary labels reveal only which Voronoi cell contains the teacher output, turning examples into inequalities rather than the directional equations available with continuous labels. The section studies how fixed vocabulary geometry shapes Bayes-optimal learning.

IV. VOCABULARY CLASSIFICATION

  • A vocabulary removes the continuous-label mechanism that makes learning noiseless regression: each example constrains weights through inequalities rather than equations.The label identifies the teacher output’s Voronoi cell, not its position within that cell.Source: §IV · p. 14
  • With two antipodal vocabulary tokens, the vocabulary classification problem reduces to the classical perceptron’s sign nonlinearity.The contrast is with continuous labels, where each example fixes the teacher-output direction.Source: §IV · p. 14
  • The vocabulary’s Voronoi-cell geometry changes which teacher directions are easier or harder to learn, and the theory depends explicitly on the embedding matrix Z.The study evaluates unit vectors sampled uniformly from the sphere, while the theory is described as applicable to more complicated or real-word embeddings.Source: §IV · p. 14
  • The authors target the statistically optimal predictor and formulate the Bayes-optimal computation from the outset, rather than focusing on the loss function as in the preceding case.The posterior partition function is the volume of students that classify the training set correctly, under the spherical teacher prior.Source: §IV · pp. 14–15
  • The replica calculation takes the quenched average in the N → ∞ limit, where finite vocabulary changes the order-parameter structure from the continuous case.The vocabulary-dependent input has a nonzero mean and covariance that need not be proportional to the identity.Source: §IV · p. 15

A. Analytical overview using a spectral anisotropic approximation

The analysis approximates vocabulary-dependent learning by resolving modes in the eigenbasis of the centered vocabulary covariance. Its learning curves compare the Bayes-optimal prediction with cross-entropy-trained students and identify where the approximation is expected to work best.

A. Analytical overview using a spectral anisotropic approximation

  • The spectral anisotropic approximation diagonalizes the order parameters in the eigenbasis of the centered vocabulary covariance, retaining one overlap parameter per positive spectral mode.The retained covariance rank is min(d, D − 1).Source: A. Analytical overview using a… · p. 16
  • For each mode, q_m measures overlap between replicas’ fluctuating fields: 0 means only the common bias is shared, while 1 means their fields coincide.The saddle is a vector q over the vocabulary covariance modes.Source: A. Analytical overview using a… · p. 16
  • The energetic term is the negative conditional entropy of the token distribution given shared fields, measuring residual uncertainty about the output token.At the saddle, the Bayes predictor chooses the token with the largest posterior probability.Source: A. Analytical overview using a… · pp. 16–17
  • The Bayes-optimal analytical curves are close to those of cross-entropy-trained students; the approximation is expected to work best when d < D.The comparison averages over 16 random spherical vocabularies for each d, D pair; when d ≥ D, some numerical errors fall below predictions.Source: Figure 6 · p. 17

B. Scaling of learning curves

Rank-normalizing the data load explains the scaling of learning curves, but vocabulary geometry still shapes learning beyond that normalization. The analysis links this geometry to which directions are learned first and contrasts the vocabulary task with spherical embeddings.

B. Scaling of learning curves

  • Learning curves rescale with eα = α/rZ², but fixed-eα differences remain because HZ depends on vocabulary geometry.The normalization saturates when d exceeds D, where rZ = min(d, D).Source: B. Scaling of learning curves · pp. 17–18

C. The role of vocabulary embeddings geometry

  • High-variance vocabulary directions are learned first, then weaker directions as data increases, even for an isotropic teacher.This filtering is stronger for more anisotropic vocabularies and becomes more evident as d approaches D.Source: C. The role of vocabulary embeddings… · pp. 18–19

V. DISCUSSION

  • For spherical embeddings with d > 1, teacher recovery is exact above αrec = d²/(d −1), and spherical cross-entropy reaches this threshold at large inverse temperature.At d = 1, the corresponding learning rule approaches a maximum-margin classifier whose generalization error remains finite at finite load.Source: §V · p. 19
  • Random-association storage capacity in the spherical setting increases proportionally to d², similarly to factual recall.The paper also contrasts this scaling with associative memories in a close setting.Source: §V · p. 20
  • Vocabulary classification makes generalization a competition among D cells whose geometry is fixed by the vocabulary, rather than alignment with only a continuous teacher direction.Posterior uncertainty can shrink at different rates across embedding modes depending on how well they distinguish vocabulary items.Source: §V · p. 20
  • The model differs from real LLM settings through explicit positional dependence, fixed rather than jointly learned embeddings, and supervised classification without sequence temporal structure.The omitted temporal structure is needed to model how LLMs learn next-token prediction in a sequence.Source: §V · p. 20

Appendix A: Replica computation for continuous embeddings

The appendix sets up the replica calculation for continuous embeddings by specifying the spherical teacher–student problem and expressing its averaged free energy through overlaps. The replicated partition function separates weight-configuration entropy from the loss contribution.

Appendix A: Replica computation for continuous embeddings

  • Inputs are independently sampled from the unit sphere, while teacher and student weights are constrained to the same sphere in the teacher–student setting.The model uses P = αN input patterns and a general empirical loss focused on the student pre-activation and teacher output.Source: Appendix A · p. 24
  • For spherical cross-entropy, the appendix derives the negative log-likelihood and divides by positive λ without changing its minimizer for fixed λ > 0.The loss is written using h = u · y and ρ = ∥u∥².Source: Appendix A · p. 25

1. Replicated partition function and order parameters

  • The averaged free energy is computed by replicating the student weights and introducing their overlaps with one another and with the teacher.The replica identity is used to average the logarithm of the partition function.Source: §1 · pp. 25–26
  • As N grows, pre-activation fields become jointly Gaussian; rotational invariance motivates a replica-symmetric isotropic overlap ansatz.The spherical constraint fixes the overlap structure used in the calculation.Source: §1 · p. 26
  • The replicated partition function separates into an entropic term counting weights compatible with given overlaps and an energetic term carrying the loss contribution.Independence of the P patterns makes the energetic contribution appear as a P-th power.Source: §1 · pp. 26–27

2. Entropic contribution

The replica calculation evaluates entropy from the volume of weight configurations at fixed overlaps, then computes the loss contribution using Gaussian cavity variables. Rotating coordinates along the teacher output separates longitudinal and orthogonal components for the energetic calculation.

2. Entropic contribution

  • The entropic contribution is the log-volume of weight configurations at fixed replica and student–teacher overlaps, yielding a term proportional to log det of the overlap matrix.The volume result follows from the Gram matrix of the collected weight vectors.Source: §2 · p. 27

3. Energetic contribution

  • For the energetic contribution, pre-activation fields are rewritten using shared and replica-specific cavity variables, which reproduce the required covariances.Conditioning on the shared variables makes the student replicas independent for the loss average.Source: §3 · p. 28
  • Aligning coordinates with the normalized teacher output decomposes each field into a longitudinal component h and an orthogonal norm s.The orthogonal component is (d −1)-dimensional, with a non-central chi-square distribution used in the calculation.Source: §3 · pp. 28–29

4. Zero-temperature limit

At zero temperature, the replica calculation focuses on minimum-loss configurations and uses a proximal minimization to evaluate the energetic term. Stationarity equations determine the overlaps, while the large-d analysis concentrates the cavity variables and uses rescaled quantities.

4. Zero-temperature limit

  • As β → ∞, only minimum-energy student configurations contribute, with the replica overlap scaled under the assumption of a unique potential minimizer.When the proximal problem has a unique minimizer, Laplace’s method yields the minimal training loss per input dimension.Source: §4 · p. 30

5. Saddle-point equations

  • The stationary conditions for Φ0 determine χ and r, and can be solved numerically together with the inner minimization over (h⋆, t⋆).For spherical cross-entropy, the minimizing pair is parallel to the specified field vector and its norm satisfies a scalar equation.Source: §5 · p. 31

6. Large embedding dimension

  • At large d, cavity variables concentrate and the saddle-point equations admit finite rescaled solutions for finite eα.The equations depend on d only through rescaled variables, confirming the scalings of α, χ, and λ with d.Source: §6 · pp. 31–32

7. Generalization error

The section relates student–teacher overlap to generalization error by considering the Gaussian pre-activation fields on a new input. It expresses the error through the angle between correlated isotropic Gaussian vectors and recovers the classical scalar perceptron dependence.

  • On a new pattern, the student and teacher pre-activation fields are jointly Gaussian, providing the basis for relating their overlap to generalization performance.Source: §7 · p. 32
  • The generalization error is written as a function of student–teacher overlap using the first moment of the angle between correlated isotropic Gaussian vectors.Source: §7 · p. 33
  • The resulting dependence recovers the classical scalar perceptron case.Source: §7 · p. 33

8. Memorization capacity

Memorization capacity is obtained from the volume of student weights satisfying random input–output constraints. For vector alignment, the capacity depends on the effective margin, while the capacity calculation reduces to distances from cavity points to admissible constraint sets.

  • The version-space volume counts student weights satisfying unrelated input–output constraints, so its capacity measures memorization rather than generalization.Source: §8 · pp. 33–34
  • As the load approaches capacity from below, the version-space volume shrinks until only a single solution survives and the overlap q approaches one.Source: §8 · p. 34
  • In the capacity calculation, the energetic contribution is governed by the squared distance from a cavity point to the admissible set, which vanishes when that point is already admissible.Source: §8 · p. 34
  • SA and VA storage impose different admissibility constraints: SA restricts H, whereas VA requires H ≥ κ_eff T.Source: §8 · pp. 34, 36
  • At zero effective margin, VA capacity is 2d^2.Source: §8 · p. 37
  • At κ_eff=1, VA capacity is 2d, and as κ_eff→∞ it converges to d^2/(d−1/2) ≈ d.Source: §8 · p. 37

Appendix B: Replica calculation for vocabulary classification

The appendix derives a replica calculation for vocabulary classification within a spectral closure. It represents posterior weight configurations by mode-wise overlaps, then uses the saddle point to obtain the Bayes decision and its error.

  • The posterior volume is built using the decision-cell indicator for the observed vocabulary label.Source: Appendix B · p. 37
  • The closure approximates output-space covariance as diagonal in a chosen basis, with one overlap per mode; replica symmetry and the Nishimori identity equate teacher–student and inter-replica overlaps.Source: Appendix B · p. 37
  • The replica calculation separates the weight-volume entropy from the data contribution, whose pattern average is −H_Z, and extremizes the resulting free energy over teacher–student and replica overlaps.Source: Appendix B · pp. 37–38
  • Input tokens probe only weight components in the span of the embedding matrix, and the overlap matrix includes the teacher and posterior replicas.Source: Appendix B · p. 38
  • At the physical saddle, the vocabulary geometry enters through H_Z, and the resulting conditional token probabilities define the Bayes decision and its error within the spectral closure.Source: Appendix B · p. 39

Related papers

Loading 2610.32788v1…