Recommender Systems (Collaborative Filtering) using Flax (JAX)¶

Recommender systems are used to predict the ratings given by a particular user to a particular item. Recommendation systems are commonly used by online websites like Amazon, Netflix, Hotstar, etc. These systems try to understand the taste of the user and then recommend items (movies, books, shopping items, etc.) based on it. Collaborative filtering is one of the ways of implementing recommender systems. In collaborative filtering, we represent the user with a vector of floats and items also with a vector of floats of pre-decided size. Each individual user and item are assigned different vectors. Then, we train our neural network to update these vectors such that when we multiply the user vector by item vector, we get a rating result that is near to the actual rating given by the user to that item. We try to bring predictions near to every user's rating of a particular movie. During this process, we learn about user taste as well as item attributes. Once, we have updated the vector after training the user then we can use it to recommend new items to users that he/she has not rated yet. During training, we update vectors of many users and items together hence generally users with the same taste will be recommended the same items.

As a part of this tutorial, we'll explain how we can implement collaborative filtering using Flax framework of Python. Flax is a deep learning framework designed on top of JAX. Deep learning frameworks generally refers to the vectors that we discussed earlier as embeddings and provides embedding layers that maps user/item ids to their vectors. We expect that reader has a background in Flax/JAX. Please feel free to explore the below links if you want to refresh some Flax/JAX content as it'll help with this tutorial.

• Flax: Framework to Create Neural Networks using JAX
• JAX - (Numpy + Automatic Gradients) on Accelerators (GPUs/TPUs)

We'll be using the Japanese Anime dataset available from kaggle as a part of this tutorial. Please feel free to download it to follow along.

• Anime Recommendations Database

Below, we have listed important sections of our tutorial to give a summary of the content covered in this guide.

Important Sections of Tutorial¶

• Load and Prepare Data
  • Load Ratings Data
  • Load Anime Data
  • Merge Ratings and Anime Data
  • Few Checks
  • Separate Users With Only One Review and Remove Them From Data
  • Shuffle Merged (Rating + Anime) Data
  • Pivot Merged Data to Analyze User-Rating Relation
  • Split Data into Train/Test Sets Based on User IDs
  • Few Checks
• Model 1: Embeddings Only
  • Create Model
  • Why We Are Treating Rating Prediction Task As A Regression Task?
  • Define Loss Function
  • Train Model
  • Predict Ratings
  • Recommend New Items
  • How to Recommend Items to New Users?
  • How to Recommend New Items to Users?
• Model 2: Embeddings + Bias
• Model 3: Embeddings + Dense Layers
• Model 4: Embeddings + Dense Layers

Below, we have imported the necessary libraries of this tutorial and printed the versions of them that we have used in this tutorial.

Load and Prepare Data ¶

In this section, we'll be loading ratings and related datasets available from kaggle. We have downloaded the anime dataset from kaggle that we'll be using in our tutorial. We'll load and prepare a dataset for the collaborative filtering task.

• Anime Recommendations Database

The dataset has 2 CSV files.

• rating.csv - It has information about ratings given by users to animes.
• anime.csv - It has information about individual anime like their name, genre, average rating, etc.

Load Ratings Data¶

Below, we have loaded rating data available from rating.csv file. It has user id, anime id, and rating given by the user to that anime. There are entries in the dataframe where the user has watched anime but has not given a rating hence those entries have ratings as -1. We have removed those entries from our data as we want only entries where the user has given a rating.

Load Anime Data¶

Below, we have loaded anime data from anime.csv file. It has a mapping from anime id to anime name, genre, type, episodes, average rating, and members. We have printed the first few rows to show a glimpse of the dataset.

Merge Ratings and Anime Data¶

In this section, we have merged the rating dataset with the anime dataset using the pandas merge() method. We have merged datasets on column anime_id which is common between the two dataframes. After merging, the dataset has two columns giving rating information hence we have renamed one from the rating dataframe as user rating and one from the anime dataframe as average anime rating.

Few Sanity Checks¶

In this section, we are calculating counts of unique items and user ids. We are also checking whether there are any Null values present in the user id, anime id, or rating.

