Towards Model Agnostic Federated Learning Using Knowledge Distillation
Authors: Andrei Afonin
Sai Praneeth Karimireddy
Venue: ICLR 2022 Poster (2022)
Keywords: Federated Learning
Knowledge Distillation
Model Agnostic Communication
Kernel Regression

Abstract: Is it possible to design an universal API for federated learning using which an ad-hoc group of data-holders (agents) collaborate with each other and perform federated learning? Such an API would necessarily need to be model-agnostic i.e. make no assumption about the model architecture being used by the agents, and also cannot rely on having representative public data at hand. Knowledge distillation (KD) is the obvious tool of choice to design such protocols. However, surprisingly, we show that most natural KD-based federated learning protocols have poor performance.
    
    To investigate this, we propose a new theoretical framework, Federated Kernel ridge regression, which can capture both model heterogeneity as well as data heterogeneity. Our analysis shows that the degradation is largely due to a fundamental limitation of knowledge distillation under data heterogeneity. We further validate our framework by analyzing and designing new protocols based on KD. Their performance on real world experiments using neural networks, though still unsatisfactory, closely matches our theoretical predictions. 

Published as a conference paper at ICLR 2022
TOWARDS MODEL- FEDERATED LEARNING USING KNOWLEDGE DISTILLATION
Andrei Afonin
Sai Praneeth Karimireddy' EPFL, UC Berkeley
EPFL
andrei.afoninee www.aaneafnnll...mmm
sp- karkmii wwddyrrerreeeeeyyy..ooomoo
ABSTRACT
Is it possible to design an universal API for federated learning using which an ad-hoc group of data-holders (agents) collaborate with each other and perform federated learning? Such an API would necessarily need to be model-agnostic i.e. make no assumption about the model architecture being used by the agents, and also cannot rely on having representative public data at hand. Knowledge distillation (KD) the obvious tool of choice to design such protocols. However, surprisingly, we show that most natural KD-based federated learning protocols have poor performance.
To investigate this, we propose a new theoretical framework, Federated Kernel ridge regression, which can capture both model heterogeneity as well as data heterogeneity. Our analysis shows that the degradation is largely due to a funda- mental limitation of knowledge distillation under data heterogeneity. We further validate our framework by analyzing and designing new protocols based on KDD. Their performance real world experiments using neural networks, though still unsatisfactory, closely matches our theoretical predictions.
INTRODUCTION
"I speak and speak, but the listener retains only the words he is expecting... is not the voice that commands the story: it is the ear." Invisible Cities, Italo Calvino.
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 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 FedAvg or its variants all proceed in rounds (Wang et al., 2021) 1). 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 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 con-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. overcome such limitations, we would need to take 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 design model agnostic federated learning algorithms which would allow each agent to train their model of choice on the combined dataset?
"Corresponding author.

---

Published as a conference paper at ICLR 2022
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 ie. predict. Our goal is to be able to collaborate with and learn from any agent which provides these two functionalities.
Simple algorithms. 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 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 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 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 (FI.). 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 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-peet communications (Nedic, 2020). A large body of work has designed algorithms for FL under the identical model setting where we either learn 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 al., 2020; Grimberg 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 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 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. KD has also been proposed to combine personalization with model compression (Ozkara et al., 2021), but our focus is for the agents to learn on the combined data. In the closely related codistillation setting (Zhang et al., 2018; Anil et al., 2018; 8; 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 theoretical analysis.
2

---

Published as a conference paper at ICLR 2022
KD analysis. Despite the empirical success of KD, is poorly understood with very little theoretical analysis. Phuong & Lampert (2021) explore generalization bound for distillation-tran trrineedddde linear model, and Tang et al. (2021) conclude that by using KD one re- weigghtt 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 kernel regression setting i.e. the model is retrained using its own predictions on the training data. They show that iterative self-distillation induces 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.
FRAMEWORK AND SETUP
Notation. We denote set A, matrix A, and vector A is the (i. j)-th element of matrix A, a/i denotes the i'th element of vector a. ||a|| denotes the &g norm of vector a.
Centralized kernel regression (warmup). Consider, training dataset 'D S IR4 x R. That is, D = uf {(æ;, }:) where In E X S IR4 and yn E ) C R. Given training set D, our aim to find best function f F X + Y. To find f we solve the following regularized optimization problem:
warm-up, the centralized setting with
f arg min I E(((en) yyn)² cRa(f), with JEF N
Rn) = 1 u(xxxx))) f
Here, F defined to be the space of all functions such that (2) is finite, e is the regularization parameter, and u(x is kernel function. That is, is symmetric u(x, x') u(x', x) and positive with Ra(f) = only when f f= 0 and Ru > o.w. Further, let k(x, t) be the function s.t.
J ula x' (k(,', t)dx'= 6(x t), where 6(-) is the Dirac delta function.
(3)
Now, we can define the positive definite matrix K E NRN and vector hi E RVV as:
Kii := k(x,, 2,) and k := k(x for I D.ViE
(4)
Note that ke is actually vector valued function which takes any X' input, and both kn and K depend on the training data D. We can then derive a closed form solution for f*.
Proposition I (Schölkopf al. (2001)). The f* which minimizes (1) is given by
(x) IIT (cI K)- for yli = ys E [N).
Note that on the training data X E Nxd with X X|i, = 2, we have ('X) K (cI + K)- Kernel regression for an input outputs weighted average the training {yx}. These weights are computed using a learned measure of distance between the input æ and the training {x,). Intuitively, the choice of the kernel u(æ, x²) creates an inductive bias and corresponds the choice of model in deep learning, and the regularization parameter acts similarly to tricks like early stopping, large learning rate, etc. which help generalization. When e = O, we completely the training data K (KK)
and the predictions exactly recover the labels with f* predictions (x) = K(cI + K) y + 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).
When e > 0 the
Federated kernel regression (our setting). We have two agents, with agent having dataset D, -N ((o).vt) and agent 2 with dataset DD2 Ug (ax v/)}. Agent aims find the best.
3

---

