Flax (JAX): Conv1D For Text Classification Tasks¶

Recurrent Neural Networks (RNNs) and their variants like LSTM and GRU are the most commonly preferred network for ML tasks that involves some kind of sequence data (time-series data, text data, speech data, etc). The RNNs are very good at capturing the long sequences of the data. Though RNNs are better at capturing sequences, it takes time to train RNN networks. Recent studies have shown that CNNs consisting of 1D convolution layers are also good at capturing sequences. Hence, we can use a network of 1D Convolutions for NLP tasks like text classification, text generation, etc. The main benefit of using CNNs is that it has fewer parameters to train compared to RNNs and gets trained faster.

As a part of this tutorial, we have designed CNNs with 1D convolution layers using Flax for text classification task. Flax is a Python deep learning library built on top of JAX for designing deep neural networks. We have used AG NEWS dataset which has text documents for 4 different categories of news. We have tried different approaches to using Conv1D layers for solving text classification tasks. For encoding text data before giving it to the convolution layer, we have used word embeddings approach. After training networks, we have also evaluated their performance by calculating various ML metrics and explained their predictions using LIME algorithm.

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 Load Datasets
   • 1.2 Vectorize Text Data
2. Approach 1: Single Conv1D Layer Network (Max Tokens=50, Embedding Length=128, Conv Output Channels=32)
   • Define Network
   • Define Loss Function
   • Train Network
   • Evaluate Performance Of Network
   • Explain Network Predictions using LIME Algorithm
3. Approach 2: Multiple Conv1D Layers Network (Max Tokens=50, Embedding Length=128, Conv Output Channels=32,32)
4. Results Summary And Further Recommendation

Below, we have imported the necessary Python libraries that we have used in our tutorial and printed the versions of them as well.

1. Prepare Data ¶

In this section, we are preparing data to be given to the neural network. As our raw data is text and neural network works on real-valued data. We need to convert our text data to real-valued data.

In this tutorial, we are going to use word embeddings approach where we break text into a list of tokens (words, punctuation marks, etc) and assign a real-valued vector to each token of text. We'll achieve this in two steps.

1. Map each token of text to a unique integer index.
2. Map integer index to real-valued vector.

We'll complete the first step in this section. The second step will be implemented in the neural network where we include Embedding Layer as the first layer which is responsible for mapping the integer index of tokens to their respective embeddings. These embeddings get updated during training of the network to better understand the meaning of token.

1.1 Load Datasets¶

In this section, we have loaded AG NEWS dataset that we are going to use for our case. It has text documents for 4 different news categories (["World", "Sports", "Business", "Sci/Tech"]) which our trained network will try to classify. The dataset is already divided into train and test sets.

1.2 Vectorize Text Data¶

In this section, we are vectorizing our text data.

First, we have initialized Tokenizer object available from keras. After initializing the tokenizer, we have trained it by calling fit_on_texts() method on it with train and test text examples. The call to this method will populate the vocabulary inside of Tokenizer object which will have all unique tokens (words, punctuation marks, etc). The vocabulary is a simple mapping from a token to a unique integer index. Each token is assigned a unique integer index starting from 1. Position 0 is reserved for unknown tokens encountered in the future.

After vocabulary is populated with tokens, we can call texts_to_sequences() method on the tokenizer object with a list of text documents. It'll tokenize them and return indexes of tokens for each text example. We have called texts_to_sequence() method on train and test text documents. We have also called pad_sequences() method around it. This method is used to bring tokens of all text examples to the same length. The text for different documents can have a different number of tokens. But we have decided to keep a maximum of 50 tokens per text example. All tokens beyond 50 will be truncated. The examples that have less than 50 tokens will be appended with 0s.

After vectorizing text data to indexes, we have also printed the first few examples with their indexes.

Below, we have explained with a simple example how text documents will be mapped to indexes.

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]

