TOWARDS MODEL AGNOSTIC FEDERATED LEARNING USING KNOWLEDGE DISTILLATION
Year: 2021

Abstract:
An often unquestioned assumption underlying most current federated learning algorithms is that all the participants use identical model architectures. In this work, we initiate a theoretical study of model agnostic communication protocols which would allow data holders (agents) using different models to collaborate with each other and perform federated learning. We focus on the setting where the two agents are attempting to perform kernel regression using different kernels (and hence have different models). Our study yields a surprising result—the most natural algorithm of using alternating knowledge distillation (AKD) imposes overly strong regularization and may lead to severe under-fitting. Our theory also shows an interesting connection between AKD and the alternating projection algorithm for finding intersection of sets. Leveraging this connection, we propose a new algorithm which improves upon AKD, but at the cost of using multiple models. Our theoretical predictions also closely match real world experiments using neural networks. Thus, our work proposes a rich yet tractable framework for analyzing and developing new practical model agnostic federated learning algorithms.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 FRAMEWORK AND SETUP
text: 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?”?

heading: 4 ALTERNATING KNOWLEDGE DISTILLATION
text: In this section, we describe a popular iterative knowledge distillation algorithm and analyze its updates in our framework. Our analysis leads to some surprising connections between KD and projection onto convex sets, and shows some limitations of the current algorithm.
Algorithm. Denote the data on agent 1 as D1 = (X1,y1) where X1[i, :] = x1n and y1[i] = y1i . Correspondingly, we have D2 = (X2,y2). Now starting from ŷ10 = y1, in each rounds t and t+ 1:
a. Agent 1 trains their model on dataset (X1, ŷ1t ) to obtain g 1 t . b. Agent 2 receives g1t and uses it to predict labels ŷ 2 t+1 = g 1 t (X
2). c. Agent 2 trains their model on dataset (X2, ŷ2t+1) to obtain g 1 t+1. d. Agent 1 receives a model g1t+1 from agent 2 and predicts ŷ 1 t+2 = g 1 t+1(X 1).
Thus the algorithm alternates between training and knowledge distillation on each of the two agents. We also summarize the algorithm in Figure 1a. Importantly, note that there is no exchange of raw data but only of the trained models. Further, each agent trains their choice of a model on their data with agent 1 training {g1t } and agent 2 training {g1t+1}. Superscript 1 means we start AKD from agent 1.

heading: 4.1 THEORETICAL ANALYSIS
text: Similar to (3), let us define functions k1(x,x′) and k2(x,x′) such that they satisfy∫ X ua(x,x ′)ka(x ′, t)dx′ = δ(x− t) for a ∈ {1, 2} . For such functions, we can then define the following positive definite matrixL ∈ R(N1+N2)×(N1+N2):
L = ( L11 L12 L21 L22 ) , La,b[i, j] =
1
N1 +N2 k1(x
a i ,x b j) for a, b ∈ {1, 2} and i ∈ [N1], j ∈ [N2] .
Note L is symmetric (with L>12 = L21) and is also positive definite. Further, each component La,b measures pairwise similarities between inputs of agent a and agent b using the kernel k1. Correspondingly, we defineM ∈ R(N1+N2)×(N1+N2) which uses kernel k2:
M = ( M11 M12 M21 M22 ) , Ma,b[i, j] =
1
N1 +N2 k2(x
a i ,x b j) and i ∈ [N1], j ∈ [N2] .
We can now derive the closed form of the AKD algorithm repeatedly using Proposition I. Proposition II. The model in round 2t learned by the alternating knowledge distillation algorithm is
g12t(x) = l > x (cI +L11) −1 (M12(cI +M22)−1L21(cI +L11)−1)t y1 , where lx ∈ RN1 is defined as lx[i] = 1N1+N2 k1(x,x 1 i ). Further, for any fixed x we have
lim t→∞
g1t (x) = 0 .
First, note that if agents 1 and 2 are identical with the same data and same model, we haveM12 = M22 = L11 = L12. This setting corresponds to self-distillation where the model is repeatedly retrained on its own predictions. Proposition II shows that after 2t rounds of self-distillation, we
obtain a model is of the form g12t(x) = l > x (cI +L11) −1 ( L11 (cI +L11) −1 )2t y. Here, the effect
of c is amplified as t increases. Thus, this shows that repeated self-distillation induces a strong regularization effect, recovering the results of (Mobahi et al., 2020).
Perhaps more strikingly, Proposition II shows that not only does AKD fail to converge to the actual optimal solution g1? as defined in (5), it will also slowly degrade and eventually converges to 0. We next expand upon this and explain this phenomenon.

heading: 4.2 DEGRADATION AND CONNECTION TO PROJECTIONS
text: While mathematically, Proposition II completely describes the AKD algorithm, it does not provide much intuition. In this section, we rephrase the result in terms of projections and contractions which provides a more visual understanding of the method.
Train Train
2 rounds
(a) Alternating KD starting from agent 1. We predict and train using predictions, alternating between the two agents.
0
(b) Alternating oblique projections starting from x converges to 0.
Oblique projections. A projection operator P linearly maps (projects) all inputs onto some linear subspace A = Range(A). In general, we can always rewrite such a projection operation as
Px = min y∈Range(A)
(y − x)>W (y − x) = A(A>WA)−1A>Wx .
Here, A defines an orthonormal basis of the space A and W is a positive definite weight matrix which defines the geometry 〈x,y〉W := x>Wy. WhenW = I , we recover the familiar orthogonal projection. Otherwise, projections can happen ‘obliquely’ following the geometry defined byW .
Contractions. A contraction is a linear operator C which contracts all inputs towards the origin: ‖Cx‖ ≤ ‖x‖ for any x .
Given these notions, we can rewrite Proposition II as follows. Proposition III. There exist oblique projection operators P>1 and P>2 , contraction matricesC1 and C2, and orthonormal matrices V1 and V2 such that the model in round 2t learned by the alternating knowledge distillation algorithm is
g12t(x) = l > x (cI +L11) −1V >1 ( C1P > 2 C2P > 1 )t V1y 1 .
In particular, the predictions on agent 2’s inputs are
g12t(X 2) = V >2 P > 1 ( C1P > 2 C2P > 1 )t V1y 1 .
Further, the projection matrices satisfy P>1 P > 2 x = x only if x = 0 .
The term ( C1P > 2 C2P > 1 )t is the only one which depends on t and so captures the dynamics of the algorithm. Two rounds of AKD (first to agent 1 and then back to 2) correspond to a multiplication with C1P>2 C2P > 1 i.e. alternating projections interspersed by contractions. Thus, the dynamics of AKD is exactly the same as that of alternating projections interspersed by contractions. The projections P>1 and P > 2 have orthogonal ranges whose intersection only contains the origin 0. As
Fig. 1b shows, non-orthogonal alternating projections between orthogonal spaces converge to the origin. Contractions further pull the inputs closer to the origin, speeding up the convergence.
Remark 1. To understand the connection to projection more intuitively, suppose that we had a 4-way classification task with agent 1 having data only of the first two classes, and agent 2 with the last two classes. Any model trained by agent 1 will only learn about the first two classes and its predictions on the last two classes will be meaningless. Thus, no information can be transferred between the agents in this setting. More generally, the transfer of knowledge from agent 1 to agent 2 is mediated by its dataX2. This corresponds to a projection of the knowledge of agent 1 onto the data of agent 2. If there is a mismatch between the two, knowledge is bound to be lost.
Speed of degradation. The alternating projection algorithm converges to the intersection of the subspaces corresponding to the projection operators (their range) (Boyd & Dattorro, 2003). In our particular case, the fixed point of P>1 and P > 2 is 0, and hence this is the point the algorithm will converge to. The contraction operations C1 and C2 only speed up the convergence to the origin 0 (also see Fig. 1b). This explains the degradation process notes in Proposition II. We can go further and examine the rate of degradation using known analyses of alternating projections (Aronszajn, 1950). Proposition IV (Informal). The rate of convergence to g1t (x) to 0 gets faster if:
• a stronger inductive bias is induced via a larger regularization constant c,
Train
Train
1 round
(a) AvgKD scheme
01 23 4
(b) Intuition behind ensemble scheme. In each round, AKD alternates between overfitting the data of agent 1 or of agent 2. We can construct an ensemble out of these models to correct for this bias and quickly converge to the true optimal.
• the kernels k1(x,y) and k2(x,y) are very different, or • the difference between the datasets D1 and D2 as measured by k1(x,y) increases.
In summary, both data and model heterogeneity may speed up the degradation defeating the purpose of model agnostic FL. All formal proofs and theorem statements are moved to the Appendix.

