Closed-Form Training Dynamics Reveal Learned Features and Linear Structure in Word2Vec-like Models
This work provides theoretical insights into self-supervised learning in language modeling, which is incremental but clarifies how semantic concepts emerge during training.
The authors tackled the problem of understanding the training dynamics and learned representations in word2vec-like models by deriving closed-form solutions for gradient flow and final embeddings, revealing that models learn orthogonal linear subspaces incrementally, each representing interpretable topics, and enabling analogy completion via vector addition.
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.