Published as a conference paper at ICLR 2022
approximation mapping g¹ E FFi y using kernel (xx x²) and objective:
1 (g(x))) y NI + N2 Eig(af) v²a2² + cRu,, (g), with geti N1 + N2
gt arg min
(5)
Ru, (g) = Sx u1 x' (g()))((( dddd''
(6)
Note that the objective of agent defined using its individual kernel un (zx, x'), ), but over the joint dataset (D1, D2). Correspondingly. agent also uses its own kernel function uz(2, æ') ee') to define the regularizer Ru,($ (g²) over the space of functions g2 E F2 and optimum g²* Thus, our setting has model heterogeneity (different kernel functions ui and u2) and data heterogeneity (D1 and D2 are not i.i.d.). Given that the setting is symmetric between agents and can focus solely error in terms of agent I's objective (5) without loss of generality.
Proposition I can used to derive a closed form form for function g¹ minimizing objective (5). However, computing this requires access to the datasets of both agents. Instead, ask "can design an iterative federated learning algorithm which can approximate g !"?
ALTERNATING KNOWLEDGE DISTILLATION
In this section, we deseribe a popular iterative knowledge distillation algorithm and analyze its updates in our framework. Our analysis leads some surprising connections between KD and projection onto convex sets, and shows some limitations of the current algorithm.
Algorithm. Denote the data on agent as Dj = (x', where X [i,: æt and y = y Correspondingly, have D = (X 2 2)). Now starting from D y² in each rounds and /
a. Agent trains their model on dataset (X1 v/ to obtain gf
b. Agent receives of and uses it to predict labels 8241 9 (x²).
Agent trains their model on dataset (x2, 2i+1) obtain gl+1
d. Agent receives model g++1 from agent and predicts úl-2 g(+1(x1)
Thus the algorithm alternates between training and knowledge distillation on each of the two agents. We also summarize the algorithm in Figure la. Importantly, note that there is exchange of raw data but only of the trained models. Further, each agent trains their choice of a model on their data with agent training {9) } and agent training {9l+1]- Superscript means we start AKD from agent
4.1 THEORETICAL ANALYSIS
Similar to (3), let us define functions kn (ax, x²) and (2(2r, x²) such that they satisfy
Jx uz x' Daa tt)dx' = 5(æ -) for a E {1,2}.
For such functions, we can then define the following positive definite matrix L E (R(M++N)xXMM
L1
L12) L22)
Lablis.] = Ni + N₂ ki (2, a²)) for a, b E {1,2} and e [N1),j E [N2].
L21
Note L is symmetric (with LT2 = L21) and also positive definite. Further, each component Lo,b measures pairwise similarities between inputs of agent and agent b using the kernel k1.
Correspondingly, define M
which uses kernel kg:
M M12
M = M21
Madli. = N1 N₂ k2(2, 2²) and i [Ni.j e [N2].
M22
We 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
22 (x) -1 (cT + L1 (M1z(cI + M22) -'
=1)) yt,
4

---

Published as a conference paper at ICLR 2022
where La E RN is defined as L N k1 (1 x.).
Further; for any fixed a we have
(x) =0. 11000
lim
First, note that if agents I and 2 are identical with the same data and same model, we have M12 = M2 Ln L This setting corresponds to iiffdiitiilation where the model is repeatedly retrained on its own predictions. Proposition shows that after 21 rounds of self-distillation, we
obtain a model is of the form d (ax) = 1o (cI + L11) (Ln (cI + L,1)-t) 2
y. Here, the effect
of e is amplified as increases. Thus, this shows that repeated self-distillation induces strong regularization effect, recovering the results of (Mobahi et al., 2020).
Perhaps more strikingly, Proposition shows that not only does AKD fail to converge to the actual optimal solution g as defined in (5), it will also slowly degrade and eventually converges to O. We next expand upon this and explain this phenomenon.
4.2 DEGRADATION AND CONNECTION TO PROJECTIONS
While mathematically, Proposition II completely describes the AKD algorithm, it does not provide much intuition. In this section, we rephrase the result terms of projections and contractions which provides a more visual understanding of the method.
Oblique projections. A projection operator P linearly maps (projects) all inputs onto some linear subspace A Range (A). In general, can always rewrite such projection operation as
P
min (y x) W x) = A(A AAAAAAAAAAA A) 1 A WW. BaRage
A)
Here, A defines orthonormal basis of the space A and W is positive definite weight matrix which defines the geometry (ar, y)w == x W y. When W = I, we recover the familiar orthogonal projection. Otherwise, projections can happen 'obliquely' following the geometry defined by W.
Contractions. contraction is linear operator C which contracts all inputs towards the origin:
||Cz|l < [x|| for any æ.
Given these notions, we can rewrite Proposition II as follows.
Proposition III. There exist oblique projection operators P1 and Ps, contraction matrices C1 and Cz and orthonormal matrices V1 and V22 such that the model in round 2t learned by the alternating knowledge distillation algorithm is
g2. = _ (cI - (CIP..C2PTT) Viy².
In particular, the predictions on agent 2's inputs are
(x Vg P.T (CiP) C2P Viy²
Further, the projection matrices satisfy P. P x = x only if x x==0 0.
The term (C.P. CTPT)' is the only one which depends t and so captures the dynamics of the algorithm. Two rounds of AKD (first to agent and then back to 2) correspond to multiplication with C P2 C, 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,T and P27 have orthogonal ranges whose intersection only contains the origin 0. As Fig. lb 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 projection more intuitively. suppose that had a 4-way classification task with agent having data only of the first two classes, and agent 2 with the last two classes. Any model trained by agent I 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 agent 2 is mediated by its data X2 This corresponds to a projection of the knowledge of agent onto the data of agent 2. If there a mismatch between the two, knowledge is bound to be lost.
5

---

Published as a conference paper at ICLR 2022
Train
a
rounds
v.
(a) Alternating KD starting from agent 1. We predict and train using predictions, alternating between the two agents.
(b) Alternating oblique projections starting from converges O.
00
X
(b) Intuition behind ensemble scheme. each round, AKD alternates between overfitting the data of agent I of agent 2. We can construct an ensemble out of these models correct for this bias and quickly converge the true optima.
round
(a) AvgKD scheme
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. and P is O, 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 (also see Fig. lb). 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 gt (x) to 0 gets faster if:
stronger inductive bias is induced via a larger regularization constant C,
the kernels ki (x, y) and (2, y) are very different,
the difference between the datasets D and D2 as measured by ki (æ, y) increases.
In summary, both data and model heterogeneity may speed up the degradation defeating the purpose of model agnostic FI.. All formal proofs and theorem statements are moved to the Appendix.
5 ADDITIONAL VARIANTS
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 g"
5.1 AVERAGED KNOWLEDGE DISTILLATION
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 information is to recombine 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 D, (X', and correspondingly agent 2 has data D2 (X2, 22) Now starting from Ø y 0e y². in each round > 0:
Agents and train their model datasets (X1 ú!) and (x2, (ff to obtain of and of

---

Published as a conference paper at ICLR 2022
b. Agents exchange their models of and gi between each other. C. Agents use exchanged models to predict labels 01+1 y +9 +x,(x)) 0i+1 = v2 (x²)
The summary the algorithm 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 y and y at every iteration, preventing degradation. We theoretically characterize its dynamies next in terms of the dfore-mentionel contraction and projection operators.
Proposition V. There exist oblique projection operators P1 and Pz, contraction matrices C1 and Cz and orthonormal matrices V1 and V2 such that the model of agent in round learned by the averaged knowledge distillation (AvgKD) algorithm is
F
C2P2 C2P )2)21
PCP
C.P.T.PPP 1')
(x)
4
4
t=U
i=0
where F = 11 (cI + L11)
VV...VV................ Further, in the limit of rounds for any fixed we have
F 2
C.P C2PT )221 FC,P.
C2P.T..PP )122.
lim gf (30) t-400
This shows that the model learned through AvgKI> does not degrade to O, unlike AKD. Instead, it converges to limit for which we can derive closed-form expressions. Unfortunately, this limit model still not the same as our desired optimal model g"
We next try to overcome this using ensembling.
5.2 ENSEMBLED KNOWLEDGE DISTILLATION
We first analyze how the limit solution of AvgKD differs from the actual optimum g². We will build upon this understanding to construct an ensemble that approximates g² For simplicity, we assume that the regularization coefficient
Understanding AvgKD. Consider AvgKD algorithm, then
Proposition VI. There exist matrices A₁ and A2 such that the minimizer g² of objective (5) predicts g² (X;) A1 + A232 for E (1,2). where Br and 32 satisfy
L11 M12)
-
y1
L21
M22
B2
y2
In contrast, for the same matrices Ai and A2, the limit models of AvgKD predict g (X.) = = 11
and da0 (X,) = A2, for E {1,2] for B1 and Ba satisfying
Lu
M12
[1
1/1
La
2 M22
B2
yy2
(8)
By comparing the equations (7) and (8), the output 2 (x) 22 (x)) is close to the output of g² (az), except that the off- off-diagonal matrices are scaled by 1 and we have Y2 on the right hand side. We need two tricks if we want to approximate g² first we need an ensemble using differences of models, and second 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. la. Let run two simultaneous runs, ones starting from agent I and another starting from agent 2, outputting models {g]}, J. and {o"] respectively. Then, instead of just using the final models, we construct the following infinite ensemble. For an input æ, we output:
(z) = Σ
( (-1) + 92(x))
1=:
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
7

---

Published as a conference paper at ICLR 2022
Trais ln
Tras oos
Trann bns
0.30
0.21
1.
Reach
Figure 3: AKD, AvgKD, and EKD methods for linear regression on synthetic data with same data (left), different data (middle), and strong regularization (right). EKD (black) eventually matches centralized performance (dashed green), whereas AvgKD (solid blue) worse than only local training (dashed blue and red). AKD (in solid red) performs the worst and degrades with increasing rounds.
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 with regularization model gradually degrades towards O, in other words, intuitively each next round obtained model adds less value to the whole sum in the EKD scheme. Although, in the lack of regularization models degradation towards 0 not always the case (see App. C). But under such an assumption we we gradually converge to the limit model, which is the point Do in the Fig. 2b. We formalize this and prove the following.
Proposition VII. The predictions of foo using (9) satisfies foc (X;) =
{1,2).
Thus, not only did we succeed in preventing degeneracy to O, but we also managed to recover the predictions of the optimal model. However, note that this comes 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 g²¹
6 EXPERIMENTS
6.1 SETUP
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).
The toy experiments solve 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 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 instead uses MLP in the 'different model' setting. Further, 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 then datasets are completely different between agents, if Alpha = 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.
split the training data randomly
6.2 RESULTS AND DISCUSSION
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 regularization used, model heterogeneity 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 well.
8