heading: 5 ADDITIONAL VARIANTS
text: In the previous section, we saw that the alternating knowledge distillation (AKD) algorithm over multiple iterations suffered slow degradation, eventually losing all information about the training data. In this section, we explore some alternative approaches which attempt to correct this. We first analyze a simple way to re-inject the training data after every KD iteration which we call averaged distillation. Then, we show an ensemble algorithm that can recover the optimal model g1?.

heading: 5.1 AVERAGED KNOWLEDGE DISTILLATION
text: As we saw earlier, each step of knowledge distillation seems to lose some information about the training data, replacing it with the inductive bias of the model. One approach to counter this slow loss of information is to recombine it with the original training data labels such as is commonly done in co-distillation (Sodhani et al., 2020).
Algorithm. Recall that agent 1 has data D1 = (X1,y1) and correspondingly agent 2 has data D2 = (X2,y2). Now starting from ŷ10 = y1, ŷ20 = y2, in each round t ≥ 0:
a. Agents 1 and 2 train their model on datasets (X1, ŷ1t ) and (X 2, ŷ2t ) to obtain g 1 t and g 2 t . b. Agents exchange their models g1t and g 2 t between each other. c. Agents use exchanged models to predict labels ŷ1t+1 = y1+g2t (X 1) 2 , ŷ 2 t+1 = y2+g1t (X 2) 2 .
The summary the algorithm is depicted in Figure 2a. Again, notice that there is no exchange of raw data but only of the trained models. The main difference between AKD and AvgKD (averaged knowledge distillation) is that we average the predictions with the original labels. This re-injects information y1 and y2 at every iteration, preventing degradation. We theoretically characterize its dynamics next in terms of the afore-mentioned contraction and projection operators.
Proposition V. There exist oblique projection operators P1 and P2, contraction matrices C1 and C2, and orthonormal matrices V1 and V2 such that the model of agent 1 in round t learned by the averaged knowledge distillation (AvgKD) algorithm is
g1t (x) = F
2 ( t−1∑ i=0 ( C1P > 2 C2P > 1 4 )i)z1 + FC1P > 2 4 ( t−2∑ i=0 ( C2P > 1 C1P > 2 4 )i)z2 .
where F = l>x (cI +L11) −1V >1 . Further, in the limit of rounds for any fixed x we have
lim t→∞
g1t (x) = F 2 (I − C1P
> 2 C2P > 1
4 )†z1 +
FC1P > 2
4 (I − C2P
> 1 C1P > 2
4 )†z2 .
This shows that the model learned through AvgKD does not degrade to 0, unlike AKD. Instead, it converges to a limit for which we can derive closed-form expressions. Unfortunately, this limit model is still not the same as our desired optimal model g1?. We next try to overcome this using ensembling.

heading: 5.2 ENSEMBLED KNOWLEDGE DISTILLATION
text: We first analyze how the limit solution of AvgKD differs from the actual optimum g1?. We will build upon this understanding to construct an ensemble that approximates g1?.
Understanding AvgKD. For simplicity, we assume that the regularization coefficient c = 0. Then,
Proposition VI. There exist matricesA1 andA2 such that the minimizer g1 ? of objective (5) predicts g1 ? (Xi) = A1β1 +A2β2 for i ∈ {1, 2}, where β1 and β2 satisfy(
L11 M12 L21 M22 )( β1 β2 ) = ( y1 y2 ) . (7)
In contrast, for the same matricesA1 andA2, the limit models of AvgKD predict g1∞(Xi) = 1 2A1β1 and g2∞(Xi) = − 12A2β2, for i ∈ {1, 2} for β1 and β2 satisfying( L11 M12 2
L21 2 M22
)( β1 β2 ) = ( y1 −y2 ) . (8)
By comparing the equations (7) and (8), the output 2(g1∞(x) − g2∞(x)) is close to the output of g1
? (x), except that the off-diagonal matrices are scaled by 12 and we have −y2 on the right hand side. We need two tricks if we want to approximate g1?: first we need an ensemble using differences of models, and second we need to additionally correct for the bias in the algorithm.
Correcting bias using infinite ensembles. Consider the initial AKD (alternating knowledge distillation) algorithm illustrated in the Fig. 1a. Let us run two simultaneous runs, ones starting from agent 1 and another starting from agent 2, outputting models {g1t }, and {g2t } respectively. Then, instead of just using the final models, we construct the following infinite ensemble. For an input x, we output:
f∞(x) = ∞∑ t=0 (−1)t(g1t (x) + g2t (x)) (9)
That is, we take the models from odd steps t with positive signs and from even ones with negative signs and sum their predictions. We call this scheme Ensembled Knowledge Distillation (EKD). The intuition behind our ensemble method is schematically visualized in 1-dimensional case in the Fig. 2b where numbers denote the t variable in the equation (9). We start from the sum of both agents models obtained after learning from ground truth labels (0-th round). Then we subtract the sum of both agents models obtained after the first round of KD. Then we add the sum of both agents models obtained after the second round of KD and so on. From the section 4.2 we know that in the AKD process model gradually degrades towards 0, in other words, each next round obtained model adds less value to the whole sum. Hence we gradually converge to the limit model, which is the point∞ in the Fig. 2b. We formalize this and prove the following.
Proposition VII. The predictions of f∞ using (9) satisfies f∞(Xi) = g1 ? (Xi) for i ∈ {1, 2}.
Thus, not only did we succeed in preventing degeneracy to 0, but we also managed to recover the predictions of the optimal model. However, note that this comes at a cost of an infinite ensemble. While we can approximate this using just a finite set of models (as we explore experimentally next), this still does not recover a single model which matches g1?.