Approach 1: Single Conv1D Layer Network (Max Tokens=50, Embedding Length=128, Conv Output Channels=32) ¶

Our first approach tries to perform a text classification task with a simple one convolution layer network. The network consists of an embedding layer, a 1D convolution layer, and a dense layer. Below, we have defined the network, trained it, evaluated the performance of the network, and tried to explain predictions made by the network as well.

Define Network¶

In this section, we have defined a network that we'll use for our text classification task in this section. The network consists of 3 layers.

• Embedding Layer
• Conv1D Layer
• Dense Layer

The first layer of the network is the embedding layer. We have created an embedding layer using Embed() constructor available from linen sub-module of Flax. We have asked it to use a length of vocabulary as a number of tokens (first parameter) and embedding length of 128. When we create this layer, it internally will create a weight matrix of shape (vocab_len, embed_len). This matrix has embeddings of length 128 for all tokens of our vocabulary. The embedding layer simply takes a list of indexes of tokens as input and returns embeddings for all those indexes from the weight matrix. The input to embedding layer is of shape (batch_size, max_tokens) = (batch_size, 50) and the output is of shape (batch_size, max_tokens, embed_len) = (batch_size, 50, 128).

The output of the embedding layer is given to Conv1D layer. We have defined Conv1D layer with 32 output channels and a kernel size of 7. We'll be treating the embedding length dimension as the channels dimension in our case. Hence, the input to convolution layer will be of shape (batch_size, max_tokens, embed_len) = (batch_size, 50, 128) and output shape will be (batch_size, max_tokens, conv_output_channels) = (batch_size, 50, 32). We have applied relu activation to the output of the convolution layer.

After applying relu activation, we have applied max() (we can also try average instead of max) function at max_tokens dimension which will return us output of shape (batch_size, 32). This output will then be given to Dense layer which has 4 output units (same as the number of target classes). The output of Dense layer is returned from our network as the prediction.

After defining the network, we initialized it, printed the shape of weights/biases of layers, and performed a forward pass through it using random data for verification purposes.

If you are someone new to Flax and want to learn how to create a neural network using it then we recommend that you go through the below tutorials that cover it in detail. It'll help us better understand Flax.

• Flax: Framework with High-level API to Create Neural Networks using JAX
• Flax: Convolutional Neural Networks

Define Loss Function¶

Below, we have defined the loss function that we'll be using for our task. We'll be using cross entropy loss for our case. The function takes network parameter, actual data, and their target labels as input. It then makes predictions on data using network parameters and one-hot encodes actual target labels. At last, it calculates cross entropy loss by giving predictions and one-hot encoded target labels to softmax_cross_entropy() function available from Optax library.

Train Network¶

In this section, we are training the neural network that we defined earlier. We have defined a simple training function that will perform training for us. 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, it performs a forward pass to make predictions, calculates loss, calculates gradients, and updates network parameters using gradients. It also records loss for each batch and prints the average loss of all batches at the end of the epoch. The function also calculates validation data accuracy at the end of each epoch.

Below, we have actually trained our network using the training routine defined in the previous cell. First, we have initialized batch size to 1024, the number of epochs to 10, and the learning rate to 0.001. Then, we have initialized our text classification network 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 value getting printed after each epoch that our model is doing a good job at the text classification task.

Evaluate Performance Of Network¶

In this section, we have evaluated the performance of our train network by calculating accuracy score, classification report and confusion matrix metrics on test predictions. We can notice from the accuracy score that the network is doing quite a good job at the classification task.

We have calculated ML metrics using functions available from scikit-learn. If you want to learn about various ML metrics available from sklearn then please feel free to check the below link.

• Scikit-Learn - Model Evaluation & Scoring Metrics

We have also created a plot of confusion matrix using Python library scikit-plot. We can notice from the plot that our model is doing a good job at classifying text documents of categories Sports and World compared to categories Business and Sci/Tech.

