Flax (JAX): GloVe Embeddings for Text Classification Tasks¶

Word embeddings is one the most preferred text encoding approach when we are trying to solve any NLP tasks using deep learning algorithms. The other encoding approaches like word frequency, one-hot, Tf-Idf, etc assign just a single scalar value to each word which has very low representation power. It can not capture the different meanings of words in different contexts as well as relationships with other words. The word embeddings on the other hand use a real-valued vector to represent each word. The size of this vector varies a lot. It can be vector of 10,15,20,50,10,200,300,etc scalars. This gives more representation power to the approach which deep learning algorithms can utilize to produce better results. We can generate embeddings for words of our datasets by training them in the network or we can use readily trained embeddings from some other network for our purpose if our dataset is small.

GloVe (Global Vectors) is an unsupervised algorithm that can generate embeddings for words. Stanford professors have already generated GloVe word embeddings for many words by training algorithms on a large corpus of data. They have generated different embeddings of different lengths from twitter and Wikipedia datasets. Glove embeddings are a very good option if your dataset is very small and you can not generate word embeddings on it by training network. Please feel free to check the below link if you are looking for detailed information on GloVe.

• GloVe: Global Vectors for Word Representation

As a part of this tutorial, we have explained how we can use GloVe word embeddings with Flax Networks for text classification tasks. Flax is a Python deep learning library designed on top of a low-level Python deep learning library JAX. We have tried different approaches at using GloVe embeddings in the tutorial and compared their results. After training networks, we have also evaluated their performance by calculating various ML metrics. We have even explained predictions made by the network using LIME (Local Interpretable Model-Agnostic Explanations) algorithm.

Please check the below link if you are looking to train your own word embeddings using Flax (JAX) Networks.

• Guide to Use Word Embeddings For Flax (JAX) Text Classification Networks

Below, we have listed important sections of Tutorial to give an overview of the material covered.

Important Sections Of Tutorial¶

1. Prepare Data
   • 1.1 Download And Unzip GloVe Embeddings (840B.300d)
   • 1.2 Load GloVe Word Embeddings in Memory
   • 1.3 Load Dataset
   • 1.4 Tokenize Examples, Populate Vocabulary And Vecorize Text Examples
   • 1.5 Create Matrix Of GloVe Embeddings for Vocabulary Tokens
2. Approach 1: GloVe Embeddings Flattened (Max Tokens=50, Embedding Length=300)
   • Define Network
   • Define Loss Function
   • Train Network
   • Evaluate Network Performance
   • Explain Network Predictions using LIME Algorithm
3. Approach 2: GloVe Embeddings Averaged (Max Tokens=50, Embedding Length=300)
4. Approach 2: GloVe Embeddings Summed (Max Tokens=50, Embedding Length=300)
5. Results Summary And Further Suggestions

Below, we have imported the necessary Python libraries that we have used in our tutorial and printed their versions for reference purposes.

1. Prepare Data ¶

In this section, we'll prepare data for the neural network. As we are going to word embeddings approach for encoding text data and we have decided to use GloVe word embeddings for our purpose, we'll be following the below steps for encoding text data.

1. Loop through all text examples, tokenize them and build a vocabulary of unique tokens (words). A vocabulary is a simple mapping from tokens to their integer indexes. Each token (word) is assigned a unique index starting from 0.
2. Loop through each text example, tokenize them and retrieve indexes of those tokens using vocabulary created in the previous step. After performing this step, we'll have a list of token indexes for each text example. The output of this step will be given to the neural network for training.
3. Next, we need to loop through tokens of our vocabulary and retrieve their GloVe embeddings. We need to stack glove embeddings for tokens of vocabulary in a matrix. We'll refer to this matrix as the embeddings matrix.
4. Now, retrieve GloVe embeddings for each token index created in the 2nd step by indexing the embeddings matrix created in the 3rd step.

So basically, we first convert our text examples to a list of token indexes and then retrieve GloVe embeddings for them by indexing the embedding matrix.

The first 3 steps mentioned above will be performed in this section. The 4th step will be implemented at the embedding layer in the neural network which will map token indexes to their embeddings. We'll set the embedding matrix as the weight matrix of the embedding layer.

Below, we have included an image giving an idea about word embeddings.

1.1 Download And Unzip GloVe Embeddings (840B.300d)¶

In this section, we have simply downloaded GloVe embeddings from the Stanford website and unzipped them. We have downloaded GloVe 840B.300d embeddings that have word embeddings of length 300 for 2.2 Million tokens.