---

Published as a conference paper at ICLR 2022
Test Acc. of sgents In MD Same mode Samt Wreeetedee Sart
Sare model Satte dita Ang.
Sorw rede Defheett 0013
1.7
1eands
Teands
Dounss
Figure 4: Test accuracy of centralized (dashed green), and AKD on MNIST using model starting from agent (blue) and agent 2 (red) with varying amount of regularization, model heterogeneity, and data heterogeneity. In all cases, performance degrades with increasing rounds with degradation speeding up with the increase in regularization, model heterogeneity, or data heterogeneity.
agents Diresed Same cala
Seme rodel Samee deb.
Some neee, 2eteree cato
1.9
Agents
25
Figure 5: Test accuracy of AvgKD on MNIST using model starting from agent (blue) and agent 2 (red) with varying model heterogeneity. and data heterogeneity. During training regularization is used. In all cases, there is no degradation of performance, though the best accuracy is obtained by agent 1 in round with only local training.
Noha
Apps
Alsha
Codtdioo
Aperi1
056
Fashng
Sounde
Sounde
Sourcc
Figure 6: Test accuracy on MNIST with varying data heterogeneity the setting of 'same model'. In case of high data heterogeneity (small Alphas), both agents benefit from the AvgKD, PKD and EKD schemes. Moreover, AvgKD and EKD consistently outperform FedAvg scheme.
EKD works but needs large ensembles. In the synthetic setting (Fig. 3), in 250 rounds (ensemble of 500 models) even matches the centralized model. However, in the real world setting (Fig. 6) the improvement is slower and it does not match centralized performance. This might be due to the small number of rounds run (only 20). EKD is also the only method that improves with subsequent rounds. Finally, observed that increasing regularization actually speeds up the convergence of EKD in the synthetic setting (Fig. 3).
Data heterogeneity is the main bottleneck. In all our experiments, both data and model hetero- geneity degraded the performance of AKD, PKD (Parallel KD introduced in App. E) and AvgKD, However, data heterogeneity has a much stronger effect. This confirms 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. 6 shows how all schemes behave dependence on data heterogeneity. Indeed, the higher data heterogeneity the faster speed of degradation for the AKD and PKD schemes. 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.

---

Published as a conference paper at ICLR 2022
6.3 ADDITIONAL EXTENSIONS AND EXPERIMENTS
In App. G. we extend our algorithms M agents and show experimentally in App. H.5 for AvgKD algorithm that the same trends hold there as well. Our conclusions also hold for the cross-entropy loss (Figs. 11, 12), for highly heterogeneous model case with MLP and Random Forests (Figs. 13, 14), as well on other datasets and models (VGG on CIFARIO in Figs. 9, 10). In the latter, we see the same trends as on MNIST for all the schemes except EKD which is probably due the use of cross-entropy loss function and, result, models being further from the kernel regime. Moreover, speed of degradation higher if there is the model heterogeneity (Figs. 8, 10) and EKD does not help even MNIST dataset (Fig. 8). Finally, all the schemes are compared to the more standard FedAvg that not applicable the "different model' setting and outperformed by the AvgKD scheme highly data heterogeneous regimes. That is, AvgKD consistently outperforms all the methods at highly data heterogencous regimes indicating is the most promising variant.
7 CONCLUSION
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 formalized notions of both model and data heterogeneity. Using this, we analyzed different knowledge distillation schemes and came the conclusion that data heterogeneity poses 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 design 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 as 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 complex datasets. We believe there is great potential further exploring our results and framework, especially in investigating how to mitigate the effect of data heterogeneity in knowledge distillation.
ACKNOWLEDGEMENTS
We are really grateful to Martin Jaggi for his insightful comments and support throughout this work. SPK is partly funded by an SNSF Fellowship and AA is funded by research scholarship from MLO lab, EPFL headed by Martin Jaggi. SPK also thanks Celestine Dünner for conversations inspiring this project.
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning, arXiv 2012.09816, 2020.
Rohan Anil, Gabricl Pereyra, Alexandre Passos, Robert Ormandi, George E Dahl, and Geoffrey Hin- ton. Large scale distributed neural network training through online distillation. arXiv 1804.03235. 2018.
Nachman Aronszajn. Theory of reproducing kernels. Transactions of the American mathematical society, 68(3):337-404 1950.
Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth. Practical secure aggregation for federated learning on user-held data. arXiv 1611.04482, 2016.
Stephen Boyd and Jon Dattorro. Alternating projections. 01 2003.
Cristian Bucilua, Rich Caruana, and Alexandru Niculescu-Mizil Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 535-541, 2006.
10

---