Separate Users With Only One Review and Remove Them From Data¶

In this section, we are just looping through all unique user ids and separating user who has given a rating to only one anime from those who has given a rating to more than one animes. We have then filtered our final dataframe and kept only user ids who have given ratings to more than one anime. The reason behind performing this step is that, when we'll divide the dataset into train and test sets, we'll be doing separation based on user ids. If we keep user ids with only ratings then we'll have to keep them in the train set and there won't be an entry for them in the test set. As we have removed entries with one rating, we'll have each user id present in both train and test sets. It's a choice that needs to be made. If the reader wants to keep user ids with only one review then it's fine.

Shuffle Merged (Rating + Anime) Data¶

In this section, we are simply shuffling the dataframe using the pandas' data frames shuffle() method. We have also created dictionaries that map anime ids to index and their title as they'll be useful in the future. We have then also performed checks to see how many user ids, anime ids, and unique rating types are present after we removed entries with a single review.

Pivot Merged Data to Analyze User-Rating Relation¶

In this section, we have just pivoted the first 10 entries of our dataframe to a tabular view of ratings given by users to particular anime.

Please make a NOTE that we are calling gc.collect() a few times in our tutorial as it initiates Python garbage collector which frees up memory.

Split Data into Train/Test Sets Based on User IDs¶

In this section, we have created a small function that divides the dataset into train and test sets. For our case, the data features are user id and anime id and target values are ratings. We'll be feeding user id and anime id to our neural network and it'll predict the rating given by that user to the anime.

In order to divide the dataset into train and test sets, we have first taken unique user ids. Then, we are looping through each user id, taking all entries of that user id (all ratings given) dividing them into the train (90% entries), and test (10%) sets. We have also handled the case where there is only one rating per user then we keep it in train set though we have removed entries with single ratings. If the reader keeps entries with single ratings then the below code will handle it as well.

Our datasets (X_train, X_test) have two entries (user id, anime id). The target values (Y_train, Y_test) are ratings.

Few Sanity Checks¶

Below, we have again performed simple checks for the count of user ids that are present in train but not in test and vice-versa. This is done to verify that the above code properly divided data into train and test sets.

Model 1: Embeddings ¶

In this section, we'll design our neural network using embeddings. Embeddings are weights that map ids (user id/ item id) to their vector values. As we had discussed earlier, in collaborating filtering, we mentioned that we keep a vector of floats for each unique user and item. Let's say there are 10 unique users and we are maintaining a vector of 50 floats per user then we'll have 10 such vectors of 50 floats, one for each user. Same way vectors will be maintained for items (movies, animes, books, etc). These vectors are generally referred to as embeddings. These embeddings for user generally represents the taste of that user (he/she likes action, adventure movies or comedy movies, or horror movies). The same way embeddings for items represent the characteristics of that item (whether the movie is horror or action or adventure.). The embedding layer for user ids in the network has embeddings for each user and it then maps the user id to its embeddings. The same goes for the embedding layer for item ids. It just maps user/item id to their embeddings (vector of floats). We generally decide the size of these vectors from the beginning and it'll be the same for all user ids and item ids.

Create Model¶

In this section, we have defined our simple network for the collaborative filtering task. We have set vector size for each user and item to 50 which means that each individual user and item will be represented with 50 floats. The number of floats is generally referred to as factors and is one of the network hyperparameters that needs to be tuned. We had tried different values for factors like 40, 50, 70, 100, and 120. Then, we decided on 50 as our final value.

We have created a model by extending flax.linen.Module class. The setup() method has layers initialized in them and call() has actual forward pass through the network. We have declared two embedding layers (one for users and one for items) using Embed layer of flax.linen. The user embedding layer has a float array of shape (n_users, n_factors) which will map the user id to its vector of floats. The item embedding layer has a float array of shape (n_items, n_factors) which will map item id to its vector of floats.

The forward pass-through network (call() method) simply multiplies user embeddings by item embeddings to generate a rating. The network will take as input tuple of (user_id, item_id), it'll then take the vector for that user id and item id, multiplies them, and sum up to predict rating given by the user to item. It can do this for a batch of user ids and item ids. We then update the weights of users and items to improve rating prediction during the training process.

