Word Embeddings in Language Models

Natural language processing (NLP) has long been a fundamental area in computer science. However, its trajectory changed dramatically with the introduction of word embeddings. Before embeddings, NLP relied primarily on rule-based approaches that treated words as discrete tokens. With word embeddings, computers gained the ability to understand language through vector space representations. In this article, you will learn about:

• How word embeddings convert words into dense vectors
• How to use pretrained word embeddings
• How to train your own word embeddings
• How word embeddings are used in modern language models

Let’s get started!

Overview

This post is divided into three parts; they are:

• Understanding Word Embeddings
• Using Pretrained Word Embeddings
• Training Word2Vec with Gensim
• Training Word2Vec with PyTorch
• Embeddings in Transformer Models

Understanding Word Embeddings

Word embeddings represent words as dense vectors in a continuous space, where semantically similar words are positioned close to each other. The core principle is that words appearing in similar contexts should have similar vector representations. This concept gained prominence through models like Word2Vec, GloVe, FastText, and ELMo.

Word embedding models are typically trained using unsupervised learning because the ideal vector representation for each word is unknown (otherwise, we could use it directly). The objective is to learn word co-occurrence patterns from the training corpus.

Word2Vec, introduced by the paper “Efficient Estimation of Word Representations in Vector Space”, pioneered this approach. It uses a neural network to predict words based on local context and comes in two variants:

1. Continuous Bag of Words (CBOW): Predicts the target word given its context
2. Skip-gram: Predicts the context words given the target word

Skip-gram generally performs better for smaller datasets and rare words, while CBOW is faster and more effective for larger datasets. Word2Vec demonstrated that computers could understand semantic relationships between words by showing that embedding vectors could satisfy equations like “king – man + woman ≈ queen”.

GloVe (Global Vectors for Word Representation) takes a different approach. Rather than using a neural network, it constructs and factorizes a word co-occurrence matrix to obtain embeddings. GloVe combines the strengths of:

• Global matrix factorization methods (like latent semantic analysis)
• Local context window methods (like Word2Vec)

The resulting embeddings capture both semantic and syntactic relationships between words and often outperform Word2Vec on tasks requiring broader semantic understanding.

FastText improved upon Word2Vec by learning vectors for character n-grams instead of whole words. This approach captures subword information, solving the out-of-vocabulary problem and providing better performance for morphologically rich languages.

ELMo, a more recent model, uses a deep bi-directional LSTM to generate context-dependent word vectors. Unlike previous models, ELMo’s word vectors are not fixed but vary based on context. While less commonly used today after the emergence of large language models, ELMo’s core idea that word meaning should depend on context forms the foundation of all modern language models.

Using Pretrained Word Embeddings

You can easily use pretrained word embeddings from popular libraries. Here’s an example using the gensim library with GloVe embeddings:

To run this code, you need to download the GloVe embeddings from https://nlp.stanford.edu/projects/glove/ and extract the file glove.6B.50d.txt from the zip file glove.6B.zip. This file contains the trained vectors of 400,000 words from a training corpus of 6 billion words.

When you run this code, you will see the following output:

The first output shows that “king” is most similar to “prince” under this embedding model. And the second output shows that “queen” is the closest word to “king + woman – man”.

Training Word2Vec with Gensim

Gensim provides a simple interface to train your own Word2Vec model. Here’s how:

Running this code will not give you any good model. For a useful embedding, you need a large corpus to train on. You probably do not want to expand the Python list sentences, but rewrite the code to read from some files on the disk.

Assume you did that, gensim will train a Word2Vec model and save it to the file word2vec.model. Once you have trained it, you can load it back and use it to get the vector for a word, as shown in the code above.

Training Word2Vec with PyTorch

You can also implement Word2Vec from scratch using PyTorch. Here’s a basic implementation:

This code will train a “skip-gram” model of Word2Vec. In this model, the training data is a window of words from the text corpus. You should do some preprocessing to make the set of vocabulary clean, for example, removal of punctuations and convert all words to lowercase. Pay attention to how the variables context and target are used. In a window, such as “the quick brown fox jumps” in the example above, the model will be fed with the center word and asked to predict any other word in the same window. The loss function for the training is the cross-entropy loss.

This example probably does not give you a good model because you need a larger corpus and more epochs to train. However, notice that the model has an embedding layer and a linear layer. The embedding layer created with nn.Embedding will be the word embedding matrix you are interested in.

Also, notice that the embedding layer is just a numerical matrix. You need a lookup table, such as vocab_to_idx in the code above, to convert the word to an index and then use the index to get the embedding vector. The lookup table should be saved together with the model because you will not be able to use it if you cannot convert words to the correct index.

Embeddings in Transformer Models

From the above example, you learned that word embeddings can be trained and you can create a nn.Embedding layer for that purpose. In fact, most modern language models use this method. Let’s look at the BERT model as an example.

When you run this, you will see:

The BERT model is sophisticated and has a lot of components. The word embedding layer is named word_embeddings. Once you have created the model, you can refer to it with model.embeddings.word_embeddings. From its parameters, you can see that it has 30522 vocabularies, and each vector has dimension 768. The second print statement will dump the embedding matrix. You should expect the matrix in the shape of (30522, 768).

In a previous post, you learned that language models need a tokenizer to split the input text into tokens. Tokenizers also assign a token ID to each token. This token ID is the row index of the embedding matrix. When you feed an input text to this model, you should feed a sequence of token IDs instead. Usually, the embedding layer is the first layer of the model. It will convert the sequence of token IDs to a sequence of embedding vectors by replacing each token ID with the corresponding rows in the embedding matrix.

Further Readings

Below are some further readings on the topic:

Summary

In this article, you learned about word embeddings and their applications. In particular, you learned that:

• Word embeddings represent words as dense vectors in a continuous space, with semantically similar words positioned nearby
• Pretrained word embeddings are readily available through popular libraries
• You can train custom word embeddings using either Gensim or PyTorch
• Modern transformer models utilize learned embeddings through the nn.Embedding layer
• Embeddings are essential for capturing semantic relationships between words