Learning Federated Representations and Recommendations with Limited Negatives

Recent years have witnessed the successes of deep retrieval models in many large-scale recommendation systems [2, 13, 19, 17, 7, 16] and natural language tasks [5, 4, 18, 1]. While these models largely improve user experience by enabling personalization and representation learning, they can also raise privacy concerns as user data typically needs to be sent to a centralized server for training.

Federated learning (FL) is a decentralized training strategy in which clients collaborate with a coordinating server to train a machine learning model [14]. It provides opportunities to leverage distributed client data to learn useful models while preserving user privacy. Bringing deep retrieval models to FL is a promising way to power recommendations and learn representations for users and items while addressing privacy concerns by reducing the centralization of user data.

In this work, we observe a challenge with training federated deep retrieval models: the insufficiency of negative examples available on a client’s device. A deep retrieval model (shown in Figure 1) learns embeddings representing users’ contexts and items like movies, songs, or websites. For the model to learn meaningful embeddings, it generally uses two types of examples: positive examples and negative examples. The positive examples pull embeddings of the training (context, item) pairs to be close together in an embedding space, while the negative examples push embeddings of unrelated pairs farther apart. Negative examples are typically required to prevent embedding space collapse, where learned representations collapse to a single point and are no longer informative.

Typically, negative examples are produced by sampling items from training data. In centralized training, training data is usually assumed to be independent and identically distributed (IID). Therefore, negatives can be sampled reliably from the overall training distribution. However, the IID data assumption does not hold for federated learning since clients generate data locally based on their circumstances. That means, on each device, negatives may not be present. Even if they are present, they may be relatively few and relatively similar. In practice, we observe that this leads to significant performance degradation when FL is applied naively, beyond the typical degradation observed for e.g., classification in the FL setting.

This work focuses on understanding and alleviating the non-IID issue for deep retrieval models. We make the following key contributions:

• •

  Observe that naive federated training of deep retrieval models causes an unusually steep performance drop. Show that performance degradation is primarily caused by sampling negatives from non-IID data, not the typical client drift issue or other aspects of federated training.

• •

  Introduce batch-insensitive losses as a class of objectives to alleviate the issue in this setting.

• •

  Perform empirical evaluation of different training objectives in a movie recommendation setting, showing that batch-insensitive losses enable performant federated deep retrieval.

• •

  Release an open-source framework to accelerate further research and practical applications of federated deep retrieval models.¹¹1https://git.io/federated_dual_encoder

Related Work: Most previous research has focused on mitigating the non-IID data issue [6, 23, 10, 11, 12, 21, 20] for improving general model convergence in federated learning settings. These works study the general problem of heterogeneity in client data causing client drift when clients perform multiple local computation steps, as in FedAvg [9, 8]. However, some of them require sharing some training data across clients [6, 23], and most of them study image classification tasks. The developed techniques are not fully relevant to the deep retrieval model, which sees unusually severe performance degradation due to explicit reliance on sampling negatives for training. Though also caused by non-IID data, this issue is orthogonal to the client drift issue: it occurs even when clients perform one local update step. We can combine the techniques explored in this work with previous techniques for addressing client drift. [21] also studies the effectiveness of hinge loss and spreadout regularizer in federated learning. However, it focuses on classification tasks and only considers extreme cases where each user has no access to negative examples. Our work focuses on a more realistic setting, aiming to understand and increase the performance of representation learning and item recommendation with deep retrieval models when each client has access to some but limited and relatively similar negative examples. More broadly, we propose batch-insensitive losses, which can be generally valuable for alleviating non-IID data issues in other federated learning tasks.

Deep retrieval models are also referred to as dual encoder, two-tower, or encoder-encoder models depending on the setting [4, 18, 1, 15, 3]. This general framework has been successfully deployed in a variety of real-world applications in embedding and recommendation learning. As illustrated in Figure 1, a deep retrieval model consists of two encoders, each of which can be a fully-connected network, convolutional neural network, Transformer, and so on, depending on the task. The input consists of (context, item) pairs encoded by the left and right encoders, respectively. For example, in a movie recommendation use-case, a context might be a sequence of previous movies a user has watched, and the item is the next movie they watch. The goals of applying a deep retrieval model in this setting would be to learn a model that can predict the next movie given an unseen sequence of previous movies and learn encoders that produce embeddings for users and movies.

