Haiku (JAX): Word Embeddings for Text Classification Tasks¶

Nowadays, many of the NLP tasks (like text classification, text generation, sentiment analysis, etc) are getting solved by developing deep neural networks. Deep neural networks are complicated mathematical functions internally that work on real-valued data to solve given tasks. In order to make deep neural networks work for text data, we need some way to represent text data as real-valued data. There are many different approaches to encoding text data before giving it to neural networks. Famous approaches are word-frequency, one-hot, Tf-IDf, embeddings, etc. Approaches like word-frequency, one-hot, and tf-idf use only one real value to represent a single text token (word/character). This has limitations as the same word can be used in a different context and a single value can not capture all meanings. To solve it, embeddings were introduced which uses a real-valued vector to represent a word/character. As more values are getting used to represent a single word/character, it gives more representation power to it.

As a part of this tutorial, we have explained how to create neural networks using Python deep learning library Haiku that uses word embeddings approach to solving text classification tasks. Haiku is a high-level deep learning library designed on top of low-level library JAX. We have explained different ways of handling embeddings through our network to get better results. Each approach is evaluated by calculating various ML metrics. The predictions are analyzed further using LIME algorithm.

Below, we have listed essential tutorial sections to give an overview of the material covered.

Important Sections Of Tutorial¶

1. Prepare Data
   • 1.1 Load Data
   • 1.2 Populate Vocabulary and Vectorize Data
2. Approach 1: Flattened Embeddings
   • 2.1 Define Network
   • 2.2 Define Loss Function
   • 2.3 Train Network
   • 2.4 Evaluate Network Performance
   • 2.5 Explain Network Predictions using LIME Algorithm
3. Approach 2: Averaged Embeddings
4. Approach 3: Summed Embeddings
5. Summary of Results
6. Further Suggestions

Below, we have imported the necessary Python libraries and printed the versions we used in our tutorial.

Haiku Installation¶

• !pip install -U dm-haiku

1. Prepare Data ¶

In this section, we are preparing data for our neural network. We'll be following the below steps to prepare data for our network.

1. Load Data.
2. Tokenize each text example and Populate the vocabulary with unique words. A vocabulary is a simple mapping from a word to an integer index. Each word is assigned a unique index starting from 0.
3. Vectorize text examples - Tokenize text examples and retrieve the index for each word from the vocabulary.

After performing the above 3 steps, we'll finally have an array of integer indexes where each index will be mapped to some word as per vocabulary.

The embedding will be implemented in a neural network as an embedding layer. This layer will have embeddings (a real-valued vector) for each word as a weight matrix. We'll retrieve embeddings for words by indexing this weight matrix using integer indexes that we retrieved from the vocabulary for each word. Don't worry if steps are not 100% clear to you as they'll become pretty clear once we implement them below.

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

1.1 Load Data¶

Below, we have loaded AG NEWS dataset available from torchtext python library. The dataset has news articles from 4 different news categories (["World", "Sports", "Business", "Sci/Tech"]). The dataset is already divided into train and test sets.

1.2 Populate Vocabulary and Vectorize Data¶

Here, we are performing steps 2 (populate vocabulary) and 3 (vectorize data) that we had discussed earlier.

First, we have created an instance of Tokenizer which is available from the Python deep learning library Keras. It'll help us tokenize data, populate vocabulary and vectorize data as well.

After initializing the tokenizer object, we have populated vocabulary by calling fit_on_texts() method on it with train and test text examples. This method will loop through all text examples one by one, tokenize them, and populate vocabulary with unique words. Tokenization is a process where we break text into a list of tokens (words). We have also printed the size of vocabulary which is available through index_word and word_index attributes of the tokenizer object. We have also printed the first few mappings from the dictionary to show its contents.

As a next step, we have vectorized data where we tokenize text examples and retrieve their index from the vocabulary. We have performed this step by calling texts_to_sequences() method on the tokenizer object with train and test text examples one by one. The output of this step will be a list of integers for each text example where the integer will represent an index of the word as per vocabulary.

The size of the text example is different for each text example. But we need to prepare a dataset that has the same size for each example for the network. In order to do this, we have decided that we'll keep only the first 50 words from the text and will keep indexes of only them. We have completed this step using pad_sequences() method. This method makes sure that each example is of length 50. The examples that have more than 50 indexes will be truncated at 50 and examples that have less than 50 tokens will be padded with 0s.

After vectorizing datasets, we have converted them to JAX arrays as required by Haiku networks.

2. Approach 1: Flattened Embeddings ¶