heading: 6 EXPERIMENTS

heading: 6.1 SETUP
text: We consider three settings corresponding to the cases Proposition IV with the agents having
• the same model architecture and close data distributions (Same model, Same data) • different model architectures and close data distributions (Different model, Same data) • the same model architecture and different data distributions (Same model, Different data).

heading: Test Acc. of agents in AKD

heading: Test Acc. of agents in AvgKD
text: The toy experiments solve a linear regression problem of the form Ax? = b. The data A and b is split between the two agents randomly in the ‘same data’ case, whereas in the ‘different data’ case the data is sorted according to b before splitting to maximize heterogeneity.
The real world experiments are conducted using Convolutional Neural Network (CNN), Multi-Layer Perceptron network (MLP), and Random Forest (RF). We use squared loss since it is closer to the theoretical setting. In the ‘same model’ setting both agents use the CNN model, whereas agent 2 instead uses an MLP in the ‘different model’ setting. Further, we split the training data randomly in the ’same data’ setting. For the ’different data’ setting, we split the data by labels and take some portion Alpha of data from each agent and randomly shuffle taken points between agents. By varying hyperparameter Alpha we control the level of data heterogeneity between agents. Notice that if Alpha = 0 then datasets are completely different between agents, if Alpha = 1 then we have the ’same data’ setting with i.i.d. split of data between two agents. All other details are presented in the Appendix A. We next summarize and discuss our results.

heading: 6.2 RESULTS AND DISCUSSION
text: AvgKD > AKD. In all settings (both synthetic in Fig. 3 and real world in Fig. 4), we see that with the increasing number of rounds the performance of AKD significantly degrades whereas that of AvgKD stabilizes (Fig. 5) regardless of regularization, model and data heterogeneity. Moreover, from the experiments on MNIST in Fig. 4, we see the following: AKD degrades faster if there is

heading: Test Acc. of agents in AvgKD
text: regularization used, model heterogeneity or data heterogeneity between agents, and the last plays the most significant role. AvgKD, on the other hand, quickly converges in a few rounds and does not degrade. However, it fails to match the centralized accuracy as well.
EKD works but needs large ensembles. In the synthetic setting (Fig. 3), in 250 rounds (ensemble of 500 models) it even matches the centralized model. However, in the real world setting (Fig. 9) the improvement is slower and it does not match centralized performance. This might be due to the small number of rounds run (only 10). EKD is also the only method that improves with subsequent rounds. Finally, we observed that increasing regularization actually speeds up the convergence of EKD, while it worsened in the cases of AKD and AvgKD.
Data heterogeneity is the main bottleneck. In all our experiments, both data and model heterogeneity degraded the performance of AKD and AvgKD. However, data heterogeneity has a much stronger effect. This confirms our theory that mismatch between agents data leads to loss of information when using knowledge distillation. Overcoming this data heterogeneity is the crucial challenge for practical model agnostic FL. Fig. 8 and 6 show how AKD and AvgKD schemes behave in dependence on data heterogeneity. Indeed, the higher data heterogeneity the faster speed of degradation for the AKD scheme. In the case of the AvgKD scheme, we see that agents do improve over their local models and this improvement is larger with more data heterogeneity. However

heading: 6.3 ADDITIONAL EXTENSIONS AND EXPERIMENTS
text: In App. G, we extend our algorithms to M agents and show experimentally in App. H.6 that the same trends hold there as well. Our conclusions also hold for the cross-entropy loss (Figs. 10, 11), for highly heterogeneous model case with MLP and Random Forests (Figs. 14, 15), as well on other datasets and models (VGG on CIFAR10 in Figs. 12, 13). In the latter, we see higher speeds of degradation which is probably due to CIFAR10 being much more complex (meaning potentially larger data heterogeneity), and also more complex models such as VGG (larger model heterogeneity). Finally, in Fig. 17, we compare AvgKD to the more standard FedAvg and yet another variant (parallel KD or PKD). AvgKD consistently outperforms these methods indicating it is most promising variant.

heading: 7 CONCLUSION
text: While the stochastic optimization framework has been very useful in analyzing and developing new algorithms for federated learning so far, it fundamentally cannot capture learning with different models. We instead introduced the federated kernel regression framework where we formalized notions of both model and data heterogeneity. Using this, we analyzed different knowledge distillation schemes and came to the conclusion that data heterogeneity poses a fundamental challenge limiting the knowledge that can be transmitted. Further, these theoretical predictions were exactly reflected in our deep learning experiments as well. Overcoming this data heterogeneity will be crucial to making KD based model agnostic federated algorithms practical.
We also utilized our framework to design a novel ensembling method motivated by our theory. However, this method could require very large ensembles (up to 500 models) in order to match the centralized performance. Thus, we view our method not as a practical algorithm, but more of a demo on how our framework can be leveraged. Similarly, our experiments are preliminary and are not on real world datasets. We believe there is great potential in further exploring our results and framework, especially in investigating how to mitigate the effect of data heterogeneity in knowledge distillation.

