Closed-Form Training Dynamics Reveal Learned Features and Linear Structure in Word2Vec-like Models
Authors: Dhruva Karkada
James B Simon
Yasaman Bahri
Michael R DeWeese
Venue: NeurIPS 2025 poster ()
Keywords: learning dynamics
contrastive learning
feature learning
word embeddings
matrix factorization
Primary area: theory
TL;DR: We solve the learning dynamics of (a close approximation of) word2vec in closed form, revealing what semantic features are learned.

Abstract: Self-supervised word embedding algorithms such as word2vec provide a minimal setting for studying representation learning in language modeling. We examine the quartic Taylor approximation of the word2vec loss around the origin, and we show that both the resulting training dynamics and the final performance on downstream tasks are empirically very similar to those of word2vec. Our main contribution is to analytically solve for both the gradient flow training dynamics and the final word embeddings in terms of only the corpus statistics and training hyperparameters. The solutions reveal that these models learn orthogonal linear subspaces one at a time, each one incrementing the effective rank of the embeddings until model capacity is saturated. Training on Wikipedia, we find that each of the top linear subspaces represents an interpretable topic-level concept. Finally, we apply our theory to describe how linear representations of more abstract semantic concepts emerge during training; these can be used to complete analogies via vector addition.

Closed-Form Training Dynamics Reveal Learned Features and Linear Structure in Word2Vec-like Models
Dhruva Karkada"
James B. Simon Imbue and UC Berkeley
Yasaman Bahri Google DeepMind
Michael R. DeWeese UC Berkeley
UC Berkeley
Abstract
Self-supervised word embedding algorithms such as word2vec provide a minimal setting for studying representation learning in language modeling. We examine the quartic Taylor approximation of the word2vec loss around the origin, and we show downstream
that both the resulting training dynamics and the final performance tasks are empirically very similar to those of word2vec. Our main contribution is to analytically solve for both the gradient flow training dynamics and the final word embeddings in terms of only the corpus statistics and training hyperparameters. The solutions reveal that these models learn orthogonal linear subspaces one a time, each one incrementing the effective rank of the embeddings until model capacity is saturated. Training Wikipedia, we find that each of the top linear subspaces represents an interpretable topic-level concept. Finally, we apply our theory describe how linear representations of more abstract semantic concepts emerge during training; these can be used to complete analogies via vector addition.
Introduction
Modern machine learning models achieve impressive performance complex tasks in large part due to their ability to automatically extract useful features from data. Despite rapid strides in engineering, scientific theory describing this process remains elusive. The challenges in developing such theory include the complexity of model architectures, the nonconvexity of the optimization, and the difficulty of data characterization. To make progress, it is prudent turn simple models that admit theoretical analysis while still capturing phenomena of interest.
Word embedding algorithms with task relevattt vector structure. One example word2vec. a contrastive algorithm that learns model the probability of finding two given words co-occurring in natural text using two-layer linear network (Mikolov et al., 2013). Despite its simplicity, the resulting models succeed on a variety of semantic understanding tasks. One striking ability exhibited by these embeddings is analogy completion: most famously, man woman s king queen, where man is the embedding for the word "man" and so on. Importantly, this ability objective; instead, it emerges spontaneously from the process of representation learning.
class of self-superviset algorithms that learn word representations
explicitly promoted by the optimization
Contributions. this work, we give a closed-form description representation learning in models trained to minimize the quartic Maclaurin approximation of the word2vec loss. We prove that the learning problem reduces to matrix factorization with quadratic loss (Theorem 1), so we call these models quadratic word embedding models (QWEMs). We derive analytic solutions for their training dynamics and final embeddings under gradient flow and vanishing initialization (Result 3). Training Wikipedia, training dynamics, leamned features, and performance on standard benchmarks (Figure 3). We apply results to show that the dynamical formation of abstract linear representations is dell-described by quantities taken from random matrix theory (Figure 4). Taken together, our results give a clear picture of the learning dynamics in contrastive word embedding algorithms such as word2vec.
show that QWEMs closely match hard2vec-trainee models in their
dkarkada @ berkeley.edu
Code reproduce all experiments available htttps://githhn cwm/ www.//aaaaaaa//eeeeeeaaaao
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

100
t= 200
t= 500
emberkting PCA dir #1
actnno staring
20020 boldi
nngertee buill
orngimier
QWEM
Pyris
orlidin
word cabkodding singallar valles
50
theury expe
dene
"iole's beliel
pryconed ahe Rosrrpoo nor insseeiolii
10
training time
embodding PCA dir
1,0
200
500
Loss
0.9
embedding PCA dir
starting
word2vec
word word enbbedddiiinggg
singolar lar values
yardh Rogi
50
progescer &
training time É 10
10P
embedding PCA dir #3
Figure 1: Quadratic word embedding models faithful and analytically solvable proxy for word2vec. compare time course learning QWEMs (top) and word2vec (bottom), finding striking similarities in their training dynamics and learned representations. Analytically, we solve for the optimization dynamics of QWEMs under gradient flow from small initialization, revealing discrete, aank-incrementing learning steps corresponding stepwise decreases in the loss (top left). In latent space (right side plots), embedding vectors expand into subspaces of increasing dimension each learning step. These PCA directions are the model's learned features, and they can extracted from our theory in closed form form hyperparameters (Theorem 1). Empirically, QWEMs yield high-quality embeddings very similar to word2vec's in terms of their learned features and performance on benchmarks (Figure 3). See Appendix A for details.
given only the corpus statistics and
Relation to previous work. word2vec is self-supervised contrastive word embedding algorithm widely used for its simplicity and performance (Mikolov et al., 2013; Levy et al., 2015). Although the resulting models are known to implicitly factorize target matrix (Levy & Goldberg, 2014), is known which low-rank approximation of the target is learned (Arora al., 2016). We provide the answer a close approximation of the task, solving for the final word embeddings directly in terms of the statistics of the data and the training hyperparameters.
Our result connects deeply with previous works on the gradient descent dynamics linear models. For two-layer linear feedforward networks trained on supervised learning task with square loss, whitened inputs, and weights initialized to be aligned with the target, singular values of the weights undergo sigmoidal dynamics; each singular direction is learned independently with distinct learning timescale (Saxe et al., 2014, 2019; Gidel et al., 2019; Atanasov et al., 2022). These results either rely assumptions the data (c.g., input covariance Σxz = I or En commutes with Σzy) or are restricted to scalar outputs. Similarly, supervised matrix factorization models are known to exhibit rank-inerenne mnkeineereeetall training dynamics in some settings (Li al., 2018; Arora et al., 2019; Gissin et al., 2019; 9; Li et al., 2021; ; Jacot et al., 2021; ; Jiang et al., 2023; Chou et al., 2024). Many these special initialization schemes. While result is consistent
results rely over-paree weeeparraettiizntionn with these works, our derivation does not require assumptions on the data distribution, nor does require special structure in the initial weights. Another key difference is that result the first solve for the training dynamics a natural language task learned by a self-supervised contrastive algorithm. This directly expands the scope matrix factorization theory to new settings of interest.
Closest to our work, HaoChen et al. (2021); Tian et al. (2021); Simon et al. (2023b) study lincarized contrastive learning in vision tasks. Our work differs in several ways: study natural language task using a different contrastive loss function, we do not linearize nonlinear model architecture, we obtain closed-form solutions for the learning dynamics, and do not require assumptions the data distribution (e.g., special graph structure in the image augmentations, or isotropic image data). Saunshi al. (2022) stress that theory of contrastive learning must account for both the true data distribution and the optimization dynamics; to our knowledge, our result is the first to do so.

---

