Robust Training Objectives Improve Embedding-based Retrieval in Industrial Recommendation Systems

In a two-stage recommendation system with a retrieval then ranking phase, the retrieval phase plays an important role filtering out the most relevant candidates to lighten the load of the ranker. Since the ranking result is largely dependent on items retrieved in the retrieval phase, a good quality retrieval model can drastically improve the final ranking. Embedding based retrieval (EBR) is a method that’s recently adopted and deployed in many content, product, and friend recommendation systems(DBLP:journals/corr/abs-2006-11632, ; 45530, ; TK:2021, ; 10.1145/3539618.3591848, ), and proved to achieve superior results. EBR transform users and items into embeddings, changing the retrieval problem into a nearest-neighbor search problem in a low-dimensional latent space. These embeddings can be determined in advance and indexed using an approximate nearest neighbor search such as FAISS (DBLP:journals/corr/JohnsonDJ17, ) and HNSW (DBLP:journals/corr/MalkovY16, ) in order to retrieve the top-$k$ most relevant items efficiently at serving.

When applying EBR to RS problems, the quality of embeddings is of upmost importance. In this paper, we use a friend recommendation system as our subject. In scenarios like friend recommendation where vast amounts of topological information relating users and items is readily available, these embeddings can be augmented with GNNs. Previous work showed that EBR for friend recommendation systems see benefits leveraging graph-aware embeddings(10.1145/3539618.3591848, ). In this setting, nodes would contain individual user features while edges map to user-user interactions. This approach compliments commonly used graph traversal approaches (eg. friend-of-friend (FoF) (Newman_2001, )), allowing for retrieval of candidates from any number of hops away from the target.

Here we describe GNNs for generating graph-aware embeddings for EBR. GNNs have demonstrated state-of-the-art performance in many problems containing rich topological information within the graph data (DBLP:journals/corr/abs-1812-08434, ), such as recommendation and forecasting. Formally, we define $G=(\mathcal{V},\mathcal{E},X)$, where $\mathcal{V}$ is the set of $n$ nodes ($|\mathcal{V}|=n)$, $\mathcal{E}$ is the set of edges ($\mathcal{E}\in\mathcal{V}\subseteq\mathcal{V}$), and $X$ is a feature matrix of dimension $d$ where $X\in\mathbb{R}^{n\times d}$. Many modern GNNs also employ a message-passing structure, consisting of an aggregation (AGG) and update (UPD) function. The goal of this paradigm is for nodes to receive information from their neighbors, collecting messages using its AGG function before updating their own messages with the UPD function, both of which are learnable and permutation-invariant. For some node $u$ at layer $k$, the next message-passing layer can be written as

where $\mathcal{N}(u)$ is the neighborhood nodes of node $u$. Different message-passing GNN models use different combinations of AGG and UPD functions. An example of a more complex GNN, Graph-attention networks (GATs) (veličković2018graph, ), use an attention mechanism for each pair of nodes $i$ and $j$

where $\mathbf{W}$ is a linear transformation applied to every node and $f_{\text{att}}$ is the attention function parameterized by a weight vector and a non-linearity function. The AGG function is then a attention-weighted sum of its neighbors features while the UPD function is implicitly defined in $\mathbf{W}$ and the non-linearity function. Typically, to generate graph-aware embeddings from GNNs, a margin based ranking loss(DBLP:journals/corr/abs-1806-01973, ; 10.1145/3539618.3591848, ) or contrastive(DBLP:journals/corr/abs-2101-01317, ) loss can be used, to encourage items that are closer in the graph to be closer in the embedding space.

Multitask learning (MTL) is an approach in machine learning where a model is trained simultaneously on several tasks. MTL has been extensively explored in recommendation as a way to improve key metrics (inproceedings1, ; 10.1145/3219819.3220007, ; ma2018entire, ; 10.1145/3383313.3412236, ). Thus, the core idea behind multitask learning is to improve the robustness of the model by leveraging the domain-specific information contained in the training signals of related tasks (NIPS2006_0afa92fc, ; Caruana1997MultitaskLearning, ). Hard parameter sharing, one of the most fundamental forms of MTL, uses a shared representation which then branches into multiple heads capable of learning task-specific information (guo2020learning, ; DBLP:journals/corr/abs-1911-12423, ; vandenhende2020branched, ).