heading: A EXPERIMENTS DETAILS
text: In all real world experiments we use the Adam optimizer with a default regularization (weight decay) of 3× 10−4, unless in the ‘no regularization’ case when it is set to 0. We split the data between 2 agents by giving a bigger part of data to agent 1 at all ’same data’ experiments. That is, for ’different data’ experiments and heterogeneous data experiments where we vary Alpha hyperparameter we split data equally between 2 agents. The non-equal split can help us to see another effect in the experiments: if agent 2 (with less data) benefits from the communication with agent 1 (with more data). In heterogeneous data experiments, we explore the significance of data heterogeneity effect on the behavior of KD scheme. That is why we design an almost ideal setting at all such experiments: the ’same model’ setting (except experiments with RF and MLP), the equal split of data in terms of its amount between agents.
Toy experiments. The toy experiments solve a linear regression problem of the form Ax? = b whereA ∈ Rn×d and x? ∈ Rd is randomly generated for n = 1.5d and d = 100. Then, the dataA and b is split between the two agents at proportion 0.6/0.4. This is done randomly in the ‘same data’ case, whereas in the ‘different data’ case the data is sorted according to b before splitting to maximize heterogeneity. This experiment with the linear kernel is supposed to show if our theory is correct, especially if the EKD scheme really works.
Real world experiments. The real world experiments are conducted using CNN and MLP networks on MNIST, MLP network and RF model on MNIST, and VGG161 (Simonyan & Zisserman, 2015) and CNN models on CIFAR10 datasets. Further, we split the training data randomly with a 70%− 30% split in the ’same data’ setting. For the ’different data’ setting, we split the data by labels: agent 1 has ’0’ to ’4’ labeled data points, agent 2 has ’5’ to ’9’. Then we take randomly from each agent some Alpha = 0.1 portion of data, combine it and randomly return data points to both agents from this combined dataset.

heading: B ALTERNATING KD WITH REGULARIZATION
text: Different models. Keeping the notation of section 4 we can construct the following matrix:
K = ( L 0 0 M ) = L11 L12 0 0L21 L22 0 00 0 M11 M12 0 0 M21 M22  . Notice that each of the sub-blocks L and M are symmetric positive semi-definite matrices. This makes matrixK symmetric positive semi-definite matrix and it has eigendecomposition form:
K = V >DV and V = ( V1V2Ṽ1Ṽ2 ) ,
whereD and V are R2(N1+N2)×2(N1+N2) diagonal and orthogonal matrices correspondingly. This means that the R2(N1+N2)×2(N1+N2) matrices V1,V2, Ṽ1, Ṽ2 are also orthogonal to each other. Then one can deduce:
Lij = V > i DVj , and Mij = Ṽ > i DṼj ∀i, j = 1, 2 .
In AKD setting only agent 1 has labeled data. The solution of learning from the dataset D1 evaluated at set X2 is the following:
g10(X2) = L21(cI +L11)−1y1 = V >2 ((V >1 (cI +D)V1)−1V >1 D)>y1 . (10)
Let us introduce notation:
P1 = V1(V > 1 (cI +D)V1) −1V >1 (cI +D) and P̃2 = Ṽ2(Ṽ > 2 (cI +D)Ṽ2) −1Ṽ >2 (cI +D) .
These are well known oblique (weighted) projection matrices on the subspaces spanned by the columns of matrices V1 and Ṽ2 correspondingly, where the scalar product is defined with Gram matrixG = cI +D. Similarly one can define P2 and P̃1.
1This model is also a convolutional neural network, but in our experiments it is bigger than CNN model.
Given the introduced notation and using the fact that V >2 V1 = 0, we can rewrite equation 10 as:
g10(X2) = V >2 P>1 V1y1 . (11)
Similarly, given X1 ⊆ D1 and agent 1 learned from Ŷ2 = g10(X2), inferred prediction Ŷ1 = g11(X1) by agent 2 can be written as:
g11(X1) = Ṽ >1 P̃>2 Ṽ2V >2 P>1 z1, where z1 = V1y1 .
At this point we need to introduce additional notation:
C1 = V1Ṽ > 1 and C̃2 = Ṽ2V > 2 .
These are matrices of contraction linear mappings.
One can now repeat the whole process again with replacement Y1 by Ŷ1. The predictions by agent 1 and agent 2 after such 2t rounds of AKD are:
g12t(X2) = V >2 P>1 ( C1P̃ > 2 C̃2P > 1 )t z1 and g12t+1(X1) = Ṽ >1 P̃>2 C̃2P>1 ( C1P̃ > 2 C̃2P > 1 )t z1 .
(12)
If one considers the case where agents have the same kernel u1 = u2 = u then one should ’remove all tildas’ in the above expressions and the obtained expressions are applied. This meansM = L,
V1 = Ṽ1 and V2 = Ṽ2. Crucial changes in this case are C1 = C̃2 → I and ( C1P̃ > 2 C̃2P > 1 )t →(
P>2 P > 1
)t . The matrix ( P>2 P > 1 )t is an alternating projection (Boyd & Dattorro, 2003) operator
after t steps. Given 2 closed convex sets, alternating projection algorithm in the limit finds a point in the intersection of these sets, provided they intersect (Boyd & Dattorro, 2003; Lewis et al., 2007). In our case matrices operators P1 and P2 project onto linear spaces spanned by columns of matrices V1 and V2 correspondingly. These are orthogonal linear subspaces, hence the unique point of their intersection is the origin point 0. This means that ( P>2 P > 1 )t −−−→ t→∞
0 and in the limit of such AKD procedure both agents predict 0 for any data point.
Speed of degradation. The speed of convergence for the alternating projections algorithm is known and defined by the minimal angle φ between corresponding sets (Aronszajn, 1950) (only non-zero elements from sets one has to consider):
|| (P2P1)t v|| ≤ (cos(φ))2t−1||v|| = (cos(φ))2t−1 as ||v|| = 1 ,
where v is one of the columns of matrix V1. In our case we can write the expression for the cosine:
cos(φ) = max v1,v2
( |v>1 (cI +D)v2|√
(v>1 (cI +D)v1) · (v>2 (cI +D)v2) ) = max v1,v2 ( |v>1 Dv2|√ (c+ v>1 Dv1) · (c+ v>2 Dv2) ) .
where v1 and v2 are the vectors from subspaces spanned by columns of matrices V1 and V2 correspondingly. Hence the speed of convergence depends on the elements of the matrix L12. Intuitively these elements play the role of the measure of ’closeness’ between data points. This means that the ’closer’ points of sets X1 and X2 the higher absolute values of elements in matrix L12. Moreover, we see inversely proportional dependence on the regularization constant c. All these sums up in the following proposition which is the formal version of the proposition IV: Proposition VIII (Formal). The rate of convergence of g1t (x) to 0 gets faster if:
• larger regularization constant c is used during the training, • smaller the eigenvalues of the matrix V1Ṽ >1 , or • smaller absolute value of non-diagonal block L12