If the process of network creation and training seems a little hard to grasp then we suggest that the reader goes through our tutorial on Flax that explains how to create neural networks as it'll help with this one as well.

• Flax: Framework to Create Neural Networks using JAX

Why We Are Treating Rating Prediction Task As A Regression Task?¶

We have treated the process of predicting rating as REGRESSION task even though ratings are unique with an integer in the range [1,10]. The reason behind treating the rating prediction task as regression is that we want a model to predict a rating that will be near the actual rating. We can then round up float prediction to an integer. The classification task treats all classes as independent of each other and it does not enforce ORDERING in target classes. It can predict any of the possible classes in case of wrong predictions whereas if we treat the task as a regression task then it'll predict a rating that is near to the actual rating even in case of wrong predictions. In our case, the ratings have ordering where a rating of 2 is better than 3, 3 is better than 4, etc. We can enforce ordering by treating the task as REGRESSION which will try to predict a rating that is near to actual.

Please make a NOTE that all of our models will approach the rating prediction task as REGRESSION task even though ratings are independent classes and seem like CLASSIFICATION task. The reason being we can't enforce an order by treating it as a classification task.

Below, we have initialized the network and printed the shape of embeddings. We have also performed a forward pass-through network for verification purposes.

Define Loss Function¶

In this section, we have defined the loss function that we have used in our case. We have used mean squared error loss as our loss which tries to reduce the distance between actual rating and predicted rating. We have also defined mean absolute error as well as that can also be used as a loss function for regression tasks.

Train Model¶

In this section, we have trained our neural network of embeddings. We have created a simple function that we have used for training. The function simply loops through train data in batches. For each batch (a bunch of (user id, item id)), it predicts ratings, calculates loss, calculates gradients, and updates weights using gradients. It also calculates the loss of our test dataset which we have provided as a validation set. The function prints both train and validation loss after each epoch. At last, the function returns updated network parameters as well.

Below, we are actually training our network using a function defined above. We have set the number of epochs to 10, batch size to 10,000, and learning rate to 0.001. We have then initialized the network and its weights. Then, We have initialized Adam optimizer to update network parameters in our case. At last, we have called our function to train network.

We can notice from the loss value getting printed after each epoch that our model seems to be improving after each epoch. The last train and validation loss are around 1.4 which indicates that our rating predictions are generally off by 1.4 from the actual rating which is near to the actual rating for the majority of samples.

Later on, we'll try more models to bring down this gap further low (below 1.4).

Predict Ratings¶

In this section, we are using our trained model to make a prediction of ratings for existing user id. We have designed a function for this. The function randomly selects one user id. It then makes predictions for items that are reviewed by this user id using our trained model. We have then merged actual rating and predicted rating in a single dataframe along with other detail which can be helpful. We have also included logic that converts float ratings to integer classes. We'll be reusing this function in our upcoming sections when we try different models.

Recommend New Items¶

In this section, we are recommending new items to existing users. As our users have not used all items, we can recommend new items to them by using our trained model which knows a little about the taste of users. We have created a small function to recommend new items to users.

The function takes as input user id, user data, recommender model, and model parameters. We then retrieve item ids that are not rated by the input user. Then, we use our trained model to predict the rating for those unrated items and sort them from high rating to low. At last, we recommend items for which our model predicted a high rating. The model tries to predict the rating that the user would have given to an item if he had rated it.

How to Recommend Items to New Users?¶

The very obvious question that can come to an individual's mind when using the collaborative filtering approach for a recommendation system is how to handle new users. In order to handle new users, we need to update our embeddings with embeddings of new users. We also need to rerun our model with embeddings of new users as initially, all embeddings will be random for them. We can train the model daily if we are getting many new users or we can train the model weekly, bi-weekly, or even monthly if the frequency of addition of new users is less. Initially, we can recommend items that have overall high ratings given by the majority of users and once the user has watched/rated other items, we can update their embeddings by training network.

How to Recommend New Items to Users?¶