For graph-aware EBR in particular, self-supervised multitask learning (SSMTL) has been proposed as a new approach to MTL, optimizing the embeddings directly to achieve desirable embedding properties without the use of positive or negative labels. In this setting, we combine several self-supervised learning (SSL) methods with a downstream retrieval task to learn both direct and indirect embedding features. Recent work (ju2023multitask, ) has shown that SSMTL can lead to improved task generalization and embedding quality on several academic benchmarks through the increasingly robust training objective. However, many of the SSL approaches used are constrained to the assumption that global graph information can be inferred within the graph structure. This is not valid in the large-scale recommendation setting, where graphs are constrained to some $K$-hop around a query user in order to fit in memory. As a result, many of these SSL methods may lead to negative transfer due to SSL task conflict with the target link prediction task, and there remains work to be done to investigate which methods perform best in this large-scale setting.

We identify two self-supervised learning approaches that are scalable and lead to improvements in the large-scale recommendation setting through a more robust training objective.

Canonical Correlation Analysis. Based on work from (DBLP:journals/corr/abs-2106-12484, ), Canonical Correlation Analysis (CCA) deploys a non-contrastive, non-discriminitive SSL method to train the GNN. The self-supervised training objective is described in Equation 3. First, given a subgraph with $n$ nodes, two augmented views of the subgraph are created and fed through the GNN, producing $\mathbf{Z_{A}}$ and $\mathbf{Z_{B}}$ where $\mathbf{Z_{A}},\mathbf{Z_{B}}\in\mathbb{R}^{n\times k}$. Each of these embeddings are fed through a task-specific head, and then are normalized so that each feature has $0$ mean and $\frac{1}{\sqrt{n}}$ standard deviation, resulting in $\mathbf{\tilde{Z}_{A}}$ and $\mathbf{\tilde{Z}_{B}}$. The loss is then computed from Equation 3. The first term in the equation seeks to minimize the distance of the same nodes between the two views. The second term enforces that the feature-wise covariance of all nodes is equal to the identity matrix.

Masked Autoencoders. Based on work from (hou2022graphmae, ), this approach leverages a graph masked autoencoder (MAE) that focuses on feature reconstruction. First, an augmented view of the subgraph is created and the features of the query users are masked out. This augmented graph is then fed through the GNN and a task-specific head. The features of the query users are then re-masked and passed through a graph convolution layer. As described in Equation 4, for all masked nodes $\mathcal{V}$, the final loss is equal to the average of the scaled cosine error between the original features $\mathbf{X}$ and generated features $\mathbf{Z}$. This approach only relies on the local neighborhood surrounding the query node, making it a good option for large-scale SSMTL.

We note that these two approaches both utilize non-contrastive methods. While experimenting with different SSL tasks, we find that contrastive SSL approaches do not perform very well in the production setting due to their assumption that global information is readily available in the original and augmented graphs. This is not necessarily true for large-scale recommendation, where subgraphs are constrained to the K-hop neighborhood surrounding each query node.

We evaluated the effectiveness of SSMTL for end-to-end friend recommendation with online A/B testing. The control group used candidates retrieved from the production model trained with retrieval baseline, while the treatment group instead used candidates retrieved with the new robust training objective in the SSMTL setting, specifically combining the previous retrieval loss with whitening decorrelation and generative reconstruction objectives.

In the A/B experimental results, we saw statistically significant improvements across several friend recommendation metrics. Specifically, we observed up to 5.45% improvements in new friends made and +1.91% new friends made with low-degree users in various markets. Overall, from these results, we see that SSMTL is able to provide improved recommendation compared with the single-task setting, in particular helping with candidate generation for low-degree users.