Published as a conference paper at ICLR 2022
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv 1461,, 2020.
Jianping Gou, Baosheng Yu, Stephen Maybank, and Dacheng Tao. Knowledge distillation: A
survey. International Journal of Computer Vision, 129(6))1789-188) 1819, 2021.
Thore Graepel, Kristin Lauter, and Michael Naehrig. MI confidential: Machine learning encrypted data. In International Conference on Information Security and Cryptology, pp- 1-21. Springer,
2012.
Felix Grimberg, Mary-Anne Hartley, Sai Praneeth Karimireddy, and Martin Jaggi. Optimal model averaging: Towards personalized collaborative learning. FL ICML workshop. 2021.
Israel Halperin. The product of projection operators. Acta Sci. (SteeeS, 23(1);6699, 1962.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv 1503.02531 2015.
Peter Kairouz. H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. arXiv 19/2.04977, 2019.
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv 2008. 03606, 2020a.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for on-device federated learning. 37th International Conference on Machine Learning (ICML), 2020b.
Adrian Lewis, Russell Luke, and Jerome Malick. Local convergence for alternating and averaged nonconvex projections. arXiv 0709.0109, 2007.
Qinbin Li, Bingsheng He, and Dawn Song. Practical one-shot federated learning for cross-silo setting. arXiv 2010.01017, 2020.
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. arXiv 2006.07242, 2020.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for
personalization with applications to federated learning. arXiv 2002. 10619, 2020.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera Arcas. Communication-efficiee learning of deep networks from decentralized data. In Proceedings of AISTATS, pp.
2017.
Aditya Krishna Menon, Ankit Singh Rawat, Sashank J. Reddi, Seungycon Kim, and Sanjiv Kumar. Why distillation helps: statistical perspective. arXiv 2005. 10419, 2020.
Hossein Mobahi, Mehrdad Farajtabar, and Peter L. Bartlett. Self-distillation amplifies regularization in hilbert space. arXiv 2002.057/5 2020.
Angelia Nedic. Distributed gradient methods for convex machine learning problems in networks: Distributed optimization. IEEE Signal Processing Mugazine, 37(3):22-10, 2020.
Kaan Ozkara, Navjot Singh, Deepesh Data, and Suhas Diggavi. Quped: Quantized personalization via distillation with applications to federated learning. Advances in Neural Information Processing Systems, 34, 2021.
Mary Phuong and Christoph H. Lampert. Towards understanding knowledge distillation. arXiv 2303, 2021.
Sashank Reddi, Zachary Charles, Manzil Zaheer. Zachary Garrett, Keith Rush, Jakub Konečný. Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiy 2003.00295, 2020.
11

---

Published as a conference paper at ICLR 2022
Felix Sattler, Arturo Marban, Roman Rischke, and Wojciech Samek. Communication-efficcin federated distillation. arXiv 2012.00632, 2020.
Bernhard Schölkopf, Ralf Herbrich, and Alex Smola. A generalized representer theorem. In International conference on computational learning theory, pp. 416-426. Springer, 2001.
Hyowoon Seo, Jihong Park, Seungeun Oh, Mehdi Bennis, and Seong-Lyun Kim. Federated knowledge distillation. arXiv 2.01.07 2020.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition, 2015.
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33, 2020.
Kennan T Smith, Donald C Solmon, and Sheldon L Wagner. Practical and mathematical aspects of the problem of reconstructing objects from radiographs. Bulletin of the American Mathematical Sociity, 83(6)12277-1271 1977.
Shagun Sodhani, Olivier Delalleau, Mahmoud Assran, Koustuv Sinha, Nicolas Ballas, and Michael Rabbat. A closer look at codistillation for distributed training. arXiv 2010.02838, 2020.
Jiaxi Tang, Rakesh Shivanna, Zhe Zhao, Dong Lin, Anima Singh, Ed H. Chi, and Sagar Jain. Understanding and improving knowledge distillation. arXiv 2002.03532, 2021.
Jianyu Wang. Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, al. field guide federated optimization. arXiv 2107.06917, 2021.
Kangkang Wang, Rajiv Mathews, Chloé Kiddon, Hubert Eichner, Françoise Beaufays, and Daniel
Ramage. Federated evaluation of on-device personalization. arXiv 10252, 2019a.
Shiqiang Wang, Tiffany Tuor, Theodoros Salonidis, Kin K. Leung, Christian Makaya, Ting He, and Kevin Chan. Adaptive federated learning in resource constrained edge computing systems. IEEE Journal on Selected Areas in Communications, 32(6): 12210, 2019b.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nathan Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods machine learning. arXiv preprint urXiv:/705.008: 2017.
Chuhan Wu, Fangzhao Wu, Ruixuan Liu, Lingjuan Lyu, Yongfeng Huang, and Xing Xie. Fedkd: Communication efficient federated learning via knowledge distillation. arXiv 2108. 13323, 2021.
Junho Yim, Donggyu Joo, Jihoon Bae, and Junmo Kim. A gift from knowledge distillation: Fast optimization, network minimization and transfer learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4133-4141, 2017.
Fuxun Yu, Weishan Zhang, Zhuwei Qin, Zirui Xu, Di Wang. Chenchen Liu, Zhi Tian, and Xiang Chen. Fed2: Feature-aligned federated learning. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery Data Mining, pp. 2066-2074, 2021.
Fuzhen Zhang. The Schur complement and its applications, volume 4. Springer Science & Business Media, 2006.
Ying Zhang, Tao Xiang, Timothy M Hospedales, and Huchuan Lu. Deep mutual learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4320-4328, 2018.
12

---

Published as a conference paper at ICLR 2022
A EXPERIMENTS DETAILS
In all real world experiments we use the Adam optimizer with default regularization (weight decay) 3 x 10 +, unless in the 'no regularization' case when it is set to O. We split the data between agents by giving a bigger part of data to agent 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 to another effect in the experiments: agent 2 (with less data) benefits from the communication with agent (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 terms of its amount between agents.
Toy experiments. The toy experiments solve linear regression problem of the form Ax" = b where E Reed and x" E Rd is randomly generated for 1.5d and = 100. Then, the data A and b is split between the two agents at proportion 0.6/0 6/0.4. This is done randomly in the 'same data' case, whereas in the 'different data' case the data is sorted according heterogeneity. This experiment with the linear kernel is supposed to show if our theory is correct, especially if the EKD scheme really works.
before splitting to maximize
Real world experiments. The real world experiments are conducted using CNN and MLP networks on MNIST, MLP network and RF model on MNIST, and VGG16 (Simonyan & Zisserman, 2015) and CNN models on CIFARIO datasets. Further, we split the training data randomly at proportion 0.7 / in the 'same data' setting. For the 'different data" setting, split the data by labels: agent 1 has 'O' 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.
B ALTERNATING KD WITH REGULARIZATION
Different models.
Keeping the notation of section 4 we can construct the following matrix:
Lu
L12
0
L
0 M
L
L22
0
0
K
0
M11
M12
M21
M
Notice that each of the sub-blocks L and M symmetrie positive semi-definite matrices. This makes matrix K symmetric positive semi-definite matrix and it has eigendecompossition form:
K V DV and V = (ViVV
where and V are R2(N Ng) diagonal and orthogonal matrices correspondingly. This means that the [R2(N + N2)) matrices V1, V2, V1, Vz are also orthogonal to each other. Then
can deduce:
Lij DV,, and M,j V. DV, vi,-
In AKD setting only agent has labeled data. The solution of learning from the dataset D evaluated X'2 the following:
(xb) = Ler(cI + L11) = ((V (CVTTcCI (cI+DDV))) -V "V
(10)
Let us introduce notation:
P Vi (cI
(cl and P (cI D)V3) (cT D)
These are well known oblique (weighted) projection matrices on the subspaces spanned by the columns of matrices V1 and V22 correspondingly, where the scalar product is defined with Gram matrix G = cl + D. Similarly one can define P2 and P
This model also convolutional neural network, but in our experiments is bigger than CNN model.
13

---

Published as a conference paper at ICLR 2022
Given the introduced notation and using the fact that V22 V1 0, rewrite equation 10 as:
(b2) 122 P V191
(11)
Similarly, given 2 S D and agent learned from Dd2 = (bt)), inferred prediction }i = gl (2i) by agent 2 can be written
( = V P V2 P 21, where 21 Vi91
At this point we need to introduce additional notation:
C = V and C2 V2V22
These are matrices of contraction linear mappings.
One can now repeat the whole process again with replacement D by Dr. The predictions by agent and agent after such 2/ rounds of AKD are:
g (b)) V221 P.P (CIP) z1 and g241( V.P.C (CIP)
(12)
If one considers the case where agents have the same kernel u2 then one should 'remove all tildas' in the above expressions and the obtained expressions are applied. This means M L,
V V12 and V2 V2. Crucial changes in this case are C₁ = C2 - and (CIPT
(P, PT)'. The matrix (P, P, is an alternating projection (Boyd & Dattorro, 2003) operator after steps. Given 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 Vi and V22 correspondingly. These are orthogonal linear subspaces, hence the unique point of their intersection is the origin point O. This means that (P.) P,
0 and in the limit of such AKD
procedure both agents predict O for any data point.
Speed of degradation. The speed of convergence for the alternating projections algorithm is known and defined by the minimal angle e between corresponding sets (Aronszajn, 1950) (only non-zero elements from sets one has to consider);
(i) vl| 20-0)
(cos()))
as ||a|| =
where v is one of the columns of matrix V1. In our case we can write the expression for the cosine:
10 (cT + D)Øzl
10 Dezl
cos(E) ( max{
max
Vo, (cI D)v,) (n) (cI (cI+D)up) D)v,)
V(c+o, D (c U. D
where 81 and U2 are the vectors from subspaces spanned by columns of matrices V1 and V2 corre- spondingly. Hence the speed of convergence depends on the elements of the matrix L Intuitively these elements play the role of the measure "closeness' between data points. This means that the "closer' points of sets Xi and Xr the higher absolute values of elements in matrix L1 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 g! (a: to 0 gets faster if:
larger regularization constant e used during the training.
smaller the eigenvalues of the matrix V1 V,V. or smaller absolute value of non- non-diagonal block L
C ALTERNATING KD WITHOUT REGULARIZATION
Before we saw that models of both agents degrade one uses regularization. A natural question to ask if models will degrade in the lack of regularization. Consider the problem (1) with 0, so
14

---

Published as a conference paper at ICLR 2022
the regularization term cancels out. In the general case, there 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), propose to analyze the minimal norm solution in the problem of alternating knowledge distillation with 2 agents each with private dataset. For simplicity let assume that agents have the same model architecture. Then the agent I solution evaluated at set 22 with all the above notation can be written as:
(o(2 K21 K| y1 = V22 DViVV (V y1
(13)
where i stands for pseudoinverse. In this section let us consider possible settings one can have:
1. Self-distillation: The datasets and models of both agents are the same.
2. Distillation with K >> 0: The datasets of both agents are different and private. Kernel matrix K' is positive definite.
3. Distillation with K 2 0: The datasets of both agents are different and private. Kernel matrix K' is positive semi-definite.
Self-distillation. Given the dataset D = X' x V, the solution of the supervised learning with evaluation at any x E Rd is:
No (z) (V DV)'y.
Then the expression for the self-distillation step with evaluation at any æ E Rd is:
o (ar) = (VT DV DV DV) y 1/ (V yv) (3)),
where property of pseudoinverse matrix was used: A'AA! = A That is, self-distillation round does not change obtained model. One can repeat self oelf-distillation step t times and there is no change in the model:
gf (V D -x)
Hence in the no regularization setting self-distillation step does not give any change in the obtained model.
Distillation with K O. In this setting have the following identity K K and results quite similar to the setting with regularization. But now the Gram matrix of scalar product in linear space D instead of cI + D in the regularized setting. By assumption on matrix K follows that D is of full rank and the alternating projection algorithm converges to the origin point O. Therefore in the limit of AKD steps predictions by models of two agents will degrade towards 0.
Distillation with K 0. In this setting kernel matrix K has least one eigenvalue and take for the analysis the minimal norm solution (13). We can rewrite this solution:
g6 (b2)) = Vg DV DV V.2 P 21.
where P VV, tV D and Viyi.
One can notice the following projection properties of matrix P:: PJ P and VV DV1)* (V, DV))
That is, P is a projection matrix with eigenspace spanned by columns of the matrix (V DVG),. Similarly, one can define Po. Then the solution for the first AKD step evaluated at set X is:
( = V, DV2 DV.) PT = VT P 21,
and after rounds we obtain for agent and agent II correspondingly:
(x,) (PT and g V, (8 (T)
In the limit of the distillation rounds operator P1,P tends to the projection on the intersection of 2 subspaces spanned by the columns of matrices V, (V, DV and V2(V, DV2)t. We should
15

---

Published as a conference paper at ICLR 2022
g
1
X²
V2
X²2
Train
1 (x)
9;
(P))))))
V1
round
(b) PKD scheme
(a) Illustrative projection 3D
highlight that in general, the intersection set in this case consists not only from the origin point but some rays that lie simultaneously in the eigenspaces of both projectors P1 and Pg. The illustrative example is shown in the Fig. 7a. We perform the alternating projection algorithm between orthogonal linear spaces that have ray the intersection and converge to some non-zero point that belongs this ray. The intersection of eigenspaces of both projectors P1 and P2 the set of points a: Span (VV, DV s.t. P Pax x.
D AVERAGED KD
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 reminder, models of both agents after round t of AvgKD algorithm follows:
(X22) 2 Lcnn + Ln) + (x1)
2 V2 P1 (z) + -2-1)
(x,) = 2
M12(cI + M22) -22)) + ye) = V. P (g--1 x2).
2
where
1
9 (cI + D)Vi + L11) + 92-2(x)),, 2
(14)
2
- (cI + D) (( Ma2) ( (X2) + +22),
2 2
(15)
t=1: 9/ P and = z2
where 21 Viy',
22 Vyy²
(16)
The illustration of this process presented the Fig. 2a. Consider the sequence of solutions for the agent 1 with evaluation at X2: Supervised Learning:
gb (X2) Vg P 21
1st round of KD:
01(x2) VgJ PT
(z)
C.P. x2)
2
2nd round of KD:
PT
C.P.T
C.P.T C,PT
2 (X2) V.
(z) +
22
21)
2
2
3rd round of KD:
P
P.P.T
C.P.P.PT.. C.P.CPPT.P...
gf = V
[1] +
22 +
21 +
22
2
2
4
16