1.2 Load GloVe Word Embeddings in Memory¶

In this section, we have simply loaded word embeddings in memory. We have created a Python dictionary whose keys are tokens and values are numpy arrays representing embeddings of those tokens.

1.3 Load Dataset¶

In this section, we have loaded the dataset that we are going to use for our text classification task. We'll be using newsgroups dataset available from scikit-learn. The dataset has text documents for 20 different news categories, though we'll be using 5 categories for our purpose. The dataset is already divided into train and test sets for our convenience.

1.4 Tokenize Examples, Populate Vocabulary, And Vecorize Text Examples¶

In this section, we have completed the first two steps of our embedding process that we had explained earlier. We have populated vocabulary and transformed text examples into a list of token indexes using this vocabulary here.

First, we have created an instance of Tokenizer() available from the Python deep learning library Keras. We have called fit_on_texts() method on tokenizer with train and test examples. This will populate the vocabulary of all unique tokens and keep it in the tokenizer object. The vocabulary can be accessed by calling word_index property of the tokenizer object. We have also later printed the size of the vocabulary to show the number of tokens.

Next, we have called texts_to_sequences() method with train and text examples separately. This will tokenize each text example into a list of tokens (words) and retrieve their indexes from the vocabulary. The output of this method will be a list of token indexes. The different text examples can have a different number of tokens but for the neural network, we need fixed input sizes. For this, we have decided to keep maximum tokens per text example at 50. To enforce this, we have called pad_sequences() method on the output of texts_to_sequences() method. This method ensures that all example has exactly 50 tokens. The examples that have more than 50 tokens will be truncated at 50 tokens and examples that have less than 50 tokens will be padded with 0s.

We have also printed the first few train examples to show how text gets encoded to a list of token indexes.

Below, we have explained with a simple example how text example is vectorized.

text = "Hello, How are you? Where are you planning to go?"

tokens = ['hello', ',', 'how', 'are', 'you', '?', 'where',
            'are', 'you', 'planning', 'to', 'go', '?']

vocab = {
    'hello': 0,
    'bye': 1,
    'how': 2,
    'the': 3,
    'welcome': 4,
    'are': 5,
    'you': 6,
    'to': 7,
    '<unk>': 8,
}

vector = [0,8,2,4,6,8,8,5,6,8,7,8,8]

1.5 Create Matrix Of GloVe Embeddings for Vocabulary Tokens¶

In this section, we have created an embedding matrix using Glove embeddings that we had described in 3rd step of our encoding process earlier.

Here, we are looping through tokens of our vocabulary and retrieving GloVe embedding for each token from the glove embedding dictionary that we had loaded earlier. We have stacked embeddings for each token to create an embedding matrix. As we know that the embedding length for GloVe 840B.300d is 300, our embedding matrix shape will be (vocab_len, 300).

We can retrieve embeddings of tokens by indexing this matrix using a token index. To explain it with an example, let's say that the index of 'how' token in our vocabulary is '10' then we can simply index embedding matrix like 'word_embeddings[10]' to retrieve the embedding of the token 'how'.

Later on, we'll set this matrix as the weight matrix of the embedding layer of the network which will take token indexes as input and index this matrix to retrieve embeddings for tokens.

Approach 1: GloVe Embeddings Flattened (Max Tokens=50, Embedding Length=300) ¶

Our approach in this section is based on the concept of flattening embeddings. The flattened embeddings are then given to dense layers for processing. After training the network in this section, we have also evaluated various ML metrics for checking the performance of the network.

Define Network¶

Below, we have defined the network that we are going to use for our text classification task. The network consists of one embedding layer and two dense layers.

The first layer of our network is the embedding layer. We have created this layer using Embed() constructor from linen module of Flax. We have provided it with vocabulary length as a number of tokens. The embedding length of 300 is set as the value of features parameter. We have initialized the weight of this layer by providing callable to embedding_init parameter. The callable simply returns the embedding matrix that we had created in the previous section. This layer is responsible for mapping token indexes to their embeddings by indexing this weight matrix. The shape of input data to this layer is (batch_size, max_tokens) = (batch_size, 50) and output shape is (batch_size, max_tokens, embed_len) = (batch_size, 50, 300).

Please make a NOTE that as we are using trained embeddings (GloVe embeddings) already, we need to prevent updates to embeddings during the training process. We have done that later by making a minor modification to the training routine that ignores updates to the embedding layer.