In this section, we have tried our first approach at using word embeddings. As we had discussed earlier, the embeddings for words will be present in the first layer of the network named Embedding layer. The embeddings for words will be retrieved by indexing them. The approach in this section flattens embeddings of single text example by stacking embeddings of words next to each other. These flattened embeddings are then given to dense layers for processing. Next, we'll check how this approach is performing on the dataset.

2.1 Define Network¶

Below, we have defined a network that we'll use for our text classification task. The network consists of 3 layers (One embedding and two linear).

The first layer of our network is the embedding layer. We have created an embedding layer using Embed() constructor available from haiku. We have given vocab size and embedding length to the constructor. The embedding length is set at 50 which means that each word index will have a real-valued vector of length 50. The constructor will internally create a weight matrix of shape (vocab_size, embed_len). As all of our examples are simple word indexes, the embedding for words will be retrieved by simply integer indexing this embedding matrix (E.g, let's say 'the' word has index 2, then embedding of 'the' will be 'weight_matrix[2]'). The input data shape to embedding layer is (batch_size, max_tokens) = (batch_size, 50) and output data shape is (batch_size, max_tokens, embed_len) = (batch_size, 50, 50).

The output of embedding layer is flattened as per our approach. This will transform shape from (batch_size, 50, 50) to (batch_size, 50 x 50) = (batch_size, 2500).

The flattened output will be given to our first linear layer which has 128 output units. The linear layer will process input data and will return the output of shape (batch_size, 128).

The output of the first linear layer is given to the second linear layer which has 4 output units (same as no of target classes). The output of the second linear layer is a prediction of our network.

After defining the network, we have transformed it to a pure JAX function (using hk.transform()) and initialized it. After initializing, we have printed the shape of weights/biases of layers of network and performed a forward pass through the network to make predictions for verification purposes.

If you are someone who is new to Haiku and want to learn how to create a neural network using it then we recommend that you go through our below tutorial. It'll get you started with haiku.

• Haiku: Guide to Create Neural Networks using JAX

2.2 Define Loss Function¶

In this section, we have defined a cross entropy loss function that we'll use for our classification task. The function takes network parameters, input data, and actual target labels as input. It then makes predictions using network parameters and input data. The actual target values are one-hot encoded. At last, loss is calculated by calling softmax_cross_entropy() function available from optax by giving predictions and one-hot encoded target values to it.

2.3 Train Network¶

In this section, we are training our network. To train the network, we have created a helper function. The function takes train data (X_train, Y_train), validation data (X_val, Y_val), number of epochs, network parameters, optimizer state, and batch size. The function executes a training loop number of epochs time. During each epoch, it loops through training data in batches. For each batch, it performs a forward pass to make predictions, calculates loss, calculates gradients, and updates network parameters. It records the loss of each batch and prints the average loss of all batches at the end of the epoch. It also calculates validation accuracy at the end of each epoch. At last, it returns updated network parameters.

Below, we are actually training network by calling our training function defined in the previous cell. We have initialized a number of epochs to 12, batch size to 1024, and learning rate to 0.001. Then, we have initialized the model and Adam optimizer from optax. At last, we have called our training routine with necessary parameters to perform training process. We can notice from the loss and accuracy values getting printed at the end of each epoch that our network is doing a good job at the text classification task as loss is reducing and accuracy is increasing.

2.4 Evaluate Network Performance¶

Below, we have evaluated the performance of our trained network by calculating metrics accuracy score, confusion matrix, and classification report (precision, recall, and f1-score) on the test dataset. We can notice from the accuracy score that our model is doing a good job at the given task. We have calculated these metrics using functions available from scikit-learn.

• Scikit-Learn - Model Evaluation & Scoring Metrics

We have also created a heatmap of confusion matrix that let us see which categories are doing better compared to others. We can notice from the chart that categories 'Sports' and 'World' are doing better compared to categories 'Business' and 'Sci/Tech'.

We have created a chart using Python library scikit-plot. It has an implementation of charts for many different ML metrics. Please feel free to check the below link if you are new to the library and want to learn about it.

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

2.5 Explain Network Predictions using LIME Algorithm¶

In this section, we have dived deeper into checking the performance of our network. We are using LIME (Local Interpretable Model-Agnostic Explanations) algorithm to interpret prediction of the network which helps us understand which words of the network are contributing to the prediction. Thee is a python library named lime that provides an implementation of the algorithm. It let us create visualization which highlights words of text document that contributed to prediction.

If you are someone who is new to LIME then we would recommend that you go through the below link to understand it in-depth.

• 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 first need to create an instance of LimeTextExplainer which we have done below.

Below, we have created a prediction function. This function takes a list of text examples and returns predictions made by the network on them. It tokenizes and vectorizes them before giving them to the network. We'll use this function in the next cell for explaining predictions.

After defining a function, we randomly took a text example from the test dataset and made predictions on it using our trained network. Our network correctly predicts the target label as 'Sci/Tech' for the selected text example. Next, we'll look at which words of text document contributed to predicting this target category.

Below, we have first created an Explanation object by calling explain_instance() method on the explainer object. We have provided a selected text example, prediction function, and actual target value to the function. The explanation object has details about words contributing to prediction.

Next, we have called show_in_notebook() method on the explanation object to create a visualization of the explanation. We can notice from the visualization that words like 'RFID', 'proponents', 'arise', 'technology', 'stores', 'consumer', etc are words contributing to predicting the target label as 'Sci/Tech' which makes sense as these are commonly used words in the tech industry. Though the network is still missing some important words which can contribute to prediction.

3. Approach 2: Averaged Embeddings ¶

In this section, we have explained another approach to using word embeddings. Our approach in this section takes embeddings of words of each text example and averages them before giving them to the linear layer. The only difference in architecture is that we are averaging embeddings at the text document level. The majority of the code in this section is the same as earlier.

3.1 Define Network¶

Below, we have defined the network that we'll use for our task. The network has the same number of layers as earlier with only a difference in forward pass. Here, we are taking the mean of embeddings before giving them to the linear layer instead of flattening them like earlier. The rest of the code is the same. We have initialized the network after defining it, printed shapes of weights/biases, and performed a forward pass for verification purposes.

3.2 Train Network¶

Here, we have trained our new network architecture using the same settings that we had used earlier. We can notice from the loss and accuracy values that our network is doing a good job at the classification task.

3.3 Evaluate Network Performance¶

Below, we have evaluated the performance of our network as usual by calculating the accuracy score, confusion matrix, and classification report metrics on test predictions. We can clearly notice from the accuracy score that our model in this section is giving better results compared to the previous approach. We have also plotted the confusion matrix for reference purposes which hints at the improvement of all categories.

3.4 Explain Predictions using LIME Algorithm¶

Here, we have explained network prediction using LIME algorithm. We have randomly selected text example and our model is correctly predicting target label as 'Sci/Tech' for it. The visualization shows that words like 'RFID', 'technology', 'arise', 'proponents', 'frequency', 'identification', 'consumer', 'questions', 'privacy', etc are contributing towards predicting target label as 'Sci/Tech'.

4. Approach 3: Summed Embeddings ¶

In this section, we are trying a third way of using word embeddings for the text classification task. Our approach in this section is almost the same as our previous approach with the only change being that we are taking sum of embeddings at text document level instead of average. The rest of the code is exactly the same as earlier.

4.1 Define Network¶

Below, we have defined the network that we'll use for our task in this section. It has the same layers as our previous networks. The only difference is in forward pass where we are taking the sum of embeddings at the text document level before giving them to the linear layer. The rest of the logic is exactly the same. After defining the network, we initialized it, printed the shape of weights/biases, and performed a forward pass for verification.

4.2 Train Network¶

Here, we have trained our network using the same settings that we have been using for all our previous approaches. We can notice from the loss and accuracy values that our network is doing a good job at the text classification task.

4.3 Evaluate Network Performance¶

Below, 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 better compared to our first approach but a little less compared to our previous approach (averaged embeddings). We have also plotted the confusion matrix for reference purposes.

4.4 Explain Predictions using LIME Algorithm¶

Below, we have checked the performance of our network using LIME algorithm. We randomly selected a text example from the test dataset and made predictions on it using our trained network. The network correctly predicts the target label as 'Sci/Tech'. Then, we created a visualization explaining the prediction of a network. We can notice from the visualization that words like 'arise', 'proponents', 'technology', 'RFID', 'identification', 'frequency', 'consumer', etc are contributing to predicting target label as 'Sci/Tech'.

5. Summary of Results ¶

6. Further Suggestions ¶

• Try training network for more epochs.
• Try different embedding lengths. We have tried embedding a length of 50.
• Try adding more dense layers.
• Try trained embeddings like GloVe, FastText, etc.
• Try other aggregation operations (max, min, etc) on embeddings. We tried averaging and summing.
• Try more max tokens per text example. We have kept only 50 tokens per text example.
• Try different activation functions (relu, tanh, etc).
• Try different weight initialization methods.
• Try regularization
• Try learning rate schedules

References¶

• Haiku: Guide to Create Neural Networks using JAX