---

Published as a conference paper at ICLR 2022
t-th round of KD:
P.
P.P.
C2P.
C2PT
P P. t-1 C2PTC.PP )0)z2
o (X2) V22
) 1/2
2 1=
The limit of KD steps:
PT
C.P
PTC.P (I
C2PTT.PP )
th (X,)
V22
(I
21 + 1/2
2
4
4
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
P.T
C.P2
C.P.
V
1/21
(17)
2
4 DV,
DV
V7T
DV2
(cI + V22 DV2) V/2 DV
2 (cl (c
y1
(18)
2
2
In a similar manner, we can deal with the second term:
PT C.P.P (I C.P.T.PP.
122
4
4
DV,
'I. DV,
V DV, yt
V DV
(cI VT DV,
(cT V/ DV-)
(cI + V22 DV2)
2
2
2
Now, we notice Schur complement expression in the equation (18):
V,T DV,
V DV. 2
(cI + V, DV
(cI V22 DV
This means that can consider the following problem:
Lu++a + cI Lx
M12
B 32
= VT DV + cl DV.
DVD 2
3
= y/
(19)
M
+ cl)
V DV2 cel
32
-
and derive that g (X2) the solution to the problem (19). Given this, we can derive the following:
V₂ DV: Br and 92 =
V,T VTT
B2, where B1, are defined as
VT DV3 + DVDVV
B2 y - -B = 20/ (2i) - (X1
(20)
2
VT DV. 2
B V.J DV,32 y2 cB2 29 (X2) (X²).
(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 matrix K and right- 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 D, and D2. And on the right-hand side, we see -22 instead of y2 which is quite 'artificial' effect. Overall these results in the fact that both limit solutions (for each agent) do not give ground truth prediction for X, and X₂ individually, which one can from equation (20) with 0. That is, we need combine the predictions of both agents in a specific way to get ground truth labels for datasets D1 and Dz2. Moreover, the way we combine from comparison of right-hand
the solutions differs between datasets Dj and D: that one 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 Vi Vi V22 Vg C₁ C2 I and (C, C2P) (P PT)
In case of 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 DV,VV, DVI' DV)) V,J DV.
17

---

Published as a conference paper at ICLR 2022
E PARALLEL KD
In this section, we theoretically analyze slight modification of the AvgKD algorithm which call Parallel KD (PKD). Keeping the notation of sections B and C, denote the data on agent as D1 (x) yy) where X [i, x/ and y [i] y Correspondingly, for agent
have
D = (x² )²). Now starting from d = y De y² in each round 2 0:
a. Agents and train their model on datasets (x+, ú!) and (x², (it) to obtain gt and sē.
b. Agents exchange g and oe between each other.
c. Agents use exchanged models to predict labels 8.+1 (x)
+o (x3) 2
The summary of the algorithm is depicted in Figure 7b. That is, we learn from the average of agents' predictions. For simplicity, we are going to analyze the scheme without regularization and we take always minimum norm solution. Notice that there is exchange of raw data but only of the trained models.
Let us analyse KD algorithm where solutions for agent I (g') and agent II (g²) obtained as:
1
(12)
K'n (Kn))n)
V P (g. + -2-1);
(22)
2
2
(xx) K12 (2)) (1)
Vi P2 (g + 9
(23)
2
2
where
1
9:1
( of (()(),
Vt 2 2
(24)
2
Φ- 1 = DV2(K2) [9) + +z-2(())
vt 2 2
(25)
2
1=1: 9 P 21
and g P₂- 22, where
2 Vaya, i 1,2.
(26)
Consider of for 1 Z 1:
1
1 P.T PT PT 2 (222 +
1 P. (gl + x(-)))
Ø , P. (ot + . 1 2-1
2
= PT (PT + +P (oi + 2222)
=..=
2' (PT P + (6)) =
2 2
21 + P. x2) = 2 (PT + PTT
1
The form of the solutions reminds the method of averaged projection (Lewis et al., 2007) with operator P1- + P2, which is similar to alternating projection converges to the intersection point of subspaces That is, similarly to AKD in the case with the regularization 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.
expect the solution to
F ENSEMBLED KD
One can consider the following problem:
Li
Mi
V DVI
V. DV2
(3)
y/1
(27)
L
M
V
DV
V DV. (3)
y/2
with the following identities:
Vi DVI + V. DV y1 and V. DV1 + V DV,32 y/2
(28)
Actually, there is explicit relation hetween alternating and averaged projections (Lewis al., 2007).
18

---

Published as a conference paper at ICLR 2022
One can find $1, P2 and deduce the following prediction by the model associated with the system (27) for = 2.2:
V. DV V. DV U | -1 -,,
V. P.
1/22 -V )
(29)
VTP
Σ C.P,, C - - Σ (.PPTPPPP
V.P.P.PP
1=0
V. P2 -CPCCR))
V C2PT DCPPPT
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 V1, V2 V22 C1 C2 I and (CIP2 C2P (P. )
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 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.
G M-AGENT SCHEMES
The natural question to ask 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 (M²2, 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 server and homomorphic encryption (Graepel et al., 2012). Homomorphic encryption allows agent 1 to compute predictions on agent 2's data fi (X2) without learning anything about X i.e. there exists a procedure Hom such that given encrypted data Enc( (x2), we compute
Hom Enc( 2))
(x
Given access to such primitive, we can use dedicated server (agent 0) to whom all mod- els f fM are sent. Let us define some weighted sum of the predictions as fx) == DM axfi(X). Then, using homomorphic encryption, each agent i can compute Enc(fa (X)))) in 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 in the case M convex sets. Suppose want to find the intersection point of M affine sets C, for i = 1, ..., M. In terms of alternating projection algorithm we can write the following extension of it (Halperin, 1962):
Prg_c, (,,,,(()) (x) = (Pess Pex.....PPP )0)
(30)
For our algorithm, it means that agent 1 passes its model to agent 2, then agent 2 passes its model to agent 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 agent as D2 (x! .y' where X : x and y [6] y for all other agents we have D (X', y'), for i = 2, ..., M. Now starting from Uh = y in each rounds I l.....++ M 1, 0:
19

---

Published as a conference paper at ICLR 2022
Agent trains their model on dataset (x!, v))) obtain g:- b. for = 2,.,, M:
b.1 Agent i receives g and uses predict labels Ui+i-1 = -i+i-2(X3)
b.2 Agent trains their model on dataset (X', --++-1 to obtain 9
C. Agent receives a model - from agent M and predicts U.+MM = g M (X
As before, there is no exchange of 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 alternating projection case 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 C₁, for i = 1,.., M then the in terms of averaged projection we have the following:
1 M
Pre e e, (z)
Σ (x) ( M
(31)
This expression easily translates into the PKD algorithm for M agents. Denote the data on all agents as D = (X', ,y)), for i = 2...,... M. Now starting from Th = y'. for i = 2, M, in each round t 12 0:
a. for i = 1,..., M:
a.l Agent trains their model on dataset (x', 9,) to obtain gi.
b. for i 1,..., M::
b. Agent receives models si for j = 1,.., M,j # from all other agents.
b.2 Agent use received models to predict أن 0i MM jss, ii of (x*)
M
As in the case of 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', for 2,... M, in each round / 2 0:
a. for = 1...,.M M:
a.l Agent trains their model on dataset (x's yt) to obtain g'.
b. for i 1,..,,M M:
b. Agent i receives models oi for j 1, ,.., M,j from all other agents.
b.2 Agent i use received models to predict أi DM D'+ 3d
of (X')
That is, there is exchange of data between agents, but only models. This scheme as well PKD requires all to all communication. That means that the scheme can not be scaled, but 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 models obtained in the process of simultancous 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 A simultancous runs of AKD algorithm starting from each agent in the multi-agent setting and summing the obtained models the process follows:
M
(x)
1-)) ((-))
gi (x))
(32)
t=0
20

---

Published as a conference paper at ICLR 2022
H ADDITIONAL EXPERIMENTS
H.1 MNIST WITH VARYING DATA HETEROGENEITY
In this section, we present the results for MNIST dataset with varying data heterogeneity in the setting of "dffferent model'. The results one can in Fig. 8. There is faster degradation trend for both AKD and PKD schemes if different models for agents are used (Fig. 8) comparing to rhe 'same model' setting (Fig. 6) at all data heterogeneity regimes. The PKD scheme is slight modification of the A scheme which is proven to degrade through rounds of distillation. We see the degradation trend for PKD scheme which suggested by our theory presented App. E. EKD does not improve with subsequent rounds in the setting of ddffferent models'. AvgKD scheme outperforms both PKD and AKD in all settings. However, its convergence not stable in extremely high heterogeneous settings, showing large oscillations. Investigating and mitigating this could be interesting future work.
Test Acc
Alpn
0.55
Central
Avg
033
Rounds
Scunds
Rounds
Rouncs
Figure 8: Test accuracy of on MNIST with varying data heterogeneity in the setting of 'different model'. Performance of PKD and AKD degrade with degradation speeding up with the increase in data heterogeneity. Performance of AvgKI) scheme converges to steady behavior at any regime of data heterogeneity. Both agents benefit from AvgKD, PKD and EKD schemes in the early rounds of communication.
H.2 CIFARIO WITH VARYING DATA HETEROGENEITY
In this section, we present the results for CIFARIO dataset with varying data heterogeneity. Note that we use cross-entropy loss function here. The results one can see in Fig. 9 and 10 that again show data heterogeneity plays key role in the behavior of all the schemes. All the trends we saw on the MNIST dataset repeated here except one: EKD does not improve in subsequent rounds in the 'same model' setting.
Test Ace
Alch=
Apri1
Agert
032
0V
Brands
Scunde
Rocrdt
Sourcce
Figure 9: Test accuracy of on CIFARIO with varying data heterogeneity in the setting of 'same model'. As on MNIST: performance of PKD and AKD degrade with degradation speeding up with the increase in data heterogeneity; performance of AvgKD scheme converges to steady behavior; both agents benefit from AvgKD, PKD and EKD schemes in early rounds of communication.
H.3
OBJECTIVE
In this section, we present the results of experiments MNIST with Cross Cross-Entropy (CE) loss for main schemes under investigation: AKD and AvgKD. In Fig. 11 and 12 one can see the results of AKD and Avg schemes correspondingly for CE loss. The results are aligned with our theory: in
21

---

Published as a conference paper at ICLR 2022
Test Ace
Acho -0
Ap - 6.0
Aphs -00
3.8
Aponl
Agert2
34
03
Rounds
Figure 10: Test accuracy of on CIFARIO with varying data heterogeneity in the setting of ddifferent model'. All the schemes behave similarly to the 'same model' setting.
Fig. 11 we see the degradation trend for AKD which is dependent the amount of the regularization, model and data heterogeneity. In Fig. 12 we see steady behavior of AvgKD scheme for both agents models: there is no degradation even if model and data different.
Test Max Acc. Same motel Same 0003
AKD D fifreer mode, Same dera
Same mide Same (V No
Same model Diffrert
3.5
Aport
333
36
24
38
24
30
1ourds
Test accuracy of AKD MNIST using CE loss and model starting from agent I (blue) and agent 2 (red) with varying amount of regularization, model heterogeneity, and data heterogeneity. In all cases, performance degrades with increasing rounds with degradation speeding up with the increase in regularization, model heterogeneity, or data heterogeneity.
Figure 11:
Test Max Acc. of agents in AvgkD Different model, Same data
Same madel, Same data
Same model, Different data
Centralzeed
Agent
Agent?
0.96
Max Accuracy
0.92
0.88
0.84
10
Rounds
Figure 12: Test accuracy of AvgKD on MNIST using CE loss and model starting from agent (blue) and agent (red) with varying model heterogeneity, and data heterogeneity. In all cases, there is no degradation of performance, though the best accuracy is obtained by agent I in round with only local training.
H.4 COLLABORATION BETWEEN MLPS AND RANDOM FORESTS
In this section, we present the results for MNIST dataset for Random Forests (RF) and MLP models with MSE loss. That is, the experiments in the setting 'different model', where agent I 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. 13 and 14 the results for accuracy
22

---

Published as a conference paper at ICLR 2022
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 degradation for AvgKD scheme, and the performance of both agents in Avg scheme highly dependent on data heterogeneity.
Test agents 2019
Same da1a
AngKD,
ANgKD,
60
30
Emnd PLF Contro kiod
Agart 1 (MLL
Agart 2 (8.7
23
25
Figure 13: Test accuracy of AKD and AvgKD MNIST using models MLP (blue) and RF (red) with varying data heterogeneity. For AKD performance degrades with increasing rounds. Degradation speeding up with the increase data heterogeneity. For AvgKD there is no degradation of performance.
Acc. of agents in AvgKD
Aphha #
Alhha 0.01
Apha 0.05
Apha 0.1
Alpha 1.0
0.72
M
Gentra VLF
Cemtraked
2
Rounds
Rounds
Rounds
Rounds
Rounds
Figure 14: Test accuracy of AvgKD on MNIST using models MLP (blue) and RF (red) with varying data heterogeneity. The increase in data heterogeneity lowers the performance of both agents without degradation trend through rounds.
H.5 AvGKD WITH M AGENTS
The AvgKD scheme does not degrade in comparison with AKD and PKD schemes that degrade already in the case agents. this section, we present results of the AvgKD scheme with M agents and use 5 agents 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. 15. We that all the agents repeat the behavior pattern in all cases of data heterogeneity. In cases of Alpha < 0.05 (high data heterogeneity) early stopping is beneficial.
the MNIST dataset in the setting of the 'same model' with varying data
Test Acc. N agents in AugkD
Agpho -
Alsha 0.01
Apha 0.05
Apha - 0.1
Alsha
Accuary
Centraloes
Agent
Apert
Apunt:
015
Rounds
Rounds
Rounds
Hounds
Rounds
Figure 15: Test accuracy of AvgKD with M agents on MNIST with varying data heterogeneity in the setting of 'same model'. All agents can benefit from the distilled knowledge early rounds of communication.
23


Reviews:
[{"review_id": "xk-Zhxbm_FY", "reviewer": "Reviewer_w5jn", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* This paper introduced the federated kernel regression framework where they formalized notions of both model and data heterogeneity, which can be useful for developing new algorithms for model agnostic federated learning.\n* The theoretical analysis for knowledge distillation, including AKD, AvgKD, and EKD is very rigorous. \nWeakness:\n* The application of the proposed method is limited. Based on the problem setting, the proposed method is only applicable for the case of two agents, while the problem of federated learning is usually for multi-agents (more than two). Can this proposed method be extended to a more agents case?\n* The analysis of the experimental section is not clear enough or even missing:\nEx1: For Figure 3, is the loss of AKD convergent or not? This result cannot be directly observed from the figure. This paper can give an analysis on this and also for the reason why the loss of AKD increase.\nEx2: The analysis for Figure 4, 5, 6 is missing, which can make the readers confused about these results. \n* Minor problem:\nSome of the notations in the paper are a bit confusing to me, especially those seem to have similar meanings. I would suggest that if have the same meaning, they can be represented by the same notation， if not, more explanations for their differences can be offered.\nEx: It seems that $h(t)$ in Algorithm AKD have similar meaning as $g_t^2$ in Algorithm AvgKD (because these two algorithms have similar schemes). ", "summary_of_the_paper": "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. ", "summary_of_the_review": "* Although this paper introduced the federated kernel regression framework which might be useful for developing new algorithms for model agnostic federated learning, however, the application of the proposed method is limited.\n* Although the theoretical analysis part is rigorous, the analysis of the experimental section is not clear enough or even missing.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VaoeLEjSa1N", "reviewer": "Reviewer_1yWX", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "NOVELTY & SIGNIFICANCE\n\nThis paper uses a simple 2-client federated regression scenario with simplified modeling choices (i.e., kernelized linear regression) with closed-form solutions to develop novel insights which can potentially be translated into new model-agnostic FL framework with strong asymptotic guarantees. The idea (as I summarized above) is novel to me. However, in terms of its practical significance, I am not convinced because of there is too much gap between the theoretic results and real-world setting, which is elaborated below.\n\n1. The theoretic setup started with 2-client setting but has not been extended to multi-agent setting. This is problematic because the result derived in 2-client setting is very specific to the iteration between the two clients and how they average between the local data & the distilled prediction of the shared models. If there are more than one shared models, how would the client set up their training output for the next iteration? Furthermore, the authors assume implicitly that all clients have the same amount of data, which is clearly not true in practical setting -- how will it impact the averaging scheme in such cases?\n\n2. Another peculiar restriction is that the presented result is tied to an obvious flaw in the way AKD is set up in this paper. From the 4 steps mentioned at the beginning of Section 4, client 2 will never get to see true output. I suspect this is a main cause that leads to AKD's degeneration. Given this, the theoretic result developed in 4.2 is kind of moot. Can the authors re-visit this result in case models at iteration t + 1 of any clients are built on distillation of models from iteration t instead -- this is so client 2 can see its true training output at iteration 1.\n\n3. The authors also proposed a data re-incorporation scheme (i.e., AvgKD) which performs better than AKD but given the 2nd point above, it is not clear if this mechanism is necessary if the peculiar flaw in AKD is fixed. My point is if this new mechanism ends up fixing only the flaw above and not anything else, then it is somewhat unnecessary because the above flaw can be fixed by allowing both client to share models concurrently rather than iteratively. Could the authors elaborate more on this?   \n\n4. The fix in 3. is later abandoned and replaced with a better EKD fix via ensembling intermediate AKD models. But again, this result is correct given the flawed setup of AKD as pointed out in 2. above. If this flaw is corrected, can the authors revisit the asymptotic convergence result for EKD?\n\n5. The result is largely based on a formulation of regression model and the result is certainly tied to this specific regression form. I am not sure if it is reasonable to anyhow impose the implication of the derived result on a very distant classification setting. In the same vein of thought, another minor restriction is that if we view this regression formulation from a probabilistic perspective then it appears the authors impose the same Gaussian likelihood across all clients and that kind of clashes with the model-agnostic motivation. \n\n6. Have the authors considered the communication cost beyond the 2-client setting? As the distillation requires access to local data, every client would have to send models to every other clients. The total communication cost is therefore N times for than the normal cost of FL and in addition, there will also be extra distillation expense, which is a lot more costly than aggregating model weights.\n\nSOUNDNESS\n\nThe results appear correct to me. But I have reservation about the theoretic implication of both AvgKD and EKD as I pointed out above: I am not sure if these proposals were meant to fix anything more than the obvious flaw in the presented AKD setup (e.g. client 2 never gets to see true training output). Furthermore, if the AKD setup were to be fixed by letting client 2 sees its true training output, would the convergence result of EKD still hold?\n\nCLARITY\n\nThe paper is well-organized with a nice narrative flow that is easy to follow. But on a minor note, there are a number of grammatical errors all over the paper.\n\nEXPERIMENT\n\nI find the experiment somewhat strange. It implies the initial local model of client 1 is already on the same level with the centralized model which means distillation does not help at all. In all 3 experiment settings, not a single one shows that distillation is being helpful. Am I missing something here? \n\nOn another note, the flaw setup on AKD is observable from the reported results. In all 3 settings, client 2 always perform worse than client 1. In fact, on the same model, same data setting, it is noticeable that client 2 is much worse than client 1 at the initial round. This is clearly because client 2 never sees its true training output and this seems to be the case that with the flaw, client 2 is initiating and leading the distillation degradation.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper presents interesting thoughts that were substantiated via a theoretical exercise on a simplified setting. But unfortunately, in both theory & practice, this is still pretty much a work in progress. Given the huge gap between the theoretic and real-world setup, it is not clear whether one can apply the insight derived from the theoretic setting to another distant real-world setting -- please see my points in 1-5 above. Thus, while this could be the beginning of an interesting theory, more work is needed to complete the idea. As a matter of fact, the current experiment implies the proposed distillation is not working. In addition, I believe the authors have not considered the compute/communication cost of this distillation scheme.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TYvO_n9WrNP", "reviewer": "Reviewer_mYLp", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper thoroughly analyzes the dynamics of co-distillation and the proposed variants in a limited setting with two clients. The local models may differ in their kernel function so that model averaging is not possible. While the constrained setup limits practical insights, it allows to clearly analyze the behavior both theoretically and empirically. To that end, the paper details the dynamics of all three approaches and shows that the straight-forward approach may degenerate, and that the ensemble approach is optimal in the limit. Moreover, it gives conditions under which the averaging approach degenerates. These theoretical results are interesting and insightful. The empirical evaluation confirms the theoretical findings and in addition shows that the behavior is similar when using neural networks instead of kernel regression. This hints at the generality of this analysis. The paper is well-written and clear, the theoretical results are sound and, to the best of my knowledge, correct.\n\nI really appreciate the theoretical contributions, but I feel that at least the empirical evaluation is too limited. While there is good reason to analyze only two clients theoretically, it would be straight-forward to apply the proposed methods on larger numbers of clients empirically. Using only two clients and only linear regression and MNIST as experiments limits the contribution. The paper would benefit greatly from a broader empirical evaluation. \n\nDetailed comments:\n- Why is only linear regression used for the kernel experiments? \n- Which kernels are used for the experiments? I guess, a linear kernel would be most suitable to the task, but then both kernel functions would be the same.\n- How does co-distillation compare to averaging kernel models [1]?\n- The approach in [2] similarly uses co-distillation. Even though they use an unlabeled reference dataset, it seems worthwhile to compare to it.\n\n[1] Kamp, Michael, et al. \"Communication-efficient distributed online learning with kernels.\" Joint European Conference on Machine Learning and Knowledge Discovery in Databases. Springer, Cham, 2016.\n[2] Bistritz, Ilai, Ariana Mann, and Nicholas Bambos. \"Distributed distillation for on-device learning.\" Advances in Neural Information Processing Systems 33 (2020).\n\n\n******** after rebuttal **********\n\nThe authors have addressed my questions to my satisfaction. While I agree with my fellow reviewers that this work is limited in its scope, I do enjoy this novel theoretical take on distributed learning with different model types via knowledge distillation. Thus, I vote for acceptance. I have updated my score accordingly.  ", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper presents an interesting theoretical analysis of co-distillation for two different kernel regression models. The empirical evaluation confirms the theoretical findings. The constrained setting on the one hand allows for strong theoretical results, but on the other limits the practical insights that can be drawn from them. Here, a broader empirical analysis would have improved the significance of the contributions.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PTaBd_6kFBi", "reviewer": "Reviewer_fTfB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n1. This work gives the first theory analysis for model agnostic FL, even though the setting is very simple. They also use the negative result to motivate the improved algorithm ensemble AvgKD. The whole story is complete and clear.\n\n2. The experiments are sufficient, and support their analysis.\n\n\nCons:\n1. Lin et al also have similar KD idea for model agnostic FL. What is difference between AvgKD and their work?\nLin, Tao, et al. \"Ensemble distillation for robust model fusion in federated learning.\" arXiv preprint arXiv:2006.07242 (2020).\n\n2. As I mentioned before, the setting is somewhat toy. In practice, there are multiple agents, and the machine learning model is way complicated, that may not have closed form solution.\n\n3. In addition, the paper only consider optimization perspective, but I am curious about the generalization ability of AvgKD. Does it provably generalize better than purely local training?", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, model agnostic FL area is lack of theory results, so this work has its own novelty. However, the setting is somewhat toy, so I am afraid it may not be very helpful towards understanding model agnostic FL in practice (e.g., multiple clients and complicated models that do not have closed form solution.) ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This manuscript proposes and analyzes a distillation approach to address heterogeneity in distributed learning. The main paper focuses on a relatively simple two-agent kernel regression setting, and the insights developed are extended (and partially analyzed) for a multiagent setting. 

There are four reviewers, all of whom agree that the method addresses an interesting and timely issue. However, reviewers are mixed on the paper score. While all reviewers agree that the setting is somewhat stylized, a subset of reviewers highlights that the results give some deep insight that might drive future analysis and implementation in the area. Other concerns raised include potential issues with the communication overhead and the simplicity of the kernel regression setting vs real-world deep learning. There are initial concerns about whether the failure case is realistic, which the authors address. Extensions to the multi-agent setting and a partial analysis are also addressed by the authors and partially satisfy the reviewers. Nevertheless, after reviews and discussion, the reviewers are mixed at the end of the discussion. 

The area chair finds, first, that the paper is much improved, and much more applicable in the updated form than in the original version, and indeed, the insights from the simple model may be informative for practice. However, the concerns raised about the distance between theory and practice are valid. The final opinion remains borderline. Authors are encouraged to address the highlighted technical concerns in any future submission of this work. In particular, the muti0agent setting should probably be central in the discussion of this work. More ambitious empirical evaluation showing that the theory translates to practice )even if there is a gap) would also help.
Decision: Accept (Poster)