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
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×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→∞ 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 αrec=d2/(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 d2 for margins O(1/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 ξ=(ξ1,…,ξN) to the field uW(ξ)=(Nd)−1/2∑i=1NWiξi . (Top) continuous embeddings, with input vectors on the unit sphere Sd−1 and prediction y^(ξ)=uW(ξ)/∥uW(ξ)∥ . (Bottom) vocabulary embeddings, with each input vector drawn from a fixed vocabulary {z1,…,zD}⊂Sd−1 . The prediction is the vocabulary vector with maximal overlap with uW(ξ) .
Figure 2: Using the rescaled load, the generalization error ϵg curves converges for large embeding dimension. We compare analytical predictions (solid lines) to numerical results (dots) for the generalization error ϵg of the spherical cross-entropy loss with λ~=λ/d=1 versus the rescaled load α~ for different values of the embedding dimensions d .
Figure 3: Spherical cross-entropy learning curves at infinite λ . Generalization error ϵg versus α=P/(Nd) for different embedding dimensions, obtained by minimizing cross-entropy at finite λ and then taking λ→∞ . Vertical dashed lines mark the exact-recovery thresholds αrec=d/(d−1) for d>1 , Eq. ( 30 ). The black dashed curve shows the d→∞ limit, with threshold αrec=1 .
Figure 4: For d>1 , generalization error decreases as λ~ increases and eventually reaches zero for λ~→∞ for α~>d−1d=α~rec(d) . Generalization error ϵg versus λ=λ/d at α=0.5 (left) and α=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 αc/d2 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→∞ 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 ϵ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 ϵg . We note that when d<D numerical results have higher ϵg then predictions as expected, while instead when d≥D some of the numerical points have lower ϵg . This is compatible with the fact that our computation is an approximation of the full d×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 α . (Left) D=16 case; we note that here d=32>D and in that case rZ saturates and the curves do not shift anymore when plotted against α , 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 HZ that depends explicitly on the vocabulary.
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 qm∗ increase monotonically with the corresponding λ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 qm∗(λ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,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,tˉ half-plane, we show the projections (crosses) on the satisfability half-plane z>dκ , indicated in green, corresponding to the points h∗,t∗ that minimise the potential on the constraint set C={h∈R,t≥0:h≥dκ} . Right panel : the minimisation for VA storage. Again, points in the z,tˉ plane (full symbols) are associated to the corresponding points h∗,t∗ that minimise the potential on the constraint set C={h∈R,t≥0:h≥κefft} (crosses). In the blue region C={h∈R,t≥0:−κefft≤h≤κefft} , to minimise the potential each point is projected onto the line z=κefft . In the orange region ( z<−κefftˉ ), all points are projected to the origin.