The output from embedding layer is flattened so that shape gets transformed from (batch_size, 50, 300) to (batch_size, 50 x 300) = (batch_size, 15000).

This flattened output is then given to our first dense layer that has 100 output units. This transforms shape from (batch_size, 15000) to (batch_size, 100). We have also applied relu activation to the output later.

The relu-activated output of the first dense layer is given to the second dense layer that has 5 output units (same as a number of target classes) for processing. This transforms shape from (batch_size, 100) to (batch_size, 5). The output of the second dense layer is returned as a prediction of our network.

After defining the network, we initialized it and printed the shape of weights and biases of layers. We have also printed a few weights of the embedding layer for verification purposes that it was initialized properly. After initializing the network, we have also performed a forward pass through the network using a few train examples for verification purposes.

Please make a NOTE that we have not described the process of network creation in detail as we expect that reader has little background on Flax library. Please feel free to go through the below link if you are new to the library as it'll help you get started with neural network creation.

• Flax: Framework to Create Neural Network using JAX

Define Loss Function¶

In this section, we have created a loss function that we'll use for our classification task. We'll be using cross entropy loss. The function takes network parameters, data, and target labels as input. It then performs a forward pass through the network to make predictions and one-hot encode target labels. Then, we calculate loss using softmax_cross_entropy() function available from Optax.

Train Network¶

In this section, we are training the network that we defined earlier. We have defined a function that will help us with the training process. The function takes train data (X, Y), validation data (X_val, Y_val), number of epochs, network parameters, optimizer state, and batch size as input. It then executes a training loop number of epochs time. For each epoch, it loops through training data in batches. For each batch of data, it performs a forward pass to make predictions, calculates loss, calculates gradients, and updates network parameters. It records loss for each batch and prints average training data loss by averaging the loss of all batches at the end of each epoch. We are also calculating the validation accuracy of the model and printing it for checking network performance on validation data.

Please make a NOTE that we are not updating the weights of the embedding layer. We have included a special line in the training routine that applies tree_map() function to updates which zeros updates to be made to the embedding layer hence preventing any updates to it.

Below, we are actually training our network using the training routine defined in the previous cell. We have first initialized bath size to 64, a number of epochs to 8, and learning rate to 0.001. Then, we have initialized our text classifier and Adam optimizer. At last, we have called our training routine with the necessary parameters to perform training. We can notice from the loss and accuracy values getting printed that our network is doing a good job at the classification task.

After training, we have verified the final weights of the embedding layer to confirm that it's not updated.

Evaluate Network Performance¶

In this section, we have evaluated the performance of our trained network by calculating accuracy score, classification report (precision, recall, and f1-score per target class) and confusion matrix metrics on test predictions. We can notice from the accuracy score that it is above average. We have calculated these ML metrics using scikit-learn.

If you want to learn about various ML metrics available from sklearn then we recommend that you go through the below link. It covers the majority of them in detail.

• Scikit-Learn: Model Evaluation and Scoring Metrics

Apart from calculations, we have also plotted confusion matrix metric using Python library scikit-plot. We can notice from the visualization that our network is doing a good job at classifying text documents from categories 'comp.graphics', 'rec.autos' and 'rec.sport.hockey' compared to categories 'alt.atheism' and 'talk.politics.mideast'.

The scikit-plot provides visualizations for many ML metrics. Please feel free to check the below link in your spare time if you want to learn about it.

• Scikit-Plot: Visualizing Machine Learning Algorithm Results & Performance Metrics

Explain Network Predictions using LIME Algorithm¶

In this section, we are trying to explain predictions made by our network using LIME (Local Interpretable Model-Agnostic Explanations) algorithm. The LIME is a very commonly used algorithm when trying to explain predictions made by any black-box ML models. We have used the Python library lime which has an implementation of the algorithm.

If you are someone who is new to the concept of LIME then we suggest that you go through the below links. It will give you in-depth knowledge of it as well as get you started with using it.

• How to Use LIME to Understand sklearn Models Predictions?
• LIME: Interpret Predictions Of Keras Text Classification Networks

In order to explain prediction using LIME, we need to create an instance of LimeTextExplainer which we have done below.

Below, we have first created a prediction function that takes a batch of text examples and returns their predicted probabilities by the network. Please make a NOTE that network outputs are not probabilities and we have converted them to probabilities by applying softmax activation to it.

Then, we randomly selected a text example from test data and made predictions it using our trained model. The model correctly predicts target label as 'rec.sport.hockey'. Next, we'll create an explanation for this prediction.

