Pytorch ELMo, trained from scratch | Towards Data Science

Embeddings from Language Model (ELMo) is a powerful contextual embedding method that finds application in a wide range of Natural Language Processing tasks. ELMo, along with others, started the trend of pretraining contextual word embeddings in NLP. The technique remains simple and intuitive, allowing itself to be easily added to existing models.
In this article, we will discuss how to train ELMo embeddings from scratch with our own text corpus and explain how it works under the hood. We will use AllenNLP, a PyTorch-based NLP framework that provides many state-of-the-art models out of the box. If you are only interested in using pre-trained ELMo embeddings, feel free to hop to the last section — Using ELMo in downstream tasks.
Training ELMo from scratchUnderstanding the implementationUsing ELMo in downstream tasks
Training ELMo from scratch
Understanding the implementation
Using ELMo in downstream tasks
I know you can’t wait so let’s dive right into the training. In the next section, we will go through the training step-by-step to understand how ELMo is implemented in AllenNLP.
First, let’s install allennlp-models. This package includes all the fancy models implemented in the AllenNLP framework. Installing this package should also find you the correct version of PyTorch and AllenNLP needed.
pip install allennlp-models=v2.0.1
Next, we get our corpus data for training. For demonstration purposes, we use the original corpus — 1 Billion Word Language Model Benchmark. You can also use any corpus of your choice.
wget http://www.statmt.org/lm-benchmark/1-billion-word-language-modeling-benchmark-r13output.tar.gztar -xzf 1-billion-word-language-modeling-benchmark-r13output.tar.gzexport BIDIRECTIONAL_LM_DATA_PATH=$PWD'/1-billion-word-language-modeling-benchmark-r13output'export BIDIRECTIONAL_LM_TRAIN_PATH=$BIDIRECTIONAL_LM_DATA_PATH'/training-monolingual.tokenized.shuffled/*'
For indexing and performance purposes, AllenNLP requires preprocessed Vocabulary, which reads from a directory with tokens.txt inside, which includes all the unique tokens on each line. (The inner working of Vocabulary is actually a bit more complicated, but we assume this for simplicity.)
Here, we download the preprocessed vocabulary. If you are using your own corpus, take a look at tokens.txt to understand its format. Pay special attention to the first three tokens — </S>, <S>, @@UNKNOWN@@. They are critical for ELMo to know the boundaries of sentences and words outside of vocabulary.
pip install awscli mkdir vocabulary export BIDIRECTIONAL_LM_VOCAB_PATH=$PWD'/vocabulary' cd $BIDIRECTIONAL_LM_VOCAB_PATH aws --no-sign-request s3 cp s3://allennlp/models/elmo/vocab-2016-09-10.txt . cat vocab-2016-09-10.txt | sed 's/<UNK>/@@UNKNOWN@@/' > tokens.txtrm vocab-2016-09-10.txt echo '*labels\n*tags' > non_padded_namespaces.txt
We have arrived at our last step of training ELMo from scratch. In this step, we use a precomposed configuration file to kick off the training.
wget https://raw.githubusercontent.com/allenai/allennlp-models/main/training_config/lm/bidirectional_language_model.jsonnetallennlp train bidirectional_language_model.jsonnet --serialization-dir output_dir/
The above steps borrowed heavily from the official AllenNLP HowTo. It might seem quite unclear if you want to tweak the underlying architecture or if you want to process your own corpus into the desired format. While offering stunning simplicity, AllenNLP sweeps a huge amount of underlying assumptions about data and architectures under its rugs.
While you are waiting for the training to complete, we will take a look at how ELMo works from the perspective of the AllenNLP framework.
To understand the ELMo architecture and how it is implemented, we take a closer look at bidirectional_language_model.jsonnet, which is essentially a configuration file that specifies everything at a very high level. It allows us to easily define how data gets processed (tokenized and embedded) and what the architecture looks like.
Note: within the JSON objects, we see many with the key "type”. This is AllenNLP’s syntax to register custom classes. By using its register decorator, we are exposing our classes to AllenNLP, enabling us to use them in our configurations.
connecting natural language to natural numbers
The first thing we are concerned with is how data flows into the model from natural language texts that sit on our disk. The data is first tokenized by a Tokenizer, splitting the sentence at white spaces. Then, individual tokens are converted into numbers by an TokenIndexer.
In many cases, it is enough to convert the words into their respective ids in a vocabulary. Since we are using a character-level CNN to capture the subword information, we need to additionally convert the words by each character.
The boundary between the model and the dataset reader lies in whether inner states are learnable. To be more concrete, the conversions from tokens to index numbers is a fixed procedure, often relying on a lookup table, therefore it belongs to the dataset reader.
connecting indices to uncontextualized embeddings
However, converting these index numbers into an uncontextualized embedding is a learnable procedure, regardless of the technique (e.g. embedding table, character convolutions, etc.)
Notice how the token_embedders share the same keys with the token_indexers. This allows the model to connect with the dataset reader internally.
In ELMo, we are only using the character-level uncontextualized embeddings. Therefore, the indices of words are discarded by the empty embedder and the indices of characters are fed into the character_encoding embedder.
Let’s take a closer look at the character embedder. Internally, it converts the character indices into character embeddings with a learning embedding matrix. Then, for each word, it uses a CNN to learn its word embedding by convolving over the character embeddings of respective words. Doing so gives us an uncontextualized word embedding that includes subword information.
Uncontextualized word embeddings have been around for quite a while. Embeddings like word2vec, GloVe have long been the standard features in NLP. A long-standing problem with such embeddings is the loss of context. With uncontextualized embeddings, the embeddings of “bank” are the same in “river bank” and “bank account”. However, we know these words have drastically different meanings. Next, we take a look at how ELMo handles this issue with its contextualizer.
nudging a word embedding by looking at the sentence
In the original paper, Peters et al. use a biLSTM as the contextualizer. The forward LSTM captures the context from words that come before; the backward LSTM captures the context from words that come after.
For example, to contextualize an embedding xt, the forward LSTM looks at x1, ..., xt−1 and creates a hidden state; the backward LSTM looks at xt+1, ..., xn and creates another hidden state. By concatenating two hidden states, we obtain a contextualized embedding that considers both the forward context and the backward context.
In a later study, Peters et al. found that while LSTM enjoys a higher accuracy in downstream tasks, transformer enjoys a much better inference speed with a little drop in performance. As a result, the implementation uses a transformer-based contextualizer. Intuitively, it achieves the same goal the LSTM tries to achieve — adding sentence context into each word embeddings.
This completes the ELMo architecture, i.e. we use the resulting contextual embeddings in our downstream tasks. One last component we have to discuss is its training objective — the bidirectional language model loss.
using previous words to predict the ____
Given a sequence of N tokens, (t1, t2, ..., tn), a forward language model computes the probability of the sequence by modeling the probability of token tk given the history (t1, ..., tk−1)
Similarly, we can derive the backward language model. Our final bidirectional language model objective is simply the sum of two log probabilities over all sentences in our corpus.
In AllenNLP, the loss is baked into the language_model model so we only have to set the parameter bidirectional to true in our configuration. We can also use sampled softmax loss to speed things up a bit by setting num_samples.
Whew! That concludes the ELMo architecture and its implementation in AllenNLP. While you are reading, your model might have finished training (ok, probably not). In the next section, we will use the trained ELMo to embed some texts!
Now that you have finished training, you should be able to see the metrics for each epoch in output_dir/, where you can also find the trained model weights model.tar.gz. Alternatively, you can grab one of the pre-trained weights and integrate it into your model.
To integrate ELMo into your AllenNLP model configuration, we simply use it as another token_embedder:
Don’t want to use a configuration file? You can also call ELMo directly in Python:
You’ve made it to the end! While the syntax of AllenNLP might still be confusing to you, we do see its power in training a PyTorch-based ELMo from scratch with 0 lines of code! It also includes many amazing NLP models out of the box. If you are interested in using them, take a look at its official guide and documentation.
Deep contextualized word representations first purposed the ELMo architecture. It obtained the state-of-the-art of many NLP tasks at its time while the architecture remained easily integrable into other models.Dissecting Contextual Word Embeddings: Architecture and Representation further studied the ELMo architecture. It tried to address how and why pre-trained contextual embeddings are so effective. In this study, a tradeoff between speed and accuracy was discovered by experiments with different contextualizers (LSTM, transformer, gated CNN). A relationship between network depth and captured information was reported, with lower layers capturing syntactic features and upper layers capturing semantic features.Long Short-Term Memory was first introduced to capture long-term dependencies in sequences and to address the vanishing gradient problem in vanilla RNNs. ELMo uses the way LSTM generates hidden states to contextualize word embeddings. For a more involved explanation of LSTM, check out colah’s blog.Attention is all you need introduced transformers into the world. It uses self-attention to capture the dependencies among tokens, making it massively parallelizable. HavardNLP dives deeply into its implementation in the annotated transformer, while Jay Alammar provides a nice intuition in his illustrated transformer.AllenNLP: A Deep Semantic Natural Language Processing Platform brings us the AllenNLP framework. It provides many common abstractions working with text data and intelligently handles things like batching and padding. I personally find it extremely useful, with its vast collection of NLP models. For more information, check out its official guide, documentation, and GitHub repo for open source models.
Deep contextualized word representations first purposed the ELMo architecture. It obtained the state-of-the-art of many NLP tasks at its time while the architecture remained easily integrable into other models.
Dissecting Contextual Word Embeddings: Architecture and Representation further studied the ELMo architecture. It tried to address how and why pre-trained contextual embeddings are so effective. In this study, a tradeoff between speed and accuracy was discovered by experiments with different contextualizers (LSTM, transformer, gated CNN). A relationship between network depth and captured information was reported, with lower layers capturing syntactic features and upper layers capturing semantic features.
Long Short-Term Memory was first introduced to capture long-term dependencies in sequences and to address the vanishing gradient problem in vanilla RNNs. ELMo uses the way LSTM generates hidden states to contextualize word embeddings. For a more involved explanation of LSTM, check out colah’s blog.
Attention is all you need introduced transformers into the world. It uses self-attention to capture the dependencies among tokens, making it massively parallelizable. HavardNLP dives deeply into its implementation in the annotated transformer, while Jay Alammar provides a nice intuition in his illustrated transformer.
AllenNLP: A Deep Semantic Natural Language Processing Platform brings us the AllenNLP framework. It provides many common abstractions working with text data and intelligently handles things like batching and padding. I personally find it extremely useful, with its vast collection of NLP models. For more information, check out its official guide, documentation, and GitHub repo for open source models.