Towards Model Agnostic Federated Learning Using Knowledge Distillation

1 INTRODUCTION . Federated learning ( and more generally collaborative learning ) involves multiple data holders ( whom we call agents ) collaborating with each other to train their machine learning model over their collective data . Crucially , this is done without directly exchanging any of their raw data ( McMahan et al. , 2017 ; Kairouz et al. , 2019 ) . Thus , communication is limited to only what is essential for the training process and the data holders ( aka agents ) retain full ownership over their datasets . Algorithms for this setting such as FedAvg or its variants all proceed in rounds ( Wang et al. , 2021 ) . In each such round , the agents first train their models on their local data . Then , the knowledge from these different models is aggregated by averaging the parameters . However , exchanging knowledge via averaging the model parameters is only viable if all the agents use the same model architecture . This fundamental assumption is highly restrictive . Different agents may have different computational resources and hence may want to use different model architectures . Further , directly averaging the model parameters can fail even when all clients have the same architecture ( Wang et al. , 2019b ; Singh & Jaggi , 2020 ; Yu et al. , 2021 ) . This is because the loss landscape of neural networks is highly non-convex and has numerous symmetries with different parameter values representing the same function . Finally , these methods are also not applicable when using models which are not based on gradient descent such as random forests . To overcome such limitations , we would need to take a functional view view of neural networks i.e . we need methods that are agnostic to the model architecture and parameters . This motivates the central question investigated in this work : Can we design model agnostic federated learning algorithms which would allow each agent to train their model of choice on the combined dataset ? Specifically , we restrict the algorithms to access the models using only two primitives ( a universal model API ) : train on some dataset i.e . fit , and yield predictions on some inputs i.e . predict . Our goal is to be able to collaborate with and learn from any agent which provides these two functionalities . A naive such model agnostic algorithm indeed exists—agents can simply transfer their entire training data to each other and then each agent can train any model of choice on the combined dataset . However , transferring of the dataset is disallowed in federated learning . Instead , we will replace the averaging primitive in federated learning with knowledge distillation ( KD ) ( Bucilua et al. , 2006 ; Hinton et al. , 2015 ) . In knowledge distillation ( KD ) , information is transferred from model A to model B by training model B on the predictions of model A on some data . Since we only access model A through its predictions , KD is a functional model-agnostic protocol . The key challenge of KD however is that it is poorly understood and can not be formulated in the standard stochastic optimization framework like established techniques ( Wang et al. , 2021 ) . Thus , designing and analyzing algorithms that utilize KD requires developing an entirely new framework and approach . Our Contributions . The main results in this work are • We formulate the model agnostic learning problem as two agents with local datasets wanting to perform kernel regression on their combined datasets . Kernel regression is both simple enough to be theoretically tractable and rich enough to capture non-linear function fitting thereby allowing each agent to have a different kernel ( hence different models ) . • We analyze alternating knowledge distillation ( AKD ) and show that it is closely linked to the alternating projection method for finding the intersection of convex sets . Our analysis reveals that AKD sequentially loses information , leading to degradation of performance . This degradation is especially severe when the two agents have heterogeneous data . • Using the connection to alternating projection , we analyze other possible variants such as averaged knowledge distillation ( AvgKD ) and attempt to construct an ‘ optimal ’ scheme . • Finally , we evaluate all algorithms on real world deep learning models and datasets , and show that the empirical behavior closely matches our insights from the theoretical analysis . This demonstrates the utility of our framework for analyzing and designing new algorithms . 2 RELATED WORK . Federated learning ( FL ) . In FL ( Kairouz et al. , 2019 ) , training data is distributed over several agents or locations . For instance , these could be several hospitals collaborating on a clinical trial , or billions of mobile phones involved in training a voice recognition application . The purpose of FL is to enable training on the union of all agents ’ individual data without needing to transmit any of the raw sensitive data . Typically , the training is coordinated by some trusted server . One can also instead use direct peer-to-peer communications ( Nedic , 2020 ) . A large body of work has designed algorithms for FL under the identical model setting where we either learn a single global model ( McMahan et al. , 2017 ; Reddi et al. , 2020 ; Karimireddy et al. , 2020b ; a ; Wang et al. , 2021 ) , or multiple personalized models ( Wang et al. , 2019a ; Deng et al. , 2020 ; Mansour et al. , 2020 ; Grimberg et al. , 2021 ) . Knowledge distillation ( KD ) . Initially , KD was introduced as a way to compress models i.e . as a way to transfer the knowledge of a large model to a smaller model ( Bucilua et al. , 2006 ; Hinton et al. , 2015 ) . Since then , it has found much broader applications such as improving generalization performance via self-distillation , learning with noisy data , and transfer learning ( Yim et al. , 2017 ) . We refer to a recent survey ( Gou et al. , 2021 ) for progress in this vast area . KD in FL . Numerous works propose to use KD to transfer knowledge from the agent models to a centralized server model ( Seo et al. , 2020 ; Sattler et al. , 2020 ; Lin et al. , 2020 ; Li et al. , 2020 ; Wu et al. , 2021 ) . However , all of these methods rely on access to some common public dataset which may be impractical . In the closely related codistillation setting ( Zhang et al. , 2018 ; Anil et al. , 2018 ; Sodhani et al. , 2020 ) , an ensemble of students learns collaboratively without a central server model . While codistillation does not need additional unlabelled data , it is only suitable for distributed training within a datacenter since it assumes all agents have access to the same dataset . In FL however , there is both model and data heterogeneity . Further , none of these methods have a theoretical analysis . KD analysis . Despite the empirical success of KD , it is poorly understood with very little theoretical analysis . Phuong & Lampert ( 2021 ) explore a generalization bound for a distillation-trained linear model , and Tang et al . ( 2021 ) conclude that by using KD one re-weights the training examples for the student , and Menon et al . ( 2020 ) consider a Bayesian view showing that the student learns better if the teacher provides the true Bayes probability distribution . Allen-Zhu & Li ( 2020 ) show how ensemble distillation can preserve their diversity in the student model . Finally , Mobahi et al . ( 2020 ) consider self-distillation in a kernel regression setting i.e . the model is retrained using its own predictions on the training data . They show that iterative self-distillation induces a strong regularization effect . We significantly extend their theoretical framework in our work to analyze KD in federated learning where agents have different models and different datasets . 3 FRAMEWORK AND SETUP . Notation . We denote a set as A , a matrix asA , and a vector as a . A [ i , j ] is the ( i , j ) -th element of matrixA , a [ i ] denotes the i ’ th element of vector a . ||a|| denotes the ` 2 norm of vector a . Centralized kernel regression ( warmup ) . Consider , as a warm-up , the centralized setting with a training dataset D ⊆ Rd × R. That is , D = ∪Ni { ( xi , yi ) } , where xn ∈ X ⊆ Rd and yn ∈ Y ⊆ R. Given training set D , our aim is to find best function f ? ∈ F : X → Y . To find f ? we solve the following regularized optimization problem : f ? : = argmin f∈F 1 N ∑ n ( f ( xn ) − yn ) 2 + cRu ( f ) , with ( 1 ) Ru ( f ) : = ∫ X ∫ X u ( x , x′ ) f ( x ) f ( x′ ) dxdx′ . ( 2 ) Here , F is defined to be the space of all functions such that ( 2 ) is finite , c is the regularization parameter , and u ( x , x′ ) is a kernel function . That is , u is symmetric u ( x , x′ ) = u ( x′ , x ) and positive with Ru ( f ) = 0 only when f = 0 and Ru ( f ) > 0 o.w . Further , let k ( x , t ) be the function s.t . ∫ X u ( x , x′ ) k ( x′ , t ) dx′ = δ ( x− t ) , where δ ( · ) is the Dirac delta function . ( 3 ) Now , we can define the positive definite matrixK ∈ RN×N and vector kx ∈ RN as : K [ i , j ] : = 1 N k ( xi , xj ) and kx [ i ] : = 1 N k ( x , xi ) , for xi ∈ D , ∀i ∈ [ N ] . ( 4 ) Note that kx is actually a vector valued function which takes any x ∈ X as input , and both kx and K depend on the training data D. We can then derive a closed form solution for f ? . Proposition I ( Schölkopf et al . ( 2001 ) ) . The f ? which minimizes ( 1 ) is given by f ? ( x ) = k > x ( cI +K ) −1y , for y [ i ] : = yi , ∀i ∈ [ N ] . Note that on the training data X ∈ RN×d with X [ i , : ] = xi , we have f ? ( X ) = K ( cI +K ) −1y . Kernel regression for an input x outputs a weighted average of the training { yn } . These weights are computed using a learned measure of distance between the input x and the training { xn } . Intuitively , the choice of the kernel u ( x , x′ ) creates an inductive bias and corresponds to the choice of a model in deep learning , and the regularization parameter c acts similarly to tricks like early stopping , large learning rate , etc . which help in generalization . When c = 0 , we completely fit the training data and the predictions exactly recover the labels with f ? ( X ) = K ( K ) −1y = y . When c > 0 the predictions f ? ( X ) =K ( cI +K ) −1y 6= y and they incorporate the inductive bias of the model . In knowledge distillation , this extra information carried by the predictions about the inductive bias of the model is popularly referred to as “ dark knowledge ” ( Hinton et al. , 2015 ) . Federated kernel regression ( our setting ) . We have two agents , with agent 1 having dataset D1 = ∪N1i { ( x1i , y1i ) } and agent 2 with dataset D2 = ∪ N2 i { ( x2i , y2i ) } . Agent 1 aims to find the best approximation mapping g1 ? ∈ F1 : X → Y using a kernel u1 ( x , x′ ) and objective : g1 ? : = arg min g∈F1 1 N1 +N2 ∑ n ( g ( x1n ) − y1n ) 2 + 1 N1 +N2 ∑ n ( g ( x2n ) − y2n ) 2 + cRu1 ( g ) , with ( 5 ) Ru1 ( g ) : = ∫ X ∫ X u1 ( x , x ′ ) g ( x ) g ( x′ ) dxdx′ . ( 6 ) Note that the objective of agent 1 is defined using its individual kernel u1 ( x , x′ ) , but over the joint dataset ( D1 , D2 ) . Correspondingly , agent 2 also uses its own kernel function u2 ( x , x′ ) to define the regularizer Ru2 ( g 2 ) over the space of functions g2 ∈ F2 and optimum g2 ? . Thus , our setting has model heterogeneity ( different kernel functions u1 and u2 ) and data heterogeneity ( D1 and D2 are not i.i.d. ) . Given that the setting is symmetric between agents 1 and 2 , we can focus solely on error in terms of agent 1 ’ s objective ( 5 ) without loss of generality . Proposition I can be used to derive a closed form form for function g1 ? minimizing objective ( 5 ) . However , computing this requires access to the datasets of both agents . Instead , we ask “ can we design an iterative federated learning algorithm which can approximate g1 ? ” ?

Summary: This paper focuses on the setting of federated learning where the two agents are attempting to perform kernel regression using different kernels (and hence have different models). Their study yields an algorithm of using alternating knowledge distillation (AKD) imposes overly strong regularization and may lead to severe under-fitting. Their theory also shows an interesting connection between AKD and the alternating projection algorithm for finding the intersection of sets. Leveraging this connection, they propose an algorithm that improves upon AKD.