Coordinated Replay Sample Selection for Continual Federated Learning (2023)

Abstract: Continual Federated Learning (CFL) combines Federated Learning (FL), the decentralized learning of a central model on a number of client devices that may not communicate their data, and Continual Learning (CL), the learning of a model from a continual stream of data without keeping the entire history. In CL, the main challenge is forgetting what was learned from past data. While replay-based algorithms that keep a small pool of past training data are effective to reduce forgetting, only simple replay sample selection strategies have been applied to CFL in prior work, and no previous work has explored coordination among clients for better sample selection. To bridge this gap, we adapt a replay sample selection objective based on loss gradient diversity to CFL and propose a new relaxation-based selection of samples to optimize the objective. Next, we propose a practical algorithm to coordinate gradient-based replay sample selection across clients without communicating private data. We benchmark our coordinated and uncoordinated replay sample selection algorithms against random sampling-based baselines with language models trained on a large scale de-identified real-world text dataset. We show that gradient-based sample selection methods both boost performance and reduce forgetting compared to random sampling methods, with our coordination method showing gains early in the low replay size regime (when the budget for storing past data is small).

## introduction the ubiquity of personal devices with a network connection, such as smart phones, watches, and home devices, offer a rich source of data for learning problems such as language modeling or facial recognition. the conventional approach is to collect all the data into one location and use dedicated hardware to learn a model; however, the privacy * work done while the author was an intern at amazon risk associated with communicating personal data makes this approach unsuitable for many applications. federated learning (fl) offers a solution by learning a central model via distributed training across user-owned devices, without communicating any data to the central server. in addition, the devices may produce a continual stream of data and, due to storage constraints and/or privacy restrictions, be able to keep only a limited amount of data at a time. thus continual federated learning (cfl) has recently emerged as a prominent topic in machine learning research. cfl incorporates methods from continual learning (cl), where a model is periodically fine-tuned on new data. the main challenge for cl is catastrophic forgetting, a phenomenon where fine-tuning on new data causes a reduction of performance on past data. this is harmful to long-term generalization, especially when different time periods comprise different tasks, or when the data distribution shifts over time or presents seasonality. among various methods, episodic replay, wherein a small, fixed-size replay buffer of past data is kept and used for fine-tuning along with new data, has proven to be among the most effective strategies to reduce forgetting and improve performance of the final model in both cl @xcite and cfl @xcite . however, only basic replay sample selection strategies, including random sampling and icarl @xcite , have been applied to cfl @xcite . to bridge this gap, we adopt the selection objective from gradient-based sample selection (gss) @xcite , a more recent approach that selects replay samples based on the diversity of their gradients. we propose a new relaxation-based selection method that results in selections closer to optimal compared to methods from prior work. any replay sample selection method from cl can be used for cfl by applying it independently at each client. however, cfl presents a yetunexplored opportunity for the central server to coordinate the selection of replay samples across clients, that is, choose samples such that the union of all clients' replay buffers, rather than each individual buffer, is optimal. the main challenge is that, to ensure privacy, the data cannot be communicated to the server, so selection techniques from cl can not be applied directly. building on our relaxation-based selection approach, we propose the first server-coordinated replay sample selection approach for cfl. by introducing auxiliary variables that make the objective of the relaxation separable across clients, we enable an alternating minimization (more generally called block coordinate descent) process whereby the optimization alternates between the server and the clients in parallel, all while maintaining communication volume and privacy very similar to standard fl training. our novel contributions are 1) a relaxation-based approach to select replay samples that maximize loss gradient diversity; 2) a practical algorithm for coordinated selection of replay samples to maximize gradient diversity jointly across many clients without sacrificing privacy or substantially increasing communication or computation cost; and 3) an empirical analysis of the effect of these strategies on performance and forgetting on a language modeling problem using real-world voice assistant data with heterogeneity across clients and time periods. ## related work fedavg @xcite is a standard fl algorithm wherein the server sends an initial model to a random sample of clients, each client in parallel fine-tunes the model with its local data and sends it back to the server, and the server averages their weights to get a new central model. this is repeated for a number of rounds. if the clients are heterogeneous (have non-i.i.d. data distributions), then the weight averaging results in client drift. as a result, convergence rates of algorithms based on fedavg generally get worse with client heterogeneity @xcite @xcite . several variations of fedavg have been proposed to address challenges such as client drift @xcite @xcite @xcite . the replay sample selection strategies proposed in this paper are orthogonal to the particulars of the fl algorithm; for our evaluation, we use standard fedavg. "continual learning" can refer to several related problems, but in this work, we consider the problem of learning a single task without forgetting from a continual stream of data, usually by periodic fine-tuning, with some limitations such as hardware capacity precluding the retention of the full history of data. the distribution of data may shift over time. common approaches to reduce forgetting are to apply regularization penalizing the difference in weights between the current model previous models @xcite ; keep a small set of historical data and project loss gradients such that they do not increase the loss on these historical data @xcite @xcite ; or keep a small set of historical data to include during training @xcite @xcite . the last approach, called episodic replay or rehearsal, has been shown to be especially effective to reduce forgetting in both cl @xcite and cfl @xcite . in particular, gradient-based sample selection (gss) @xcite ) is an episodic replay strategy that chooses replay samples to maximize the diversity of the loss gradients. it is shown to outperform other strategies and is the foundation for our proposed cfl methods. continual federated learning (cfl) is a setting where each client receives a continual stream of data and federated learning is periodically applied to update a central model. this setting faces challenges of both heterogeneity across clients, as in fl, and heterogeneity across time steps, as in cl. cfl works that focus on improving performance by reducing forgetting, like this one, include the following: @xcite applies model regularization methods from cl to fl, but focuses on improving generalization of fl by reducing client drift; @xcite proposes a general cfl framework with convergence analysis and applies cl techniques including model regularization, generative data augmentation, and episodic replay strategies including naive random sampling and icarl @xcite , finding that episodic replay outperforms the other cl strategies by a wide margin, with the naive method being superior; @xcite uses a distillation strategy with both central and past local models as teachers for new local models; @xcite uses parameter masking to preserve and reuse knowledge; and @xcite proposes a different take on cfl using lightweight models with ensemble methods, focusing mainly on practical limitations of low-power devices, but also discussing applicability to single-task cl problems with distribution shift. to the best of our knowledge, we are the first to apply gradient-based replay sample selection methods to cfl and the first to propose a servercoordinated approach. other cfl works focus on fl challenges such as client interference @xcite or variable sampling rate, device capabilities, latency, and availability issues @xcite . ## problem formulation in fl, each client @xmath0, and we aim to find a model w that solves the optimization problem with @xmath1,t n m,t the total number of samples. since data is generated sequentially and that userowned devices typically have limited storage, at time period t each client only has access to the data generated during t and a small subset of the past data. thus in a cfl setting, we learn a series of models w 1 , . . . , w t , with the goal that w t minimizes (2); each w t , for @xmath2, is trained on x 1,t , . . . , x m,t using federated learning with initialization from w t-1 , except w 1 , which is initialized randomly or pre-trained, e.g., on publicly available data. ## episodic replay strategies for each t, w t is trained on x 1,t , . . . , x m,t , so we may expect that w t minimizes m nm,t nt ℓ(x m,t ; w t ); however, it is not necessarily true that w t minimizes m n m,t n t ℓ(x m,t ; w t ) for @xmath3. episodic replay is a simple and effective remedy whereby, at each time period t, each client m has a replay buffer r m,t containing at most n m data from x m,1 , . . . , x m,t-1 , where n m is the replay buffer size for client m. then w t is trained on x 1,t ∪ r 1,t , . . . , x m,t ∪ r m,t using federated learning. the purpose of the replay buffer is to alleviate forgetting and ultimately result in a good performing model across time periods, and it has been shown in numerous cl @xcite @xcite and cfl @xcite works that episodic replay is effective in accomplishing that. the defining feature of an episodic replay strategy is how r m,t+1 is selected from x m,t ∪ r m,t . we next describe several such sample selection strategies, which we call uncoordinated if the selection is made independently at each client, or coordinated if the selection is made jointly across clients. ## experiments we run experiments to demonstrate the quality of our relaxation-based sample selection and the performance of models trained using cfl with the proposed sample selection strategies. additional experimental details and results are in appendix c. ## discussion we proposed a new relaxation for gradient-based selection of replay samples in continual learning. based on this, we proposed the first algorithm for coordinated replay sample selection in continual federated learning, which converges to the optimal selection under our relaxation while maintaining privacy and low communication cost. our experiments show that, compared to random sampling, the gradient-based selection of replay samples improves performance of the final model for various replay buffer sizes, and coordinated selection improves for small buffer sizes. ## limitations the reproducibility of this work is limited because the data used for some experiments is not public. moreover, training language models in a large cfl setting is extremely demanding of both time and computational resources. this contradicts the optimality of