heading: C ALTERNATING KD WITHOUT REGULARIZATION
text: Before we saw that models of both agents degrade if one uses regularization. A natural question to ask if models will degrade in the lack of regularization. Consider the problem (1) with c → 0, so
the regularization term cancels out. In the general case, there are many possible solutions as kernel matrix K may have 0 eigenvalues. Motivated by the fact that Stochastic Gradient Descent (SGD) tends to find the solution with minimal norm (Wilson et al., 2017), we propose to analyze the minimal norm solution in the problem of alternating knowledge distillation with 2 agents each with private dataset. For simplicity let us assume that agents have the same model architecture. Then the agent I solution evaluated at set X2 with all the above notation can be written as:
g?0(X2) =K21K † 11y1 = V > 2 DV1(V > 1 DV1) †y1, (13)
where † stands for pseudoinverse. In this section let us consider 3 possible settings one can have:
1. Self-distillation: The datasets and models of both agents are the same. 2. Distillation withK > 0: The datasets of both agents are different and private. Kernel matrix K is positive definite.
3. Distillation withK ≥ 0: The datasets of both agents are different and private. Kernel matrix K is positive semi-definite.
Self-distillation. Given the dataset D = X × Y , the solution of the supervised learning with evaluation at any x ∈ Rd is:
g?0(x) = l > x (V >DV )†y.
Then the expression for the self-distillation step with evaluation at any x ∈ Rd is:
g?1(x) = l > x (V >DV )†V >DV (V >DV )†y = l>x (V >DV )†y = g?0(x),
where property of pseudoinverse matrix was used: A†AA† = A†. That is, self-distillation round does not change obtained model. One can repeat self-distillation step t times and there is no change in the model:
g?t (x) = l > x (V >DV )†y = g?0(x),
Hence in the no regularization setting self-distillation step does not give any change in the obtained model.
Distillation withK > 0. In this setting we have the following identityK† =K−1 and results are quite similar to the setting with regularization. But now the Gram matrix of scalar product in linear space isD instead of cI +D in the regularized setting. By assumption on matrixK it follows that D is of full rank and the alternating projection algorithm converges to the origin point 0. Therefore in the limit of AKD steps predictions by models of two agents will degrade towards 0.
Distillation withK ≥ 0. In this setting kernel matrixK has at least one 0 eigenvalue and we take for the analysis the minimal norm solution (13). We can rewrite this solution:
g?0(X2) = V >2 DV1(V >1 DV1)†y1 = V >2 P̂>1 z1,
where P̂1 = V1(V >1 DV1) †V >1 D and z1 = V1y1.
One can notice the following projection properties of matrix P̂1:
P̂ 21 = P̂1 and P̂1V1(V > 1 DV1) † = V1(V > 1 DV1) † .
That is, P̂1 is a projection matrix with eigenspace spanned by columns of the matrix V1(V >1 DV1) †. Similarly, one can define P̂2. Then the solution for the first AKD step evaluated at set X1 is:
g?1(X1) = V >1 DV2(V >2 DV2)†V >2 P̂>1 z1 = V >1 P̂>2 P̂>1 z1,
and after t rounds we obtain for agent I and agent II correspondingly:
g?2t(X2) = V >2 P̂>1 (P̂>2 P̂>1 )tz1 and g?2t+1(X1) = V >1 (P̂>2 P̂>1 )t+1z1.
In the limit of the distillation rounds operator (P̂2P̂1)t tends to the projection on the intersection of 2 subspaces spanned by the columns of matrices V1(V >1 DV1) † and V2(V >2 DV2) †. We should
0
(a) Illustrative projection 3D
Train
Train
1 round
(b) PKD scheme
highlight that in general, the intersection set in this case consists not only from the origin point 0 but some rays that lie simultaneously in the eigenspaces of both projectors P̂1 and P̂2. The illustrative example is shown in the Fig. 7a. We perform the alternating projection algorithm between orthogonal linear spaces that have a ray in the intersection and we converge to some non-zero point that belongs to this ray. The intersection of eigenspaces of both projectors P̂1 and P̂2 is the set of points x ∈ Span(V1(V >1 DV1)†) s.t. P̂1P̂2x = x.

heading: D AVERAGED KD
text: In this section we present the detailed analysis of the algorithm presented in the section 4.2 keeping the notation of the section B. As a reminder, models of both agents after round t of AvgKD algorithm are as follows:
g1t (X2) = 1
2 L21(cI +L11)
−1(y1 + g 2 t−1(X1)) =
1 2 V >2 P > 1 (z1 + g 2 t−1),
g2t (X1) = 1
2 M12(cI +M22)
−1(g1t−1(X2) + y2) = 1
2 Ṽ >1 P̃ > 2 (g 1 t−1 + z2).
where
g1t−1 = 1
2 (cI +D)V1(cI +L11)
−1(y1 + g 2 t−2(X1)), ∀t ≥ 2 (14)
g2t−1 = 1
2 (cI +D)Ṽ2(cI +M22)
−1(g1t−2(X2) + y2), ∀t ≥ 2 (15)
t = 1 : g10 = P > 1 z1 and g 2 0 = P̃ > 2 z2, where z1 = V1y 1, z2 = Ṽ2y 2 . (16)
The illustration of this process is presented in the Fig. 2a.
Consider the sequence of solutions for the agent 1 with evaluation atX2:
• Supervised Learning: g10(X2) = V > 2 P > 1 z1
• 1st round of KD:
g11(X2) = V > 2 P>1 2 (z1 +C1P̃ > 2 z2)
• 2nd round of KD:
g12(X2) = V > 2 P>1 2 (z1 + C1P̃ > 2 2 z2 + C1P̃ > 2 C̃2P > 1 2 z1)
• 3rd round of KD:
g13(X2) = V > 2 P>1 2 (z1 + C1P̃ > 2 2 z2 + C1P̃ > 2 C̃2P > 1 4 z1 + C1P̃ > 2 C̃2P > 1 C1P̃ > 2 4 z2)
• t-th round of KD:
g1t (X2) = V > 2 P>1 2 ( t∑ i ( C1P̃ > 2 C̃2P > 1 4 )i)z1 + V > 2 P>1 C1P̃ > 2 4 ( t−1∑ i ( C̃2P > 1 C1P̃ > 2 4 )i)z2
• The limit of KD steps:
g1∞(X2) = V > 2 P>1 2 (I − C1P̃ > 2 C̃2P > 1 4 )†z1 + V > 2 P>1 C1P̃ > 2 4 (I − C̃2P > 1 C1P̃ > 2 4 )†z2,
where † stands for pseudoinverse.
Let us analyze the limit solution and consider the first term of its expression. One can deduce the following identity: 2
V >2 P̃>1 2 (I − C1P̃ > 2 C̃2P > 1 4 )†z1 = (17) V >2 DV1 2 (cI + V >1 DV1 − Ṽ >1 DṼ2 2 (cI + Ṽ >2 DṼ2) −1V > 2 DV1 2 )†y1 (18)
In a similar manner, we can deal with the second term:
V >2 P>1 C1P̃ > 2 4 (I − C̃2P > 1 C1P̃ > 2 4 )†z2 = V >2 DV1 2 (cI + V >1 DV1 − Ṽ >1 DṼ2 2 (cI + Ṽ >2 DṼ2) −1V > 2 DV1 2 )† Ṽ >1 DṼ2 2 (cI + Ṽ >2 DṼ2) −1y2
Now, we notice Schur complement expression in the equation (18):
(cI + V >1 DV1 − Ṽ >1 DṼ2
2 (cI + Ṽ >2 DṼ2) −1V > 2 DV1 2 ) .
This means that we can consider the following problem:( L11 + cI M12 2
L21 2 M22 + cI
)( β1 β2 ) = ( V >1 DV1 + cI
Ṽ >1 DṼ2 2
V >2 DV1 2 Ṽ > 2 DṼ2 + cI
)( β1 β2 ) = ( y1 −y2 ) , (19)
and derive that g1∞(X2) = V >2 DV1 2 β1 and g 2 ∞(X1) = − Ṽ >1 DṼ2 2 β2, where β1,β2 are defined as the solution to the problem (19). Given this, we can derive the following:
V >1 DV1β1 + Ṽ >1 DṼ2
2 β2 = y1 − cβ1 = 2g1∞(X1)− g2∞(X1), (20)
−V > 2 DV1 2 β1 − Ṽ >2 DṼ2β2 = y2 + cβ2 = 2g2∞(X2)− g1∞(X2). (21)
As one can see, there is a strong relation between the limit KD solutions and the solution of a linear system of equations with modified matrixK and right-hand side. Mainly, we take the kernel matrix and divide its non-diagonal blocks by 2, which intuitively shows that our final model accounts for the reduction of the ’closeness’ between datasets D1 and D2. And on the right-hand side, we see −y2 instead of y2 which is quite a ’artificial’ effect. Overall these results in the fact that both limit solutions (for each agent) do not give a ground truth prediction forX1 andX2 individually, which one can see from equation (20) with c = 0. That is, we need combine the predictions of both agents in a specific way to get ground truth labels for datasets D1 and D2. Moreover, the way we combine the solutions differs between datasets D1 and D2 that one can see from comparison of right-hand sides of expressions (20) and (21). Given all the above, to predict optimal labels we need to change we way we combine models of agents in dependence on a dataset, but an usual desire is to have one model that predicts ground truth labels for at least both training datasets.
To obtain the expressions for the case of identical models one should ’remove all tildas’ in the above expressions and by setting V1 = Ṽ1, V2 = Ṽ2, C1 = C̃2 → I and (C1P̃>2 C̃2P>1 )t → (P>2 P>1 )t
2In case of c = 0, the whole analysis can be repeated by replacing inverse sign with † sign and using the following fact for positive semidefinite matrices (Zhang, 2006):
V >2 DV1(V > 1 DV1) †(V >1 DV1) = V > 2 DV1 .