If you are new to library scikit-plot and want to learn various ML metrics plots it provides then please do check the below link.

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

Explain Network Predictions using LIME Algorithm¶

In this section, we have tried to explain predictions made by our network using LIME algorithm. We'll be using an implementation of the algorithm available through lime python library. It let us generate visualization showing which words from our text example contributed to predicting a particular target label.

Please feel to check the below links if you are new to LIME and want to learn about it in-depth.

• How to Use LIME to Understand sklearn Models Predictions?
• LIME: Interpret Predictions Of Flax (JAX) Text Classification Networks

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

Now, we have randomly selected 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.

Apart from this, we have also defined the prediction function in the below cell. The function takes a list of text examples as input and returns their prediction probabilities using our model. The function tokenizes data and then gives it to the network to make predictions. At last, softmax activation is applied to the output of the network to convert it to probabilities and returned from the function.

Below, we have first generated an Explanation object by calling explain_instance() method on LimeTextExplainer object. We have given text examples, prediction function, and the target label to the method. This explanation object has details about words contributing to prediction.

Then, we have called show_in_notebook() method on Explanation object to generate a visualization showing words contributing to predicting the target label as Sci/Tech.

We can notice from the visualization that words like 'RFID', 'privacy', 'frequency', 'identification', etc are used for predicting the target label as Sci/Tech which makes sense as they are commonly used in the field.

Approach 2: Multiple Conv1D Layers Network (Max Tokens=50, Embedding Length=128, Conv Output Channels=32,32) ¶

Our approach in this section uses multiple convolution layers. The majority of the code in this section is the same as our code from the previous section except for the definition of neural network which uses two 1d convolution layers this time instead of one.

Define Network¶

In this section, we have defined a network that we'll use for our text classification task in this section. The network consists of one embedding layer, two convolution layers, and one dense layer. We have again used the embedding length of 128 in this section as well. The two convolution layers have 32 output channels and a kernel size of 7. The output of the embedding layer is given to the convolution layer. The relu activation is applied to the output of the first convolution layer and then max-pooling operation is performed on it. The output of max-pooling is given to the second convolution layer. Then, we have applied relu activation to the output of the second convolution layer and given it to the dense layer. The output of the dense layer is returned as a prediction of the network.

After defining the network, we initialized it, printed the shape of weights/biases of layers, and performed a forward pass through it using random data for verification purposes.

Train Network¶

Here, we have trained our network using exactly the same settings that we have used in the previous section. We can notice from the loss and accuracy values getting printed after each epoch that our model is doing a good job at the text 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 test accuracy that it's a little low compared to the previous approach which is surprising as we had expected that trying more convolution layers might improve accuracy. We have also plotted the confusion matrix for reference.

Explain Network Predictions using LIME Algorithm¶

In this section, we have again tried to explain predictions made by our network using LIME algorithm. Our network correctly predicts the target label as Sci/Tech for a selected random test example. From the visualization created using LIME, we can notice that words like 'RFID', 'technology', 'frequency', 'identification', 'privacy', etc are contributing to predicting the target label as Sci/Tech.

4. Results Summary And Further Recommendation ¶

Below, we have listed a summary of various approaches we tried above.

Further Suggestions¶

• Try different convolution layer output channels.
• Try different kernel sizes.
• Try a different number of maximum tokens.
• Try different embedding lengths.
• Add more dense layers at the end.
• Try different weight initialization methods.
• Try different optimizers (RMSProp, Adagrad, etc).
• Add recurrent layers (LSTM, GRU) after the convolution layer.
• Try learning rate schedulers.

This ends our small tutorial explaining how we can design CNNs of 1D Convolution layers using Flax (JAX) framework for solving text classification tasks. Please feel free to let us know your views in the comments section.

References¶

• Guide to Use Word Embeddings with Flax (JAX) for Text Classification Tasks
• Text Classification Using Flax (JAX) Networks