summaries: To overcome statistical heterogeneity among data shards in collaborative federated learning, this paper proposes a novel personalization schema which only requires smaller memory footprint in clients and possibly is less susceptible to catastrophic forgetting. The main idea is to split the trainable parameters into  shared and personalized parameters where, unlike existing personalization schema, only the shared model is exchanged with the server in communication rounds to be aggregated- thus partial personalization. In regression, this corresponds to learning the residual error of shared model via personalized model and in the classification setting corresponds to output averaging (unlike interpolation based personalization methods that do parameter mixing). The authors propose two algorithms FedSim (simultaneous updating of shared and personalized models locally) and FedAlt (alternative updating of shared and personalized models locally), to learn in this setting, theoretically analyze the convergence rates in non-convex settings and conduct empirical studies on image classification and next-word prediction to evaluate the proposed methods.  This paper studies personalization in federated setting, i.e., instead of collaboratively training a global model, personalizing a model for each client.  This paper proposed to personalize only part of the model parameters instead of the full model, and studied two algorithms FedSim and FedAlt. In local client updates, FedSim will simultaneously train the shared and personalized parameters, while FedAlt will train the personalized model first, then train the shared parameter.  This paper proposed a personalized FL framework with partial model personalization. It separates the model parameters into two parts, shared model and personalized model, and optimize them in an interleaving manner. The authors proposed two optimization algorithm, named as FedSim and FedAlt, with partial clients participation. They also analyzed the algorithms' convergence rate on general smooth nonconvex function. In experiments, they consider 3 different model split methods: 1. input layer as personalized model, the rest as shared model; 2. output layer as personalized model, the rest as shared model; 3. adding adapter as personalized model.  The experiments conduct on NLP or vision tasks also demonstrate that the proposed algorithm outperforms other personalization method. The paper discusses using partial personalization objective function defined in equation (3) to achieve personalized models in federated learning. While the idea of partial personalization and the objective of (3) have been widely studied in previous literatures, the main contributions of the paper are the convergence analysis of two proposed algorithms FedSim and FedAlt in nonconvex case. The authors also did extensive experiments to compare partial personalized models with fully personalized models by using various model structures and on different datasets.
abstractText: We propose and analyze a general framework of federated learning with partial model personalization. Compared with full model personalization, partial model personalization relies on domain knowledge to select a small portion of the model to personalize, thus imposing a much smaller on-device memory footprint. We propose two federated optimization algorithms for training partially personalized models, where the shared and personal parameters are updated either simultaneously or alternately on each device, but only the shared parameters are communicated and aggregated at the server. We give convergence analyses of both algorithms for minimizing smooth nonconvex functions, providing theoretical support of them for training deep learning models. Our experiments on real-world image and text datasets demonstrate that (a) partial model personalization can obtain most of the benefit of full model personalization with a small fraction of personalized parameters, and, (b) the alternating update algorithm often outperforms the simultaneous update algorithm.
references.authors: Martı́n Abadi Andy Chu Ian J. Goodfellow H. Brendan McMahan Ilya Mironov Kunal Talwar Li Zhang
references.title: Deep Learning with Differential Privacy
references.venue: In Proc. of ACM SIGSAC,
references.authors: Alekh Agarwal John Langford Chen-Yu Wei
references.title: Federated Residual Learning
references.venue: arXiv Preprint,
references.authors: Manoj Ghuhan Arivazhagan Vinay Aggarwal Aaditya Kumar Singh Sunav Choudhary
references.title: Federated Learning with Personalization Layers
references.venue: arXiv Preprint,
references.authors: Duc Bui Kshitiz Malik Jack Goetz Honglei Liu Seungwhan Moon Anuj Kumar Kang G. Shin
references.title: Federated User Representation Learning
references.venue: arXiv Preprint,
references.authors: Gregory Cohen Saeed Afshar Jonathan Tapson André van Schaik
references.title: EMNIST: an extension of MNIST to handwritten letters
references.venue: arXiv Preprint,
references.authors: Liam Collins Hamed Hassani Aryan Mokhtari Sanjay Shakkottai
references.title: Exploiting Shared Representations for Personalized Federated Learning
references.venue: In Proc. of ICML,
references.authors: James W. Demmel
references.title: Applied Numerical Linear Algebra
references.venue: SIAM, Philadelphia,
references.authors: Jia Deng Wei Dong Richard Socher Li-Jia Li Kai Li Fei-Fei Li
references.title: ImageNet: A large-scale hierarchical image database
references.venue: In Proc. of CVPR,
references.authors: Yuyang Deng Mohammad Mahdi Kamani Mehrdad Mahdavi
references.title: Adaptive Personalized Federated Learning
references.venue: arXiv Preprint,
references.authors: Canh T. Dinh Nguyen Tran Josh Nguyen
references.title: Personalized Federated Learning with Moreau Envelopes
references.venue: In Proc. of NeurIPS,
references.authors: Alireza Fallah Aryan Mokhtari Asuman E. Ozdaglar
references.title: Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach
references.venue: In Proc. of NeurIPS,
references.authors: Filip Hanzely Boxin Zhao Mladen Kolar
references.title: Personalized Federated Learning: A Unified Framework and Universal Optimization Techniques
references.venue: arXiv Preprint,
references.authors: Kaiming He Xiangyu Zhang Shaoqing Ren Jian Sun
references.title: Deep Residual Learning for Image Recognition
references.venue: In Proc. of CVPR,
references.authors: Neil Houlsby Andrei Giurgiu Stanislaw Jastrzebski Bruna Morrone Quentin de Laroussilhe Andrea Gesmundo Mona Attariyan Sylvain Gelly
references.title: Parameter-Efficient Transfer Learning for NLP
references.venue: In Proc. of ICML,
references.authors: Kevin Hsieh Amar Phanishayee Onur Mutlu Phillip B. Gibbons
references.title: The Non-IID Data Quagmire of Decentralized Machine Learning
references.venue: In Proc. of ICML,
references.authors: Tzu-Ming Harry Hsu Hang Qi Matthew Brown
references.title: Federated Visual Classification with Real-World Data Distribution
references.venue: In Proc. of ECCV,
references.authors: Sai Praneeth Karimireddy Satyen Kale Mehryar Mohri Sashank Reddi Sebastian Stich Ananda Theertha Suresh
references.title: SCAFFOLD: Stochastic controlled averaging for federated learning
references.venue: In Proc. of ICML,
references.authors: James Kirkpatrick Razvan Pascanu Neil Rabinowitz Joel Veness Guillaume Desjardins Andrei A Rusu Kieran Milan John Quan Tiago Ramalho Agnieszka Grabska-Barwinska Demis Hassabis Claudia Clopath Dharshan Kumaran Raia Hadsell
references.title: Overcoming catastrophic forgetting in neural networks
references.venue: Proceedings of the National Academy of Sciences,
references.authors: Anastasia Koloskova Nicolas Loizou Sadra Boreiri Martin Jaggi Sebastian Stich
references.title: A Unified Theory of Decentralized SGD with Changing Topology and Local Updates
references.venue: In Proc. of ICML,
references.authors: Tian Li Shengyuan Hu Ahmad Beirami Virginia Smith
references.title: Ditto: Fair and Robust Federated Learning Through Personalization
references.venue: In Proc. of ICML,
references.authors: Paul Pu Liang Terrance Liu Ziyin Liu Ruslan Salakhutdinov Louis-Philippe Morency
references.title: Think Locally, Act Globally: Federated Learning with Local and Global Representations
references.venue: In NeurIPS Workshop on Federated Learning,
references.authors: Yishay Mansour Mehryar Mohri Jae Ro Ananda Theertha Suresh
references.title: Three Approaches for Personalization with Applications to Federated Learning
references.venue: arXiv Preprint,
references.authors: Michael McCloskey Neal J. Cohen
references.title: Catastrophic Interference in Connectionist Networks: The Sequential Learning Problem. volume 24 of Psychology of Learning and Motivation, pp. 109–165
references.authors: Brendan McMahan Eider Moore Daniel Ramage Seth Hampson Blaise Agüera y Arcas
references.title: Communication-Efficient Learning of Deep Networks from Decentralized Data
references.venue: In Proc. of AISTATS,
references.authors: Tomas Mikolov Geoffrey Zweig
references.title: Context dependent recurrent neural network language model
references.venue: In IEEE SLT,
references.authors: Sylvestre-Alvise Rebuffi Hakan Bilen Andrea Vedaldi
references.title: Learning multiple visual domains with residual adapters
references.venue: In Proc. of NeurIPS,
references.authors: Sashank J. Reddi Zachary Charles Manzil Zaheer Zachary Garrett Keith Rush Jakub Konečný Sanjiv Kumar Hugh Brendan McMahan
references.title: Adaptive Federated Optimization
references.venue: In Proc. of ICLR,
references.authors: Karan Singhal Hakim Sidahmed Zachary Garrett Shanshan Wu Keith Rush Sushant Prakash
references.title: Federated reconstruction: Partially local federated learning
references.venue: In Proc. of NeurIPS,
references.authors: Virginia Smith Chao-Kai Chiang Maziar Sanjabi Ameet S Talwalkar
references.title: Federated Multi-Task Learning
references.venue: In Proc. of NeurIPS,
references.authors: Iulia Turc Ming-Wei Chang Kenton Lee Kristina Toutanova
references.title: Well-read students learn better: On the importance of pre-training compact models
references.venue: arXiv Preprint,
references.authors: Ashish Vaswani Noam Shazeer Niki Parmar Jakob Uszkoreit Llion Jones Aidan N. Gomez Lukasz Kaiser Illia Polosukhin
references.title: Attention is All you Need
references.venue: In Proc. of NeurIPS,
references.authors: Tobias Weyand André Araujo Bingyi Cao Jack Sim
references.title: Google Landmarks Dataset v2 - A Large-Scale Benchmark for Instance-Level Recognition and Retrieval
references.venue: In Proc. of CVPR,
references.authors: Houlsby
references.title: 2019) and residual adapters proposed by Rebuffi et al
references.authors: Reddi
references.title: 2021), we restrict the vocabulary to the top 10000 most frequently occurring words in the dataset. We pad and truncate each sequence of each client to length 20 and consider at most 1000 training sequences on each client. Model. We use a transformer model (Vaswani et al., 2017) commensurate in size with BERT
references.venue: Mini (Turc et al.,
sections.heading: 1 INTRODUCTION
sections.text: Federated Learning (McMahan et al., 2017) has emerged as a powerful paradigm for distributed and privacy-preserving machine learning over a large number of edge devices (see Kairouz et al., 2021, and references therein). We consider a typical setting of Federated Learning (FL) with n devices (also called clients), where each device i has a training dataset of Ni samples zi,1, · · · , zi,Ni . Let w ∈ Rd represent the parameters of a (supervised) learning model and fi(w, zi,j) be the loss of the model on the training example zi,j . Then the loss function associated with device i is Fi(w) = (1/Ni) ∑Ni j=1 fi(w, zi,j). A common objective of FL is to find model parameters that minimize the weighted average loss across all devices (without transferring the datasets):
  minimize w n∑ i=1 αiFi(w), (1)
  where weights αi are nonnegative and satisfy ∑n i=1 αi = 1. A common practice is to choose the
  weights as αi = Ni/N where N = ∑n k=1Nk, which corresponds to minimizing the unweighted
  average loss across all samples from the n devices: (1/N) ∑n i=1 ∑Ni j=1 fi(w, zi,j).
  The main motivation for minimizing the average loss over all devices is to leverage their collective statistical power for better generalization, because the amount of data on each device can be very limited. This is especially important for training modern deep learning models with large number of parameters. However, this argument assumes that the datasets from different devices are sampled from the same, or at least very similar, distributions. Given the diverse characteristics of the users and increasing trend of personalized on-device services, such an i.i.d. assumption may not hold in practice. Thus, the one-model-fits-all formulation in (1) can be less effective and even undesirable.
  Several approaches have been proposed for personalized FL, including ones based on multi-task learning (Smith et al., 2017), meta learning (Fallah et al., 2020), and proximal methods (Dinh et al., 2020; Li et al., 2021). A simple formulation that captures their main idea is
  minimize w0,{wi}ni=1 n∑ i=1 αi ( Fi(wi) + λi 2 ‖wi − w0‖2 ) , (2)
  where wi for i = 1, . . . , n are personalized model parameters at the devices, w0 is a reference model maintained by the server, and the λi’s are regularization weights that control the extent of personalization. A major disadvantage of the formulation (2), which we call full model personalization, is that it requires twice the memory footprint of the model, wi and w0 at each device, which severely limits the size of trainable models. On the other hand, the flexibility of full model personalization can be unnecessary. Modern deep learning models are composed of many simple functional units and are typically organized into layers or a more general interconnected architecture. Personalizing the “right” components, selected with domain knowledge, may result in a substantial benefit with only a small increase in memory footprint. In addition, partial model personalization can be less susceptible to “catastrophic forgetting” (McCloskey & Cohen, 1989), where a large model finetuned on a small local dataset forgets the original (non-personalized) task, leading to a degradation of test performance.
  We propose a framework for FL with partial model personalization. Specifically, we partition the model parameters into two groups: the shared parameters u ∈ Rd0 and the personal parameters vi ∈ Rdi for i = 1, . . . , n. The full model on device i is denoted as wi = (u, vi), and the local loss function is Fi(u, vi) = (1/Ni) ∑Ni i=1 fi ( (u, vi), zi,j ) . Our goal is to solve the optimization problem
  minimize u, {vi}ni=1 n∑ i=1 αiFi(u, vi). (3)
  Notice that the dimensions of vi can be different across the devices, allowing the personal components of the model to have different number of parameters or even different architecture.
  We investigate two FL algorithms for solving problem (3): FedSim, a simultaneous update algorithm and FedAlt, an alternating update algorithm. Both algorithms follow the standard FL protocol. During each round, the server randomly selects a subset of the devices for update and broadcasts the current global version of the shared parameters to devices in the subset. Each selected device then performs one or more steps of (stochastic) gradient descent to update both the shared parameters and the personal parameters, and sends the updated shared parameters to the server for aggregation. The updated personal parameters are kept local at the device to serve as the initial states when the device is selected for another update. In FedSim, the shared and personal parameters are updated simultaneously during each local iteration. In FedAlt, the devices first update the personal parameters with the received shared parameters fixed and then update the shared parameters with the new personal parameters fixed. We provide convergence analysis and empirical evaluation of both methods.
  The main contributions of this paper are summarized as follows:
  • We propose a general framework of FL with partial model personalization, which relies on domain knowledge to select a small portion of the model to personalize, thus imposing a much smaller memory footprint on the devices than full model personalization. This framework unifies existing work on personalized FL and allows arbitrary partitioning of deep learning models.
  • We provide convergence guarantees for the FedSim and FedAlt methods in the general (smooth) nonconvex setting. While both methods have appeared in the literature previously, they are either used without convergence analysis or with results on limited settings (assuming convexity or full participation) Our analysis provides theoretical support for the general nonconvex setting with partial participation. The analysis of FedAlt with partial participation is especially challenging and we develop a novel technique of virtual full participation to overcome the difficulties.
  • We conduct extensive experiments on image classification and text prediction tasks, exploring different model personalization strategies for each task, and comparing with several strong baselines. Our results demonstrate that partial model personalization can obtain most of the benefit of full model personalization with a small fraction of personalized parameters, and FedAlt often outperforms FedSim.
  • Our experiments also reveal that personalization (full or partial) may lead to worse performance for some devices, despite improving the average. Typical forms of regularization such as weight decay and dropout do not mitigate this issue. This phenomenon has been overlooked in previous work and calls for future research to improve both performance and fairness.
  Related work. Specific forms of partial model personalization have been considered in previous works. Liang et al. (2019) propose to personalize the input layers to learn a personalized representation per-device (Figure 1a), while Arivazhagan et al. (2019) and Collins et al. (2021) propose to personalize the output layer while learning a shared representation with the input layers (Figure 1b). Both FedSim and FedAlt have appeared in the literature before, but the scope of their convergence analysis is limited. Specifically, Liang et al. (2019), Arivazhagan et al. (2019) and Hanzely et al. (2021) use FedSim, while Collins et al. (2021) and Singhal et al. (2021) proposed variants of FedAlt. Notably, Hanzely et al. (2021) establish convergence of FedSim with full device participation in the convex and non-convex cases, while Collins et al. (2021) prove the linear convergence of FedAlt for a two-layer linear network where Fi(·, vi) and Fi(u, ·) are both convex for fixed vi and u respectively. We analyze both FedSim and FedAlt in the general nonconvex case with partial device participation, hence addressing a more general and practical setting.
  While we primarily consider the problem (3) in the context of partial model personalization, it can serve as a general formulation that covers many other problems. Hanzely et al. (2021) demonstrate that various full model personalization formulations based on regularization (Dinh et al., 2020; Li et al., 2021), including (2), as well as interpolation (Deng et al., 2020; Mansour et al., 2020) are special cases of this problem. The rates of convergence we prove in §3 are competitive with or better than those in previous works for full model personalization methods in the non-convex case.
sections.heading: 2 PARTIALLY PERSONALIZED MODELS
sections.text: Modern deep learning models all have a multi-layer architecture. While a complete understanding of why they work so well is still out of reach, a general insight is that the lower layers (close to the input) are mostly responsible for feature extraction and the upper layers (close to the output) focus on complex pattern recognition. Depending on the application scenarios and domain knowledge, we may personalize either the input layer(s) or the output layer(s) of the model; see Figure 1.
  In Figure 1c, the input layers are split horizontally into two parts, one shared and the other personal. They process different chunks of the input vector and their outputs are concatenated before feeding
  Algorithm 1 Federated Learning with Partial Model Personalization (FedSim / FedAlt)
  Input: initial states u(0), {v(0)i }ni=1, number of rounds T , number of devices per round m 1: for t = 0, 1, · · · , T − 1 do 2: server randomly samples m devices as S(t) ⊂ {1, . . . , n} 3: server broadcasts u(t) to each device in S(t) 4: for each device i ∈ S(t) in parallel, do 5: ( u
  (t+1) i , v (t+1) i
  ) = LocalSim / LocalAlt ( u(t), v
  (t) i
  ) . v
  (t+1) i = v (t) i if i /∈ S(t)
  6: send u(t+1)i back to server 7: server updates u(t+1) = 1m ∑ i∈S(t) u (t+1) i
  to the upper layers of the model. As demonstrated in (Bui et al., 2019), this partitioning can help protect user-specific private features (input 2 in Figure 1c) as the corresponding feature embedding (through vi) are personalized and kept local at the device. Similar architectures have also been proposed in context-dependent language models (e.g., Mikolov & Zweig, 2012).
  A more structured partitioning is illustrated in Figure 2a, where a typical transformer layer (Vaswani et al., 2017) is augmented with two adapters. This architecture is proposed by Houlsby et al. (2019) for finetuning large language models. Similar residual adapter modules are proposed by Rebuffi et al. (2017) for image classification models in the context of multi-task learning. In the context of FL, we treat the adapter parameters as personal and the rest of the model parameters as shared.
  Figure 2b shows a generalized additive model, where the outputs of two separate models, one shared and the other personalized, are fused to generate a prediction. Suppose the shared model is h(u, ·) and the personal model is hi(vi, ·). For regression tasks with samples zi,j = (xi,j , yi,j), where xi,j is the input and yi,j ∈ Rp is the output, we let Fi(u, vi) = (1/Ni) ∑Ni j=1 fi ( (u, vi), zi,j ) with
  fi ( (u, vi), zi,j ) = ‖yi,j − h(u, xi,j)− hi(vi, xi,j)‖2 .
  In this special case, the personal model fits the residual of the shared model and vice-versa (Agarwal et al., 2020). For classification tasks, h(u, ·) and hi(vi, ·) produce probability distributions over multiple classes. We can use the cross-entropy loss between yi,j and a convex combination of the two model outputs: θh(u, xi,j) + (1− θ)hi(vi, xi,j), where θ ∈ (0, 1) is a learnable parameter. Finally, we can cast the formulation (2) of full model personalization as a special case of (3) by letting
  u← w0, vi ← wi, Fi(u, vi)← Fi(vi) + (λi/2)‖vi − u‖2. Many other formulations of full personalization can be reduced to (3); see Hanzely et al. (2021).
sections.heading: 3 ALGORITHMS AND CONVERGENCE ANALYSIS
sections.text: In this section, we present and analyze two FL algorithms for solving problem (3). To simplify presentation, we denote V = (v1, . . . , vn) ∈ Rd1+...+dn and focus on the case of αi = 1/n, i.e.,
  minimizeu, V F (u, V ) := 1 n ∑n i=1 Fi(u, vi). (4)
  This is equivalent to (3) if we scale Fi by nαi, thus does not lose generality. Moreover, we consider more general local functions Fi(u, vi) = Ez∼Di [fi((u, vi), z)], where Di is the local distribution. The FedSim and FedAlt algorithms share a common outer-loop description given in Algorithm 1. They differ only in the local update procedures LocalSim and LocalAlt, which are given in Algorithm 2 and Algorithm 3 respectively. In the two local update procedures, ∇̃u and ∇̃v represent stochastic gradients with respect to w and vi respectively. In LocalSim (Algorithm 2), the personal variables vi and local version of the shared parameters ui are updated simultaneously, with their (stochastic) partial gradients evaluated at the same point. In LocalAlt (Algorithm 3), the personal parameters are updated first with the received shared parameters fixed, then the shared parameters are updated with the new personal parameters fixed. They are analogous to the classical Jacobi update and Gauss-Seidel update in numerical linear algebra (e.g., Demmel, 1997, §6.5).
  In order to analyze the convergence of the two algorithms, we make the following assumptions.
  Algorithm 2 LocalSim ( u, vi ) Input: number of steps τ , step sizes γv and γu
  1: initialize vi,0 = vi 2: initialize ui,0 = u 3: for k = 0, 1, · · · , τ − 1 do 4: vi,k+1 = vi,k − γv∇̃vFi ( ui,k, vi,k
  ) 5: ui,k+1 = ui,k − γu∇̃uFi ( ui,k, vi,k
  ) 6: update v+i = vi,τ 7: update u+i = ui,τ 8: return ( u+i , v + i )
  Algorithm 3 LocalAlt ( u, vi ) Input: number of steps τv, τu, step sizes γv, γu
  1: initialize vi,0 = vi 2: for k = 0, 1, · · · , τv−1 do 3: vi,k+1 = vi,k − γv∇̃vFi ( u, vi,k ) 4: update v+i = vi,τv and initialize ui,0 = u 5: for k = 0, 1, · · · , τu−1 do 6: ui,k+1 = ui,k − γu∇̃uFi ( ui,k, v + i
  ) 7: update u+i = ui,τu 8: return ( u+i , v + i
  ) Assumption 1 (Smoothness). The function Fi is continuously differentiable for each i = 1, . . . , n, and there exist constants Lu, Lv , Luv and Lvu such that for each i = 1, . . . , n, it holds that
  • ∇uFi(u, vi) is Lu-Lipschitz with respect to u and Luv-Lipschitz with respect to vi; • ∇vFi(u, vi) is Lv-Lipschitz with respect to vi and Lvu-Lipschitz with respect to u.
  Due to the definition of F (u, V ) in (4), it is easy to verify that∇uF (u, V ) has Lipschitz constant Lu with respect to u, Luv/ √ n with respect to V , and Luv/n with respect to any vi. We also define
  χ := max{Luv, Lvu} /√
  LuLv, (5) which measures the relative cross-sensitivity of ∇uFi with respect to vi and ∇vFi with respect to u. Assumption 2 (Bounded Variance). The stochastic gradients in Algorithm 2 and Algorithm 3 are unbiased and have bounded variance. That is, for all u and vi,
  E [ ∇̃uFi(u, vi) ] = ∇uFi(u, vi), E [ ∇̃vFi(u, vi) ] = ∇vFi(u, vi) .
  Furthermore, there exist constants σu and σv such that E [∥∥∇̃uFi(u, vi)−∇uFi(u, vi)∥∥2] ≤ σ2u , E[∥∥∇̃vFi(u, vi)−∇vFi(u, vi)∥∥2] ≤ σ2v .
  We can view ∇uFi(u, vi), when i is randomly sampled from {1, . . . , n}, as a stochastic partial gradient of F (u, V ) with respect to u. The following assumption imposes a variance bound. Assumption 3 (Partial Gradient Diversity). There exist δ ≥ 0 and ρ ≥ 0 such that for all u and V ,
  1 n ∑n i=1 ∥∥∇uFi(u, vi)−∇uF (u, V )∥∥2 ≤ δ2 + ρ2∥∥∇uF (u, V )∥∥2 . With ρ = 0, this assumption is similar to a constant variance bound on the stochastic gradient ∇uFi(u, vi); with ρ > 0, it allows the variance to grow with the norm of the full gradient.
  Throughout this paper, we assume F is bounded below by F ? and denote ∆F0 = F ( u(0), V (0) ) −F ?. Further, we use shorthand V (t) = (v(t)1 , . . . , v (t) n ) and
  ∆ (t) u = ∥∥∇uF (u(t), V (t))∥∥2 , and ∆(t)v = 1n∑ni=1∥∥∇vFi(u(t), v(t)i )∥∥2 . For smooth and nonconvex loss functions Fi, we obtain convergence in expectation to a stationary point of F if the expected values of these two sequences converge to zero.
  We first present our main result for FedSim (Algorithm 1 with LocalSim), proved in Appendix A.2. Theorem 1 (Convergence of FedSim). Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedSim are chosen as γu = η/(Luτ) and γv = η/(Lvτ) with
  η ≤ min {
  1 12(1+χ2)(1+ρ2) ,
  √ m/n
  196(1−τ−1)(1+χ2)(1+ρ2)
  } .
  Then, ignoring absolute constants, we have
  1 T ∑T−1 t=0 ( 1 Lu E [ ∆ (t) u ] + mnLv E [ ∆ (t) v ]) ≤ ∆F0ηT + η(1 + χ 2) ( σ2u+δ 2(1−mn ) mLu + mσ2v nLv ) + η2(1− τ−1)(1 + χ2) ( σ2u+δ 2
  Lu + σ2v Lv
  ) . (6)
  The left-hand side of (6) is the average over time of a weighted sum of E [ ∆ (t) u ] and E [ ∆ (t) v ] . The right-hand side contains three terms of order O(1/(ηT )), O(η) and O(η2) respectively. We can minimize the right-hand side by optimizing over η. By considering special cases such as σ2u = σ 2 v = 0 and m = n, some terms on the right-hand side disappear and we can obtain improved rates. Table 1 shows the results in several different regimes along with the optimal choices of η.
  Challenge in Analyzing FedAlt. We now turn to FedAlt. Note that the personal parameters are updated only for the m selected devices in S(t) in each round t. Specifically,
  v (t+1) i =
  { v
  (t) i − γv ∑τv k=0 ∇̃vFi ( u(t), v (t) i,k ) if i ∈ S(t), v (t) i if i /∈ S(t).
  Consequently, the vector V (t+1) of personal parameters depends on the random variable S(t). This makes it challenging to analyze the u-update steps of FedAlt because they are performed after V (t+1) is generated (as opposed to simultaneously in FedSim). When we take expectations with respect to the sampling of S(t) in analyzing the u-updates, V (t+1) becomes a dependent random variable, which prevents standard proof techniques from going through (see details in Appendix A.3).
  We develop a novel technique called virtual full participation to overcome this challenge. Specifically, we define a virtual vector Ṽ (t+1), which is the result if every device were to perform local v-updates. It is independent of the sampling of S(t) and we can derive a convergence rate for related quantities. We carefully translate this rate from the virtual Ṽ (t+1) to the actual V (t) to get the following result. Theorem 2 (Convergence of FedAlt). Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedAlt are chosen as γu = η/(Luτu) and γv = η/(Lvτv), with
  η ≤ min {
  1 24(1+ρ2) , m 128χ2(n−m) , √ m χ2n } .
  Then, ignoring absolute constants, we have
  1 T ∑T−1 t=0 ( 1 Lu E [ ∆ (t) u ] + mnLv E [ ∆ (t) v ]) ≤ ∆F0
  ηT + η
  ( σ2u+δ
  2(1−mn ) mLu + σ2v Lv m+χ2(n−m) n ) + η2 ( σ2u+δ 2
  Lu (1− τ−1u ) + σ2vm Lvn (1− τ−1v ) + χ2σ2v Lv
  ) .
  The proof of Theorem 2 is given in Appendix A.3. Similar to the results for FedSim, we can choose η to minimize the above upper bound to obtain the best convergence rate, as summarized in Table 1.
  Comparing FedSim and FedAlt. Table 1 shows that both FedSim and FedAlt exhibit the standard O(1/ √ T ) rate in the general case. Comparing the constants in their rates, we identify two regimes in terms of problem parameters. The regime where FedAlt dominates FedSim is characterized by σ2v Lv ( 1− 2mn ) < σ2u+δ 2(1−m/n) mLu . A practically relevant scenario where this is true is σ2v ≈ 0 and σ2u ≈ 0 from using large or full batch on a small number of samples per device. Here, the rate of FedAlt is better than FedSim by a factor of (1 + χ2), indicating that the rate of FedAlt is less affected by the coupling between the personal and shared parameters. Our experiments in §4 corroborate the practical relevance of this regime.
  The rates from Table 1 also apply for full personalization schemes without convergence guarantees in the nonconvex case (Agarwal et al., 2020; Mansour et al., 2020; Li et al., 2021). Our rates are better than those of (Dinh et al., 2020) for their pFedMe objective.
sections.heading: 4 EXPERIMENTS
sections.text: In this section, we experimentally compare different model personalization schemes using FedAlt and FedSim as well as no model personalization. Details about the experiments, hyperparameters and additional results are provided in the appendices. The code to reproduce the experimental results will be publicly released.
  Datasets, Tasks and Models. We consider three learning tasks; they are summarized in Table 2.
  (a) Next-Word Prediction: We use the StackOverflow dataset, where each device corresponds to the questions and answers of one user on stackoverflow.com. This task is representative of mobile keyboard predictions. We use a 4-layer transformer model (Vaswani et al., 2017).
  (b) Visual Landmark Recognition: We use the GLDv2 dataset (Weyand et al., 2020; Hsu et al., 2020), a large-scale dataset with real images of global landmarks. Each device corresponds to a Wikipedia contributor who uploaded images. This task resembles a scenario where smartphone users capture images of landmarks while traveling. We use a ResNet-18 (He et al., 2016) model with group norm instead of batch norm (Hsieh et al., 2020) and images are reshaped to 224× 224.
  (c) Character Recognition: We use the EMNIST dataset (Cohen et al., 2017), where the input is a 28 × 28 grayscale image of a handwritten character and the output is its label (0-9, a-z, A-Z). Each device corresponds to a writer of the character. We use a ResNet-18 model, with input and output layers modified to accommodate the smaller image size and number of classes.
  All models are trained with the cross entropy loss and evaluated with top-1 accuracy of classification.
  Model Partitioning for Partial Personalization. We consider three partitioning schemes.
  (a) Input layer personalization: This architecture learns a personalized representation per-device by personalizing the input layer, while the rest of the model is shared (Figure 1a). For the transformer, we use the first transformer layer in place of the embedding layer.
  (b) Output layer personalization: This architecture learns a shared representation but personalizes the prediction layer (Figure 1b). For the transformer model, we use the last transformer layer instead of the output layer.
  (c) Adapter personalization: In this architecture, each device adds lightweight personalized adapter modules between specific layers of a shared model (Figure 2a). We use the transformer adapters of Houlsby et al. (2019) and for ResNet-18, the residual adapters of Rebuffi et al. (2017).
  Algorithms and Experimental Pipeline. For full model personalization, we consider three baselines: (i) Finetune, where each device finetunes (using SGD locally) its personal full model starting from a learned common model, (ii) Ditto (Li et al., 2021), which is finetuning with `2 regularization, and, (iii) pFedMe (Dinh et al., 2020) which minimizes the objective (2). All methods, including FedSim, FedAlt and the baselines are initialized with a global model trained with FedAvg.
sections.heading: 4.1 EXPERIMENTAL RESULTS
sections.text: Partial personalization nearly matches full personalization and can sometimes outperform it. Table 3 shows the average test accuracy across all devices of different FL algorithms. We see that on the StackOverflow dataset, output layer personalization (25.05%) makes up nearly 90% of the gap between the non-personalized baseline (23.82%) and full personalization (25.21%). On EMNIST, adapter personalization exactly matches full personalization. Most surprisingly, on GLDv2, adapter personalization outperforms full personalization by 3.5pp (percentage points).
  This success of adapter personalization can be explained partly by the nature of GLDv2. On average, the training data on each device contains 25 classes out of a possible 2028 while the testing data contains 10 classes not seen in its own training data. These unseen classes account for nearly 23% of all testing data. Personalizing the full model is susceptible to “forgetting” the original task (Kirkpatrick et al., 2017), making it harder to get these unseen classes right. Such catastrophic forgetting is worse when finetuning on a very small local dataset, as we often have in FL. On the other hand, personalizing the adapters does not suffer as much from this issue (Rebuffi et al., 2017).
  Partial personalization only requires a fraction of the parameters to be personalized. Figure 3 shows that the number of personalized parameters required to compete with full model personalization is rather small. On StackOverflow, personalizing 1.2% of the parameters with adapters captures 72% of the accuracy boost from personalizing all 5.7M parameters; this can be improved to nearly 90% by personalizing 14% of the parameters (output layer). Likewise, we match full personalization on EMNIST and exceed it on GLDv2 with adapters, personalizing 11.5-12.5% of parameters.
  The best personalized architecture is model and task dependent. Table 3 shows that personalizing the final transformer layer (denoted as “Output Layer”) achieves the best performance for StackOverflow, while the residual adapter achieves the best performance for GLDv2 and EMNIST. This shows that the approach of personalizing a fixed model part, as in several past works, is suboptimal. Our framework allows for the use of domain knowledge to determine customized personalization.
  Finetuning is competitive with other full personalization methods. Full finetuning matches the performance of pFedMe and Ditto on StackOverflow and EMNIST. On GLDv2, however, pFedMe outperforms finetuning by 0.07pp, but is still 3.5pp worse than adapter personalization.
  FedAlt outperforms FedSim for partial personalization. If the optimization problem (3) were convex, we would expect similar performance from FedAlt and FedSim. However, with nonconvex optimization problems such as the ones considered here, the choice of the optimization algorithm often affects the quality of the solution found. We see from Table 4 that FedAlt is almost always better than FedSim by a small margin, e.g., 0.08pp for StackOverflow/Adapter and 0.3pp for GLDv2/Input Layer. FedSim in turn yields a higher accuracy than simply finetuning the personalized part of the model, by a large margin, e.g., 0.12pp for StackOverflow/Output Layer and 2.55pp for GLDv2/Adapter.
sections.heading: 4.2 EFFECTS OF PERSONALIZATION ON PER-DEVICE GENERALIZATION
sections.text: Personalization hurts the test accuracy on some devices. Figure 4 shows the change in training and test accuracy of each device, compared with a non-personalized model trained by FedAvg. We see that personalization leads to an improvement in training accuracy across all devices, but a reduction in test accuracy on some of the devices over the non-personalized baseline. In particular, devices whose testing performance is hurt by personalization are mostly on the left side of the plot, meaning that they have relatively small number of training samples. On the other hand, many devices with the most improved test accuracy also appear on the left side, signaling the benefit of personalization. Therefore, there is a large variation of results for devices with few samples.
  Additional experiments (see Appendix C) show that using `2 regularization, as in (2), or weight decay does not mitigate this issue. In particular, increasing regularization strength (less personalization) can reduce the spread of per-device accuracy, but only leads to a worse average accuracy that is close to using a common model. Other simple strategies such as dropout also do not fix this issue.
  An ideal personalized method would boost performance on most of the devices without causing a reduction in (test) accuracy on any device. Realizing this goal calls for a sound statistical analysis for personalized FL and may require sophisticated methods for local performance diagnosis and more structured regularization. These are very promising directions for future research.
sections.heading: 5 DISCUSSION
sections.text: In addition to a much smaller memory footprint than full model personalization and being less susceptible to catastrophic forgetting, partial model personalization has other advantages. For example, it reduces the amount communication between the server and the devices because only the shared parameters are transmitted. While the communication saving may not be significant (especially when the personal parameters are only a small fraction of the full model), communicating only the shared parameters may have significant implications for privacy. Intuitively, it can be harder to infer private information from partial model information. This is especially the case if the more sensitive features of the data are processed through personal components of the model that are kept local at the devices. For example, we speculate that less noise needs to be added to the communicated parameters in order to satisfy differential privacy requirements (Abadi et al., 2016). This is a very promising direction for future research.
sections.heading: REPRODUCIBILITY STATEMENT
sections.text: For theoretical results, we state and discuss the assumptions in Appendix A. The full proofs of all theoretical statements are also given there.
  For our numerical results, we take multiple steps for reproducibility. First, we run each numerical experiment for five random seeds, and report both the mean and standard deviation over these runs. Second, we only use publicly available datasets and report the preprocessing at length in Appendix B. Third, we give the full list of hyperparameters used in our experiments in Table 8 in Appendix B. Finally, we will publicly release the code to reproduce the our experimental results.
sections.heading: ETHICS STATEMENT
sections.text: The proposed framework for partial model personalization is immediately applicable for a range of practical federated learning applications in edge devices such as text prediction and speech recognition. One of key considerations of federated learning is privacy. Partial model personalization maintains the all the privacy benefits of current non-personalized federated learning systems. Indeed, our approach is compatible with techniques to enhance privacy such as differential privacy and secure aggregation. We also speculate that partial personalization has the potential for further reducing the privacy footprint — an investigation of this subject is beyond the scope of this work and is an interesting direction for future work.
  On the flip side, we also observed in experiments that personalization (both full or partial) leads to a reduction in test performance on some of the devices. This has important implications for fairness, and calls for further research into the statistical aspects of personalization, performance diagnostics as well as more nuanced definitions of fairness in federated learning.
sections.heading: Appendix
sections.heading: Table of Contents
sections.heading: A Convergence Analysis: Full Proofs 14
sections.text: A.1 Review of Setup and Assumptions . . . . . . . . . . . . . . . . . . . . . . . . 14 A.2 Convergence Analysis of FedSim . . . . . . . . . . . . . . . . . . . . . . . . . 15 A.3 Convergence Analysis of FedAlt . . . . . . . . . . . . . . . . . . . . . . . . . 22 A.4 Technical Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29
sections.heading: B Experiments: Detailed Setup and Hyperparameters 31
sections.text: B.1 Datasets, Tasks and Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 B.2 Experimental Pipeline and Baselines . . . . . . . . . . . . . . . . . . . . . . . 34 B.3 Hyperparameters and Evaluation Details . . . . . . . . . . . . . . . . . . . . . 35
sections.heading: C Experiments: Additional Results 36
sections.text: C.1 Ablation: Final Finetuning for FedAlt and FedSim . . . . . . . . . . . . . . . . 36 C.2 Effect of Personalization on Per-Device Generalization . . . . . . . . . . . . . . 37 C.3 Partial Personalization for Stateless Devices . . . . . . . . . . . . . . . . . . . 38
sections.heading: A CONVERGENCE ANALYSIS: FULL PROOFS
sections.text: We give the full convergence proofs here. The outline of this section is:
  • §A.1: Review of setup and assumptions; • §A.2: Convergence analysis of FedSim and the full proof of Theorem 1; • §A.3: Convergence analysis of FedAlt and the full proof of Theorem 2; • §A.4: Technical lemmas used in the analysis.
sections.heading: A.1 REVIEW OF SETUP AND ASSUMPTIONS
sections.text: We consider a federated learning system with n devices. Let the loss function on device i be Fi(u, vi), where u ∈ Rd0 denotes the shared parameters across all devices and vi ∈ Rdi denotes the personal parameters at device i. We aim to minimize the function
  F (u, V ) := 1
  n n∑ i=1 Fi(u, vi) , (7)
  where V = (v1, · · · , vn) is a concatenation of all the personalized parameters. This is a special case of (3) with the equal per-device weights, i.e., αi = 1/n. Recall that we assume that F is bounded from below by F ?.
  For convenience, we reiterate Assumptions 1, 2 and 3 from the main-paper as Assumptions 1′, 2′ and 3′ below respectively, with some additional comments and discussion. Assumption 1′ (Smoothness). For each device i = 1, . . . , n, the objective Fi is smooth, i.e., it is continuously differentiable and,
  (a) u 7→ ∇uFi(u, vi) is Lu-Lipschitz for all vi, (b) vi 7→ ∇vFi(u, vi) is Lv-Lipschitz for all u, (c) vi 7→ ∇uFi(u, vi) is Luv-Lipschitz for all u, and, (d) u 7→ ∇vFi(u, vi) is Lvu-Lipschitz for all vi. Further, we assume for some χ > 0 that
  max{Luv, Lvu} ≤ χ √ LuLv .
  The smoothness assumption is a standard one. We can assume without loss of generality that the cross-Lipschitz coefficients Luv, Lvu are equal. Indeed, if Fi is twice continuously differentiable, we can show that Luv, Lvu are both equal to the operator norm ‖∇2uvFi(u, vi)‖op of the mixed second derivative matrix. Further, χ denotes the extent to which u impacts the gradient of vi and vice-versa.
  Our next assumption is about the variance of the stochastic gradients, and is standard in literature. Compared to the main paper, we adopt a more precise notation about stochastic gradients. Assumption 2′ (Bounded Variance). Let Di denote a probability distribution over the data space Z on device i. There exist functions Gi,u and Gi,v which are unbiased estimates of ∇uFi and ∇vFi respectively. That is, for all u, vi:
  Ez∼Di [Gi,u(u, v, z)] = ∇uFi(u, vi), and Ez∼Di [Gi,v(u, v, z)] = ∇vFi(u, vi) .
  Furthermore, the variance of these estimators is at most σ2u and σ 2 v respectively. That is,
  Ez∼Di‖Gi,u(u, v, z)−∇uFi(u, vi)‖ 2 ≤ σ2u ,
  Ez∼Di‖Gi,v(u, v, z)−∇vFi(u, vi)‖ 2 ≤ σ2v .
  In practice, one usually has Gi,u(u, vi, z) = ∇ufi((u, vi), z), which is the gradient of the loss on datapoint z ∼ Di under the model (u, vi), and similarly for Gi,v . Finally, we make a gradient diversity assumption. Assumption 3′ (Partial Gradient Diversity). There exist δ ≥ 0 and ρ ≥ 0 such that for all u and V ,
  1
  n n∑ i=1 ‖∇uFi(u, vi)−∇uF (u, V )‖2 ≤ δ2 + ρ2‖∇uF (u, V )‖2 . (8)
  Algorithm 4 FedSim: Simultaneous update of shared and personal parameters
  Input: Initial iterates u(0), V (0), Number of communication rounds T , Number of devices per round m, Number of local updates τ , Local step sizes γu, γv .
  1: for t = 0, 1, · · · , T − 1 do 2: Sample m devices from [n] without replacement in S(t) 3: for each selected device i ∈ S(t) in parallel do 4: Initialize v(t)i,0 = v (t) i and u (t) i,0 = u (t) 5: for k = 0, · · · , τ − 1 do . Update all parameters jointly 6: Sample data z(t)i,k ∼ Di 7: v(t)i,k+1 = v (t) i,k − γvGi,v(u (t) i,k, v (t) i,k, z (t) i,k) 8: u(t)i,k+1 = u (t) i,k − γuGi,u(u (t) i,k, v (t) i,k, z (t) i,k) 9: Update v(t+1)i = v (t) i,τ and u (t+1) i = u (t) i,τ
  10: Update u(t+1) = ∑ i∈S(t) αiu (t+1) i∑
  i∈S(t) αi at the server with secure aggregation
  11: return u(T ), v(T )1 , · · · , v (T ) n
  This assumption is analogous to the bounded variance assumption (Assumption 2′), but with the stochasticity coming from the sampling of devices. It characterizes how much local steps on one device help or hurt convergence globally. Similar gradient diversity assumptions are often used for analyzing non-personalized federated learning (Koloskova et al., 2020; Karimireddy et al., 2020). Finally, it suffices for the partial gradient diversity assumption to only hold at the iterates (u(t), V (t)) generated by either FedSim or FedAlt.
sections.heading: A.2 CONVERGENCE ANALYSIS OF FEDSIM
sections.text: We give the full form of FedSim in Algorithm 4 for the general case of unequal αi’s but focus on αi = 1/n for the analysis. In order to simplify presentation, we denote V (t) = (v (t) 1 , . . . , v (t) n ) and define the following shorthand for gradient terms
  ∆(t)u = ∥∥∥∇uF (u(t), V (t))∥∥∥2 , and ∆(t)v = 1n n∑ i=1 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 . For convenience, we restate Theorem 1 from the main paper. Theorem 1 (Convergence of FedSim). Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedSim are chosen as γu = η/(Luτ) and γv = η/(Lvτ) with
  η ≤ min {
  1 12(1+χ2)(1+ρ2) ,
  √ m/n
  196(1−τ−1)(1+χ2)(1+ρ2)
  } .
  Then, ignoring absolute constants, we have
  1 T ∑T−1 t=0 ( 1 Lu E [ ∆ (t) u ] + mnLv E [ ∆ (t) v ]) ≤ ∆F0ηT + η(1 + χ 2) ( σ2u+δ 2(1−mn ) mLu + mσ2v nLv ) + η2(1− τ−1)(1 + χ2) ( σ2u+δ 2
  Lu + σ2v Lv
  ) . (6)
  Before proving the theorem, we give the following corollary with optimized learning rates. Corollary 3. Consider the setting of Theorem 1 and let ε > 0 be given. Suppose we set the learning rates γu = η/(τLu) and γv = η/(τLv), where (ignoring absolute constants),
  η = ε(
  δ2 Lu ( 1− mn ) + σ2u Lu + σ2vm Lun ) (1 + χ2)
  ∧ ε( δ2
  Lu ∨ σ 2 u Lu ∨ σ 2 v Lv
  ) (1− τ−1)(1 + χ2)
  1/2
  ∧ 1 (1 + χ2)(1 + ρ2) ∧( m/n (1− τ−1)(1 + ρ2)(1 + χ2) )1/2 .
  We have,
  1
  T T−1∑ t=0
  ( 1 Lu E ∥∥∥∇uF (u(t), V (t))∥∥∥2 + m Lvn2 n∑ i=1 E ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 ) ≤ ε
  after T communication rounds, where, ignoring absolute constants,
  T ≤ ∆F0(1 + χ 2)
  ε2
  ( σ2u + δ 2 ( 1− mn ) mLu + mσ2v nLv )
  + ∆F0
  √ (1− τ−1)(1 + χ2)
  ε3/2
  ( σu + δ√ Lu + σv√ Lv ) +
  ∆F0 ε (1 + χ2)(1 + ρ2)
  ( 1 + √ (1− τ−1)n
  m
  ) .
  Proof. The choice of the constant η ensures that each of the constant terms in the bound of Theorem 1 is O(ε). The final rate is now O ( ∆F0/(ηε) ) ; plugging in the value of η completes the proof.
  We now prove Theorem 1.
  Proof of Theorem 1. The proof mainly applies the smoothness upper bound to write out a descent condition with suitably small noise terms. We start with some notation.
  Notation. Let F (t) denote the σ-algebra generated by ( u(t), V (t) ) and denote Et[·] = E[·|F (t)]. For all devices, including those not selected in each round, we define virtual sequences ũ(t)i,k, ṽ (t) i,k as the SGD updates in Algorithm 4 for all devices regardless of whether they are selected. For the selected devices k ∈ S(t), we have ( u
  (t) i,k, v (t) i,k
  ) = ( ũ
  (t) i,k, ṽ (t) i,k ) . Note now that the random variables ũ(t)i,k, ṽ (t) i,k
  are independent of the device selection S(t). The updates for the devices i ∈ S(t) are given by
  v (t+1) i = v (t) i − γv τ−1∑ k=0 Gi,v ( ũ (t) i,k, ṽ (t) i,k, z (t) i,k ) ,
  and the server update is given by
  u(t+1) = u(t) − γu m ∑ i∈S(t) τ−1∑ k=0 Gi,u ( ũ (t) i,k, ṽ (t) i,k, z (t) i,k ) . (9)
  Proof Outline. We use the smoothness of Fi, more precisely Lemma 16, to obtain F ( u(t+1), V (t+1) ) − F ( u(t), V (t) ) ≤ 〈 ∇uF (u(t), V (t)), u(t+1) − u(t)
  〉 ︸ ︷︷ ︸
  T1,u
  + 1
  n n∑ i=1 〈 ∇vFi(u(t), v(t)i ), v (t+1) i − v (t) i 〉 ︸ ︷︷ ︸
  T1,v
  + Lu(1 + χ
  2)
  2 ∥∥∥u(t+1) − u(t)∥∥∥2︸ ︷︷ ︸ T2,u + 1 n n∑ i=1 Lv(1 + χ 2) 2 ∥∥∥v(t+1)i − v(t)i ∥∥∥2︸ ︷︷ ︸ T2,v .
  (10)
  Our goal will be to bound each of these terms to get a descent condition from each step of the form
  Et
  [ F ( u(t+1), V (t+1) ) − F ( u(t), V (t) )] ≤ −γuτ
  8 ∥∥∥∇uF (u(t), V (t))∥∥∥2 − γvτm 8n2 n∑ i=1 ∥∥∥∇vFi(u(t), v(t)i )∥∥∥2 +O(γ2u + γ2v) ,
  where the O(γ2u + γ 2 v) terms are controlled using the bounded variance and gradient diversity assumptions. Telescoping this descent condition gives the final bound.
  Main Proof. Towards this end, we prove non-asymptotic bounds on each of the terms T1,v, T1,u, T2,v and T2,u, in Claims 4 to 7 respectively. We then invoke them to get the bound
  Et
  [ F ( u(t+1), V (t+1) ) − F ( u(t), V (t) )] ≤ −γuτ
  4 ∆(t)u −
  γvτm
  4n ∆(t)v
  + Lu(1 + χ
  2)γ2uτ 2
  2
  ( σ2u + 12δ2
  m (1−m/n)
  ) + Lv(1 + χ 2)γ2vτ 2σ2vm
  2n
  + 2
  n n∑ i=1 τ−1∑ k=0 Et ∥∥∥u(t)i,k − u(t)∥∥∥2 (L2uγu + mn χ2LuLvγv) + 2
  n n∑ i=1 τ−1∑ k=0 Et ∥∥∥v(t)i,k − v(t)∥∥∥2 (mn L2vγv + χ2LuLvγu) . (11)
  Note that we simplified some constants appearing on the gradient norm terms using
  γu ≤ ( 12Lu(1 + χ 2)(1 + ρ2)τ )−1 and γv ≤ ( 6Lv(1 + χ 2)τ )−1 .
  Our next step is to bound the last two lines of (11) with Lemma 8 and invoke the gradient diversity assumption (Assumption 3′) as
  1
  n n∑ i=1 ∥∥∥∇uFi(u(t), v(t)i )∥∥∥2 ≤ δ2 + (1 + ρ2)∥∥∥∇uF (u(t), V (t))∥∥∥2 . This gives, after plugging in the learning rates and further simplifying the constants,
  Et
  [ F ( u(t+1), V (t+1) ) − F ( u(t), V (t) )] ≤− c∆ (t) u
  8Lu − cm∆
  (t) v
  8Lvn + c2(1 + χ2)
  ( σ2u
  2Lu + mσ2v nLv + 6δ2 Lum
  ( 1− m
  n )) + c3(1 + χ2)(1− τ−1) ( 24δ2
  Lu + 4σ2u Lu + 4σ2v Lu
  ) .
  Taking full expectation, telescoping the series over t = 0, · · · , T − 1 and rearranging the resulting terms give the desired bound in Theorem 1.
  Claim 4 (Bounding T1,v). Let T1,v be defined as in (10). We have,
  Et[T1,v] ≤ − γvτm
  2n2 n∑ i=1 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γvm
  n n∑ i=1 τ−1∑ k=0 Et [ χ2LuLv ∥∥∥ũ(t)i,k − u(t)∥∥∥2 + L2v∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2] .
  Proof. Define T1,v,i to be contribution of the ith term to T1,v. For i /∈ St, we have that T1,v,i = 0, since v(t+1)i = v (t) i . On the other hand, for i ∈ S(t), we use the unbiasedness of the gradient estimator
  Gi,v and the independence of z (t) i,k from u (t) i,k, v (t) i,k to get
  Et [T1,v,i] = −γv τ−1∑ k=0 Et 〈 ∇vFi ( u(t), v (t) i ) ,∇vFi ( u (t) i,k, v (t) i,k )〉 = −γv
  τ−1∑ k=0 Et 〈 ∇vFi ( u(t), v (t) i ) ,∇vFi ( ũ (t) i,k, ṽ (t) i,k )〉 =− γvτ
  ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 − γv
  τ−1∑ k=0 Et 〈 ∇vFi ( u(t), v (t) i ) ,∇vFi ( ũ (t) i,k, ṽ (t) i,k ) −∇vFi ( u(t), v (t) i )〉 ≤ −γvτ
  2 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γv2 τ−1∑ k=0 Et ∥∥∥∇vFi (ũ(t)i,k, ṽ(t)i,k)−∇vFi (u(t), v(t)i )∥∥∥2 . (12) For the second term, we add and subtract∇vFi ( u(t), ṽ (t) i,k ) and use smoothness to get
  ∥∥∥∇vFi (ũ(t)i,k, ṽ(t)i,k)−∇vFi (u(t), v(t)i )∥∥∥2 ≤ 2χ2LuLv∥∥∥ũ(t)i,k − u(t)∥∥∥2 + 2L2v∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2 . (13)
  Since the right hand side of this bound is independent of St, we get,
  Et[T1,v] = m
  n Et  1 m ∑ i∈S(t) T1,v,i  = m n2 n∑ i=1 Et[T1,v,i] ,
  and plugging in (12) and (13) completes the proof.
  Claim 5 (Bounding T1,u). Consider T1,u defined in (10). We have the bound,
  Et[T1,u] ≤ − γuτ
  2 ∥∥∥∇uF (u(t), V (t))∥∥∥2 + γu n n∑ i=1 τ−1∑ k=0 Et [ L2u ∥∥∥ũ(t)i,k − u(t)∥∥∥2 + χ2LuLv∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2] .
  Proof. Due to the independence of S(t) from ũ(t)i,k, ṽ (t) i,k, we have,
  Et
  [ u(t+1) − u(t) ] = −γuEt  1 m ∑ i∈S(t) τ−1∑ k=0 ∇uFi ( u (t) i,k, v (t) i,k ) = −γuEt  1 m ∑ i∈S(t) τ−1∑ k=0 ∇uFi ( ũ (t) i,k, ṽ (t) i,k
  ) = −γu
  n n∑ i=1 τ−1∑ k=0 Et [ ∇uFi ( ũ (t) i,k, ṽ (t) i,k )] ,
  where the last equality took an expectation over S(t), which is independent of ũ(t)i,k, ṽ (t) i,k. Now, using the same sequence of arguments as Claim 4, we have,
  Et 〈 ∇uF ( u(t), V (t) ) , u(t+1) − u(t) 〉 = −γu
  τ−1∑ k=0 Et
  〈 ∇uF ( u(t), V (t) ) , 1
  n n∑ i=1 ∇uFi ( ũ (t) i,k, ṽ (t) i,k
  )〉
  ≤ −γuτ 2 ∥∥∥∇uF (u(t), V (t))∥∥∥2 + γu 2 τ−1∑ k=0 Et ∥∥∥∥∥ 1n n∑ i=1 ∇uFi ( ũ (t) i,k, ṽ (t) i,k ) −∇uF ( u(t), V (t) )∥∥∥∥∥ 2
  (∗) ≤ −γuτ
  2 ∥∥∥∇uF (u(t), V (t))∥∥∥2 + γu 2n n∑ i=1 τ−1∑ k=0 Et ∥∥∥∇uFi (ũ(t)i,k, ṽ(t)i,k)−∇uFi (u(t), v(t)i )∥∥∥2 ≤ −γuτ
  2 ∥∥∥∇uF (u(t), V (t))∥∥∥2 + γu n n∑ i=1 τ−1∑ k=0 Et [ L2u ∥∥∥ũ(t)i,k − u(t)∥∥∥2 + L2uv∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2] , where the inequality (∗) follows from Jensen’s inequality as∥∥∥∥∥ 1n n∑ i=1 ∇uFi ( ũ (t) i,k, ṽ (t) i,k ) −∇uF ( u(t), V (t) )∥∥∥∥∥ 2 ≤ 1 n n∑ i=1 ∥∥∥∇uFi (ũ(t)i,k, ṽ(t)i,k)−∇uFi (u(t)i,k, v(t))∥∥∥2 .
  Claim 6 (Bounding T2,v). Consider T2,v as defined in (10). We have the bound,
  Et[T2,v] ≤ 3Lv(1 + χ
  2)γ2vτ 2m
  2n2
  n∑ i=1 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + Lv(1 + χ2)γ2vτ2mσ2v2n + 3Lv(1 + χ 2)γ2vτm
  2n2
  n∑ i=1 τ−1∑ k=0 Et [ L2v ∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2 + χ2LuLv∥∥∥ũ(t)i,k − u(t)∥∥∥2] . Proof. We start with
  Et ∥∥∥ṽ(t)k,τ − v(t)∥∥∥2 = γ2vEt ∥∥∥∥∥ τ−1∑ k=0 Gi,v ( ũ (t) i,k, ṽ (t) i,k, z (t) i,k )∥∥∥∥∥ 2
  ≤ γ2vτ τ−1∑ k=0 Et ∥∥∥Gi,v (ũ(t)i,k, ṽ(t)i,k, z(t)i,k)∥∥∥2 ≤ γ2vτ2σ2v + γ2vτ
  τ−1∑ k=0 Et ∥∥∥∇vFi (ũ(t)i,k, ṽ(t)i,k)∥∥∥2 ≤ γ2vτ2σ2v + 3γ2vτ2
  ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + 3γ2vτ
  τ−1∑ k=0 Et [ L2v ∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2 + χ2LuLv∥∥∥ũ(t)i,k − u(t)∥∥∥2] . Using (a) v(t+1)i = ṽ (t) i,τ for i ∈ S(t), and, (b) S(t) is independent from ũ (t) i,k, ṽ (t) i,k, we get,
  Et[T2,v] = Lv(1 + χ
  2)m
  2n Et  1 m ∑ i∈S(t) ∥∥∥ṽ(t)i,τ − v(t)i ∥∥∥2 
  ≤ Lv(1 + χ 2)m
  2n2
  n∑ i=1 Et ∥∥∥ṽ(t)i,τ − v(t)i ∥∥∥2 Plugging in the bound Et ∥∥∥ṽ(t)i,τ − v(t)∥∥∥2 completes the proof.
  Claim 7 (Bounding T2,u). Consider T2,u as defined in (10). We have,
  Et[T2,u] ≤ Lu(1 + χ
  2)γ2uτ 2
  2m
  ( σ2u + 12δ 2 (
  1− m n )) + 3Lu(1 + χ 2)γ2uτ 2(1 + ρ2)
  ∥∥∥∇uFi (u(t), V (t))∥∥∥2 + 3Lu(1 + χ 2)γ2uτ
  2n
  n∑ i=1 τ−1∑ k=0 Et [ L2u ∥∥∥ũ(t)i,k − u(t)∥∥∥2 + χ2LuLv∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2] . Proof. We proceed with the first two inequalities as in the proof of Claim 6 to get
  Et ∥∥∥u(t+1) − u(t)∥∥∥2 ≤ γ2uτ2σ2u m + γ2uτ τ−1∑ k=0 Et ∥∥∥∥∥∥ 1m ∑ i∈S(t) ∇uFi ( ũ (t) i,k, ṽ (t) i,k )∥∥∥∥∥∥ 2
  ︸ ︷︷ ︸ =:T3,j
  .
  For T3,j , (a) we add and subtract ∇uF (u(t), V (t)) and ∇uFi(u(t), ṽ(t)i,k), (b) invoke the squared triangle inequality, and, (c) use smoothness to get
  T3,j = 6Et ∥∥∥∥∥∥ 1m ∑ i∈S(t) ∇uFi ( u(t), v (t) i ) −∇uF ( u(t), V (t) )∥∥∥∥∥∥ 2 + 6 ∥∥∥∇uF (u(t), V (t))∥∥∥2
  + 3Et  1 m ∑ i∈S(t) ( L2u ∥∥∥ũ(t)i,k − u(t)∥∥∥2 + χ2LuLv∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2) 
  For the first term, we use the fact that S(t) is obtained by sampling without replacement to apply Lemma 17 together with the gradient diversity assumption to get
  Et ∥∥∥∥∥∥ 1m ∑ i∈S(t) ∇uFi ( u(t), v (t) i ) −∇uF ( u(t), V (t) )∥∥∥∥∥∥ 2
  ≤ 1 m ( n−m n− 1 ) 1 n n∑ i=1 ∥∥∥∇uFi (u(t), v(t)i )−∇uF (u(t), V (t))∥∥∥2 ≤ 1 m ( n−m n− 1 )( δ2 + ρ2
  ∥∥∥∇uF (u(t), V (t))∥∥∥2) . Therefore,
  T3,j = 12δ2
  m
  ( 1− m
  n
  ) + 6(1 + ρ2) ∥∥∥∇uF (u(t), V (t))∥∥∥2 + 3
  n n∑ i=1 Et [ L2u ∥∥∥ũ(t)i,k − u(t)∥∥∥2 + χ2LuLv∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2] , where we also used the independence between S(t) and (ũ(t)i,k, ṽ (t) i,k). Plugging this into the expression for Et‖u(t+1) − u(t)‖2 completes the proof.
  Lemma 8. Let Fi satisfy Assumptions 1′-3′, and consider the iterates
  uk+1 = uk − γuGi,u(uk, vk, zk) , and, vk+1 = vk − γvGi,v(uk, vk, zk) ,
  for k = 0, · · · , τ − 1, where zk ∼ Di. Suppose the learning rates satisfy γu = cu/(τLu) and γv = cv/(τLv) with cu, cv ≤ 1/ √ 6 max{1, χ−2}. Further, define,
  A = γuL 2 u + fχ 2γvLuLv , and, B = fγvL2v + χ 2γuLuLv ,
  where f ∈ (0, 1] is given. Then, we have the bound,
  τ−1∑ k=0 E [ A‖uk − u0‖2+B‖vk − u0‖2 ] ≤ 4τ2(τ − 1) ( γ2uσ 2 uA+ γ 2 vσ 2 vB )
  + 12τ2(τ − 1) ( γ2uA‖∇uFi(u0, v0)‖2 + γ2vB‖∇vFi(u0, v0)‖2 ) .
  Proof. If τ = 1, there is nothing to prove, so we assume τ > 1. Let ∆k := A‖uk − u0‖2 +B‖vk − v0‖2 and denote by Fk the sigma-algebra generated by (wk, vk). Further, let Ek[·] = E[·|Fk]. We use the inequality 2αβ ≤ α2/δ2 + δ2β2 for reals α, β, δ to get,
  Ek‖uk+1 − u0‖2 ≤ ( 1 + 1
  τ − 1
  ) ‖uk − u0‖2 + τγ2uEk‖Gi,u(uk, vk, zk)‖ 2
  ≤ ( 1 + 1
  τ − 1
  ) ‖uk − u0‖2 + τγ2uσ2u + τγ2u‖∇uFi(uk, vk)‖ 2
  ≤ ( 1 + 1
  τ − 1
  ) ‖uk − u0‖2 + τγ2uσ2u + 3τγ2u‖∇uFi(u0, v0)‖ 2
  + 3τγ2uL 2 u‖uk − u0‖2 + 3τγ2uLuv‖vk − v0‖2 ,
  where the last inequality followed from the squared triangle inequality (from adding and subtracting ∇uFi(u0, vk) and∇uFi(u0, v0)) followed by smoothness. Together with the analogous inequality for the v-update, we get,
  Ek[∆k+1] ≤ ( 1 + 1
  τ − 1
  ) ∆k +A ′‖uk − u0‖2 +B′‖vk − v0‖2 + C ,
  where we have
  A′ = 3τ(γ2uL 2 uA+ γ 2 vχ 2LuLvB), and, B′ = 3τ(γ2vL 2 vB + γ 2 uχ 2LuLvA) and,
  C ′ = τγ2uσ 2 uA+ τγ 2 vσ 2 vB + 3τγ 2 uA‖∇uFi(u0, v0)‖2 + 3τγ2vB‖∇vFi(u0, v0)‖2 .
  Next, we apply Lemma 20 to get that A′ ≤ A/τ and B′ ≤ B/τ under the assumed conditions on the learning rates; this allows us to write the right hand side completely in terms of ∆k and unroll the recurrence. The intuition behind Lemma 20 is as follows. Ignoring the dependence on τ, Lu, Lv, χ for a moment, if γu and γv are both O(η), then A′, B′ are both O(η3), while A and B are O(η). Thus, making η small enough should suffice to get A′ ≤ O(A) and B′ ≤ O(B). Concretely, Lemma 20 gives
  Ek[∆k+1] ≤ ( 1 + 2
  τ − 1
  ) E[∆k] + C ,
  and unrolling this recurrence gives for k ≤ τ − 1
  E[∆k] ≤ k−1∑ j=0 ( 1 + 2 τ − 1 )j C ≤ τ − 1 2 ( 1 + 2 τ − 1 )k C
  ≤ τ − 1 2
  ( 1 + 2
  τ − 1
  )τ−1 C ≤ e 2
  2 (τ − 1)C ,
  where we used (1 + 1/α)α ≤ e for all α > 0. Summing over k and using the numerical bound e2 < 8 completes the proof.
  Remark 9. We only invoked the partial gradient diversity assumption (Assumption 3) at iterates (u(t), V (t)); therefore, it suffices if the assumption only holds at iterates (u(t), V (t)) generated by FedSim, rather than at all (u, V ).
  Algorithm 5 FedAlt: Alternating updates of shared and personalized parameters
  Input: Initial iterates u(0), V (0), Number of communication rounds T , Number of devices per round m, Number of local updates τu, τv , Local step sizes γu, γv ,
  1: for t = 0, 1, · · · , T − 1 do 2: Sample m devices from [n] without replacement in S(t) 3: for each selected device i ∈ S(t) in parallel do 4: Initialize v(t)i,0 = v (t) i 5: for k = 0, · · · , τv − 1 do . Update personalized parameters 6: Sample data z(t)i,k ∼ Di 7: v(t)i,k+1 = v (t) i,k − γvGi,v(u(t), v (t) i,k, z (t) i,k) 8: Update v(t+1)i = v (t) k,τv 9: Initialize u(t)i,0 = u (t)
  10: for k = 0, · · · , τu − 1 do . Update shared parameters 11: u(t)i,k+1 = u (t) i,k − γuGi,u(u (t) i,k, v (t+1) i , z (t) i,k) 12: Update u(t+1)i = u (t) i,τu
  13: Update u(t+1) = ∑ i∈S(t) αiu (t+1) i∑
  i∈S(t) αi at the server with secure aggregation
  14: return u(T ), v(T )1 , · · · , v (T ) n
sections.heading: A.3 CONVERGENCE ANALYSIS OF FEDALT
sections.text: We give the full form of FedAlt in Algorithms 5 for the general case of unequal αi’s but focus on αi = 1/n for the analysis. For convenience, we reiterate Theorem 2 below. Recall the definitions
  ∆(t)u = ∥∥∥∇uF (u(t), V (t+1))∥∥∥2 , and, ∆(t)v = 1n n∑ i=1 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 . Theorem 2 (Convergence of FedAlt). Suppose Assumptions 1, 2 and 3 hold and the learning rates in FedAlt are chosen as γu = η/(Luτu) and γv = η/(Lvτv), with
  η ≤ min {
  1 24(1+ρ2) , m 128χ2(n−m) , √ m χ2n } .
  Then, ignoring absolute constants, we have
  1 T ∑T−1 t=0 ( 1 Lu E [ ∆ (t) u ] + mnLv E [ ∆ (t) v ]) ≤ ∆F0
  ηT + η
  ( σ2u+δ
  2(1−mn ) mLu + σ2v Lv m+χ2(n−m) n ) + η2 ( σ2u+δ 2
  Lu (1− τ−1u ) + σ2vm Lvn (1− τ−1v ) + χ2σ2v Lv
  ) .
  Before proving the theorem, we have the corollary with optimized learning rates.
  Corollary 10. Consider the setting of Theorem 2 and fix some ε > 0. Suppose we set γu = η/(τLu) and γv = η/(τLv) such that, ignoring absolute constants,
  η = ( σ2v εLv (m n + χ2(1−m/n) ))−1∧(σ2u + δ2(1−m/n) mLuε )−1∧(σ2u + δ2 Luε (1− τ−1u ) )−1/2
  ∧( σ2vm Lvnε (1− τ−1v ) )−1/2∧ 1 1 + ρ2 ∧ m χ2(n−m) ∧√ m χ2n .
  Then, we have,
  1
  T T−1∑ t=0
  ( 1 Lu E ∥∥∥∇uF (u(t), Ṽ (t))∥∥∥2 + m Lvn2 n∑ i=1 E ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 ) ≤ ε
  after T communication rounds, where, ignoring absolute constants,
  T ≤ ∆F0 ε2
  ( σ2u + δ 2 ( 1− mn ) mLu + σ2v Lv (m n + χ2 ( 1− m n )))
  + ∆F0 ε3/2 ( σu + δ√ Lu √ 1− τ−1u + σv√ Lv √ 1− τ−1v ) +
  ∆F0 ε
  ( 1 + ρ2 + χ2 ( n m − 1 ) + √ χ2n m ) .
  Proof. We get the bound by balancing terms from the bound of Theorem 2. The choice of η ensures that all the O(η) and O(η2) terms are at most O(ε). Finally, the smallest number of communication rounds to make the left hand side of the bound of Theorem 2 smaller than ε is ∆F0/(ηε).
  We are now ready to prove Theorem 2.
  Proof of Theorem 2. The proof mainly applies the smoothness upper bound to write out a descent condition with suitably small noise terms. We start with some notation.
  We introduce the notation ∆̃(t)u as the analogue of ∆ (t) u with the virtual variable Ṽ (t+1):
  ∆̃(t)u = ∥∥∥∇uF (u(t), V (t))∥∥∥2 .
  Notation. Let F (t) denote the σ-algebra generated by ( u(t), V (t) ) and denote Et[·] = E[·|F (t)]. For all devices, including those not selected in each round, we define virtual sequences ũ(t)i,k, ṽ (t) i,k as the SGD updates in Algorithm 5 for all devices regardless of whether they are selected. For the selected devices i ∈ S(t), we have v(t)i,k = ṽ (t) i,k and u (t) i,k = ũ (t) i,k. Note now that the random variables ũ (t) i,k, ṽ (t) i,k are independent of the device selection S(t). Finally, we have that the updates for the selected devices i ∈ S(t) are given by
  v (t+1) i = v (t) i − γv τv−1∑ k=0 Gi,v ( u(t), ṽ (t) i,k, z (t) i,k ) ,
  and the server update is given by
  u(t+1) = u(t) − γu m ∑ i∈S(t) τu−1∑ k=0 Gi,u ( ũ (t) i,k, ṽ (t) i,τv , z (t) i,k ) .
  Proof Outline and the Challenge of Dependent Random Variables. We start with F ( u(t+1), V (t+1) ) − F ( u(t), V (t) ) =F ( u(t), V (t+1) ) − F ( u(t), V (t) ) + F ( u(t+1), V (t+1) ) − F ( u(t), V (t+1) ) .
  (14)
  The first line corresponds to the effect of the v-step and the second line to the u-step. The former is easy to handle with standard techniques that rely on the smoothness of F ( u(t), · ) . The latter is more challenging. In particular, the smoothness bound for the u-step gives us
  F ( u(t+1), V (t+1) ) − F ( u(t), V (t+1) ) ≤ 〈 ∇uF ( u(t), V (t+1) ) , u(t+1) − u(t) 〉 + Lu 2
  ∥∥∥u(t+1) − u(t)∥∥∥2 . The standard proofs of convergence of stochastic gradient methods rely on the fact that we can take an expectation w.r.t. the sampling S(t) of devices for the first order term. However, both V (t+1) and
  u(t+1) depend on the sampling S(t) of devices. Therefore, we cannot directly take an expectation with respect to the sampling of devices in S(t).
  Virtual Full Participation to Circumvent Dependent Random Variables. The crux of the proof lies in replacing V (t+1) in the analysis of the u-step with the virtual iterate Ṽ (t+1) so as to move all the dependence of the u-step on S(t) to the u(t+1) term. This allows us to take an expectation; it remains to carefully bound the resulting error terms.
  Finally, we will arrive at a bound of the form
  1
  T T−1∑ t=0 (γuτu 8 E[∆̃(t)u ] + γvτvm 16n E[∆(t)v ] ) ≤ ∆F0 T +O(γ2u + γ 2 v) .
  Next, we translate this bound from gradient E[∆̃(t)u ] of the virtual Ṽ (t+1) to E[∆ (t) u ], which is the gradient computed at the actual iterate V (t). A careful analysis shows that we only incur a lower order term of O(γuγ2v) in this translation. Choosing γu and γv small enough will give us the final result.
  Analysis of the u-Step with Virtual Full Participation. We introduce the virtual iterates Ṽ (t+1) into the analysis of the u-step as follows:
  F ( u(t+1), V (t+1) ) − F ( u(t), V (t+1) ) ≤ 〈 ∇uF ( u(t), V (t+1) ) , u(t+1) − u(t) 〉 + Lu 2
  ∥∥∥u(t+1) − u(t)∥∥∥2 = 〈 ∇uF ( u(t), Ṽ (t+1) ) , u(t+1) − u(t) 〉 + Lu 2
  ∥∥∥u(t+1) − u(t)∥∥∥2 + 〈 ∇uF ( u(t), V (t+1) ) −∇uF ( u(t), Ṽ (t+1) ) , u(t+1) − u(t)
  〉 ≤ 〈 ∇uF ( u(t), Ṽ (t+1) ) , u(t+1) − u(t) 〉 + Lu
  ∥∥∥u(t+1) − u(t)∥∥∥2 + 1
  2Lu ∥∥∥∇uF (u(t), V (t+1))−∇uF (u(t), Ṽ (t+1))∥∥∥2 ≤ 〈 ∇uF ( u(t), Ṽ (t+1) ) , u(t+1) − u(t)
  〉 ︸ ︷︷ ︸
  T1,u
  +Lu ∥∥∥u(t+1) − u(t)∥∥∥2︸ ︷︷ ︸ T2,u + χ2Lv 2n n∑ i=1 ∥∥∥ṽ(t+1)i − v(t+1)i ∥∥∥2︸ ︷︷ ︸ T3,u .
  The last two inequalities follow from Young’s inequality and Lipschitzness of V 7→ ∇uF (u, V ) respectively.
  We have now successfully eliminated the dependence of the first-order term T1,u on V (t+1). The virtual iterates Ṽ (t+1) are now independent of S(t). This allows us to take an expectation w.r.t. the sampling S(t) of the devices.
  We bound each of these terms in Claims 11 to 13 below to get
  Et
  [ F ( u(t+1), V (t+1) ) − F ( u(t), V (t+1) )]
  ≤ −γuτu 4 Et[∆̃ (t) u ] +
  2γuL 2 u
  n n∑ i=1 τu−1∑ k=0 Et ∥∥∥ũ(t)i,k − u(t)∥∥∥2︸ ︷︷ ︸ =:T ′2,u +4γ2vτ 2 vLvσ 2 vχ 2(1−m/n)
  + Luγ
  2 uτ 2 u
  m
  ( σ2u + 3δ 2 (
  1− m n
  )) + 8γ2vτ 2 vLvχ 2(1−m/n)∆(t)v .
  Note that we used the fact that 24Luγuτu(1 + ρ2) ≤ 1 to simply the coefficients of some of the terms above. The second term has also been referred to as client drift in the literature; we bound it with
  Lemma 18 and invoke the assumption on gradient diversity (Assumption 3′) to get
  T ′2,u ≤ 16γ3uL 2 uτu(τu − 1) n n∑ i=1 Et ∥∥∥∇uFi (u(t), ṽ(t+1)i )∥∥∥2 + 8γ3uL2uτ2u(τu − 1)σ2u ≤ 16γ 3 uL
  2 uτu(τu − 1) n
  ( δ2 + ρ2Et ∥∥∥∇uF (u(t), Ṽ (t+1))∥∥∥2)+ 8γ3uL2uτ2u(τu − 1)σ2u . Plugging this back in, we get,
  Et
  [ F ( u(t+1), V (t+1) ) − F ( u(t), V (t+1) )]
  ≤ −γuτu 8 Et[∆̃ (t) u ] +
  Luγ 2 uτ 2 u
  m
  ( σ2u + 2δ 2(1−m/n) ) + 4γ2vτ 2 vLvσ 2 vχ 2(1−m/n)
  + 8γ2vτ 2 vLvχ 2(1−m/n)∆(t)v + 8γ2uL3uτ2u(τu − 1)(σ2u + 2δ2u) .
  Note that we used 128γ2uL 2 uτu(τu − 1)ρ2 ≤ 1, which is implied by 24Luγuτu(1 + ρ2) ≤ 1.
  Bound with the Virual Iterates. We plug this analysis of the u-step and Claim 14 for the v-step into (14) next. We also simplify some coefficients using 128γvτvLvχ2(n/m− 1) ≤ 1. This gives us
  Et
  [ F ( u(t+1), V (t+1) ) − F ( u(t), V (t) )] ≤ − γuτu
  8 Et[∆̃
  (t) u ]−
  γvτvm
  16n Et[∆
  (t) v ] + 4γ 2 vLvτ 2 vσ 2 v (m n + χ2(1−m/n) )
  + γ2uLuτ 2 u
  m
  ( σ2u + 2δ 2(1−m/n) ) + 8γ3uL 2 uτ 2 u(τu − 1)(σ2u + 2δ2) + 4γ3vL 2 vτ
  2 v (τv − 1)σ2vm n .
  Taking an unconditional expectation, summing it over t = 0 to T − 1 and rearranging this gives
  1
  T T−1∑ t=0 (γuτu 8 E[∆̃(t)u ] + γvτvm 16n E[∆(t)v ] ) (15)
  ≤ ∆F0 T + 4γ2vLvτ 2 vσ 2 v (m n + χ2(1−m/n) ) + γ2uLuτ 2 u m ( σ2u + 2δ 2(1−m/n) )
  + 8γ3uL 2 uτ 2 u(τu − 1)(σ2u + 2δ2) +
  4γ3vL 2 vτ 2 v (τv − 1)σ2vm n .
  This is a bound in terms of the virtual iterates Ṽ (t+1). However, we wish to show a bound in terms of the actual iterate V (t).
  Obtaining the Final Bound. It remains now to relate ∆̃(t)u with ∆(t)u . Using the Cauchy-Schwartz inequality and smoothness, we have,
  Et ∥∥∥∇uF (u(t), V (t))−∇uF (u(t), Ṽ (t+1))∥∥∥2 ≤ 1 n n∑ i=1 Et
  ∥∥∥∇uFi (u(t), v(t)i )−∇uFi (u(t), ṽ(t+1)i )∥∥∥2 ≤ χ
  2LuLv n n∑ i=1 Et ∥∥∥ṽ(t+1)i − v(t)i ∥∥∥2 ≤ χ
  2LuLv n n∑ i=1 ( 16γ2vτ 2 v ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + 8γ2vτ2vσ2v) = 8γ2vτ 2 vσ 2 vχ 2LuLv + 16γ 2 vτ 2 vχ 2LuLv∆ (t) v ,
  where the last inequality followed from Lemma 19. Using∥∥∥∇uF (u(t), V (t))∥∥∥2 ≤ 2∥∥∥∇uF (u(t), V (t))−∇uF (u(t), Ṽ (t+1))∥∥∥2+2∥∥∥∇uF (u(t), Ṽ (t+1))∥∥∥2 ,
  we get, E[∆(t)u ] ≤ 2E[∆̃(t)u ] + 16γ2vτ2vσ2vχ2LuLv + 32γ2vτ2vχ2LuLv E[∆(t)v ] . Therefore, we get, γuτu 16 E[∆(t)u ] + γvτvm 32n E[∆(t)v ]
  ≤ γuτu 8 E[∆̃(t)u ] + γvτvm 16n
  ( 1
  2 +
  32η2χ2m
  n
  ) E[∆(t)v ] + γuτuγ 2 vτ 2 vσ 2 vχ 2LuLv
  ≤ γuτu 8 E[∆̃(t)u ] + γvτvm 16n E[∆(t)v ] + γuτuγ 2 vτ 2 vσ 2 vχ 2LuLv ,
  where we used 32η 2χ2m n ≤ 1/2, which is one of the conditions we assume on η.
  Summing this up and plugging in (15) gives
  1
  T T−1∑ t=0 (γuτu 16 E[∆(t)u ] + γvτvm 32n E[∆(t)v ] ) ≤ 1 T T−1∑ t=0 (γuτu 8 E[∆̃(t)u ] + γvτvm 16n E[∆(t)v ] ) + γuτuγ 2 vτ 2 vσ 2 vχ 2LuLv
  ≤ ∆F0 T + 4γ2vLvτ 2 vσ 2 v (m n + χ2(1−m/n) ) + γ2uLuτ 2 u m ( σ2u + 2δ 2(1−m/n) )
  + 8γ3uL 2 uτ 2 u(τu − 1)(σ2u + 2δ2) +
  4γ3vL 2 vτ 2 v (τv − 1)σ2vm n + γuτuγ 2 vτ 2 vσ 2 vχ 2LuLv .
  Plugging in γu = η/(Luτu) and γv = η/(Lvτv) completes the proof.
  The analysis of each of the terms in the u-step is given in the following claims.
  Claim 11 (Bounding T1,u). We have,
  Et [T1,u] ≤ − γuτu
  2 Et ∥∥∥∇uF (u(t), Ṽ (t+1))∥∥∥2 + γuL2u n n∑ i=1 τu−1∑ k=0 Et ∥∥∥ũ(t)i,k − u(t)∥∥∥2 . Proof. For i ∈ S(t), we have that ũ(t)i,k = u (t) i,k. Therefore, we have,
  Et[T1,u] = −γuEt 〈 ∇uF ( u(t), Ṽ (t+1) ) , 1
  m ∑ i∈S(t) τu−1∑ k=0 ∇uFi ( ũ (t) i,k, ṽ (t+1) i )〉 .
  Using that ũ(t)i,k is independent of S (t), we get,
  Et[T1,u] = −γuEt 〈 ∇uF ( u(t), Ṽ (t+1) ) , 1
  n n∑ i=1 τu−1∑ k=0 ∇uFi ( ũ (t) i,k, ṽ (t+1) i
  )〉
  = − γuτuEt ∥∥∥∇uF (u(t), Ṽ (t+1))∥∥∥2
  − γu τu−1∑ k=0 Et
  〈 ∇uF ( u(t), Ṽ (t+1) ) , 1
  n n∑ i=1 ∇uFi ( ũ (t) i,k, ṽ (t+1) ) −∇uFi ( u(t), ṽ(t+1) )〉 Invoking 〈x, y〉 ≤ ‖x‖2/2+‖y‖2/2 for vectors x, y followed by smoothness completes the proof.
  Claim 12 (Bounding T2,u). We have, Et [T2,u] ≤ 3Luγ2uτ2u ( 1 + 2ρ2
  m (1−m/n)
  ) Et ∥∥∥∇uF (u(t), Ṽ (t+1))∥∥∥2 + 3L2uγ 2 uτu
  n
  n∑ i=1 τu−1∑ k=0 Et ∥∥∥ũ(t)i,k − u(t)∥∥∥2 + 6Luγ2uτ2uδ2m (1−m/n) .
  Proof. We use E‖z‖2 = ‖E[z]‖2 + E‖z −E[z]‖2 for a random vector z to get
  Et[T2,u] ≤ Luγ
  2 uτ 2 uσ 2 u m + Luγ 2 uτu τu−1∑ k=0 Et ∥∥∥∥∥∥ 1m ∑ i∈S(t) ∇uFi ( ũ (t) i,k, ṽ (t+1) i )∥∥∥∥∥∥ 2
  ︸ ︷︷ ︸ =:T ′k
  .
  We break the term T ′k as
  T ′k ≤ 3 ∥∥∥∥∥∥ 1m ∑ i∈S(t) ( ∇uFi ( ũ (t) i,k, ṽ (t+1) i ) −∇uFi ( u(t), ṽ (t+1) i ))∥∥∥∥∥∥ 2
  + 3 ∥∥∥∥∥∥ 1m ∑ i∈S(t) ∇uFi ( u(t), ṽ (t+1) i ) −∇uF ( u(t), Ṽ (t+1) )∥∥∥∥∥∥ 2 + 3 ∥∥∥∇uF (u(t), Ṽ (t+1))∥∥∥2 .
  For the first term, we use Jensen’s inequality to take the squared norm inside the sum, then use smoothness and take an expectation over the sampling of devices to get
  Et ∥∥∥∥∥∥ 1m ∑ i∈S(t) ( ∇uFi ( ũ (t) i,k, ṽ (t+1) i ) −∇uFi ( u(t), ṽ (t+1) i ))∥∥∥∥∥∥ 2 ≤ L 2 u n n∑ i=1 Et ∥∥∥ũ(t)i,k − u(t)∥∥∥2 . For the second term, we use the fact that S(t) was sampled without replacement (cf. Lemma 17) and invoke the gradient diversity assumption (Assumption 3′) to get,∥∥∥∥∥ 1m ∑
  i∈S(t) ∇uFi
  ( u(t), ṽ
  (t+1) i ) −∇uF ( u(t), Ṽ (t+1) )∥∥∥∥∥ 2
  ≤ ( n−m n− 1 ) 1 mn n∑ i=1 ∥∥∥∇uFi (u(t), ṽ(t+1)i )−∇uF (u, Ṽ (t+1))∥∥∥2 ≤ 2 m ( 1− m n )( δ2 + ρ2Et
  ∥∥∥∇uF (u(t), Ṽ (t+1))∥∥∥2) . To complete the proof, we plug these terms back into the definition of T ′k and Et[T2,u] to complete the proof.
  Claim 13 (Bounding T3,u). We have, Et [T3,u] ≤ 8γ2vτ2vLvχ2 (
  1− m n
  ) ∆(t)v + 4χ 2γ2vτ 2 vLvσ 2 v ( 1− m
  n
  ) .
  Proof. Since v(t+1)i = ṽ (t+1) i for i ∈ S(t), we have that
  T3,u = χ2Lv
  2n ∑ i/∈S(t) ∥∥∥ṽ(t+1)i − v(t)i ∥∥∥2 . Since
  ∥∥∥ṽ(t+1)i − v(t)i ∥∥∥2 is independent of S(t), we can take an expectation to get Et[T3,u] =
  χ2Lv 2n n∑ i=1 P(i /∈ S(t))Et ∥∥∥ṽ(t+1)i − v(t)i ∥∥∥2
  = χ2Lv
  2n
  ( 1− m
  n ) n∑ i=1 Et ∥∥∥ṽ(t+1)i − v(t)i ∥∥∥2 . Plugging in Lemma 19 completes the proof.
  The analysis of the v-step is given in the next result. Claim 14. Consider the setting of Theorem 2 and assume that γvτvLv ≤ 1/8. We have,
  Et
  [ F ( u(t), V (t+1) ) − F ( u(t), V (t) )] ≤ −γvτvm∆ (t) v
  8n + γ2vτ 2 vLvσ 2 vm 2n + 4γ3vL 2 vτ 2 v (τv − 1)σ2vm n .
  Proof. From smoothness, we get,
  Fi
  ( u(t), ṽ
  (t+1) i ) − Fi ( u(t), v (t) i ) ≤ 〈 ∇vFi ( u(t), v (t) i ) , ṽ (t+1) i − v (t) i 〉 ︸ ︷︷ ︸
  T1,v
  + Lv 2 ∥∥∥ṽ(t+1)i − v(t)i ∥∥∥2︸ ︷︷ ︸ T2,v .
  We bound the first term as
  Et[T1,v] = −γvEt 〈 ∇vFi ( u(t), v (t) i ) , τv−1∑ k=0 ∇vFi ( u(t), ṽ (t) i,k )〉
  =− γvτv ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2
  − γv τv−1∑ k=0 Et 〈 ∇vFi ( u(t), v (t) i ) ,∇vFi ( u(t), ṽ (t) i,k ) −∇vFi ( u(t), v (t) i )〉 ≤ −γvτv
  2 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γv2 τv−1∑ k=0 Et ∥∥∥∇vFi (u(t), ṽ(t)i,k)−∇vFi (u(t), v(t)i )∥∥∥2 ≤ −γvτv
  2 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γvL2v2 τv−1∑ k=0 ∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2 . Next, we observe that Ez‖Gi,v(u, vi, z)‖2 = ‖∇vFi(u, vi)‖2+Ez‖Gi,v(u, vi, z)−∇vFi(u, vi)‖2 ≤ ‖∇vFi(u, vi)‖2+σ2v . We invoke this inequality to handle the second term as
  Et[T2,v] ≤ γ2vLvτv
  2
  τv−1∑ k=0 Et ∥∥∥Gi,v (u(t), ṽ(t)i,k, z(t)i,k)∥∥∥2 ≤ γ 2 vLvτ 2 vσ 2 v
  2 + γ2vLvτv 2 τv−1∑ k=0 Et ∥∥∥∇vFi (u(t), ṽ(t)i,k)∥∥∥2 ≤ γ 2 vLvτ 2 vσ 2 v
  2 + γ2vLvτ 2 v ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γ2vLvτv
  τv−1∑ k=0 Et ∥∥∥∇vFi (u(t), ṽ(t)i,k)−∇vFi (u(t), v(t)i )∥∥∥2 ≤ γ 2 vLvτ 2 vσ 2 v
  2 + γ2vLvτ 2 v ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γ2vL3vτv τv−1∑ k=0 Et ∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2 . Plugging these bounds for T1,v and T2,v into the initial smoothness bound and using γvLvτv ≤ 1/4 gives
  Et [ Fi ( u(t), ṽ (t+1) i ) − Fi ( u(t), v (t) i ) ] ≤
  − γvτv 4 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γvL2v τv−1∑ k=0 ∥∥∥ṽ(t)i,k − v(t)i ∥∥∥2 + γ2vLvτ2vσ2v2 . We invoke Lemma 18 to bound the ∑ k Et‖ṽ (t) i,k − v (t) i ‖2 term, which is also known as client drift. We simplify some coefficients using 8γvτvLv ≤ 1 to get Et [ Fi ( u(t), ṽ (t+1) i ) − Fi ( u(t), v (t) i ) ] ≤
  − γvτv 8 ∥∥∥∇vFi (u(t), v(t)i )∥∥∥2 + γ2vLvτ2vσ2v2 + 4γ3vLvτ2v (τv − 1)σ2v .
  It remains to invoke that S(t) is a uniformly random sample of m devices from {1, · · · , n} and that ṽ
  (t+1) i is independent of S (t). To this end, note that
  Et
  [ F ( u(t), V (t+1) ) − F ( u(t), V (t) )] = m
  n Et  1 m ∑ i∈S(t) Fi ( u(t), ṽ (t+1) i ) − Fi ( u(t), v (t) i ) ≤ m n2 n∑ i=1 Et [ Fi ( u(t), ṽ (t+1) i ) − Fi ( u(t), v (t) i )] .
  Plugging in the previous bound completes the proof.
  Remark 15. We only invoked the partial gradient diversity assumption (Assumption 3) at (virtual) iterates (u(t), Ṽ (t+1)); therefore, it suffices if the assumption only holds at iterates (u(t), Ṽ (t+1)) generated by FedAlt, rather than at all (u, V ).
sections.heading: A.4 TECHNICAL LEMMAS
sections.text: The first lemma involves smoothness of two blocks of variables; we use this in the proof of FedSim. Lemma 16 (Block Smoothness). Suppose Fi : Rd × Rdi satisfy Assumption 1′. Then, it holds that
  Fi(w ′, v′i)− Fi(w, vi) ≤〈∇wFi(w, vi), w′ − w〉+ 〈∇vFi(w, vi), v′i − vi〉
  + Lw 2 (1 + χ2)‖w′ − w‖2 + Lv 2 (1 + χ2)‖v′i − vi‖2 .
  Proof. Using the Lw-smoothness of F (·, v′i) and the Lv-smoothness of F (w, ·), we have
  Fi(w ′, v′i)− Fi(w, v′i) ≤ 〈∇wFi(w, v′i), w′ − w〉+ Lw 2 ‖w′ − w‖2,
  Fi(w, v ′ i)− Fi(w, vi) ≤ 〈∇wFi(w, vi), v′i − vi〉+ Lv 2 ‖v′i − vi‖2.
  Summing the above two inequalities together gives
  Fi(w ′, v′i)− Fi(w, vi) ≤〈∇wFi(w, v′i), w′ − w〉+ 〈∇vFi(w, vi), v′i − vi〉
  + Lw 2 ‖w′ − w‖2 + Lv 2 ‖v′i − vi‖2 . (16)
  We can bound the first inner product term on the right-hand side of the above inequality as
  〈∇wFi(w, v′i), w′ − w〉 = 〈∇wFi(w, vi), w′ − w〉+ 〈∇wFi(w, v′i)−∇wFi(w, vi), w′ − w〉 ≤ 〈∇wFi(w, vi), w′ − w〉+ ‖∇wFi(w, v′i)−∇wFi(w, vi)‖‖w′ − w‖ ≤ 〈∇wFi(w, vi), w′ − w〉+ Lwv‖v′i − vi‖‖w′ − w‖
  ≤ 〈∇wFi(w, vi), w′ − w〉+ χ √ LwLv‖v′i − vi‖‖w′ − w‖
  ≤ 〈∇wFi(w, vi), w′ − w〉+ χ2 Lv 2 ‖v′i − vi‖2 + χ2 Lw 2 ‖w′ − w‖2,
  where the first inequality is due to Cauchy-Schwarz, the second inequality is due to Lwv-Lipschitz property of∇wFi(w, ·), the third inequality is due to the definition of χ in (5), and the last inequality is due to Young’s inequality. Substituting the above inequality into (16) yields the desired result.
  Next, we have the variance of sampling without replacement. Note the correction factor of (n − m)/(n− 1) over sampling with replacement. We include the elementary proof for completeness. Lemma 17 (Sampling Without Replacement). Let a1, · · · , an ∈ Rd be given. Let S be a uniformly random sample of size m from this collection, where the sampling is without replacement. Denoting the mean ā = ∑n i=1 ai/n, we have,
  ES ∥∥∥∥∥ 1m∑ i∈S ai − ā ∥∥∥∥∥ 2 ≤ ( n−m n− 1 ) 1 m ( 1 n n∑ i=1 ‖ai − ā‖2 ) .
  Proof. The statement is trivially true if m = 1 or m = n. Therefore, we assume now that 2 ≤ m ≤ n− 1. Further, without loss of generality, we assume that ā = 0. Finally, let S denote the set of all subsets of [n] of size m. Note that |S| = ( n m ) . We now have,
  ES ∥∥∥∥∥ 1m∑ i∈S ai ∥∥∥∥∥ 2 = 1 m2 ( n m ) ∑ S∈S ∑ i∈S ‖ai‖2 + ∑ i,j∈S: i 6=j 〈ai, aj〉  . For the first term, we have,∑
  S∈S ∑ i∈S ‖ai‖2 = n∑ i=1 ∑ S∈S : i∈S ‖ai‖2 = ( n− 1 m− 1 ) n∑ i=1 ‖ai‖2 .
  Likewise, for the second term, we use ∑ j 6=i aj = −ai to get,∑
  i,j∈S: i6=j 〈ai, aj〉 = n∑ i=1 ∑ j 6=i ∑ S∈S : i,j∈S 〈ai, aj〉 = ( n− 2 m− 2 ) n∑ i=1 ∑ j 6=i 〈ai, aj〉 = − ( n− 2 m− 2 ) n∑ i=1 ‖ai‖2 .
  Therefore, we get,
  ES ∥∥∥∥∥ 1m∑ i∈S ai ∥∥∥∥∥ 2 = ( n−1 m−1 ) − ( n−2 m−2 ) m2 ( n m ) n∑ i=1 ‖ai‖2 = ( n−2 m−1 ) m2 ( n m ) n∑ i=1 ‖ai‖2 = n−m mn(n− 1) n∑ i=1 ‖ai‖2 .
  The next two lemmas are about the effect of the local updates in the local SGD literature. The first lemma has also appeared in (Karimireddy et al., 2020); we give the proof for completeness.
  Lemma 18. Consider f : Rd → R which is L-smooth and fix a w(0) ∈ Rd. Define the sequence (w(t)) of iterates produced by stochastic gradient descent with a fixed learning rate γ starting from w(0): w(t+1) = w(t) − γg(t) , where g(t) is an unbiased (and independent of w) estimator of ∇f(w) with bounded variance σ2. Fix a number τ of steps. If γ ≤ ( √ 2τL)−1, we have the bound
  τ−1∑ t=0 ‖w(t) − w(0)‖2 ≤ 8γ2τ2(τ − 1)‖∇f(w(0))‖2 + 4γ2τ2(τ − 1)σ2 .
  Proof. If τ = 1, we have nothing to prove. Assume now that τ ≥ 2. Let F (t) be the sigma-algebra generated by w(t) and denote Et[·] = E[· |F (t)]. We will use the inequality
  Et ∥∥∥g(t)∥∥∥2 = Et∥∥∥g(t) −∇f(w(t))∥∥∥2 + ∥∥∥∇f(w(t))∥∥∥2 ≤ σ2 + ∥∥∥∇f(w(t))∥∥∥2 . (17) We now successively deduce,
  Et‖w(t+1) − w(0)‖2 = ‖w(t) − w(0) − γg(t)‖2
  (a) ≤ ( 1 + 1
  τ − 1
  ) ‖w(t) − w(0)‖2 + γ2τEt‖g(t)‖2
  (b) ≤ ( 1 + 1
  τ − 1
  ) ‖w(t) − w(0)‖2 + 2γ2τ‖∇f(w(t))−∇f(w(0))‖2 + 2γ2τ‖∇f(w(0))‖2 + γ2τσ2
  (c) ≤ ( 1 + 1
  τ − 1 + 2γ2τL2
  ) ‖w(t) − w(0)‖2 + 2γ2τ‖∇f(w(0))‖2 + γ2τσ2
  (d) ≤ ( 1 + 2
  τ − 1
  ) ‖w(t) − w(0)‖2 + 2γ2τ‖∇f(w(0))‖2 + γ2τσ2 .
  Above, we used (a) the inequality 2αβ ≤ α2/δ2 + δ2β2 for reals α, β, δ, (b) Eq. (17), (c) Lsmoothness of f , and, (d) the condition on the learning rate.
  Let C = 2γ2τ‖∇f(w(0))‖2 + γ2τσ2. Unrolling the inequality and summing up the series gives for all t ≤ τ − 1
  ‖w(t) − w(0)‖2 ≤ C t−1∑ j=0 ( 1 + 2 τ − 1 )j ≤ C 2 (τ − 1) ( 1 + 2 τ − 1 )t
  ≤ C 2
  (τ − 1) ( 1 + 2
  τ − 1
  )τ−1 ≤ C
  2 (τ − 1)e2 ,
  where we used the bound (1 + 1/α)α ≤ e for all α > 0. Summing over t and using the numerical bound e2 < 8 completes the proof.
  Lemma 19. Consider the setting of Lemma 18. If γ ≤ (2τL)−1, we have the bound
  ‖w(τ) − w(0)‖2 ≤ 16γ2τ2‖∇f(w(0))‖2 + 8γ2τ2σ2 .
  Proof. Proceeding similar to the last proof (expect using δ = τ ) gives us
  Et ∥∥∥w(t+1) − w(0)∥∥∥2 ≤ (1 + 2 τ )∥∥∥w(t) − w(0)∥∥∥2 + 4γ2τ∥∥∥∇f(w(0))∥∥∥2 + 2γ2τσ2 . Unrolling and summing up the sequence completes the proof, similar to that of Lemma 18.
  Finally, the last lemma is about constants. Lemma 20. Let γu, γv, Lw, Lv, χ, f ∈ R+ and a natural number τ be given. Denote
  A := γuL 2 u + fγvχ 2LuLv , and, B := fγvL2v + γuχ 2LuLv .
  Suppose γu = cu/(τLu) and γv = cv/(τLv) with cu, cv > 0 satisfying
  cu, cv ≤ 1√ 6 max{1, χ−2} .
  Then, we have that
  γ2vχ 2LuLvB + γ 2 uL 2 uA ≤ A/(3τ2) , and, γ2uχ2LuLvA+ γ2vL2vB ≤ B/(3τ2) .
  Proof. Note that it suffices to show
  3τ2χ2γ2vLuLvB ≤ A/2 , and, 3τ2χ2γ2uLuLvA ≤ B/2 . Plugging in γu, γv , these are equivalent to
  6χ2fc3v + 6χ 4c2vcu ≤ χ2fcv + cu and, 6χ2c3u + 6χ4fcvc2u ≤ fcv + χ2cu .
  The assumption on cv implies that 6χ2fc3v ≤ χ2fcv and 6χ4c2vcu ≤ cu. Therefore, the first condition holds. Similarly, the second condition holds too.
sections.heading: B EXPERIMENTS: DETAILED SETUP AND HYPERPARAMETERS
sections.text: We conduct our experiments on three datasets from two modalities, namely images and text. The datasets contain a natural, non-iid split of data which is reflective of data heterogeneity encountered in federated learning. We describe in detail the experimental setup and hyperparameters. The code to reproduce the experimental results will be publicly released.
  The outline of this section is:
  • §B.1 describes the tasks and their associated datasets and metrics. • §B.2 describes the experimental pipeline as well as the baselines we compare to. • §B.3 presents the hyperparameters of all the algorithms.
  As discussed in §1, we take the weight αk to be proportional to the number of datapoints available on the device.
sections.heading: B.1 DATASETS, TASKS AND MODELS
sections.text: We consider three tasks motivated by real-world applications of federated learning. The tasks are summarized in Table 2 of the main paper and the distribution of data across the clients is visualized in Figure 5.
  For each model, we consider three partial personalization architectures:
  (a) Input layer personalization: Motivated by Liang et al. (2019), this architecture places the first layer on-device to learn a personalized representation per-client, while the rest of the model is shared. For the transformer model, we use the first transformer layer in place of the word embedding layer owing to its large size.
  (b) Output layer personalization: Motivated by Collins et al. (2021), this architecture learns a shared global representation but personalizes the prediction layer. For the transformer model, we use the last transformer layer in place of the last prediction layer owing to its large size.
  (c) Adapter personalization: We also consider a novel partial personalization architecture, where the full model is shared among all clients, while each client adds personalized adapter modules, which are lightweight modules added between layers of the shared model. We use the transformer adapters proposed by Houlsby et al. (2019) and residual adapters proposed by Rebuffi et al. (2017).
sections.heading: B.1.1 STACKOVERFLOW FOR NEXT WORD PREDICTION
sections.text: Dataset. The StackOverflow dataset comprises of questions and answers from the programming question-answer website stackoverflow.com. The goal of the next word prediction task is to predict the next word given a partial sequence of words in a question or answer. This task is a good opensource benchmark for next word predictions in mobile keyboards. We use the StackOverflow dataset provided by TensorFlow Federated.
  Client Distributions. Each client corresponds to one user on Stack Overflow; the data on the client corresponds to the questions and answers posted by this user. We only consider clients with at least 100 training sequences and 10 testing sequences, where a sequence refers to either a question or an answer. We use a fixed subsample of 1000 of them. Following Reddi et al. (2021), we restrict the vocabulary to the top 10000 most frequently occurring words in the dataset. We pad and truncate each sequence of each client to length 20 and consider at most 1000 training sequences on each client.
  Model. We use a transformer model (Vaswani et al., 2017) commensurate in size with BERT Mini (Turc et al., 2019). It has with 4 transformer blocks and 4 attention heads in each self-attention layer with a transformer hidden dimension of 256 and a fully-connected hidden dimension of 1024. The output layer is a causal language modeling head, i.e., a fully connected layer which assigns
  a score for each possible vocabulary item, including the special tokens. The model has 6 million parameters, which require around 23 megabytes of memory.
  Partial Personalization Architecture. The partial personalization architectures used are summarized in Table 5.
  Loss Function and Evaluation Metric. We train the model with the causal language modeling objective. That is, for each partial sequence, we treat the prediction of the next word as a multiclass classification problem to minimize the multinomial logistic loss, also known as cross entropy loss. For evaluation, we use the top-1 accuracy of predicting words in the proper 10000-word vocabulary (i.e., ignoring special tokens such as padding, out-of-vocabulary, and beginning/end of sequence).
sections.heading: B.1.2 GLDV2 FOR VISUAL LANDMARK RECOGNITION
sections.text: Dataset. GLDv2 stands for Google Landmarks Dataset v2 (Weyand et al., 2020), which is a largescale image dataset. It contains images of popular landmarks from around the world taken and uploaded by Wikipedia contributors. While the images vary in size, the most common image size is 800× 600 pixels. The goal of the visual landmark recognition task is to identify the landmark from its image. This task resembles a scenario where smartphone users take photos of natural and architectural landmarks while traveling. We use the federated version of the GLDv2 dataset introduced by Hsu et al. (2020) with 2028 landmarks and provided by TensorFlow Federated.
  Client Distributions. Each client corresponds to one Wikipedia user and contains all the images contributed by that user. We only all 823 clients with at least 50 datapoints. We do not use original test set from GLDv2 from evaluation as it comes from different clients. Instead, we take 50% of the data on each client as a testing set.
  Model. We use a ResNet-18 (He et al., 2016) model pretrained on ImageNet (Deng et al., 2009), with group normalization instead of batch normalization (Hsieh et al., 2020). We resize all images to 224×224. We use two data augmentations for training: a random crop from 256×256 and a random horizontal flip. The model has 12 million parameters, which require around 49 megabytes of storage.
  Partial Personalization Architecture. The partial personalization architectures used are summarized in Table 6.
  Loss Function and Evaluation Metric. We use the multinomial logistic loss, also known as cross entropy loss. We evaluate the performance of the model using its classification accuracy.
sections.heading: B.1.3 EMNIST FOR CHARACTER RECOGNITION
sections.text: Dataset. EMNIST (Cohen et al., 2017) is a character recognition dataset. The goal is to identify images of handwritten digits or letters; there are 62 possible options (a-z,A-Z, 0-9). The images are grey-scaled pictures of 28× 28 = 784 pixels. We use the EMNIST dataset provided by TensorFlow Federated.
  Client Distributions. Each client corresponds to one “writer”, i.e., the human subject who handwrote the digit/letter during the data collection process. We only use those clients with at least 100 training points and 25 testing points: there are 1114 of such clients.
  Model. We use a ResNet-18 (He et al., 2016) model with group normalization instead of batch normalization (Hsieh et al., 2020). We make two modifications to handle the smaller image size (28× 28× 1 as opposed to the 224× 224× 3 which the original ResNet was designed to accept): (a) we use a convolutional kernel of size 3× 3 rather than the original 7× 7 in the first convolution layer, and, (b) we drop the first pooling layer. The model has 11 million parameters, which require around 45 megabytes. Note that the number of parameters in this ResNet is smaller than the one for GLDv2 due to the architectural modifications we make for smaller images as well as the smaller number of classes.
  Partial Personalization Architecture. The partial personalization architectures used are summarized in Table 7.
  Loss Function and Evaluation Metric. We use the multinomial logistic loss, also known as cross entropy loss. We evaluate the performance of the model using its classification accuracy.
sections.heading: B.2 EXPERIMENTAL PIPELINE AND BASELINES
sections.text: There are three components in the training pipeline for all experiments:
  (a) Non-personalized federated training: The first step involves training a global model wg using the one-model-fits-all approach of (1) with FedAvg variants.
  (b) Personalized federated training: This optional second step involves training the shared parameters w together with the personalized parameters vk using a personalized federated learning approach. We warm-start w, vk from the non-personalized model wg from the previous step.
  (c) Final finetuning: The last step involves only finetuning the personalized parameters vk while the shared parameters w remain unchanged.
  For step (b), we initialize vk for each k to be the appropriate part of wg for input/output layer personalization. On the other hand, for adapters, we initialize vk to be equal to the same set of randomly initialized weights for each device k.
  We consider the following baselines:
  • Non-personalized: This denotes the performance of step (a) of the pipeline above, i.e., nonpersonalized federated training with FedAvg variants.
  • Full model personalization: We consider three baselines of personalization of the full model: (i) Finetune: The non-personalized model from step (a) of the pipeline above is finetuned
  locally on each client (step (c) of the pipeline). Step (b) is skipped for this baseline.
sections.heading: B.3 HYPERPARAMETERS AND EVALUATION DETAILS
sections.text: All the tuning of hyperparameters was performed on validation data, formed by holding out 20% of the training data on each device. Once the tuning was complete, we reran the experiments on the full training data, including those held out for validation.
  Evaluation Metric. Our primary evaluation metric is the weighted average of the test accuracy on each client, weighted by the number of test examples (the details of how the accuracy is computed on each dataset is given in §B.1 in the paragraph on “Loss Function and Evaluation Metric”). This corresponds to the unweighted accuracy obtained by pooling all the data locally, similar to the loss as discussed in §1. The same metric is used for hyperparameter tuning and is reported in all the tables and plots, unless explicitly noted otherwise.
  The final hyperparameters we use are given in Table 8.
  Rounds. We start with the number of communication rounds (i.e., the number of calls to secure aggregation routine for the shared parameters), which is used to measure the progress of each algorithm. For the non-personalized training, we use 1000 rounds for StackOverflow, 2500 rounds for GLDv2 and 2000 rounds for EMNIST. For the personalized training, we warm-start the model from the non-personalized one, and run the training for 500 rounds for StackOverflow and EMNIST and 600 rounds for GLDv2.
  Devices per Round. All devices are assumed to be available and selections are made uniformly at random. Following (Reddi et al., 2021; Weyand et al., 2020), we select 50 devices per round for
  StackOverflow/GLDv2 and 10 per round for EMNIST, for both the non-personalized as well as the personalized training.
  Local Updates and Minibatch Size. Each selected device locally runs 1 epoch of mini-batch stochastic gradient descent locally for non-personalized as well as personalized federated training. The final finetuning at the end of personalized training is performed for 5 epochs. We use a minibatch size of 64 for StackOverflow/GLDv2 and 32 for EMNIST for all settings.
  Server and Client Optimizer Details. We use FedAvg for EMNIST and FedAdam (Reddi et al., 2021) for StackOverflow and GLDv2. We also use a global scheduler, which applies a schedule on the client learning rates across rounds, while the client learning rate within each round is held constant. We use either a linear scheduler or an exponential scheduler (also called “stepLR” in PyTorch). A linear scheduler applies a linear warmup, if applicable, until the maximum learning rate followed by a linear decay to 0. An exponential scheduler halves the client learning rate once every fixed number of rounds. Both the client and server learning rates are tuned using the validation set.
  Regularization Coefficient for pFedMe and Ditto. We tune the regularization coefficient λk = λ for pFedMe and Ditto using the validation data from the set {10−4, 10−3, · · · , 100} of possible values. The tuned values are:
  • StackOverflow: 10−3 for Ditto and 10−4 for pFedMe, • GLDv2: 10−1 for both Ditto and pFedMe, • EMNIST: 10−1 for both Ditto and pFedMe.
  Random Seed. We report numbers averaged over 5 random seeds.
sections.heading: C EXPERIMENTS: ADDITIONAL RESULTS
sections.text: We now present the detailed experimental results.
  We start with Table 9, which shows the standard deviations of Table 4. We see that the numbers are fairly consistent across runs so the standard deviation is quite small (0.01 to 0.05 percentage points).
sections.heading: C.1 ABLATION: FINAL FINETUNING FOR FEDALT AND FEDSIM
sections.text: We now study the effect of the final finetuning (step (c) of the experimental pipeline; cf. §B.2) for FedAlt and FedSim.
  The final finetuning has a minimal impact on partial personalization. We see from Table 10 that the effect of the final finetuning is much smaller than the improvements from personalization. For instance, the improvements from finetuning are close to 0 for FedAlt on the StackOverflow dataset. For GLDv2, the finetuning accounts for < 0.5pp of improvement, whereas personalization overall accounts for 5 to 15pp.
  The final finetuning is more important to FedSim than FedAlt. Table 10 also shows that the final finetuning helps FedSim more than FedAlt. However, FedAlt still outperforms FedSim, as we saw in
  Table 9. Overall, this shows that FedAlt is a better algorithm than FedSim. The final finetuning helps FedSim make up some percentage points in accuracy, but not enough to make up its gap with FedAlt.
sections.heading: C.2 EFFECT OF PERSONALIZATION ON PER-DEVICE GENERALIZATION
sections.text: Summary of all scatter plots. All the scatter plots shown in the main paper are summarized in the violin plot of Figure 6. We see from the leftmost figure that the training accuracies on all devices improve with personalization. From the second figure, we see that the test accuracy of some of the devices reduces with personalization; this is true for both partial and full personalization.
  From the third plot of Figure 6, we see that regularization does not mitigate this overfitting. In fact, the regularization tuned for best average accuracy leads to a nearly identical distribution of test accuracies. A larger regularization reduces the spread of accuracies, but does so at the expense of a smaller median (white dot). The fourth plot of Figure 6 shows that the effect of dropout is similar. The best dropout improves the median accuracy, but it does not mitigate the issue of some devices being hurt by personalization.
  Train Accuracy plots for devices. From Figure 7, we see that personalization leads to a reduction in test accuracy on some of the devices beyond the initial non-personalized model. The corresponding train accuracy plot is given in Figure 7. We observe that the personalization always leads to an improvement in the training accuracy but not in the test accuracy. The analogous plots for GLDv2 are in Figure 8, where the trends are similar.
  Whether personalization helps a device or not depends on the random seed. We see in Figure 10 that the shaded region for some of the devices intersects the dotted line at 0. In other words, personalization sometimes helps this device and sometimes hurts it, depending on the random seed. This indicates that the best fix in practice is to use A/B testing on the deployed model to choose whether to use the personalized model or the non-personalized one.
  Regularization and dropout do not mitigate this issue. From the first row of Figure 9, we see that the weight decay with best mean accuracy exactly matches the unreguarlized case in terms of per-device statistics. Increasing the regularization weight can reduce the spread of per-device accuracy. However, this only leads to a worse mean accuracy and does not mitigate the issue of personalization hurting individual devices.
  From the second row of Figure 9, we see that the best dropout (0.3 in this case) leads to slight increase in average accuracy (0.18 pp). It also reduces the number of devices hurt by personalization from 256 out of 1000 to 193, but it does not fix this issue. Increasing dropout further only leads to a degradation of per-device statistics.
sections.heading: C.3 PARTIAL PERSONALIZATION FOR STATELESS DEVICES
sections.text: The algorithms we considered in this paper, namely FedAlt and FedSim, require the devices to maintain the personalized parameters vi’s as state across rounds. In cross-device federated learning settings, it is also interesting to consider stateless devices, which are not allowed to maintain state between training rounds.
  We give preliminary experiments in this setting. We modify the FedAlt and FedSim algorithms from the main paper so that the personalized parameters vi are reinitialized each time device i is chosen for participation. We warm-start vi from the appropriate part of the non-personalized model trained in step (a) of the pipeline. For adapters, we fix a random initialization once, and reuse it.
  FedAlt is better than FedSim for stateless devices, although the improvement is smaller. We see from Table 11 that all algorithms perform similarly for the stateless setting. Nevertheless, we see that FedAlt obtains mild improvements over both FedSim and finetuning for GLDv2, e.g., 0.24pp with adapters.
  The final finetuning is crucial for stateless devices. We see from Table 12 that the final finetuning accounts for most of improvements in the stateless case. For instance, for GLDv2, the final finetuning accounts for 11.68 and 10.42pp out of a total of 12.67 and 11.76pp for FedAlt and FedSim respectively. However, the personalized federated training (step (b) of the pipeline; cf. §B.2) still leads to an increase in accuracy of 1 to 1.34pp.