Another obvious question that comes to mind is how to handle new items that are new to the platform. In addition, new items will have random embeddings just like new users. One of the solutions can be that we keep on recommending new items randomly to users until we have few ratings on that item. We can then update embeddings of the new item with those few ratings by training our network. We can initially highlight new addition items above to the home pages of websites as done by famous platforms like Netflix, Amazon, etc. We can also send a notification to users who have used the item but have not been given a rating.

Model 2: Embeddings + Bias ¶

In this section, we have introduced bias to embeddings so that it can add one single number to each item/user that can be used to represent more information. It'll also work as a regularization term and help weights settle better.

Create Model¶

Our model for this section has much of the code the same as an earlier section with few additions. We have declared user and item bias using Embed layer in setup() method. In call() method, we have first multiplied user embeddings by item embeddings and then added user and item biases to it. At last, we are returning results.

Train Model¶

In this section, we are training our new recommendation system with bias. We are a training network for 10 epochs with a batch size of 10,000 and a learning rate of 0.001. From the loss value getting printed at the end of epochs, we can notice that this model has seem to have done little better than our earlier model.

Predict Ratings¶

In this section, we are making predictions for items that the user has seen and rated. From the results, it seems that this model has done almost the same job as our previous model.

Recommend New Items¶

In this section, we are recommending new items to the user.

Model 3: Embeddings + Dense Layers ¶

Both of our previous models used only embedding layers and operations between them. But in this example, we'll extend the concept and add dense layers to the network which can help improve the results further. This time, we haven't multiplied user and item embeddings. Instead, we have merged them before feeding them to dense layers.

Create Model¶

Our model for this section first creates embedding layers for users and items in setup() method. It then declared two dense layers of units 100 and 1. The dense layer with a single output unit will produce our output.

Inside of call() method, we are first applying users embedding layer to user ids and then items embedding layer to item ids. Then, we are merging user embeddings and item embeddings by stacking them next to each other. The resulting matrix will of shape (n_samplesx100) resulted from merging user embeddings shape (n_samplesx50) and item embeddings of shape (n_samplesx50). We have then fed these merged embeddings to our dense layer with 100 units and applied relu (rectified linear unit) activation function to the output. Then, we have fed the output of relu to dense layer with 1 unit which will be our prediction (rating prediction). Now, our model has extra weights to learn due to the addition of 2 dense layers which can further improve the results of the network.

Train Model¶

Below, we have trained our new recommendation system for 10 epochs with a batch size of 10,000 and a learning rate of 0.001. The loss value getting printed after completion of all epochs is better than both of our previous models which were merely based on embeddings.

Predict Ratings¶

In this section, we have used our new recommendation system with dense layers to make predictions. The results are overall the same as our previous systems.

Recommend New Items¶

In this section, we have recommended new items using our new recommendation system. Many of the items recommended here are almost the same as the ones recommended by previous models.

Model 4: Embeddings + Dense Layers ¶

In this section, we have tried one more model where we have used one extra dense layer compared to our previous model. This experiment was performed to check whether adding one extra dense layer is helping improve results further or not.

Create Model¶

Our model for this example has almost the same code as the previous example with the only difference that we have used on the extra dense layer. We have applied that extra dense layer of 50 units after our dense layer of 100 units.

Train Model¶

Below, we have trained our new recommendation system for 10 epochs with a batch size of 10,000 and a learning rate of 0.001. From the resulting loss values, we can notice that the model performance is almost the same as our previous model. There does not seem to be much of an improvement with the addition of one extra dense layer.

Predict Ratings¶

In this section, we have used our new model to make predictions. The results seem a little better compared to our previous models though.

Recommend New Items¶

In this section, we have recommended new items to users using our new system. The majority of items are almost the same as the ones recommended by previous models.

This ends our small tutorial explaining how we can create recommender systems using the collaborative filtering approach and Flax (JAX) framework. The results of all models seem to be almost the same but we think that model 3 which had dense layers with embeddings did a little better job compared to others. We would recommend trying different settings to improve results before settling with any approach. Please feel free to let us know your views in the comments section.

References¶

• Flax: Framework to Create Neural Networks using JAX
• JAX - (Numpy + Automatic Gradients) on Accelerators (GPUs/TPUs)
• Guide to Create Neural Networks using High-level JAX API