heading: E PARALLEL KD
text: In this section, we theoretically analyze a slight modification of the AvgKD algorithm which we call Parallel KD (PKD). Keeping the notation of sections B and C, denote the data on agent 1 as D1 = (X1,y1) where X1[i, :] = x1n and y1[i] = y1i . Correspondingly, for agent 2 we have D2 = (X2,y2). Now starting from ŷ10 = y1, ŷ20 = y2, in each round t ≥ 0:
a. Agents 1 and 2 train their model on datasets (X1, ŷ1t ) and (X 2, ŷ2t ) to obtain g 1 t and g 2 t . b. Agents exchange g1t and g 2 t between each other. c. Agents use exchanged models to predict labels ŷ1t+1 = ŷ1t+g 2 t (X 1) 2 , ŷ 2 t+1 = ŷ2t+g 1 t (X 2) 2 .
The summary of the algorithm is depicted in Figure 7b. That is, we learn from the average of 2 agents’ predictions. For simplicity, we are going to analyze the scheme without regularization and we take always a minimum norm solution. Notice that there is no exchange of raw data but only of the trained models.
Let us analyse KD algorithm where solutions for agent I (g1) and agent II (g2) obtained as:
g1t (X2) = 1
2 K21(K11)
†(g1t−1(X1) + g2t−1(X1)) = 1
2 V >2 P̂ > 1 (g 1 t−1 + g 2 t−1), (22)
g2t (X1) = 1
2 K12(K22)
†(g1t−1(X2) + g2t−1(X2)) = 1
2 V >1 P̂ > 2 (g 1 t−1 + g 2 t−1) , (23)
where
g1t−1 = 1
2 DV1(K11)
†(g1t−2(X1) + g2t−2(X1)), ∀t ≥ 2 (24)
g2t−1 = 1
2 DV2(K22)
†(g1t−2(X2) + g2t−2(X2)), ∀t ≥ 2 (25)
t = 1 : g10 = P̂ > 1 z1 and g 2 0 = P̂ > 2 z2, where zi = Viyi, i = 1, 2 . (26)
Consider g1t for t ≥ 1:
g1t = 1
2 P̂>1 (g 1 t−1 + g 2 t−1) =
1 2 P̂>1 ( 1 2 P̂>1 (g 1 t−2 + g 2 t−2) + 1 2 P̂>2 (g 1 t−2 + g 2 t−2)) =
= 1
2 P̂>1 (
1 2 (P̂>1 + P̂ > 2 )(g 1 t−2 + g 2 t−2)) = ... = 1 2t P̂>1 (P̂ > 1 + P̂ > 2 ) t−1(g10 + g 2 0) =
= 1
2t P̂>1 (P̂ > 1 + P̂ > 2 ) t−1(P̂>1 z1 + P̂ > 2 z2)
The form of the solutions reminds the method of averaged projection (Lewis et al., 2007) with operator P̂1 + P̂2, which is similar to alternating projection converges to the intersection point of 2 subspaces3. That is, similarly to AKD in the case with the regularization we expect the solution to converge to the origin point 0 in the limit of the distillation steps. As a result, after some point, one expects steady degradation of the predictions of both agents in the PKD scheme.

