Solvable Dynamics of Contrastive Word Embeddings and the Formation of Linear Representations
Dhruva Karkada, James Simon, Yasaman Bahri, Michael Deweese
Self-supervised word embedding algorithms provide a relevant and minimal setting for studying representation learning in natural language modeling. We examine the quartic approximation of word2vec and empirically show that it closely matches both the training dynamics of word2vec and its performance on downstream tasks.Our main contribution is to analytically solve for both the training dynamics and the final word embeddings in terms of only the corpus statistics and 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 concepts emerge during training; these can be used to complete analogies via vector addition.