More formally, we denote the feature vectors representing contexts and items as $\mathbf{x}$ and $\mathbf{y}$. The two encoders are denoted as two parameterized functions $f(\cdot)$ and $g(\cdot)$, with $f(\mathbf{x})$ and $g(\mathbf{y})$ mapping $\mathbf{x}$ and $\mathbf{y}$ to a shared embedding space. The model outputs the similarity score between the encoded context and item, e.g., the inner product of context and item embeddings, $s(\mathbf{x},\mathbf{y})=\big{\langle}f(\mathbf{x}),g(\mathbf{y})\big{\rangle}$. A loss function is applied to enforce that positive examples (i.e., similar context and item pairs) have high similarity, and negative examples have low similarity. Once the parameterized functions $f(\cdot)$ and $g(\cdot)$ are learned, the model can predict relevant items given a new context. The representations produced by $f(\cdot)$ and $g(\cdot)$ are general representations of user contexts and items and can also be used for other downstream applications, such as classification [3, 1].

Loss Function: The loss function plays a key role in training a deep retrieval model. The most commonly used one [3, 17, 19] is softmax cross-entropy loss over similarities : $\ell(\mathbf{X}_{i},\mathbf{Y}_{i})=-\log(e^{s(\mathbf{X}_{i},\mathbf{Y}_{i})}/\sum_{\mathbf{Y}_{j}\in\mathcal{N}}e^{s(\mathbf{X}_{i},\mathbf{Y}_{j})})$. $\mathbf{X}$ and $\mathbf{Y}$ represent all the contexts and items in a batch, and $\mathcal{N}$ is a set of negative labels used to construct negative example pairs $(\mathbf{X}_{i},\mathbf{Y}_{j})$. The model is incentivized to maximize the similarity between positive example pairs and minimize the similarity between negative example pairs. Note that if it only did the former, then the embedding space produced by $f(\cdot)$ and $g(\cdot)$ would collapse. Therefore, negative examples are important to learn a good model. A standard method for getting negatives is to use in-batch negatives [19, 17], which means given a training batch and any ($\mathbf{X}_{i}$, $\mathbf{Y}_{i}$) pair in the batch, all other items $\mathbf{Y}_{j,j\neq i}$ in the same batch are treated as negatives for $\mathbf{X}_{i}$. We refer to this as batch softmax below.

Federated training of a deep retrieval model involves three main steps in each training round. First, a central server sends the current model to several randomly sampled clients. Second, each sampled client trains on its dataset and updates its model locally. Finally, the local model updates are sent back to the server and aggregated to update the server model. The de-facto standard federated optimization method is FedAvg [14]. Each client updates its model multiple times before sending the model update back to the server to be averaged. We also later refer to FedSGD, where each client only runs a single local training step at each round, similar to standard distributed training.

As discussed in Section 4.2, naively applying federated learning to this setting produces a steep performance drop, worse than typical when comparing centralized and federated performance. We will show that training degrades due to non-IID negatives. Note that this is distinct from the client drift phenomenon discussed in other works [9], which causes slight performance deterioration when clients take multiple local steps. In contrast, the problem we observe occurs even in the FedSGD regime (see Figure 3). This work aims to characterize this problem better and propose methods that enable federated deep retrieval models to perform comparably to centralized counterparts.

We propose batch-insensitive losses as a potential solution to address the non-IID data issue.

See Appendix A for proof. Proposition 1 shows that FedSGD with $\ell_{BI}$ approximates standard large-batch SGD when the number of clients per round is large (so $\mathcal{E}$ is representative of the centralized data). This motivates using batch-insensitive losses to mitigate the non-IID data issue.

Note that batch softmax loss does not have these properties and is batch-sensitive. We now describe specific batch-insensitive losses for the deep retrieval setting.

We evaluate the efficacy of batch-insensitive losses on the non-IID negatives issue with a movie recommendation task. We train and evaluate a deep retrieval model on the MovieLens 1M dataset ²²2https://grouplens.org/datasets/movielens/1m/. The model takes in a user’s movie-watching history and predicts a relevant next movie for this user.

We are releasing a general code framework for experimenting with federated deep retrieval models built on the TensorFlow Federated library³³3https://www.tensorflow.org/federated. The code is released under Apache License 2.0. The framework enables reproduction of our experiments and provides a flexible, well-documented interface for researchers to train federated and centralized deep retrieval models with different models and losses. We provide libraries for training and evaluation for MovieLens next movie prediction, which can be easily extended for new tasks. We hope that this framework spurs further research and lowers the barrier to more practical applications.

This work investigates the effect of non-IID negatives on federated training of deep retrieval models and proposes batch-insensitive losses to alleviate the issue. We compare model performance using various loss functions and show that batch-insensitive losses produce better federated deep retrieval models that can approximately match centralized models. We also open-source our code framework to accelerate future research and applications. Note that our proposed techniques do not directly address the separate, well-studied issue of client drift when clients do multiple steps of local training–approaches addressing this issue are complementary and can be combined with our work.

We thank Warren Morningstar, Chung-Ching Chang, and Zachary Garrett and for their helpful comments and discussions. We also thank Warren Morningstar for his contribution to the federated training pipeline.