heading: F ENSEMBLED KD
text: One can consider the following problem:( L11 M12 L21 M22 ) = ( V >1 DV1 Ṽ > 1 DṼ2
V >2 DV1 Ṽ > 2 DṼ2
)( β1 β2 ) = ( y1 y2 ) , (27)
with the following identities:
V >1 DV1β1 + Ṽ > 1 DṼ2β2 = y1 and V > 2 DV1β1 + Ṽ > 2 DṼ2β2 = y2 . (28)
3Actually, there is an explicit relation between alternating and averaged projections (Lewis et al., 2007).
One can find β1,β2 and deduce the following prediction by the model associated with the system (27) for i = 1, 2:
V >i DV1β1 + Ṽ > i DṼ2β2 =
V >i P > 1 (I −C1P̃>2 C̃2P>1 )†z1 − V >i P>1 C1P̃>2 (I − C̃2P>1 C1P̃>2 )†z2+ Ṽ >i P̃ > 2 (I − C̃2P>1 C1P̃>2 )†z2 − Ṽ >i P>2 C̃2P>1 (I −C1P̃>2 C̃2P>1 )†z1 =
V >i P > 1 ∞∑ t=0 (C1P̃ > 2 C̃2P > 1 ) tz1 − V >i P>1 C1P̃>2 ∞∑ t=0 (C̃2P > 1 C1P̃ > 2 ) tz2+
Ṽ >i P > 2 ∞∑ t=0 (C̃2P > 1 C1P̃ > 2 ) tz2 − Ṽ >i P>2 C̃2P>1 ∞∑ t=0 (C1P̃ > 2 C̃2P > 1 ) tz1.
From this one can easily deduce (28) which means that for datasets of both agents this model predicts ground truth labels. To obtain the expressions for the case of identical models one should ’remove all tildas’ in all the above expressions and by setting V1 = Ṽ1, V2 = Ṽ2, C1 = C̃2 → I and (C1P̃ > 2 C̃2P > 1 ) t → (P>2 P>1 )t
The last question is how one can construct the scheme of iterative KD rounds to obtain the above expression for the limit model. From the form of the prediction, we conclude that one should use models obtained in the process of AKD. There are many possible schemes how one can combine these models to obtain the desired result. One of the simplest possibilities is presented in the section 5.2.

heading: G M-AGENT SCHEMES
text: The natural question to ask is how one could extend the discussed schemes to the setting of M agents. In this section, we address this question and explicitly provide the description of each algorithm for the setting of M agents.
Scalability and privacy. Before diving into particular algorithms we investigate some important concerns about our KD based framework. A naive implementation of our methods will require each agent sharing their model with all other agents. This will incur significant communication costs (M2), storage costs (each agent has to store M models), and is not compatible with secure aggregation (Bonawitz et al., 2016) potentially leaking information. One potential approach to alleviating these concerns is to use an server and homomorphic encryption (Graepel et al., 2012). Homomorphic encryption allows agent 1 to compute predictions on agent 2’s data f1(X2) without learning anything aboutX2 i.e. there exists a procedure Hom such that given encrypted data Enc(X2), we can compute
Hom(f1,Enc(X2)) = Enc(f1(X2)) .
Given access to such a primitive, we can use a dedicated server (agent 0) to whom all models f1, . . . , fM are sent. Let us define some weighted sum of the predictions as fα(X) :=∑M
i=1 αifi(X) . Then, using homomorphic encryption, each agent i can compute Enc(fα(X i))) in a private manner without leaking any information to agent 0. This makes the communication cost linear in M instead of quadratic, and also makes it more private and secure. A full fledged investigation of the scalability, privacy, and security of such an approach is left for future work. With this caveat out of the way, we next discuss some concrete algorithms.
AKD with M agents. To extend AKD scheme to a multi-agent setting and corresponding theory we need to start by understanding what the alternating projection algorithm is in the case ofM convex sets. Suppose we want to find the intersection point of M affine sets Ci, for i = 1, ...,M . In terms of alternating projection algorithm we can write the following extension of it (Halperin, 1962):
P∩ni=1Ci(x) = (PCMPCM−1 ...PC1) ∞(x) (29)
For our algorithm, it means that agent 1 passes its model to agent 2, then agent 2 passes its model to agent 3 and so until agent M that passes its model to agent 1, and then the cycle repeats. That is, as before, we denote the data on agent 1 as D1 = (X1,y1) where X1[i, :] = x1n and y1[i] = y1i ., for all other agents we have Di = (Xi,yi), for i = 2, ...,M . Now starting from ŷ10 = y1, in each rounds t, ..., t+M − 1, t ≥ 0:
a. Agent 1 trains their model on dataset (X1, ŷ1t ) to obtain g 1 t . b. for i = 2, ...,M : b.1 Agent i receives g1t+i−2 and uses it to predict labels ŷ i t+i−1 = g 1 t+i−2(X i).
b.2 Agent i trains their model on dataset (Xi, ŷit+i−1) to obtain g 1 t+i−1.
c. Agent 1 receives a model g1t+M−1 from agent M and predicts ŷ 1 t+M = g 1 t+M−1(X 1).
As before, there is no exchange of raw data but only of the trained models. And given all the results deduced before, all the models from some point will start to degenerate. The rate of convergence of such an algorithm is defined similarly to alternating projection in case 2 sets and can be found in Smith et al. (1977).
PKD with M agents. PKD scheme can be easily extended to the multi-agent setting analogously to how averaged projection algorithm can be extended to the multi-set setting (Lewis et al., 2007). Suppose we want to find the intersection point of M affine sets Ci, for i = 1, ...,M then the in terms of averaged projection we have the following:
P∩ni=1Ci(x) = ( 1
M M∑ i=1 PCi) ∞(x) (30)
This expression easily translates into the PKD algorithm for M agents. Denote the data on all agents as Di = (Xi,yi), for i = 2, ...,M . Now starting from ŷi0 = yi, for i = 2, ...,M , in each round t ≥ 0:
a. for i = 1, ...,M : a.1 Agent i trains their model on dataset (Xi, ŷit) to obtain g i t.
b. for i = 1, ...,M :
b.1 Agent i receives models gjt , for j = 1, ...,M, j 6= i from all other agents.
b.2 Agent i use received models to predict ŷit+1 = ŷit+
∑M j=1,j 6=i g j t (X i)
M .
As in the case of 2 agents, there is no exchange of data between agents, but only models. This scheme requires all to all communication.
AvgKD with M agents. Similarly to PKD algorithm we can extend AvgKD algorithm as follows: starting from ŷi0 = y
i, for i = 2, ...,M , in each round t ≥ 0: a. for i = 1, ...,M :
a.1 Agent i trains their model on dataset (Xi, ŷit) to obtain g i t.
b. for i = 1, ...,M :
b.1 Agent i receives models gjt , for j = 1, ...,M, j 6= i from all other agents.
b.2 Agent i use received models to predict ŷit+1 = yi+
∑M j=1,j 6=i g j t (X i)
M .
That is, there is no exchange of data between agents, but only models. This scheme as well as PKD requires all to all communication. That means that the scheme can not be scaled, but it is still useful for small numbers of agents (e.g collaboration of companies). The last is motivated by the simplicity of the scheme without any need for hyperparameter tuning, the non-degrading behavior as well as its superiority over the FedAvg scheme at highly heterogeneous data regimes.
EKD with M agents. EKD scheme in the case of 2 agents is based on models obtained in the process of 2 simultaneous runs of the AKD algorithm. This means that the extension of EKD to the multi-agent setting, in this case, is straightforward by using the M simultaneous runs of AKD algorithm starting from each agent in the multi-agent setting and summing the obtained models in the process as follows:
f∞(x) = ∞∑ t=0 (−1)t( M∑ i=1 git(x)) (31)

