I recently read an article from folks at TikTok called
Deep Retrieval: Learning A Retrievable Structure for Large-Scale Recommendations"
Recommendation systems need to be able to quickly obtain relatively relevant candidates, which are then reranked to produce the final output.
Typically, for candidate generation, an inner-product model (such as metric learning) is used, followed by an ANN (approximate nearest neighbour; a popular option is
FAISS.
In this article, however, t
he authors want to show how retrieval can be done directly through item-user interactions without making assumptions about the Euclidean nature of the space and the proximity of entities within it (which, in my opinion, isn't really a problem, considering that in metric learning we specifically train for this kind of representation).
The authors train a model with D layers. Each layer consists of an MLP with a softmax function over K nodes that outputs the probability of belonging to one of the K clusters. The input to layer D1 is an embedding of the user( embedding takes into account their previous actions, and a recurrent neural network with GRU is used to project the behaviour sequence onto a fixed dimension embedding as the input). The target is the cluster of the item with which the user has interacted with (e.g., clicked or purchased). The output of D1, let's call it K1, is then concatenated with user embedding and used as input for D2. An output K2 is then concatenated with K1 and user embedding and used as an input for D3.
Any user potentially has K^D different paths. For example, if there are 30 clusters and three layers, the model can output the following path for user X: 1-10-15, meaning cluster 1 among the first 30, cluster 10 among the next 30, and cluster 15 among the next 30. Additionally, because we have a distribution, we can go deeper and take the top-3 from each layer, resulting in n^D (27 in our case) different paths instead of the original one. Since we train on user-item interactions, we can get paths for both the user and the item.
A question arises: how can an item belong to different clusters? For example, an item related to a kebab could belong to a "food" cluster, while an item related to flowers could belong to a "gift" cluster. However, an item related to chocolate or cakes could belong to both clusters in order to be recommended to users interested in either food or gifts.
This is actually one of the advantages over tree-based deep models.
The reasonable question is, how do we determine the initial clusters? Okay, we have user embeddings and user-item interactions, but where do we get labels for K? We can randomly distribute and turn on the
EM machine.
In the first iteration, we distribute them randomly and train the model; then, we re-train the item's mapping to the cluster to maximise the model output.
How is Deep retrieval applied during inference?
We input the user's embedding -> we get N paths (the greedy algorithm outputs one path)
We gather all items that are in these paths
We run them through the intermediate reranker
There are a few additional points.
I. Despite the fact that DR (Deep Retrieval) outputs significantly fewer items than all, there are still many of them, so it is also trained with a reranker in order to output the top (this is still not the final reranker!)
II. The mapping to clusters is discrete, so it cannot be updated with gradient methods (hence the use of EM)
III. They add a penalty for adding another item (passing the same path) to the path. Otherwise, there is a risk that all items will fall into one path, and they used a penalty in the form of c^4/4 where c is the number of items in the path
IIII. They updated the model from the incoming data stream - this affected some things, such as step M in EM. They also used exponential decay with a coefficient of 0.999
Metrics