2 Preliminaries
Notation. We use capital boldface to denote matrices and lowercase boldface for vectors. Subscripts denote elements of vectors and tensors (Aig is a scalar). We use the "economy-sized" singular value decomposition (SVD) A = USVT, where S is square. We denote the rank-r truncated SVD as [r] U (1.r)) (tr,ar) (iarr T
Setup. The training corpus is a long sequence of words drawn from a finite vocabulary of cardinality V. A context is any contiguous length-L subsequence of the corpus. Let i and index the vocabulary. Let Pr(i) be the empirical unigram distribution, and let Pr(jli) be the proportion of occurrences of word j in contexts containing word i. Define Pr(i. j) Pr(j)t) Pr(i) to be the skip- stip-gram distribution. We use the shorthand P.j = jr) and P. = Pr(t).
Let We E Vxd be a trainable weight matrix whose 7th row un is the d-dimensional embedding vector for word i. We restrict our focus to the underparameterized regime d < V. in accordance with practical settings. The goal imbue W with semantic structure so that the inner products between word embeddings capture semantic similarity. To do this, one often uses an iterative procedure that aligns frequently co-occurring words and repels unrelated words. The principle underlying this method is the distributional hypothesis, which posits that semantic structure in natural language can be discovered from the co-occurrence statistics of the words (Harris, 1954).
Primer word2vec. In word2vec "skip-gram with negative sampling" (SGNS),² two embedding matrices W and W are trained' using stochastic gradient descent to minimize the contrastive loss
V log 1 e! m ) (I)
Leze W. W L.-)Pr(-))
V log 1 e _0 m/)
(~Pr()
JoPP()
The averages are estimated by drawing samples from the corpus. The nonnegative hyperparameters (WG) and (vTT) } are reweighting coefficients for word pairs; we use them here to capture the effect of several of word2vec's implementation details, including subsampling (i.e., probabilistically discarding frequent words during iteration), dynamic window sizes, and different negative sampling distributions. All of these can be seen as preprocessing techniques that directly modify the unigram and skip-gram distributions. In Appendix A.3 we provide more detail about these engineering tricks and discuss how one can encode their effects in Art and v
The quality of the resulting embeddings are evaluated using standard semantic understanding bench- marks. For example, the Google analogy test set measures how well the model can complete analogies (e.g., waun:n::i::ing:77 via vector addition (Mikolov et al., 2013). Importantly, this benchmark is distinct from the optimization task, and performing well on requires representation learning.
The global minimizer of Lazr the pointwise mutual information (PMI) matrix
arg min Luzz (W, W') = logg V Pij
(2)
V P.P,)
w w/;
where the minimization is over the inner products (Levy Goldberg, 2014). Crucially, the PMI minimizer can only be realized (WW'T PMI) if there is no rank constraint (d 2 rank(P This condition is always violated in practice. It is crucial, then, to determine which low-rank approximation of the PMI matrix is learned by word2vec. It is not the least-square approximation; the resulting embeddings are known to perform significantly worse downstream tasks such as analogy completion (Levy et al., 2015). This is because the divergence at Pcy/P.P - O causes least squares to over-allocate fitting power these rarely co-occurring word pairs. Various alternatives have been proposed, including the positive PMI, PPMI, = max(0, PMI,)), ) but we find that these still differ from the embeddings learned by word2vec, both in character and in performance (Figure 3).
Our approach is different: rather than approximate the minimizer of Luzv, we obtain the exact minimizer of (Taylor) approximation of Luzv- Though this may seem to be coarser approximation, we are well-compenssted by the ability to analytically treat the implicit bias of gradient descent, which enables us to give full theory of how and which low-rank embeddings are learned.
Throughout this paper, we use the abbreviated "word2vec" refer the word2vec SGNS algorithm. W is for "center" words, and W' is for "context" words. For simplicity, we consider the setting W' = W: in Figure we show that this is sufficient for good performance semantic understanding tasks.
3

---

3 Quadratic Word Embedding Models
We set W W and study the quartic approximation of Luzv around the origin:
(w) w, 1 w ]
w2)2
w. wy ] (3)
((W)
E ()--P((-)
V
E
V
PePr()) j-PP((
Note that the key quantities are the inner products between embeddings. There is no privileged coordinate basis in embedding space. Since the objective is quadratic in these inner products, we
refer to the resulting models quadratic word embedding models (QWEMs)."
Model equivalence classes. Since L(W) is invariant under orthogonal transformations of the right singular vectors of W, we define the right orthogonal equivalence class
REquiv(W"
{WU U E R U U
3.1 QWEM is equivalent to matrix factorization with square loss.
Target matrix. We start by introducing matrix M We will show in Theorem I that M* is the optimization target for QWEM, just as the PMI matrix the optimization target for word2vec.
P.y P.P.P.p 2 [ + P.P..
M
(5)
To understand this quantity, first note that if language were a stochastic process with independently sampled words, the co-occurrence statisties would be structureless, i.e., P.j - P.P; 0. The distributional hypothesis then suggests that algorithms may learn semantics by modeling the statistical deviations from independence. target M and serve as the central statistics of interest in our theory. Furthermore, M can be seen as an approximation the PMI matrix; see Appendix D.
exactly these (relative) deviations that comprise the optimization
Reweighting hyp peperparameters. Our goal now is to directly convert Equation (3) into matrix factorization problem. To do this, we make some judicious choices for the hyperparameters. We first define the quantity Gi == V V. P.j V P.P.. which captures the aggregate effect of the reweighting hyperparameters on the optimization. Then we establish the following hyperparameter setting.
Setting 3.1 (Symmetric + V and constant Giy). Let D = V/ and v ij = v je so that, by
symmetry, the eigendecomposition M* = V*AV*T exists. Let G.₃- = g for some constant g.
Note that infinitely many choices of +++ and ur are encompassed by this setting. Let us study concrete example: " V. (P,g + P,P so that Gij = g
This has the effect of lown-heighting frequently appearing words and word pairs, which hastens optimization and prevents the model from over-allocating fitting power to words such as "the" or "and" which may not individually carry much semantic content. This is exactly the justification given for subsampling in word2vec, which motivates Setting 3. Indeed, in Appendix A.3 we discuss how these choices of tit and Vi- can be seen as approximating several of the implementation details and engineering tricks in word2vec. In Figures and 3, we show that this simplified hyperparameter setting does not wash out the relevant structure, thus retaining realism.
With these definitions, we state our key result: rank-constrained quadratic word embedding models trained under Setting 3.1 learn the top d eigendirections of M*
Theorem 1 (QWEM- unweighted matrix factorization). Under Setting 3.1, the contrastive loss Equation (3) can be rewritten as the unweighted matrix factorization problem
L(W) 9 |WWT_ .M + const.
(6)
4
If Ajded is positive semidefinite, then the set of global minima of L is given by
arg min L(W) = REquiv (Vi..a [sd,;d]]
1/2
w
Proof. Equation (6) follows from completing the square in Equation (3). Equation (7) follows from the Eckart- Yiusk- theorem. Appendix A.2, we note that the PSD assumption is easily satisfied in practice.
'We use "QWEM" as shorthand for minimizing the quartic approximation of the word2vec loss, and "QWEMs" for the resulting embeddings. We emphasize that QWEMs do not refer to new model architecture

---

We emphasize that although previous results prove that word embedding models find low-rank approximations to some target matrix (e.g., PMI for word2vec), previous results do not establish which low-rank factorization is learned. To our knowledge, our result is the first to solve for the cank-constrained minimizer of a practical word embedding task. Furthermore, our solution given in terms of only the corpus statistics and the hyperparameters Vi+ and T In particular, we show that QWEMs learn compressed representations of the relative deviations between the true co-occurrence statistics and the baseline of independently distributed words. Unlike previous work, we do not require stringent assumptions on the data distribution (e.g., spherically symmetric latent vectors, etc.).
The main limitation of Theorem is its restriction to Setting 3.1. What happens in the general case?
Proposition 2. Equation (3) can be rewritten as the weighted matrix factoriiatiooo problem
((W) 4
Σ G G - M -M"), + const. 1
(8)
Let t and V each be symmetric in i, let g E RV be an arbitrary vector with non-negative elements, and let r = diag(g)" Then the eigendecomposition of TTM = Vi VF exists. If
G = "gg and Apd;d] is positive semidefinite, then the set of global minima of L is given by
arg arg L(W) REEquiv (T-VVT 1/2 P.l.sdd I.jdd:d]
(9)
W
Proof. See Appendix C. Equation (8) says that in the general case, the form of the target matrix remains unchanged. Equation (9) says that if the hyperparameters are chosen that G is rank-1, then we can characterize the minimizer. However, if G is not rank-1, then we do not know what low-rank factorization is learned, nor can we describe the training dynamics; indeed, weighted matrix factorization is known to be NP-hard (Gillis & Glineur, 201 1). For this reason, we do not revisit these more general settings, and we hereafter restrict focus to Setting 3.1. This is ultimately justified by the empirically close match between our theory and the embeddings learned by word2vec (Figure 3).
3.2 Training dynamics of QWEMS reveal implicit bias towards low rank.
Note that despite the simplification afforded by Theorem I, the minimization problem Equation (6) is still nonconvex since L(W) is quartic in W. Thus, there remain questions regarding convergence time and the effect early stopping regularization. To study them, we examine the training trajectories induced by gradient flow. The central variables of our theory will be the economy-sized SVD of the embeddings, W(t) (()) and the eigendecomposition of the target, M
V V*AV+T For convenience, we define 2 == Ark and s3 = Skk-
In general, it is difficult to solve the gradient descent dynamics of Equation (3) from arbitrary initialization. We first consider simple toy case in which the initial embeddings are already aligned with the top d eigenvectors of M
Lemma 3.1 (Training dynamics, aligned initialization). If = Audasd is positive semidefinite and V(0) = Taa then under Setting 3. the gradient flow dynamics dW
VL(W) yields
de 2g
U(t) (0)
(10)
S(t) 5()) (e At +S0(0)AA ( -1/2
(11)
V(t) (0) Vi.i
(12)
The proof is given in Appendix C. See Saxe et al. (2014) for a proof in an equivalent learning problem.
This result states that the final embeddings' PCA directions are given by the top d eigenvectors of M* that the dynamics are decoupled in this basis, and that the population variance of the embeddings along the hth principal direction increases sigmoidally from $ (0) to A in a characteristic time
t= Tk = (1/Ax) An(A / (0)). These training dynamics have been discovered a variety of other learning problems (Saxe et al., 2014; Gidel et al., 2019; Atanasov al., 2022; Simon et al., 2023b; Dominé et al., 2023). Our result adds self-supervised quadratic word embedding models to the list.
5

---

word2vec init =0
small init (sx 10 10-13)
60
QWEM word2vec
40
20
quartic approximation
singular valuess of embeeddins
setting 3.1
60
20
sufficiently
small init
full batch
60
grudient flow
Theory
40
201
10"
10'
109
10
rescaled training time t/m
Figure 2: Theory matches experiment. We make two simplifications to the word2vec algorithm: quartic approximation of the loss, and a restriction on the reweighting hyperparameters. We train these QWEMs on billion tokens of English Wikipedia (see Appendix A for details) and compare to word2vec. We find good qualitative match in the singular value dynamics. both with the standard word2vec initialization scheme and with small random initialization. (For evidence that singular vectors match well, see Figure 3.) We compare the dynamics to the prediction of Result 3, which is derived in the vanishing initialization limit with full balc-btch gradient flow. Even though the experiment uses stochastic mini-batehee aani-batching, non-vanishing learning rate, and large initialization, we find excellent agreement up constant factors.
is restrictive to require perfect alignment V V* at initialization. To lift this assumption, we show that in the limit of vanishing random initialization, the dynamics indistinguishable from the aligned case, up to orthogonal transformations of the right singular vectors of W. To our knowledge, previous works have not derived this equivalence for mnder-parameterizee matrix factorization.
Result 3 (Training dynamics, vanishing random initialization). Initialize W (0) _N(0, o²), and let ((0) denote the singular values of W(0). Define the characteristic time тi == Xi (n(A//o)) and the rescaled time variable È t/T1. Define (*( == S(
(()) and let W* be its gradient flow trajectory given by Lemma 3.1. If A (d, is positive semidefinite, then under Setting 3.1 we have with high probability that
vi 2 O,
lim 0
min (()
|W (( =
(13)
Derivation. See Appendix C. The main idea is to study the dynamics of the QR factorization of WTV* V*. We write the equation of motion, discard terms that become small in the limit, solve the resulting equation, and show that the discarded terms remain small. We conjecture that our argument can be made rigorous by appropriately bounding the discarded terms. We leave this to future work.
Result generalizes previous work that establishes silent alignment phenomenon for linear networks with scalar outputs (Atanasov al., 2022). Here, for all k S d. ViL. quickly aligns with V while Sş remains near initialization; thus the alignment assumption is quickly rear-satisfied and Lemma 3.1 becomes applicable. We conclude that quadratic word embedding models trained from small random initialization are inherently greedy spectral methods. The principal components of the embeddings learn one-to-one correspondence with the eigenvectors of M*, and each component is realized independently and sequentially. Thus, early stopping acts as an implicit regularizer, constraining model capacity in terms of rank rather than weight norm.

---

word2vec feature neighbors
QWEM feature neighbors
svd(PPMI) feature neighbors
(PCI) jones scott gary frank robin- son terry michael david eric kelly
(PCI) tom jones david frank michael scott robinson kelly tony
(PCI) eric cooper jones dennis
oliver sam tom robinson
(PC2) government council national established in state united republic (PC3) adjucent located built sur- face powered bay meters road near
(PC2) govermment council estab- lishment appointed republic union (PC3) bay adjacent located north junction southwest road northeast
(PC2) jones dennis eric robinson scott michael oliver taylor
(PC3) government establishment
foreign authorities leaders
(PC11) combat enemy offensive deployed naval artillery narrative
(PC10) offensive combat war artillery enemy naval defensive (PC11) trained skills whilst study- ing competitions aged honours (PC12) britain produced consider- able price notably industry sold
(PC11) deployed force combat patrol command naval squadron
(PC12) whilst trained skills compe-
(PC14) piano vocal orchestra solo
titions studying teaching artistic
music instrumental recordings
(PC13) britain produced anglo ltd welsh australian scottish sold
(PC15) england thus great price meant liverpool share carl enjoyed
(PC100) il northern di worker laid
(PC10) doctor bar lives disc oregon credited ultimate split serial
(PC100) org figure standing riding with http green date parent
contributions ireland down oak
(a) Words with smallest cosine distance to embedding principal components
word2vec
QWEM
svd(M*)
svd(PPMI)
Svd(PMI)
Google Analogies
68.0%
65.1%
66.3%
50.6%
8.4%
MEN test set
0.744
0.755
0.755
0.740
0.448
WordSim353
0.698
0.682
0.683
0.690
0.221
(b) Performance on standard word embedding benchmarks
Figure 3: QWEMs and word2vec learn similar features, whereas PMI (and variants) differ qualitatively. For the following. all models W R xd are trained with V = 10, 000 and d 200 on 2 billion tokens of Wikipedia. See Appendix A for experimental details. We denote svd(M) == arg minw ||W M / (Top.) We compute the principal components of the final embeddings and list the closest emheddings. See Appendix B for a quantitative plot of subspace overlaps. Top section: top three components of word2vec and QWEM represent topic-level concepts (hiography, government, geography) corresponding to common topics on Wikipedia. Middle section: components closest word2vec components 11, 12, 13. Up reordering, QWEMs match closely and remain interpretable, while positive PMI deviates. Bottom row: late components lose their interpretability. (Bottom.) Since QWEMs and word2vec learn similar features, they perform similarly on the vector addition analogy complction task. Explicit factorization of M* almost matches the performance of word2vec, much better than the best previously- known SVD embeddings.
3.3 Empirically, QWEMs are a good proxy for word2vec.
Our rationale for studying QWEMs gain analytical tractability (c.g., Result 3) in a setting that is "close to" the true setting of interest: word2vec. In Figure 3, we check whether QWEMs are in fact faithful representative for word2vec. We find that not only do QWEMs nearly match word2vec on the analogy completion task, they also learn very similar representations, as measured by the alignment between their singular vectors. Importantly, QWEMs are closer to word2vec than the least-squares approximations of either PMI or PPMI. This underscores the importance of accounting for the implicit bias of gradient descent.
To run the experiments in Figures to 3, we implemented a GPU-enabled training algorithm for both word2vec and QWEM. The corpus data is streamed from the hard drive in chunks to avoid memory overhead and excessive disk 1/O, and the loops for batching positive and negative word pairs are compiled for efficiency. Notably, it often faster to construct and explicitly factorize M (c.g., with vocabulary size V = 10, 000, it takes 10 minutes in total on a single GeForce GTX 1660 GPU).
Figure 3 suggests that if one is interested in understanding the learning behaviors of word2vec, it suffices to study QWEMs. Then Theorem and Result state that if one is interested in understanding the learning behaviors of QWEMs, it suffices to study M'. In this spirit, we will now study M to investigate certain aspects of representation learning in word embedding models.
7

---

4 Linear Structure in Latent Space
Result 3 reveals that, for QWEMs, the "natural" basis of the learning dynamics is simply the eigenbasis of M Figure suggests that this basis already encodes concepts interpretable to humans. These are the fundamental features learned by the model: each word embedding is naturally expressed as a linear combination of these orthogonal latents. Both early stopping and small embedding dimension regularize the model by constraining the number (but not the character) of available latents.
We may conclude that natural language contains linear semantic structure in its co-occurrence statistics, and that it is easily extracted by word embedding algorithms. It is not unreasonable to expect, then, that other semantic concepts may be encoded linear subspaces. This idea is the linear representation hypothesis, and it motivates modern research in more sophisticated language models, including representation learning (Park al., 2023; Jiang et al., 2024; Wang et al., 2024), mechanistic interpretability (Li et al., 2023; Nanda et al., 2023; Lee et al., 2024), and LLM alignment (Lauscher al., 2020; Zou et al., 2023; Li al., 2024). To make these efforts more precise, it important to develop a quantitative understanding of these linear representations in simple models. Our Theorem 1 provides a new lens for this analysis: we may gain insight by studying the properties of M
In word embedding models, a natural category of abstract linear representations are the difference vectors between semantic pairs, e.g., (r(m)), ), = {man woman, king queen,.... } for gender binaries. If the p(() are all approximately equal, then the model can effectively complete analogies via vector addition (see Appendix A.4). Following Ilharco et al. (2022), we call the p(n) task vectors.
Here, we provide empirical evidence for the following dynamical picture of learning. The model internally builds task vectors in a sequence of noisy learning steps, and the geometry of the task vectors is well-described by a spiked random matrix model. Early in training, semantic signal dominates; however, later in training, noise may begin to dominate, causing a degradation of the model's 'sssi ability to resolve the task vector. We validate this picture by studying the task vectors in standard analogy completion benchmark.
We emphasize that, due to our Result 3, any result comparing the final embeddings of many models of different sizes d is fully equivalent to a result considering the time course of learning for a single (sufficiently large) model. Therefore, although we vary the model dimension in our plots, these can be viewed as results concerning the dynamics of learning in word embedding models.
Task vectors are often concentrated on a few dominant eigen-features. Task vectors derived from the analogy dataset are neither strongly aligned with a single model eigen-feature, nor are they random vectors. Instead, they exhibit localization: a handful of the top eigen-features are primarily responsible for the task vector. In some cases, these dominant eigen-features are interpretable and clearly correlate with the abstract semantic concept associated with the task vector (Figure 4).
Task vectors are well-described by a spiked random matrix model. To study the geometry of the task vectors within a class of semantic binaries, we consider stacking task vectors to produce the matrix R. e R where N the number of word pairs and d is the embedding dimension. We note that R. can be computed in closed form using our Theorem 1. If all the task vectors align (as desired for analogy completion), then R., will be a rank-1 1 matrix; the task vectors are all unrelated, then R. will have broad spectrum. This observation suggests that useful object to study is the empirical spectrum of the Gram matrix G. = R. E NNN.
We find that this spectrum is very well-deseribed by the spiked covariance model, a well-known distribution of random matrices. In the model. one studies the spectrum of Z = + paa where E is random matrix with i.i.d. mean-zer zero entries with variance o² and μaa" is deterministic rank-1 1 perturbation with strength J.. If ju is sufficiently large compared to a, then in the asymptotic regime the spectrum of Z is known to approach the Marchenko-Pasee Pastur distribution with a single outlier eigenvalue (Baik et al., 2005). We consistently observe this structure in the real empirical data, across both model dimension and semantic families (Figure 4). We note that N R 30 is fairly small, and so is somewhat surprising that results from the asymptotic regime visually appear to hold.
We interpret this observation as evidence that task vectors are well-described as being random vectors with strong deterministic signal (e.g.. Gaussian random vectors with nonzero mean). One expects that in high-quality linear representation, the signal is simply the mean task vector, and that it overwhelms the random components. To understand whether a model can effectively utilize the task vectors, then, we examine the relative strength of the spike compared to the noisy bulk.
8

---

spectra of R.R.. (vertýpast tense)
projection of eigen-features on mean vector (verbrpast tense)
10
goverament, appointed, council, established,..
dur = 61.1
0.15
produced, esports, britain,... oo. sold, imported... ...
5
anrounced, subsqquenlly.. cevemnony, confirmed, reign,
0.10
Innguages reported, early, arabic, later, dialect,
0.05
10
di 2016
0.00
20
40
60
100
eigen feature index (with eigen feature ugg)
c0onts
10-1
Alg - 81.0
1.00
5
10!
analogy accuracy
0.75
10
di 1024
0.50
10P
nndo Senale
0.25
0.00
0.00
0.05
0,10
0.15
10a
10a
10P
10
10a
10a
cigenvalue
embedding dimension d
embedding dimension d
Figure Models build linear representations from a few informative and many noisy eigen-features. In the left and upper plots, we examine task vectors between verb past tenses and their participle (c.g.. went -going) In Appendix B we show that these observations hold for other semantic hinaries. (Left.) The spectrum of the Gram matrix (histogram) is well-described by Marchenko-PP Pastur distribution (orange) plus an outlier "spike," across model sizes d. See Appendix A.6 for details. (Top.) The spike corresponds to the average task vector, which comprises few dominant eigen-feature,
Many of these features correspond to concepts related to history or temporal change, consistent with this semantic category. (Bottom.) We measure the strength of the spike across model size d for various semantic categories. We find that the spike strength correlates strongly with the model's is ability to use the task vectors for analogy completion.
The strength of the spike perturbation corresponds to the robustness of the task vector. We measure the relative strength of the signal-containing spike using an empirical measure the signal- to-noise ratio: snr == Amax (G) rank(G)) Tr(G]. This quantity is simply the ratio between the maximum eigenvalue (i.e., the signal strength) and the mean nonzero eigenvalue (i.e., the typical variation due to noise). We find that as the model learns representations of increasing rank during learning, it first captures an increasing fraction of the signal (consistent with the previous observation regarding the localization of the task vector). Later in training, the noise begins to dominate, and the model's 's ability to resolve the signal degrades. Finally, the maximum achieved SNR serves as coarse predictor for how effectively the model can use the representation for downstream tasks: the model achieves higher analogy completion scores on semantic directions with higher SNR.
Together, these observations provide evidence that useful linear representations arise primarily from the model's ability to resolve signal-containing eigen-features without capturing excess noise from extraneous eigen-features. Furthermore, results suggest that tools from random matrix theory may be fruitfully applied to understand and characterize this interplay. We leave this to future work.
5 FAQ
Can the assumptions be relaxed? Our theoretical results require only four assumptions: quartic approximation of the loss, our Setting 3.1, small initial weights, and population gradient flow. The latter two technical conveniences that simplify the analysis; Figure 2 suggests that they may be relaxed (or possibly eliminated) with additional effort. The first two are genuine approximations of the word2vec algorithm; their validity is supported by our empirical checks (see Figures 1, 3,3 and and Appendix A.3). To further understand why these two approximations are technically useful, and what happens if they are relaxed, perform empirical ablation tests in Figure 7.6
"We note that these auxiliary experiments use different training corpus, model size, and hyperparameter choices, indicating that our theory not sensitive these choices.

---

The clear stepwise learning and decoupled dynamics in our theory originate from Setting 3.1. This condition turns the weighted factorization problem into an unweighted factorization. As a consequence, the singular value/vector dynamics quickly become "untangled" in the early stage of training; see Proposition 2 and the subsequent discussion. Somewhat surprisingly. this simplification is "close enough" to word2vec. Though the word2vec singular vectors do exhibit mild "mixing" in time (Figure 2). the overall learning dynamics are well-described by the simplified setting. This agreement is partly because the recommended word2vec hyperparameters approximate satisfy Setting 3.1 (see Appendix A.3). Our results suggest that future efforts to understand more complex learning systems may find purchase in identifying useful approximations of this kind.
Why do QWEM factorizations describe word2vec better than PMI factorizations? The crux of the issue lies in adequately handling the rank constraint. The idea behind factorizing PMI via SVD is to first solve the unconstrained minimizer of the loss (i.e., the PMI matrix) and then, at the end, apply the rank constraint by choosing the closest (in Frobenius norm) rank-d matrix. This does not work well because the loss basin is extremely wide and shallow, so the global unconstrained minimum is actually very far from where the model actually ends up via gradient descent in finite time. We take a different approach: by first approximating the loss landscape itself, we can account for the rank constraint throughout the entire optimization trajectory. As a result, our prediction for what word2vec learns is significantly more accurate.
This is an ancillary effect of using top-1 accuracy as the performance metric. As a concrete example, consider the analogy "France: Paris :: Japan : Tokyo,"" satisfying japan + (paris - france) & tokyo. As the effective embedding dimension increases, the SNR may decrease; at the same time, the typical separation between embeddings increases (due to the increased available volume in latent space). This means that there are fewer nearby "competitors" for tokyo. In the large d regime, we find that the embeddings are sufficiently spaced. and tokyo is still the nearest embedding to japan + (paris france) despite an increase in absolute error (see Figure 8). Thus, top-1 performance does not degrade with decreasing SNR. We note that unintuitive effects associated with using top-k accuracy have been observed in LLMs as well (Schaeffer al., 2024).
Why does decreasing SNR sometimes still yield high performance in Figure 4?
How might apply these results to other self-supervised learning tasks? Since our results distribution agnostic, Theorem 1 can be extended to establish an explicit equivalence between self- supervised contrastive learning and supervised matrix factorization. In particular, if the contrastive objective has the functional form L(W)- = p pairs (f+ (wo w')] + Emegative pains [f (w"" ww') for some diffcrentiable f (.) and f (-), then one can simply take the quadratic Taylor approximation, complete the square, and obtain a target matrix in terms of the positive and negative distributions of the learning problem. See Appendix D.2 for a demonstration of this idea for the SimCLR loss. To obtain closed-form learning dynamics, one needs to find equivalent of our Setting 3.1. This may be challenging depending on the form of the input data; for instance, when the imputs are not one- hot, it may require a whitened data covariance.
What do these results tell us about feature learning? Two operational notions of feature learning are 1) optimization trajectories must escape the local vicinity of the (typically small) initialization (Yang & Hu, 2021; Jacot et al., 2021; Zhu et al., 2022; Atanasov et al., 2024; Kunin et al., 2025), and 2) learned weight matrices must project data onto target-relevant subspaces (Damian et al., 2022; Radhakrishnan et al.. 2024). Our message complements this line of research by offering a practical yet solvable setting in which both behaviors appear: word2vec escapes its near-isotropic initialization region to learn a dense compression of the relative excess co-occurrence probability M thereby aligning the embedding geometry with the most salient semantic linear subspaces. Our result is thus step forward in the broader scientific project of obtaining quantitative, predictive descriptions of learning in practical algorithms.
Limitations. Our results are limited to the symmetric setup (tied encoder/decoder weights, i.e., W = W') We did not train and evaluate at scales that are considered state-of-the-ard nor did we compare against other embedding models such as GloVe (Pennington et al., 2014).
Author contributions. DK developed the analytical results, ran all experiments, and wrote the manuscript with input from all authors. JS proposed the initial line of investigation and provided insight at key points in the analysis. YB and MD helped shape research objectives and gave oversight and feedback throughout the project's execution.
10

---

Acknowledgements. DK is grateful to Daniel Kunin and Carl Allen for early conversations about learning linear analogy structure from small initialization. DK thanks Google Deep.Mnnd and BAIR for funding support and compute access. MRD thanks the U.S. Army Research Laboratory and the U.S. Army Research Office for supporting this work under Contract No. W911NF-20-1-0151.
References
Arora, S., Li, Y. Liang, Y., Ma, T., and Risteski, A. A latent variable model approach to pmi-based word embeddings. Transactions of the Association for Computational Linguistics, 4:385-399, 2016. Cited page 2.
Arora, S., Cohen, N., Hu, W., and Luo, Y. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019. Cited on page 2.
Atanasov, A., Bordelon, B., and Pehlevan, C. Neural networks as kernel learners: The silent alignment
effect. In International Conference Learning Representations, 2022. Cited pages 2, 5, and 6,
Atanasov, A., Meterez, A., Simon, J. B., and Pehlevan, C. The optimization landscape of sgd across the feature learning strength. arXiv preprint urXiv:2410006642, 2024. Cited on page 10.
Baik, J., Ben Arous, G., and Péché, Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices. Annals of Probability, 2005. Cited on page 8.
Baldi, P. and Hornik, K. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58. 1989. Cited on page 34.
Bordelon, B., Canatar, A., and Pehlevan, C. Spectrum dependent learning curves in kernel regression and wide neural networks. In International Conference on Machine Learning, pp. 1034. PMLR, 2020. Cited on page 34.
Bradbury, J., Frostig, R., Hawkins, P.. Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //gpth/b//hubbb com/ jax ml/ jax. Cited on page 23.
Bruni, E., Tran, N.-K., and Baroni, M. Multimodal distributional semantics. Journal of artificial intelligence research, 49:1-47, 2014. Cited on page 24.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1559--1077 PMLR, 2020. Cited on page 33.
Chizat, L., Oyallon, E., and Bach, F. On lazy training in differentiable programming. Advances in neural information processing systems, 32, 2019. Cited page 34.
Chou, H.-H., Gieshoff, C., Maly, J., and Rauhut, H. Gradient descent for deep matrix factorization: Dynamics and implicit bias towards low rank. Applied and Computational Harmonic Analysis, 68: 101595, 2024. Cited page 2.
Church, K. and Hanks, P. Word association norms, mutual information, and lexicography. Computa- tional linguistics, 16(1):22-29, 1990. Cited on page 33.
Damian, A., Lee, J., and Soltanolkotabi, M. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pp. 1413-5452, PMLR, 2022. Cited page 10.
Dominé, C. C., Braun, L., Fitzgerald, J. E., and Saxe, A. M. Exact learning dynamics of deep linear networks with prior knowledge. Journal of Statistical Mechanics: Theory and Experiment, 2023 (11):114004, 2023. Cited on page 5.
Finkelstein, I.., Gabrilovich, E.,, Matias, Y., Rivlin, E., Solan, Z.,, Wolfman, G., and Ruppin, E. Placing search in context: The concept revisited. In Proceedings of the 10th international conference on World Wide Web, pp. 406-414, 2001. Cited on page 24.
Gidel, G., Bach, F., and Lacoste-Julien, S. Implicit regularization of discrete gradient dynamics in linear neural networks. Advances in Neural Information Processing Systems, 32, 2019. Cited on pages and 5.

---

Gillis, N. and Glineur, F. Low-rank matrix approximation with weights or missing data is np-hard. SIAM Journal on Matrix Analysis and Applications, 32(4):1149-1163 2011. Cited on page
Gissin, D., Shalev-Shwartz, S., and Daniely, A. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv: 1909. /2051, 2019. Cited on page 2.
HaoChen, J. Z.. Wei, C., Gaidon, A., and Ma, T. Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in Neural Information Processing Systems, 34: 5000-5011, 2021. Cited page 2.
Harris, Z. S. Distributional structure, 1954. Cited on page 3.
Huang. J. and Yau. H.-T. Dynamics of deep neural networks and neural tangent hierarchy. In International conference on machine learning, pp. 4542-4551, PMLR, 2020. Cited on page 34.
Ilharco, G., Ribeiro, M. T., Wortsman, M., Gururangan, S., Schmidt, L., Hajishirzi, H., and Farhadi, A. Editing models with task arithmetic. arXiv preprini 2rXiv:2212.04005 2022. Cited on page 8.
Jacot, A., Gabriel, F,, and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018. Cited on page 34.
Jacot, A., Ged, F., Şimşek, B., Hongler, C.,, and Gabriel, Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprini arXiv:21066 15933, 2021. Cited on pages 2, 10, and 34.
Jiang, L.. Chen, Y., and Ding, L. Algorithmic regularization in model-free overparametrized asymmetric matrix factorization. SIAM Journal on Mathematics of Data Science, 5(3):723-744, 2023. Cited page 2.
Jiang, Y., Rajendran, G., Ravikumar, P., Aragam, B., and Veitch, V. On the origins of linear representations in large language models. In Proceedings of the 41st International Conference on Machine Learning. 2024. Cited on page 8.
Karkada, D. The lazy (ntk) and rich (Hp) regimes: a gentle tutorial. arXiv preprint arXi:2244.119719 2024. Cited on page 34.
Kunin, D., Marchetti, G. L., Chen, F., Karkada, D., Simon, J. B., DeWeese, M. R., Ganguli, S., and Miolane, N. Alternating gradient flows: A theory of feature learning in two-layer neural networks. arXiv preprint arXiv:2506.0644
2025. Cited on page 10.
Lauscher, A.,
Glavaš, G., Ponzetto, S. P., and Vulić, I. A general framework for implicit and explicit debiasing of distributional word vector spaces. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 8131-8138, 2020. Cited on page 8.
Lee, A., Bai, X., Pres. L. Wattenberg, M., Kummerfeld, J. K., and Mihalcea, R. A mechanis- tic understanding of alignment algorithms: A case study on dpo and toxicity. arXiv preprint arXiv:2401.00 019667 2024. Cited page 8.
Lee, J., Xiao, L., Schoenholz, S., Bahri, Y., Novak, R., Sohl-Dickstein, J., and Pennington, J. Wide neural networks of any depth evolve linear models under gradient descent. Advances in neural information processing systems, 32, 2019. Cited on page 34.
Levy, O. and Goldberg, Y. Neural word embedding as implicit matrix factorization. Advances in neural information processing systems, 27, 2014. Cited on pages 2, 3, and 33.
Levy, O., Goldberg, Y., and Dagan, Improving distributional similarity with lessons learned from word embeddings. Transactions of the association for computational linguistics, 3:211-225, 2015, Cited on pages and
Li, K., Patel, O., Viégas, F., Pfister, H., and Wattenberg, M. Inference-time intervention: Eliciting truthful answers from a language model. Advances Neural Information Processing Systems, 36, 2024. Cited page 8.
12

---

Li, Y., Ma, T., and Zhang, H. Algorithmic regularization in over parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018. Cited on page 2.
Li, Y. Li, Y.. and Risteski, A. How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning, pp. 19689-19729, PMLR, 2023. Cited on page
Li, Z., Luo, Y., and Lyu, K. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2021. Cited page 2.
Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., and Dean, Distributed representations of words and phrases and their compositionality. Advances in neural information processing systems, 26, 2013. Cited on pages 1, 2, 3, 23, and 24.
Nanda, N., Lee, A., and Wattenberg, M. Emergent linear representations in world models of self- supervised sequence models. In Proceedings of the 6th BlackboxNI.P Workshop: Analyzing and Interpreting Neural Networks for NLP. Association for Computational Linguistics, 2023. Cited on page
Park, K., Choe, J.J., and Veitch, V. The linear representation hypothesis and the geometry of large language models. arXiv preprint arXiv:231 1.03658, 2023. Cited on page 8.
Pennington, J., Socher, R., and Manning, C. D. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532-1543, 2014. Cited page 10.
Radhakrishnan, A., Beaglehole, D., Pandit, P., and Belkin, M. Mechanism for feature learning in neural networks and backpropagation-fre machine learning models. Science, 36696)))) 1461-1467, 2024. Cited on page 10.
Rumelhart, D. E. and Abrahamson, A. A. A model for analogical reasoning. Cognitive Psychology, 5 (1):1-2,, 1973, Cited on page 24.
Saunshi, N., Ash, J., Goel, S., Misra, D., Zhang, C., Arora, S., Kakade, S., and Krishnamurthy, A. Understanding contrastive learning requires incorporating inductive biases. In International Conference Machine Learning, pp. 19250-19286. PMLR, 2022. Cited on page 2.
Saxe, A., McClelland, J., and Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In Proceedings of the International Conference Learning Represenatations 2014. International Conference on Learning Represenatations 2014, 2014. Cited on pages 2, 5, and 29.
Saxe, A. M., McClelland, J. L., and Ganguli, S. A mathematical theory of semantic development in deep neural networks. Proceedings of the National Academy of Sciences, 1 116(23):1153--11546 2019. Cited on page 2.
Schaeffer, R., Miranda, B., and Koyejo, S. Are emergent abilities of large language models a mirage? Advances in Neural Information Processing Systems, 36, 2024. Cited on page 10.
Simon, J. B., Dickens, M., Karkada, D., and Deweese, M. The eigenleaming framework: A conservation law perspective on kernel ridge regression and wide neural networks. Transactions on Machine Learning Research, 2023a. Cited on page 34.
Simon, B., Knutins, M., Ziyin, L., Geisz, D., Fetterman, J., and Albrecht, J. On the stepwise nature of self-supervised learning. In International Conference on Machine Learning, pp. 31852- 31876. PMLR, 2023b. Cited on pages 2, 5, and 34.
Tian, Y., Chen, X., and Ganguli, S. Understanding self-supervised learning dynamics without contrastive pairs. International Conference on Machine Learning, pp. 10268-10277 PMLR, 2021. Cited page 2.
13

---

Vyas, N., Atanasov, A., Bordelon, B., Morwani, D., Sainathan, S., and Pehlevan, C. Feature- learning networks are consistent across widths at realistic scales. Advances Neural Information Processing Systems, 36: 1036--0000 2023. Cited on page 34.
Wang, Z. Gui, L., Negrea, J., and Veitch, V. Concept algebra for (score-based) text-controlled generative models. Advances in Neural Information Processing Systems, 36, 2024. Cited on page 8.
Yang, G. and Hu, J. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pp. 11727-11737. PMLR, 2021. Cited on page 10.
Zhu, L., Liu, C., Radhakrishnan, A., and Belkin, M. Quadratic models for understanding catapult dynamics of neural networks. arXiv preprint arXiv:2205. 11787, 2022. Cited on page 10.
Zou, A., Phan, L., Chen, S., Camphell, J., Guo, P., Ren, R., Pan, A., Yin, X., Mazeika, M., Dombrowski, A.-.., et al. Representation engineering: A top-down approach to ai transparency. arXiv preprint 2r 2023. Cited on page 8.
14

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], , [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] provvdddd proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA]" is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and intro state the contributions of the paper and frame them in relation to the larger goal of developing theory for machine learning.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
15

---

Answer: [Yes]
Justification:
The limitations are discussed in the final section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
model well-specificatie
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-ted system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Proofs can be found in the appendices.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
16

---

Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Experimental details be found in the appendices. Code is provided in supplemental material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source datasct or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The code to run experiments is provided the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (tttps://nips.cd public/guie httpa///giiiee//ddeelaiiissiioio oliyy) for more details.
17

---

While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (httpa: cttpsssccpppblcccc gttg///iiee//e// (ooeeSumiissioo/PP
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
The experimental details can be found in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: The main contribution of this paper is theoretical; the experiment simply confirms the analytic theory. Our sanity checks revealed that our results not sensitive to particular initializations or data subsets.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
18

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eттоr rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Given in appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https http///eurupp..eeee htt.:rrpp.cc.puu.iiittttiiiieeee cc/pppblii/ttthhccccueellellllll
Answer: [Yes]
Justification:
No ethics violations.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
This paper is primarily theoretical.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
19

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: No high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We train Wikipedia, which has CC 4.0 license.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
20

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/dttasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: No human subjects.
21

---

Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
22

---

A
Experimental details
All our implementations use jax (Bradbury al., 2018). all reported experiments we use a vocabulary size V= = 10 model dimension d = 200, and context length I, = 32 neighbors for each word. These were chosen fairly arbitrarily and our robustness checks indicated that our main empirical results were not sensitive to these choices.
A.1 Training corpus
We train all word embedding models on the November 2023 dump of English Wikipedia, downloaded from
ct/cco/////oo /ad/aaaats///iiiieeeaaaaa hettss/wwwiieiiii / We tokenize by replacing all non-alphabetical characters (including numerals and punctuation) with whitespace and splitting by whitespace. The full corpus contains 6.4 million articles and 4.34 billion tokens total. The number of tokens per article follows long-tailed distribution: the (50%, 90%, 95%, 99%) quantiles of the article token counts are (364, 1470, 2233, 5362) tokens. We discard all articles with fewer than 500 tokens, leaving training corpus of 1.58 million articles with 2.00 billion tokens in total.
We use vocabulary consisting of the V 10* most frequently appearing words. We discard oft-of-vocabulle words from both the corpus and the benchmarks. Our robustness checks indicated that long as the corpus is sufficiently large (as is the case here), it does not matter practically whether out-of-vocabulary words are removed or simply masked. This choice of V retains 87% of the tokens in the training corpus and 53% of the analogy pairs.
A.2 Optimization
We optimize using stochastic gradient descent with no momentum nor weight decay. Each minibatch consists of 100,000 word pairs (50,000 positive pairs and 50,000 negative pairs). For benchmark evaluations (Figure 3) we use a stepwise learning rate schedule in which the base learning rate is decreased by 90% at = 0. max- We found that this was very beneficial, especially for QWEMs, which appear more sensitive to large learning rates. The finite-stepsize gradient descent dynamics of matrix factorization problems remains an interesting area for future research.
We directly train the tensor of embedding weights W E RdXV One potential concern is that since the model WWT positive semidefinite, cannot reconstruct the eigenmodes of M E gVXV with negative eigenvalue. However, we empirically find that with V = 10 M* has 4795 non-negative eigenvalues. Therefore, in the underparameterizee regime d < V, the model lacks the capacity to attempt fitting the negative eigenmodes. Thus, the PSD-ness of our model poses no problem.
A.3 Reweighting hyperparameters
Taking from the original word2vec implementation, we use the following engineering tricks to improve performance.
Dynamic context windows. Rather than using a fixed context window length L,, word2vec uni- formly samples the context length between and L each center word. In aggregate, this has the effect of more frequently sampling word pairs with less separation. Let dej be the mean distance between words and j. when found co-occurring in contexts of length L. (Thus, d, is small for the words "phase" and "transition" since they are a linguistic collocation, but large for "proved" and "derived" since verbs are typically separated by many words.) Then not hard calculate the effect of mniformly-distriuur dynamic context window lengths: it is equivalent to setting
dij
(14)
Ev (L
deg Pes
Subsampling frequent words. Mikolov et al. (2013) suggest discarding very frequent words with reqquency-dependen discard probability. This is akin to rejection sampling, with the desired unigram distribution flatter than the true Zipfian distribution. The original word2vec implementation uses an acceptance probability of
10-3
10
10
Pana min
1,
+
(15)
P
P.
This is equivalent to setting the hyperparameters
V
Vij
Paee Pace
(16)
23

---

Different negative sampling rate. Finally, Mikolov et al. (2013) draw negative samples from different distribution (empirically, they find that pa is particularly performant) and upweight the negative sampling term by hyperparameter (they recommend & = 5 for large corpora). This is equivalent to setting
kPj 1/4
(17)
E, P
Note that this is asymmetric choice, it is not covered by Setting 3.1. However, both this and the previous subsampling technique are methods for modifying and flattening the word distributions, as seen by the training algorithm. Our Setting 3.1 accomplishes the same thing. We speculate that these manipulations accomplish for language data what spectral whitening does for image data.
Taken together, these give prescription for cleanly and concisely capturing several of the implemen- tation details word2vec in a set of training hyperparameters. word2vec, use all of these settings they can be combined by multiplying the different y and likewise for V For QWEMs, we use " = Vij = (Pi + P.PP) and modify V by combining with the setting for dynamic context windows. For fair comparison, we use k for the SGNS experiments in the main text.
our experiments, when training
A.4 Benchmarks
We use the Google analogies described in Mikolov et al. (2013) for the analogy completion benchmark, https://a thub. com/ com/tmikee amikolov' compute the analogy accuracy using
words. txt. We then
acc (W)
Σ
1(b)
arg max (aa.bba)
aa' + (6 a))
(18)
D
(a.b.a' (a.b.a'.b)ED
1
1 [(b)
arg min (a..ba.
à- b a' + 2
(19)
D
(D
where the 4-tuple of embeddings (a, b, a', b') constitute an analogy from the benchmark data D, hats (c.g., w) denote normalized unit vectors, 1 the indicator function, and {W} the set containing the word embeddings. The first expression measures cosine alignment between embeddings and the "expected" representation obtained by summing the query word and the task vector. The second expression measures the degree to which four embeddings form closed parallelogram (Rumelhart & Abrahamson, 1973). The two forms of the accuracy are equivalent since
arg min à b à' + uò = arg min (2w T (a b - a') const..) 111
(20)
10
arg max u T (à' + à).
(21)
w
To understand the role of normalization, we compared this metric with one in which only the candidate embedding is normalized:
winner
arg max
à (a' + (b (b-a)) a)).
(22)
WE W (aa,ba,
We found that the accuracy metric using this selection criterion yields performance measurements that are nearly identical to those given by the standard tully-normalized accuracy metric. We conclude that the primary role of embedding normalization is to prevent longer w's from "winning" the arg max just by virtue of their length. The lengths of a, b, a' are only important if there is significant angular discrepancy between (á' + b à) and (a' + b a); in the high-dimensional regime with relatively small variations in embedding length, we expect such discrepancies be negligible.
For semantic similarity benchmarks. we use the MEN dataset (Bruni et al., 2014) and the WordSim353 dataset (Finkelstein al., 2001). These datasets consist of of N word pairs ranked by humans in order of increasing perceived semantic similarity. The model rankings are generated by computing the inner product between the embeddings each pair and sorting them. The model is scored by computing the Spearman's 's rank correlation coefficient between its rankings and the human rankings:
p(W)= == Corr(r (1), Fmodel (+)
(23)
where r(i) is the rank of pair i and Corr is the standard Pearson's 's correlation coefficient.
24

---

A.5
Computational resources
Our implementations are relatively lightweight. The models can be trained on a home desktop computer with 17 4-core processor, 32GB RAM, and consumer-grade NVIDIA GTX 1660 graphics card, in about hours. For the experiments in Figure 3, we train for 12 hours with a much lower learning rate. Our code is publicly available at https: / /githutt.t..///g com/ dkarkada/quem.
A.6 Empiries of task vectors
We construct the theoretical embeddings W by constructing M' according to Equation (5), diago- nalizing it, and applying Theorem 1. Due to right orthogonal symmetry, we are free to apply any right orthogonal transformation; we choose the identity the right singular matrix, so that the components of each embedding in the canonical basis are simply its projections on the rigen-features. With this setup, it is easy to extract the embeddings resulting from smaller model: simply truncate the extrancous columns of W.
We construct the task vectors by first collecting a dataset of semantic binaries. These can, for example, be extracted from the analogy dataset. The task vectors are then simply the difference vectors between the embeddings in each binary. In our reported results, we do not normalize either the word embeddings nor the task vectors. However, our robustness checks indicated that choosing to normalize does not qualitatively change our results.
Fitting the empirical spectra with the Marchenko- Pasttur distribution. The Marchenko-Paiee distribution is limiting empirical spectral distribution given in terms of aspect ratio 7 N/d and an overall scale o². In particular, if N < d and =\t E ggNxd is random matrix with i..d. elements of mean zero and variance o², then as N, d + oo with N/d = y fixed, the MP law for d d- is
1)(1 - Yp
du((A) 2
(24)
7t
A = o1=
(25)
where the support is A_ < A S A4, and the density is zero outside. Thus, the MP law gives the expected distribution of eigenvalues for large noisy covariance (or Gram) matrix.
To fit the MP law to the Gram matrix of task vectors, we first subtract off the mean task vector. Then compute the population variance of the elements of the centered matrix use this to set the o² parameter. We manually fit the aspect ratio 7, and we find that the best-fitting 7 = N/datt corresponds effective dimension which is often much lower than the true embedding dimension d. We hypothesize that deif < d due to anisotropic noise in the task vectors. We report the dea resulting from the fit in Figure
Composition of mean task vector. We measure the degree to which the spike captures the mean task vector by computing the ratio
R $
signal in mcan N Amax RRa R
(26)
where 1 is the ones-vector. If this quantity is (close to) 1, then the mean task vector is a large component of the spike, and vice versa. We find that this quantity is consistently greater than 0.9 for all the semantic categories in the analogy benchmark.
25

---

B
Additional figures
1.0
word2vec eigen-feeee featfeetuuess V3
10
10
0.8
20
20
0.6
30
30
0.4
40
50
0.2
60
60
10
20
30
40
50
60
0 10
20
30
40
50 &
60
QWEM eigen-features V
PPMI eigen-features
Figure 5: word2vec eigen-features align very closely with QWEMs' and less with PPML. The heatmap indicates the squared overlaps between the latent feature vectors (i.c., the left singular vectors) between the three models considered. This is quantitative version of the table in Figure 3.
spectra of R.R !
peojection of eigen-features
features on mean task vector
d= 255
0.075
dig
male/ female
0.050
0.025
0.000
d= 256
adjective! adverb
di UK.4
0.2
0.1
0.0
10
266
0.075
Ag 52.8
noun! plural counts
0.050
0.025
0.000
adjectivel comparative
dig 83.4
0.2
0.1
0.00
0.10
0.20
80
100
oigenvalue
eigen-feature indox (with eigen- u,)
Figure 6: The empirical observations regarding linear representations extend across the semantic classes in the analegy dataset. We show that across analogy categories, the corresponding task vectors exhibit the geometric structure discussed in Figure 4.
26

---

regular SGNS
Quartic loss only
40
20
singular valuees of embeddings
0-
Setting 3.1 only
Quartic loss Setting 3.1 (QWEM)
10P
10"
rescaled training time t/T)
Figure 7: Ablation tests. Starting from vanilla word2vec (top left), we separately add both the quartic approximation (top right) and the reweighting condition Setting 3.1 (bottom left). We include QWEM on the bottom right (i.e., both effeets). The quartic approximation hardly changes the singular value dynamics all: the clean mode separation due to the reweighting. However, is common with logistic losses, weights begin diverge at late times. Using the quartic loss prevents this. The experiments here use 41-million-token mixture of Corpus Contemporary American English and the News on We use vocabulary size of 20,000; the embedding dimension is 150; 2 negative samples per positive sample; and context length of 16.
dataset.
1.2
1.2
1.0
1.0
0.8
0.8
cos w)
000 w)
IIN,, cve(v,
0.6
0.6
country-nnioonaliy male-female
0.4
0.4
verb-pasttense
0.2
adj-adverb
0.2
adj antonym
0.0
0.0
10a
102
103
10
10a
10a
embedding dimension d
embedding dimension d
Figure 8: Although the analogy geometry degrades with d, relative scores remain high. We plot two different measures analogy geometry across d illustrate subtleties associated with the top-1 accuracy metric. For an analogy a b :: c: w. we denote á = c+ (b a) the prediction for the fourth word. We plot the average cosine similarity between & and the true w over the analogy family. The prediction degrades dramatically as the embedding dimension increases, complementing the observation in Figure 4 that the SNR decreases with d. increases:
However, this geometric breakdown fails capture the fact that all embedding vectors separate we normalize by the maximum cosine similarity between w and non-analogy embeddings, the score remains roughly stable at large This provides an explanation for why top-1 accuracy often remains high despite the breakdown of geometric analogical structure.
27

---

C Proofs
Theorem 1 (QWEM = unweighted matrix factorization). Under Setting 3.1, the contrastive loss Equation (3) can be rewritten the unweighted matrix factorization problem
((W) 9 WW M 1 + const.
-
(6)
If (oted] is positive semidefinite, then the set of global minima of given by
arg min L(W) -REquiv (VE.nng A 1/2 (sd.id)
W
Proof. Define M = WW!
Rewriting Equation (3) and plugging in Equation (5) and Setting 3.1,
1 M MMs) + P.P.V ("MS, + Mij
<(W)
Σ
Pij V
(27)
4
My 4(P,) V.J P.P,
(28)
1 Σ Gij
P. P.P.V. Mij
MR - 2 [[, V i.P.
(29)
4
M
1 Σ
Gy (M²) 2M My + +M² Mi;
(30)
4
V
9 (IIM 1 MM)
(31)
Since Pij. P.Py, Wij. Vij all real symmetric, so is M so it has eigendecomposition. By the Yckart-Young-No Mirsky theorem, the loss-mininizing M must be the truncated SVD Maj whose symmetric factors exactly given by Equation (7).
Proposition 2. Equation (3) can be rewritten as the weighted
problem
((W) 1 Σ G. MM + const.
(8)
Let * and V each be symmetric in i, let g RV be an arbitrary vector with non-negative elements, and let r = diag(g) Then the eigendecomposition of T V Vr exists. If
G gg and Aidad] is positive semidefinite, then the of global minima of L is given by
arg min L(W) -REquiv VI W
(9)
Proof. The formulation as a weighted matrix factorization follows from Equation (30). In the case that is rank substituting r, we find
L(W) I M
(32)
After distributing factors and invoking the Eckart- Young- theorem, we conclude that the rank-d minimizer is
Mmin r (TM'T) r
(33)
whose symmetric factors are exactly given by Equation (9).
28

---

Lemma 3.1 (Training dynamics, aligned initialization). If == Apdcdd is positive semidefiite and
V(0)= T.ag then under Setting 3.1, the gradient flow dynamics
dW di 2g
VL(W) yields
((t) U())
(10)
S(t (0) (e - 552 (I-c-Ar)) (S-(0)---T--- -1/2
(11)
(t)) = ( (ird
(12)
Proof. By Theorem we write the loss (Equation (3)) as
((W) 9 Tr [(WW]- -M) M") (WWT- -M)].
(34)
We neglect constant terms since they do not affect the gradient descent dynamics. Under the stated gradient flow, the equation of motion is
dw
1
VL(W)
(35)
dt
2g
(2(WW - (2W))
(36)
8
5 (M-
(37)
From here on we adopt the dot notation for time derivatives. We use the spectral decompositions
" = (t) and M = V*AV*T Then the above can be written
W VSU V VSU (V V SUT
(38)
deft-multiplying by and right hightmmultipliing by U, we obtain
1 (A S
VTS+ SÜ
(39)
where we now use the alignment assumption, VTV = Note that since we use the economy SVD in our notation, use non-standard notation where I < RaxV is rectangular matrix with ones on the main diagonal and zeros elsewhere. Now, since the RHS is diagonal, we must have that
(VTVS + SÜ SÜTU); =0 for i
(40)
Furthermore, since V and U are orthogonal,
V V V and ÚT must both be antisymmetric (since for any orthogonal matrix Q. de Q Q Q Q Q Q = 0). It follows that, for all i i j.
(V 1888 + 0) = (V 8: + (ÜTU) O.
(41)
Isolating (V )ij. we that this can only hold if s; 8j or if VV )is (T,)), 0. The former is ruled out by level repulsion in Gaussian random matrices; with probability we have that S contains distinct singular values. We conclude that Ú = 0 and V =0
Returning to Equation (39), we have that
$ 5= (A- SA)
(42)
These are precisely the dynamics studied in Saxe al. (2014). These dynamics are now decoupled, so we may solve each component separately. The solution to this equation is
()) Au Avt
(43)
A (0) (e -1)
Thus, the each singular direction of the embeddings is realized in a characteristic time
1
Tk
In
(44)
29

---

Result 3 (Training dynamics, vanishing random initialization). Initialize W ~ NN (") and let S(0) denote the singular values of (0). Define the characteristic time == X (n and the rescaled time variable i i= T/ Define (0) Vi ((0) and let (* be its gradient flow trajectory given by Lemma 3.1. If (ddad positive semidefinite, then under Setting 3.1 we have with high probability that
vi 2 O.
lim o² o2
min W'EREqui(W((n))
|
W (fn)))l = O.
(13)
Derivation. Before starting the main derivation, we give qualitative argument for why one expects the result of Lemma 3.1 to hold in the small initialization limit despite a lack of initial alignment.
Warmup. We begin with the equation of motion for M == WWT:
M WWT + WWT = (MM* + M 2
(45)
We again consider the dynamics in terms of the spectral decompositions (t) (tv
and M V*AVV* Note that here V,S VVV are square. We define the eigenbasis overlap
O = V VV V. After transforming coordinates to the target eigenbasis, we find
v y-Tvv+ V (VS'VV + V(SS)V + +VS'VV
(46)
O + 20S$00 + T
(47)
O OSSO o
(48)
For clarity, we rotate coordinates again into the O basis and find
0'0 + +25 O AO + O AOS?
S*.
(49)
2
Since O is orthogonal, satisfies O O I). Therefore the first two terms the LHS of Equation (49), which concern the singular vector dynamics, have zero diagonal; the third term, which concerns singular value dynamics, has zero off-diagonal. This implies
+ T Ó 0 (this follows from differentiating the identity
255 (O (TAO S², S
(50)
where diag(-) is the diagonal matrix formed from the diagonal of the argument. From Equation (50), we see that initialization 25$ scales with the initialization size o² since (²) o². On the other hand, from the off off-diagonal of Equation (49), we see that Ô is order (since the scale of O is fixed by orthonormality). Therefore, in the limit of small initialization, we expect the model to align quickly compared the dynamics of S2. This motivates the silent alignment ansatz, which informally posits that with high probability, the top d x d submatrix of O converges towards the identity matrix well before 52 reaches the scale of A. As O + I, the dynamics decouple and enter a regime well-described by Lemma 3.1. We formalize this idea in our concrete derivation of Result 3.
Main derivation.
We are interested in showing that
V
0,
lim -0
min (()
I WT 0,
(51)
where we express the statement in terms of rescaled time i -= L/T; since we anticipate that 71 - 0o as 0² + 0. For notational clarity, though, let us switch back to the original time variable. Then
min (C))
min [[W'(t) - W (t)llrr-= (()eRrEqii((Wee))))
( - (
(52)
min
(U((U(()) (()))(((T)-)))) (1) -0)-5 S* ||F (53)
min U'EeE
(U()))(())(())))) (() - (r))l
(54)
where we again define the eigenbasis overlap 0 = TV
30

---

Motivated by the expectation that we will see sequential learning dynamics starting from the top mode and descending into lower modes, we seek a change of variables in which the dynamics are expressed in an upper-triangular matrix. We can achieve this reparameterization using a QR factorization: USO QR, where Q is orthogonal and R. is upper triangular. Then
min (())
(W) W (f)llr = min (I)(()(()) S (1)llr
(55)
U'itEE
(IU( (t))) U'(D)E
min
(56)
where the second equation follows since U' is a variable to be minimized, and it can simply absorb Q We emphasize that this convenience follows from the right orthogonal symmetry of the embeddings. If we can show that R(t)) - (t)lle + 0 for all t in the vanishing initialization limit, then we will have completed the derivation.
Our starting point will be transpose of Equation (37), right-multiilied by V - QR (A- R
1 -UOO (A- =
dt (USOT)= = 2
(57)
Left-multiplying by 2Q' T and rearranging, we find
2R = R (A- R R) - 2Q QR
(58)
RA RRR' =2Q Q)
(59)
RA RR.
(60)
where we define R. RR + 2QQTQ Q. Note that R must be upper triangular: since the other terms in the equation are upper triangular, so must be RR; and since R and R.R are both upper triangular, then R must be upper triangular.
In fact, this is enough to fully determine the elements of R. We know that Q Q is antisymmetric (since QTQ Q- I by orthogonality. Q Q Q Q 0). Additionally using the fact that RR is symmetric and imposing upper-triangull on the sum, we have that
2 RR Ky (RR
if i<j
R.g-
Da
if i i=j.
(61)
0
if i>j
Here, we take a moment to examine the dynamics in Equation (60). Treating the initialization scale scaling variable, expect that R.j ~ Thus, in the small initialization limit, we expect the second term (which scales like σ²) to contribute negligibly until late times; initially, we will see an exponential growth in the elements of R with growth rates given by A. Later, R will (roughly speaking) reach the scale of A 1/2, and there will be competitive dynamics between the two terms. We now write out the element- dynamics of R to see this precisely.
2Rg RujAy - Σ
Rys
(62)
jzk2i
-Rid Σ
8t Rut
(63)
12R21 EZk
- -RidA- Σ
R Ruj
Σ
Σ
Rut Ru Rkj
(64)
2
12k 12k
-Rijd-
Σ
R Raj - Σ
Rh - 2
Σ
Rae Raj
(65)
E21
j2k2i
f2kzifs12k
A₂
R
R - Σ
R
R,j - 2
Ra f2
Σ
(66)
j2k5i
We have separated the dynamics of Ray into part that is explicitly linear in R.j and a part which has no explicit dependence on Raj- (Of course, there is coupling between all the elements of R and Ry through their own dynamical equations.) So far, everything we have done is exact.
31

---

Our main approximation is to argue that at all times, only the diagonal elements of R contribute non-negligibly to the dynamics. This holds if the off-diagonal elements are initialized vanishingly small and if they remain vanishingly small throughout. In this case, we may discard any terms that include couplings between off-diagonal elements.
With this approximation, we may discard the entire second term on the RHS, as well as some of the summands in the first prefactor for Rij. This leaves the approximate dynamics
2R = (2g R 2(1 big R3 Raj-
(67)
We may now directly solve for the diagonal dynamics. Letting Th = RE
02(0) d exat
tx = 2 A2 r Fk
(t) A + (0) (Et.t -1
(68)
We obtain the same sigmoidal dynamics as in Lemma 3.1. If we show that the off-diagonal elements remain vanishingly small in the limit o² -> O, then: a) approximation is justified, and b) it follows that || (t) - 5* (t)lle + 0 for all t, completing the derivation.
To do this, we examine the dynamics of the off-diagonals and show that the maximum scale they achieve (at any time) decays to zero as x² > For <j j we have
2Rg (Λ3 r - 2r R.j
(69)
This is linear first-order homogeneous ODE with a time-dependent coefficient, and thus it can be solved exactly:
2
X2 r)) (e (1) (
A; d + r²( (ettt
(t)
ext
(70)
R0 (1) (t)) (0) (0)
-e- -A,+t,))
(71)
This product contains two factors with sigmoidal dynamics of different timescales, and one factor with an exponential decay to the dynamics. Thus, as + co, the first two factors saturate, and the decay drives the off-diagonal elements R.j to zero. We now show that max; (t) vanishes as ²² 0. Focusing on the scaling w.r.t. the initialization scale, we may approximate R o² for all Discarding (aat) terms and solving R.j = O, we find
A:
Ar/de
log d:4s 02(Ai -1y)
max R
RJ
-2/-2 when t t= )
(72)
Therefore, for i < j. assuming A + Xjj
lim max R. =
(73)
02
and we conclude that R(() - S - 0 desired.
In fact, under these approximations, as long as the initialization scale satisfies
X log
A Agy
log 0²
(74)
X
for all and j. the off-diagonal elements will remain much smaller than the diagonal elements, and we may view the diagonal element dynamics simply being the singular value dynamics. This follows from Weyl's inequality for matrix perturbations. Thus we may expect that our results hold in the small-but but-finite initialization regime (e.g., the regime accessed by our experiments).
32

---

D Relation between QWEMs and known algorithms
D.1
Relation to PMI
Early word embedding algorithms obtained low-dimensional embeddings by explicitly constructing some target matrix and employing dimensionality reduction algorithm. One popular choice was the pointwise mutual information (PMI) matrix (Church & Hanks, 1990), defined
Ps P log P.P
(75)
Later, Levy Goldberg (2014) showed that PMI is the rank unconstrained minimizer of Luzv. To see the relation between the QWEM target M and PMI, let us write
Psj P.P
= 1+ A
(76)
where the function (xx) yields the fractional deviation away from independent word statistics, in terms of some small parameter x of our choosing (so that A()) = =0). This setup allows us to Taylor expand quantities of interest around
0. A judicious choice will produce terms that cancel the - -2 that arises from the Taylor expansion of log(1 + A), leaving only third-order corrections.
One such example is A (z) 2z (2 z). which yields
Ps P.P,
Jij
[[,, + P.PP)
(77)
and
2x 2
3
PMT = log 1
=
x
(78)
12
80
This calculation reveals that M learns very close approximation to the PMI matrix; the leading correction is third order. However, M is much friendlier least-squares approximation, since bounded (-2 < Mi < 2).
D.2 Relation to SimCLR
SimCLR is widely-used contrastive learning algorithm for learning visual representations (Chen et al., 2020). It uses a deep convolutional encoder to produce latent representations from input images. Data augmentation is used to construct positive pairs; negative pairs are drawn uniformly from the dataset. The encoder is then trained using the normalized cumperatur-scaales cross entropy loss:
exp (Bf(()(())) T So(z;))
((fo)
E (.j-Pt(-.-)
log
(79)
Σke3 exp (fo(z,) fo(
where Pr(s.)) is the positive pair distribution, fox; is the learned representation of x;, B is an inverse temperature hyperparameter, and B is the batch size. Defining (i fa())) fo(#;), in the limit of large batch size, we can Taylor expand this objective function around the origin:
C(Se)
E ()-Pr(+))
BSo(1.)) + log
E (-PP())
(xp((S + log B
(80)
E i)-Pr(-))
(So(i.))+ + E k
(exp(BSo(i,,)) - + log B
(81)
1 8 S k)
E ()-PP((s))
BSo(i.))
+ iP(())
+ (i, + 2
1 log
(82)
k-Pr(-)
3 (S(..00 So(i.)) +
3
E
So(i.))
E i-Pr(.) j-Pr(.)
const.
(83)
P.()P((.))
(84)
33

---

If the model is in linearized regime, we may approximate Sp(i.) = f; for some linearized feature vectors f. Then the loss can be written as an unweighted matrix factorization problem using exactly the same argument as in Theorem 1. Thus, we expect that vision models trained under the SimCLR loss in a linearized regime will undergo stepwise sigmoidal dynamics. This provides an explanation for the previously unresolved observation in Simon et al. (2023b) that vision models trained with SimCLR from small initialization exhibit stepwise learning.
D.3 Relation to next-token prediction.
Word embedding targets are order-2 tensors M that captures two-token (skip-gram) statistics. These two-token statistics are sufficient for coarse semantic understanding tasks such as analogy completion. To perform well on more sophisticated tasks, however, requires modeling more sophisticated language distributions.
The current LLM paradigm demonstrates that the next-token distribution is largely sufficient for most downstream tasks of interest. The next-token prediction (NTP) task aims to model the probability of finding word given preceding window of context tokens of length L Therefore, the NTP target is an order-L tensor that captures the joint distribution of length- contexts. NTP thus generalizes the word embedding task. Both QWEM and LLMs are underparameterized models that learn internal representations with interpretable and task-relevant vector structure. Both are trained using self-supervised gradient descent algorithms, implicitly learning a compression of natural language statistics by iterating through the corpus.
Although the size of the NTP solution space is exponential in L (i.e., much larger than that of QWEM), LLMs succeed because the sparsity of the target tensor increases with L. We conjecture, then, that dynamical description of learning sparse sigh-dimensional tensors is necessary for general scientific theory of when and how LLMs succeed on reasoning tasks and exhibit failures such hallucinations or prompt attack vulnerabilities.
D.4 Relation to neural tangent kernel.
Our result describing implicit bias towards low rank directly contrasts the well-studied neural tangent kernel training (NTK) regime (Jacot et al., 2018; Chizat et al., 2019; Karkada, 2024). Here compare the two learning regimes.
Learning with NTK.
Learning dynamics and generalization performance can be solved (Lee et al., 2019; Bordelon et al., 2020; Simon al., 2023a).
Extreme wver-paaameterization required; large finite-width corrections at practical widths (Huang & Yau, 2020).
Model weights do not learn task-relevant features.
Optimization remains in a locally convex region of the loss landscape (Chizat et al., 2019).
Learning from small initialization.
Learning dynamics are generally complicated; can be solved in very simple cases (linear networks, special data distributions, etc.).
Behavior is qualitatively consistent across network widths, with only moderate finite-width corrections (Vyas et al., 2023).
Model weights learn task-relevant features.
Optimization tends to pass near sequence of saddle points. (Baldi & Hornik, 1989; Jacot et al., 2021).
34


Reviews:
[{"review_id": "9BfjugKOHv", "reviewer": "Reviewer_WyQ8", "summary": "The authors present an analysis of the word2vec skip-gram negative sampling algorithm. They derive an exact minimizer of a Taylor approximation to the skip-gram objective (QWEM), and show that the optimal solution of QWEM is related to the PMI matrix describing word-word co-occurrences. Stochastic gradient optimization of QWEM yields a sequence of discrete rank-incrementing learning steps that correspond to substantial decreases in loss; the semantics of the additional dimensions added at a step $t$ correspond to dimensions added in a similar learning step in word2vec. QWEM can be optimized in either this stochastic way or solved for in closed-form; the resulting word vectors perform similarly to those from word2vec. A final section studies the geometry underlying linear vector analogy in the QWEM embedding space.", "questions": "- How do you account for the dissociation between SNR and accuracy in higher dimensions (the right halves of the bottom-right plots in Figure 4; compare e.g. the blue curves, where the SNR curve peaks and then decreases while the accuracy curve appears to monotonically increase)? If it is the case that this lone \"spike\" eigenvalue is the task-relevant direction, why is it the case that decreasing SNR at higher dimensions still yields higher performance?\n- Why do we not see the same clear discrete rank-incremental performance in the word2vec loss function?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "This is an extremely clear and interesting paper studying what is now a fundamental algorithm of the field, combining theoretical work with thorough empirical validation. This is an important result and could serve as the foundation both for further analysis of word2vec and for intuition about the function of more complex models.\n\nThe relationship between Section 4 and the rest of the paper is unclear to me — it seems like this is mostly a free-standing study, in principle independent of the theory motivated in earlier sections. Perhaps the authors want to link the claimed training dynamics (rank-incrementing steps of training with interpretable semantics) to the notion of task vectors operant in analogy evaluations — but from what I can see in fig. 4 top right, these task vectors don't have a strong one-to-one alignment with the model eigenfeatures; instead they combine many of those dimensions.", "was_revised": true, "final_justification": "Some other reviewers questioned whether this theoretical study of word2vec is useful in an age when word2vec has been replaced by LLMs as the primary object of study. I agree with the authors' rebuttal here: word2vec still establishes a relatively simple setting for preliminary theoretical analysis, and there is good chance (given some convergent findings about linear geometry in both spaces) that some of the insights could generalize to other more modern word embedding methods.\n\nI will keep my score.", "rebuttal": "Thank you very much for your review! We’re glad you found the paper clear and interesting. Some comments follow.\n\n**Why does decreasing SNR sometimes still yield high performance?** This is a very interesting question and we spent some time exploring this. Our leading hypothesis (corroborated by several empirical checks we did) is that this is an ancillary effect of the benchmark using top-1 accuracy as the error metric. As a concrete example, consider the analogy “France : Paris :: Japan : Tokyo”, satisfying japan+(paris-france) ≈ tokyo. As the effective embedding dimension increases, the SNR may decrease; at the same time, all the embeddings become more spaced out (due to the increased available volume in latent space), meaning there are fewer nearby “competitors” for tokyo. In the large $d$ regime, we empirically found that the embeddings are sufficiently spaced, and tokyo is still the nearest embedding to japan+(paris-france) despite an increase in absolute error. Thus, the measured top-1 performance does not degrade with the increased (relative) noise. Note that unintuitive effects associated with using top-k accuracy have been observed in LLMs as well (see https://arxiv.org/abs/2304.15004). We plan to include this discussion in the final version of the manuscript.\n\n**Where is the clear, discrete, rank-incrementing behavior in bona fide word2vec?** You correctly observe that the rank-incrementing steps in word2vec do not appear clearly separated (especially early in training). The reason that clear stepwise learning occurs in QWEMs is that our Setting 3.1 turns the weighted factorization problem into an unweighted factorization – as a consequence, the singular value dynamics become decoupled (i.e., “untangled”). Without our Setting 3.1, there is “mixing” between the singular directions, even when the corresponding singular values have grown to be $O(1)$ in magnitude. Exactly characterizing this mixing is NP hard (see Proposition 2 and the subsequent discussion). The somewhat surprising fact is that word2vec is “close enough” to the unweighted factorization – though the singular value dynamics do exhibit some mixing (as you observe), the overall learning dynamics are well-described by the simplified setting. This agreement is partly because the hyperparameters chosen in word2vec coarsely approximate our Setting 3.1 (see Appendix A.3).\n\n**Regarding the link between Section 4 and the theoretical results.** Though the empirical results in Section 4 do not rely on the correctness of the preceding theory, the two are related. Our first result (left and top right of Fig. 4) suggests that it may be possible to obtain a *sparse* effective theory of semantic linear representations by studying their decomposition in terms of the semantic basis identified by our theory. Our second result (bottom right) identifies a potential benefit of early stopping in large models – it may prevent the realization of noisy semantic directions in the latents. This insight follows from the duality between model size and training time identified by our theory (i.e., the low-rank bias). In general, all the empirics in Section 4 were relatively easy to iterate on because they only involved factorizing a large matrix once, rather than re-training many word embedding models. This computational saving follows directly from the theory (we discuss it in Section 3.3).\n\nLet us know if you have further questions!"}, {"review_id": "3dBPnIYLrM", "reviewer": "Reviewer_7VDM", "summary": "This paper formulates a quadratic approximation (QWEM) to the contrastive objective of word2vec. QWEM's training dynamics can be solved for in closed form when making some extra assumptions on the parameters (e.g., W = W'). The authors demonstrate the trained QWEM solutions empirically well-approximate the trained word2vec vectors, and better than other proposals (Levy et al 2014). A key insight (Result 3) is that the \"natural basis\" of the gradient flow dynamics are the eigenvectors of the matrix M* (Eq 5). What follows from this result is that the components are learned sequentially, which makes sense in light of early stopping as regularization. A nice follow-up observation is that task vectors (e.g. the \"man\"-\"woman\" vector) empirically appear rank 1, supporting the linear representation hypothesis. Overall, QWEMs appear a useful approximation of word2vec, in a way that not only approximates the end solution (as in Levy and Goldberg 2014) but crucially its training dynamics.", "questions": "1. Figure 4: (Bottom.) took me a while to parse. For readability I would add (\"SNR\") to spike strength and (\"analogy accuracy\") to \"the model's ability to use task vectors for analogy completion\". \n\n2. l251 State the benchmark in main text\n\n3. Is it enough for the $R^d$ to be rank-1 (line 261)? If the model is good at analogy completion, then shouldn't the task vectors be exactly equal rather than rank-1, where rank-1 instead implies they are the same up to scaling? If I understand correctly, rank-1 is a necessary but not sufficient condition for good analogy completion-- perhaps it'd be useful to comment on this in the manuscript.\n\n4. Figure 1: Missing discussion in the main text (only reference is line 217). For the left plots, I assume each line is one singular value? It seems like in QWEM, the SVs are learned more \"sequentially\", but each one is learned earlier than in the word2vec plot. In the word2vec plot, it seems there is more \"mixed information\" being learned earlier, but each SV individually converges later. Does this imply that QWEM promotes sparsity? And is that a good thing?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Disclaimer: I read the paper to the best of my ability, but I'm not well-versed in this literature / the methods. I don't envision changing my score or confidence score as I lack the background to judge the contribution.\n\n## Strengths\n1. I found this paper highly readable and clear, even not coming from this background. In particular, I liked that the assumptions were very clearly stated (Setting 3.1 and Lemma 3.1), and later justified empirically with closeness to the word2vec solutions.\n2. The theory is backed up with experiments that clearly show how rank increases in the representations with QWEM, and how that mirrors the learning dynamics of word2vec.\n\n## Weaknesses\nHad some minor clarity issues, see Questions.", "was_revised": true, "final_justification": "Overall nice paper! I disagree with other reviewers that word2vec is trivial and that the insights need to have an application for LLMs. Word2vec is a model that stands on its own, and besides, it is good to have some diversity when the field is already inundated by LLMs.", "rebuttal": "Thank you very much for your review! We’re glad you found the paper readable. Thank you for your suggestions regarding clarity – they are very useful to us and we’ll make the requisite changes. Some comments follow regarding your other questions.\n\n**Does early mixing and delayed convergence imply that QWEM promotes sparsity?** Both QWEMs and word2vec exhibit a low-rank bias in their learning dynamics, although the effect is more pronounced in QWEMs, as you observe. We don’t believe that this low-rank bias is necessarily good or bad; that depends on the downstream task of interest. Our theory is simply descriptive of the learning dynamics.\n\nThe reason that clear stepwise learning occurs in QWEMs is that our Setting 3.1 turns the weighted factorization problem into an unweighted factorization – as a consequence, the singular value dynamics become decoupled (i.e., “untangled”). Without our Setting 3.1, there is “mixing” between the singular directions (as you observe), even when the corresponding singular values have grown to be $O(1)$ in magnitude. Exactly characterizing this mixing is NP hard (see Proposition 2 and the subsequent discussion). The somewhat surprising fact is that word2vec is “close enough” to the unweighted factorization – though the singular value dynamics do exhibit some early mixing, the overall learning dynamics are well-described by the simplified setting. This agreement is partly because the hyperparameters chosen in word2vec coarsely approximate our Setting 3.1 (see Appendix A.3).\n\nAs for the slightly delayed convergence, we don’t think this is an important effect. The two algorithms optimize different objectives, so the magnitudes of the gradients are not exactly the same; one would not expect their learning timescales to match exactly. Our theory predicts the convergence time with relatively small error – this is already much better than previous results, which only give the scaling behavior (or loose bounds) for the convergence time, rather than a direct and accurate prediction.\n\n**Shouldn't the task vectors be exactly equal rather than rank-1?** Strictly speaking, for exact analogy recovery, you’re right that rank-1 task vectors are necessary but not sufficient. In practice though, when performance is measured with top-1 accuracy, the length of the task vectors is rather unimportant in the high-dimensional regime. The angular information is far more relevant. (One can intuitively see this by considering the toy setting of Gaussian random vectors in the high-dimensional limit – their lengths concentrate and only angles matter.) This is especially true since the embeddings are normalized to be unit norm before evaluating the analogy task. We discuss this subtlety in Appendices A.4 and A.6. We will clarify this in the main text.\n\nLet us know if you have further questions!"}, {"review_id": "7UC9LZ3xxO", "reviewer": "Reviewer_sHUx", "summary": "This paper formulates a quadratic model that closely approximates what is learned from Word2Vec using a Taylor approximation of the loss.  \nThey show the resulting embedding learned on Wikipedia closely matches what Word2Vec does, and works almost as well.", "questions": "If I wanted to apply this quadratic approximation to other semi-supervised learning tasks (for instance is other domains like genomics or finance), how would I do that?  What could I learn from this to apply to other areas where semi-supervised methods are still being developed (and transformer based methods may be overkill).", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The paper is (for the most part) very well written.  And I found it very easy to read and understand.  \n\n\nStrengths:  \n  - The word2vec model is an incredibly influential model, so it's a great case study.  \n\n  - The proposed model shows to work much better than the Levy+Goldberg 2015 linear approximation model\n\n  - The results on the one data set are quite convincing that it works well.  \n\n  - as mentioned, the paper is very clearly written.\n\n\nWeaknesses:\n  - The word2vec model is nice, but is now outdated, as it thoroughly replaced with contextual embeddings (e.g., BERT) and their derivatives.  The analysis also does not seem to extend to other contemporary models like GloVe or FastText.  \n\n  - The embedding seems to do slightly worse than word2vec in tests.  It's similar, but not better.  \n \n  - I do not see what the use of this analysis is.  It probably would have been very useful in 2016 or so, but I am less sure now.", "was_revised": true, "final_justification": "The paper provides a new mathematical/optimization perspective on how the classic word embeddings word2vec form.  This would have been very useful and important 10 years ago, when they were SoTA.  I am convinced it is related to modern LLMs via discussion -- however, this analysis is still not of those models, and do not yet apply there.  Its interesting, but not essential.  I'll be fine if accepted and I lean that way, but also fine if not accepted.", "rebuttal": "Thanks very much for your review, we appreciate it. We’re glad you found the result clear and convincing. Some comments regarding your concerns and questions:\n\n**What’s the point of this analysis? Isn’t word2vec outdated?** Though the algorithm itself has now been superseded by larger context-aware embeddings, the core mechanistic idea behind word2vec is back in vogue! Many recent papers empirically report linear structure in the latent representations in LLMs – a direct analog of the word2vec embedding geometry. (For example, see https://arxiv.org/abs/2305.16130, https://arxiv.org/pdf/2303.04245, and the references we mention in Section 4).\n\nTrying to understand this behavior directly in large language models is an admirable goal to strive towards. However, we believe it is too difficult to tackle directly (for now). Previous efforts in this direction have not gotten very far. A useful way for science to proceed in this situation is to focus on first understanding the mysterious behavior in the simplest possible system that exhibits it. Doing this is one of our primary contributions.\n\nIn general, one of the main outstanding questions of deep learning theory is: *what is the mechanism and character of representation learning?* Theorists have very few solvable cases, and these analyses tend to rely on unrealistic assumptions on either the data or the learning task (or both). We believe our result is powerful because it answers this question in a relevant natural language task under very few, relatively weak assumptions.\n\n(Also, it’s a bit surprising to lack a theory of learning dynamics for one of the foundational algorithms of the field! Even if it *weren’t* directly connected to understanding modern models, for posterity’s sake we should try to solve and deeply understand the small handful of truly breakthrough algorithms in the history of neural learning.)\n\n**What about GloVe and FastText?** We expect the main idea of our proofs to hold for both these models. FastText is essentially the same architecture and loss as word2vec, and its extra engineering tricks will not change the core story much. GloVe is very similar to word2vec; while word2vec uses a contrastive loss, GloVe instead performs *explicit* least-squares factorization. Incorporating something like our Setting 3.1 in the analysis of the GloVe algorithm would yield results exactly analogous to the ones presented here. We plan to include this discussion and some supporting experiments in the final version of the manuscript.\n\n**How can I apply these results to other self-supervised learning tasks?** Since our results are distribution agnostic, they carry over to pretty much any linearized contrastive algorithm. To be specific, if your contrastive loss has the form $\\\\mathcal{L}(W) = \\\\mathbb{E}\\_\\\\text{positive pairs}[f^{+}(w^T w’)] + \\\\mathbb{E}\\_\\\\text{negative pairs}[f^{-}(w^T w’)]$, then you can simply take the quadratic Taylor approximation of your $f^{+}(w^Tw’)$ and $f^{-}(w^T w’)$, complete the square, and obtain a target matrix in terms of the positive and negative distributions for your learning problem of interest. In this sense, the idea behind our Theorem 1 is actually very general. See Appendix D.2 for an example using the SimCLR loss. To get closed-form learning dynamics, you’d need to find an equivalent of our Setting 3.1. This may be difficult depending on the form of your input data; for instance, it may require a whitening data preprocessing step. The idea is to choose these hyperparameters in such a way that the weighted matrix factorization becomes unweighted. For more details, see Proposition 2 and the subsequent discussion.\n\nLet us know if you have further questions."}, {"review_id": "sdr76I4WHt", "reviewer": "Reviewer_TGpe", "summary": "This paper provides an analytical study of self-supervised word embedding learning by introducing Quadratic Word Embedding Models (QWEM), derived via a quartic Taylor approximation of the word2vec loss function. The authors analyze both the training dynamics and the final solution under gradient flow, revealing a rank-incremental learning process that matches empirical observations in standard word2vec training. They demonstrate that QWEM closely approximates word2vec in terms of learned features and benchmark performance, while being analytically tractable. The study also explores the emergence of linear representations in latent space and proposes a random matrix theory perspective on task vectors in analogy tasks.", "questions": "Is there any possible relations between your focus and some more modern methods?\n\nIs this possible to give a more explicit explanation of why gradient descent in Word2vec leads to a reasonable embedding(or similar to QWEM rather than directly approximating PMI)?\n\nIn figure 4, SNR and task performance don't really match, and some tasks suggest that bigger d , better performance. I think this part is not really convincing.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The story of theory part is clear, and authors are indeed trying to provide experimental evidence showing that theory matches practice. The idea of taylor approximating the loss function is simple but seemingly effective of understanding behaviors of word2vec. Introducing M, which is element-wise bounded and hence able to solve the inherent problem of PMI, shows a strong ability to capture the concrete matrix factorization which word2vec learns.\n\nHowever, word2vec is simple and far from SOTA models that could be very complex, leading to only limited insight into understanding the success of recent LLMs. As a result, the significance and utility of this paper could be somehow limited. Also, there is no theoretical guarantee that QWEM would always behave similarly to original word2vec. Though the paper claims to provide a closed-form solution of training dynamics, this solution is under a lot of assumptions unverified or only verified by single evidence. The similarity of behaviors could be dependent on datasets and models with different sizes and features, which is partially admitted in the paper.", "was_revised": true, "final_justification": "Based on the authors’ rebuttal and the strengths I mentioned in my initial review, I have decided to keep my score.", "rebuttal": "Thank you for your assessment and review, we appreciate it. We are glad that you found our results clear and our main idea effective. Some comments regarding your concerns and questions:\n\n**What’s the point of this analysis? Isn’t word2vec outdated?** Actually, the core mechanistic idea behind word2vec is back in vogue! Many recent papers empirically report linear structure in the latent representations in LLMs – a direct analog of the word2vec embedding geometry. (For example, see https://arxiv.org/abs/2305.16130, https://arxiv.org/pdf/2303.04245, and the references we mention in Section 4).\n\nTrying to understand this behavior directly in large language models is an admirable goal to strive towards. However, we believe it is too difficult to tackle directly (for now). Previous efforts in this direction have not gotten very far. A useful way for science to proceed in this situation is to focus on first understanding the mysterious behavior in the simplest possible system that exhibits it. Doing this is one of our primary contributions.\n\nIn general, one of the main outstanding questions of deep learning theory is: *what is the mechanism and character of representation learning?* Theorists have very few solvable cases, and these analyses tend to rely on unrealistic assumptions on either the data or the learning task (or both). We believe our result is powerful because it answers this question in a relevant natural language task under very few, relatively weak assumptions.\n\n(Also, it’s a bit surprising that we lack a theory of learning dynamics for one of the foundational algorithms of the field! Even if it *weren’t* directly connected to understanding modern models, for posterity’s sake we should try to solve and deeply understand the small handful of truly breakthrough algorithms in the history of neural learning.)\n\n**On our assumptions.** We respectfully disagree that our solution relies on a lot of unverified assumptions. On the contrary, our result relies on rather few assumptions, especially compared to what is par for the course in theory. We tried to emphasize this point throughout the manuscript. For example, our solution does not rely on any distributional assumptions on the data, which is uncommon in theoretical literature. We clearly list the set of approximations and assumptions in Figure 2 – there are only four in total. The first two are approximations to word2vec which are supported by our empirical checks (see Figure 1, Figure 3, the discussion in Appendices A.3, and Figure 5 in Appendix B). The other two are technical conveniences that simplify the analysis; we expect that they may be relaxed (or possibly eliminated completely) with some additional effort, as is suggested by Figure 2.\n\nIn particular, we do *not* believe that the similarity between QWEMs and word2vec is sensitive to the choice of dataset or the model size. (Please do point us to where in the paper this was unclear, as we’d like to clarify this in the writing.) We have done some empirical checks confirming this, as we mention in Appendix A; we will include more details in the final manuscript.\n\nAs a side note, since our results are data distribution agnostic, we expect them to carry over to a broad class of linearized contrastive algorithm. For example, if your contrastive loss has the form $\\\\mathcal{L}(W) = \\\\mathbb{E}\\_\\\\text{positive pairs}[f^{+}(w^T w’)] + \\\\mathbb{E}\\_\\\\text{negative pairs}[f^{-}(w^T w’)]$, then you can simply take the quadratic Taylor approximation of your $f^{+}(w^Tw’)$ and $f^{-}(w^T w’)$, complete the square, and obtain a target matrix in terms of the positive and negative distributions for your learning problem of interest. In this sense, the idea behind our Theorem 1 is actually very general. See Appendix D.2 for an example using the SimCLR loss.\n\n**Why does gradient descent in word2vec lead to embeddings similar to QWEM rather than directly approximating PMI?** This is a good question, and the answer is a key insight of our result. The crux of the issue is in adequately handling the rank constraint. (In general, rank-constrained optimization is NP-hard, so handling this properly is a subtle business.) Let us compare the two approaches.\n* The idea behind SVD-approximating PMI is to say “Let us first solve the unconstrained minimizer of the loss (i.e., the PMI matrix) and then, at the end, apply the rank constraint by choosing the closest (in Frobenius norm) rank-$d$ matrix.” This fails. The reason is because the loss basin is extremely wide and shallow, so the global unconstrained minimum is actually *very far* from where the model actually ends up via gradient descent in finite time.\n* Our insight is to change the approach: we say “Let us first approximate the loss landscape itself; then, we can account for the rank constraint *throughout the entire optimization process.*” With the quartic approximation of the loss, this turns out to be solvable. As a result, our prediction for what word2vec learns is significantly more accurate.\nWe discuss this at the end of Section 2. We will include more details to clarify this point, as it is central to our result.\n\n**Why does decreasing SNR sometimes still yield high performance?** Although the relationship between SNR and task performance is qualitative rather than quantitative, it’s still quite interesting that the relative ordering of analogy categories is preserved. In other words, SNR (which is computed from the corpus statistics and hyperparameters alone) is a useful measure of how difficult a particular class of analogies is to learn. To our knowledge, this is the first time that a measurable *a priori* proxy for analogy difficulty has been proposed.\n\nThis still leaves the question of why, within an analogy category, larger dimension yields better performance despite the degradation in SNR. Our leading hypothesis (corroborated by several empirical checks we did) is that this is an ancillary effect of the benchmark using top-1 accuracy as the error metric. As a concrete example, consider the analogy “France : Paris :: Japan : Tokyo”, satisfying japan+(paris-france) ≈ tokyo. As the effective embedding dimension increases, the SNR may decrease; at the same time, all the embeddings become more spaced out (due to the increased available volume in latent space), meaning there are fewer nearby “competitors” for tokyo. In the large $d$ regime, we empirically found that the embeddings are sufficiently spaced, and tokyo is still the nearest embedding to japan+(paris-france) despite an increase in absolute error. Thus, the measured top-1 performance does not degrade with the increased (relative) noise. Note that unintuitive effects associated with using top-k accuracy have been observed in LLMs as well (see https://arxiv.org/abs/2304.15004). We plan to include this discussion in the final version of the manuscript.\n\nThanks for the thoughtful comments on our manuscript. Let us know if you have further questions."}]

Author rebuttal:

Meta-review: This paper analyzes word2vec by deriving a quartic approximation, which exhibits training dynamics and downstream performance closely aligned with the original model.

The reviewers recognized the solid technical contributions and found the analysis interesting and valuable, as reflected in their positive evaluations. No major concerns were raised, and the author rebuttal addressed the reviewers’ questions satisfactorily.

As noted by the reviewers, word2vec has become a classic approach for obtaining word embeddings. While the work may not appear super exciting, it nonetheless provides a rigorous analysis. Given its solid contribution to word embedding, this paper represents an important and worthwhile addition to NeurIPS.
Decision: Accept (poster)