heading: H ADDITIONAL EXPERIMENTS

heading: H.1 AKD WITH VARYING DATA HETEROGENEITY
text: In this section, we present the results for AKD scheme with varying data heterogeneity on MNIST dataset. We vary Alpha hyperparameter in the experiments with AKD scheme on MNIST. The results one can see in Fig. 8 which shows that the increase in data heterogeneity imposes faster degradation through rounds of AKD scheme which is aligned with the proposition IV.

heading: Test Acc. of agents in AKD

heading: H.2 EKD ON MNIST
text: In this section, we present the results of Ensembled KD (EKD) scheme on MNIST dataset. The results are presented in Fig. 9 for MSE loss and accuracy in comparison with centralized model and separate agents models. From the results, we see that EKD consistently improves with the increase in the number of models used.

heading: Test Loss

heading: Test Accuracy

heading: H.3 CROSS-ENTROPY OBJECTIVE
text: In this section, we present the results of experiments on MNIST with Cross-Entropy (CE) loss for 2 main schemes under investigation: AKD and AvgKD. In Fig. 10 and 11 one can see the results of AKD and AvgKD schemes correspondingly for CE loss. The results are aligned with our theory: in Fig. 10 we see the degradation trend for AKD which is dependent on the amount of the regularization, model and data heterogeneity. In Fig. 11 we see steady behavior of AvgKD scheme for both agents models: there is no degradation even if model and data are different.

heading: Test Max Acc. of agents in AKD

heading: Test Max Acc. of agents in AvgKD

heading: H.4 PERFORMANCE OF AVGKD AND AKD ON CIFAR10
text: In this section, we present the results for CIFAR10 dataset for AKD scheme in Fig. 12 and 13 for MSE and CE losses correspondingly. One can see the same results for CIFAR dataset: degradation trend in AKD scheme for both MSE and CE losses and dependence on model and data heterogeneity of this trend.

heading: H.5 COLLABORATION BETWEEN MLPS AND RANDOM FORESTS
text: In this section, we present the results for MNIST dataset for Random Forests (RF) and MLP models with MSE loss. That is, the experiments are in the setting ’different models’, where agent 1 has MLP model and agent 2 has RF model. These experiments can show how AKD and AvgKD schemes behave in the setting of fundamentally different models. In the Figs. 14 and 15 the results for accuracy are presented. We see the alignment of these results with theory and other experiments with deep learning models. Mainly, there is a degradation trend for AKD scheme which is speeding up with the increase in data heterogeneity, there is no degradation for AvgKD scheme, and the performance of both agents in AvgKD scheme is highly dependent on data heterogeneity.

heading: H.6 AVGKD WITH M AGENTS
text: The AvgKD scheme does not degrade in comparison with AKD and PKD schemes that degrade already in the case of 2 agents. Let us experiment if AvgKD keeps its steady behavior in the case of M agents. In this section, we present the results of AvgKD scheme with M agents where each round for each agent we average its ground truth labels with the predictions of all other actors. Here

heading: Test Max Acc. of agents in AKD

heading: Test Max Acc. of agents in AKD

heading: Test Acc. of agents
text: we use 5 agents on MNIST dataset in the setting of the ’same model’ and vary data heterogeneity. Such experiments show how AvgKD model can perform in a more practical setting with a bigger number of agents. Moreover, we can understand if there is a dependence on data heterogeneity. In case of full data heterogeneity (Alpha = 0) we assign the labels (2(i−1), 2i−1) to the actor number i, for i = 1...5. The results are presented in the Fig. 16. We see that all the agents repeat the behavior

heading: Test Acc. of agents in AvgKD
text: of each other and in all cases of data heterogeneity can benefit from the AvgKD scheme: in cases of Alpha < 0.05 (high data heterogeneity) early stopping is beneficial.

heading: Test Acc. of N agents in AvgKD

heading: H.7 COMPARISOIN WITH PKD AND FEDAVG
text: In this section, we present the results of PKD scheme defined and analyzed in section E with 2 agents. The PKD scheme is a slight modification of the AvgKD scheme which is proven to degrade through rounds of distillation. The comparison results of AvgKD and PKD are presented in the Fig. 17. We see the degradation trend for PKD scheme which is suggested by our theory presented in section E. Moreover, AvgKD scheme outperforms both PKD and FedAvg in nearly all settings. However, its convergence is less stable in extremely high heterogeneous settings, showing large oscillations. Investigating and mitigating this could be interesting future work.

heading: Test Acc. of agents in PKD

Review summaries:
The paper analyzes the dynamics of optimizing two kernel regression models via co-distillation in a distributed setup where local models may differ in the kernel used. In each round, each local client uses the other client's model to produce novel labels for its local dataset and then retrains its local model using these novel labels. The paper analyzes three variants of this approach, the vanilla variant, a variant where novel labels are an average of the actual label and the one predicted by the other client's model, and an ensembling approach that uses all model iterations to produce predictions. The approaches are analyzed theoretically and empirically.

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. 

This paper analyze knowledge distillation based model agnostic federated learning. They consider simple two agent kernel regression scenario, where each agent has its own dataset and predicting function. They propose to train agent 1 model on dataset 1, and then use agent 1's model to make predictions on dataset 2, and agent 2 will train model using these predictions. They analyze the dynamics of agent 1's model, and show that Alternating Knowledge distillation will degrade the model prediction to 0. They provide another algorithm ensemble AvgKD which can actually avoid this issue and converge to optimal solution. The experiments also show that AvgKD can converge to zero loss while other approaches do not work.

This paper introduce an interesting perspective on federated learning via knowlegde distillation, which allows participating clients to have their own choices of models. In particular, the paper develops theoretical results for a two-client federated regression scenario, which demonstrates (a) the degeneration of an alternating knowledge distillation where iterative distillations (i.e., a client model at each iteration is re-trained based on unlabeled input + prediction of the other client's model of the previous iteration) gradually lose information over the iterations and eventually converge towards a vacuous model; and (b) a new ensembling technique that aggregates intermediate models produced by both clients over the iterations such that in the limit (i.e., when the no. of iterations tends to infinity), the aggregated model is the same as the (oracle) centralized model. This developed intuition on ensembling intermediate models is then applied to more realistic 2-client federated classification scenarios on MNIST and CIFAR-10 datasets.