Below, we have first created an Explanation object by calling explain_instance() method on LimeTextExplainer instance. We have provided the method with a selected text example, prediction function, and target label. The explanation object has details about words contributing to prediction. Next, we have visualized the explanation by calling show_in_notebook() method on it.

The visualization highlights that words like 'leafs', 'games', 'goals', 'goalie', etc contributed to predicting target label as 'rec.sport.hockey'. This makes sense as these are commonly used words in the hockey world.

Approach 2: GloVe Embeddings Averaged (Max Tokens=50, Embedding Length=300) ¶

Our approach in this section has a minor change in the way we handle word embeddings. In this section, we have taken the average of word embeddings across tokens of each example. The majority of the code is exactly the same as earlier with a minor change in the network architecture.

Define Network¶

Below, we have defined the network that we'll use for our text classification task. The network has 3 layers like earlier. The only difference is in forward pass logic. Earlier, we had flattened embeddings whereas here, we have taken the average of embeddings across tokens of text examples using mean() function. This will transform data shape from (batch_size, 50, 300) to (batch_size, 300). This output will then be given to a dense layer for processing.

After defining the network, we initialized it and performed a forward pass using a few train examples to make predictions for verification purposes.

Train Network¶

In this section, we have trained the network using the same settings that we had used in our previous section. We can notice from the loss and accuracy values getting printed after each epoch that our network is doing a good job at the classification task.

Evaluate Network Performance¶

In this section, we have evaluated the performance of our trained network by calculating accuracy score, classification report and confusion matrix metrics on test predictions. We can notice from the accuracy score that it is quite better compared to our previous approach. We have also plotted a confusion matrix for reference purposes which also hints at improvements in classifying text documents.

Explain Network Predictions using LIME Algorithm¶

In this section, we have explained predictions made by our trained network using LIME algorithm. The network correctly predicts the target label as 'rec.sport.hockey'. The visualization shows that words like 'leafs', 'games', 'goalie', 'goals', 'defense', etc are contributing to predicting target label as 'rec.sport.hockey'.

Approach 3: GloVe Embeddings Summed (Max Tokens=50, Embedding Length=300) ¶

Our approach in this section is the same as our previous approach with a minor change in that we are taking the sum of embeddings across tokens this time instead. The majority of the code is the same as earlier with only a change in network architecture.

Define Network¶

Below, we have defined the network architecture that we'll use for our text classification task. The network has the same layer as earlier. The only difference is in forward pass logic which now uses sum() function to take the sum of embeddings instead of the average.

After defining the network, we initialized it and performed a forward pass to make predictions for verification purposes.

Train Network¶

In this section, we have trained the network using exactly the same settings that we have used for all our approaches. We can notice from the loss and accuracy score that the network is doing a good job at the given task.

Evaluate Network Performance¶

In this section, we have evaluated the performance of our network by calculating ML metrics like accuracy score, classification report and confusion matrix on test predictions. We can notice from the accuracy score that it is a little better compared to the previous approach and the highest of all our approaches. We have even plotted a confusion matrix for reference purposes.

Explain Network Predictions using LIME Algorithm¶

In this section, we have explained predictions made by the network using LIME algorithm. The network correctly predicts the target category as 'rec.sport.hockey'. The visualization highlights that words like 'leafs', 'games', 'host', 'goalie', etc contributes to predicting target label as 'rec.sport.hockey'.

5. Results Summary And Further Suggestions ¶

Further Recommendations¶

• Try different token sizes.
• Try different GloVe embeddings like 42B, 6B, 27B, etc.
• Try different combinations of dense layers after the embedding layer.
• Try different ways to handle embedding other than flattening, averaging, and summing (max, min, std, etc.).
• Try different weight initialization methods.
• Try different optimizers.
• Train network for more epochs.
• Try learning rate schedulers

This ends our small tutorial explaining how we can use GloVe embeddings with Flax(JAX) networks for text classification tasks. Please feel free to let us know your views in the comments section.

References¶

• Text Classification using Flax (JAX) Networks
• Guide to Use Word Embeddings For Flax (JAX) Text Classification Networks
• Flax: CNNs with Conv1D for Text Classification Tasks
• LIME: Interpret Predictions Of Flax (JAX) Text Classification Networks
• Eli5.lime: Interpret Predictions Of Flax(JAX) Text Classification Networks
• Explain Flax (JAX) Text Classification Network Predictions using SHAP Values