Towards Diverse Device Heterogeneous Federated Learning via Task Arithmetic Knowledge Integration
Authors: Mahdi Morafah
Vyacheslav Kungurtsev
Hojin Matthew Chang
Chen Chen
Bill Lin
Venue: NeurIPS 2024 poster (2024)
Keywords: Federated Learning
Heterogeneous Device Prototypes
Knowledge Distillation
Task Arithmetic
Machine Learning
Primary area: other

Abstract: Federated Learning (FL) has emerged as a promising paradigm for collaborative machine learning, while preserving user data privacy. Despite its potential, standard FL algorithms lack support for diverse heterogeneous device prototypes, which vary significantly in model and dataset sizes---from small IoT devices to large workstations. This limitation is only partially addressed by existing knowledge distillation (KD) techniques, which often fail to transfer knowledge effectively across a broad spectrum of device prototypes with varied capabilities. This failure primarily stems from two issues: the dilution of informative logits from more capable devices by those from less capable ones, and the use of a single integrated logits as the distillation target across all devices, which neglects their individual learning capacities and and the unique contributions of each device. To address these challenges, we introduce TAKFL, a novel KD-based framework that treats the knowledge transfer from each device prototype's ensemble as a separate task, independently distilling each to preserve its unique contributions and avoid dilution. TAKFL also incorporates a KD-based self-regularization technique to mitigate the issues related to the noisy and unsupervised ensemble distillation process. To integrate the separately distilled knowledge, we introduce an adaptive task arithmetic knowledge integration process, allowing each student model to customize the knowledge integration for optimal performance. Additionally, we present theoretical results demonstrating the effectiveness of task arithmetic in transferring knowledge across heterogeneous device prototypes with varying capacities. Comprehensive evaluations of our method across both computer vision (CV) and natural language processing (NLP) tasks demonstrate that TAKFL achieves state-of-the-art results in a variety of datasets and settings, significantly outperforming existing KD-based methods. Our code is released at https://github.com/MMorafah/TAKFL and the project website is available at https://mmorafah.github.io/takflpage .

Towards Diverse Device Heterogeneous Federated Learning via Task Arithmetic Knowledge Integration
Mahdi Morafah" Vyacheslav Kungurtsev", Hojin Chang', Chen Chen Bill Lin
Uniiersityy of Californaa San Diego (UCSD). 'Czech Technical University in Prague, University of Central Florida (UCF)
*Correspondence: mmoraf werrafah@ucsd.do edu
Abstract
Federated Learning (FL.) has emerged as a promising paradigm for collabora- tive machine learning, while preserving user data privacy. Despite its potential, standard FL. algorithms lack support for diverse heterogeneous device prototypes, which vary significantly in model and dataset sizes-from small IoT devices large workstations. This limitation is only partially addressed by existing knowl- edge distillation (KD) techniques, which often fail to transfer knowledge effec- tively across broad spectrum of device prototypes with varied capabilities. This failure primarily stems from two issues: the dilution of informative logits from more capable devices by those from less capable ones, and the use of single integrated logits as the distillation target across all devices, which neglects their individual learning capacities and and the unique contributions of each device. To address these challenges, we introduce TAKFL, novel KD-based frame- work that treats the knowledge transfer from each device prototype's ensemble as separate task, independently distilling each to preserve its unique contribu- tions and avoid dilution. TAKFL also incorporates a KD-based self-regularization technique to mitigate the issues related to the noisy and unsupervised ensemble distillation process. To integrate the separately distilled knowledge, intro- duce an adaptive task arithmetic knowledge integration process, allowing each student model to customize the knowledge integration for optimal performance. Additionally, we present theoretical results demonstrating the effectiveness task arithmetic in transferring knowledge across heterogeneous device prototypes with varying capacities. Comprehensive evaluations of our method across both computer vision (CV) and natural language processing (NLP) tasks demonstrate that TAKFL achieves state-of-the-ae results in a variety of datasets and settings, significantly outperforming existing KD-based methods. Our code is released at tattpp://gt //github........ thub. com/MMorafah/TAKK
Introduction
Federated Learning (FI.) has rapidly gained traction as a promising approach train machine learn- ing models collaboratively across multiple devices, while preserving the privacy of user data. Stan- dard federated learning methods, such as FedAvg [33], however, are primarily designed for un- realistic device-homogeneous scenarios, where all devices are assumed to have identical compute resource and can train the same neural network architecture [28, 33, 49, 21, 27, 48, 31J. Therefore, standard FL cannot support the participation of heterogeneous devices, all of which could signif- icantly contribute to model training due to their unique and invaluable local datasets. To address this gap. knowledge distillation (KD) techniques have emerged a promising approach to establish federation among heterogeneous device prototypes and facilitate knowledge transfer between them. In this approach, locally updated client models from different device prototypes, collectively termed as ensembles, serve as teachers to distill their knowledge into each device prototype's 's server student model using unlabeled public dataset.
Project Website: https:/ //moorafah.. github. io/takfllage..
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Despite their success, however, existing KD-based methods for device heterogeneous FL are primar- ily designed for scenarios where device prototypes are in the same-size with similar capabilities, i.e. same model and dataset sizes. However, in practice, device capabilities vary widely, ranging from small devices like loT's with small models and small datasets to large devices like workstations with large models and large datasets. This diversity, often overlooked in the existing literature, results in device prototypes with varying strengths and information qualities. Unfortunately, existing methods struggle to establish effective knowledge transfer in these challenging, real-world device heteroge- neous settings, primarily due to two reasons: S Existing methods often disregard the individual strengths and information quality of each device prototype's ensembles and integrate their logits into a single distillation target. This approach dilutes the richer, more informative logits from larger; more capable devices with less informative logits from smaller. less capable ones. ? Additionally, these methods employ this single integrated distillation target to transfer knowledge across all dif- ferent size student models. This one-size-fit--de approach fails to provide customized knowledge integration based on the unique learning capacities of each student and the specific helpfulness of each device prototype's 's ensembles.
Moreover, the heterogeneous ensemble distillation process can inadvertently lead student models into erroneous learning directions, causing them to forget their self-knowledge acquired through averaged locally updated parameters. This issue arises primarily due to two reasons: ① The distilla- tion process introduces noise, as the ensembles' logits are inferred distinct from their original training data. Additionally, the presence of data heterogeneity and the insufficient training of some ensembles, due to computational constraints, can further exacerbate this noise. 2 The distillation process lacks supervision from the actual private datasets, which are the ultimate learning objectives. Consequently, these factors, combined with the limitations outlined earlier, result in suboptimal knowledge transfer in device heterogeneous settings. This underscores the urgent need for more effective knowledge transfer framework.
unfamiliar public dataset,
In this paper, we introduce TAKFL, a novel "Task Arithmetic Knowledge Transfer Integration for Federated Learning" framework, designed to overcome the fundamental limitations in the existing methods and improve knowledge transfer in scenarios where device prototypes vary in size-bor both model and dataset-and device prototype's 's ensembles as separate tasks, distilling them independently to ensure that each prototype's 's unique contributions are accurately distilled without interference. To tackle the chal- lenges associated with noisy and unsupervised ensemble distillation, we incorporate a KD-based reff-regularization technique into this individual knowledge transfer process. Subsequently, to lectively integrate the separately distilled knowledge from heterogeneous prototypes' ensembles, we introduce an adaptive task arithmetic knowledge integration method by extending the notion of task vectors from centralized learning to federated learning. Our approach enables the student model to strategically customize the knowledge integration process based on the quality of knowledge from each prototype's 's ensembles and its intrinsic capacity, aiming to achieve optimal performance. We present theoretical results, grounded on the established theoretical learning properties of over- parametrized neural networks, that conceptualize knowledge distillation as the allocation of device prototypes' capacities to accurately fit the chosen logits. These results demonstrate the advantages of employing task arithmetic for knowledge transfer in terms of overall accuracy, coverage, and effi- ciency, as well as the adaptive knowledge integration based on the capacity of the student prototype. Furthermore, we comprehensively evaluate our method across both computer vision (CV) and nat- ural language processing (NLP) tasks, utilizing various datasets and architectures, and demonstrate that TAKFL consistently achieves ofate-of-the-art (SOTA) performance.
consequently, in strength. TAKFL treats knowledge transfer from each
The contribution of our paper is as follows:
1. We formalize and review the important considerations of the problem statement of federated learning with heterogeneous device prototypes.
2. We introduce TAKFL, a novel KD-based method designed to overcome the fundamental limita- tions of existing approaches, effectively facilitating knowledge transfer across diverse heteroge- neous device prototypes with varying capabilities.
We present the first theoretical model for device heterogeneous KD, and demonstrate the effec- tiveness and efficiency of TAKFL compared to the standard alternatives that do not adapt to the student's self-knowledge quality and available learning capacity.
4. Our comprehensive experimental evaluations on both CV and NLP tasks, spanning various datasets and architectures, reveal that TAKFL consistently achieves SOTA performance, out- performing existing KD-based methods.
2

---

Medium
Small
Entembler of Diverte Device Protohypas
Ensembles
Student
Studert
Student
Loga
Medium Ensembles
Avgroged Wool
Aweraged Modid
Model
Unkibeied F Dute
Small Ensembiee
disstittd
e avv
avsa
e distited
1s
Student
distitted
avg
A T.
@ merggd
aang
A Tu
O distitted
O ass
0 merged
0 ave + (ReTT AuT +2gTs)
(a) Vanilla Ensemble Distillation
(b) Overview of TAKEL
Figure I: Overview of our approach and its distinction from prior works. (a) This figure il- lustrates the vanilla ensemble distillation process, where logits from ensembles of various sizes averaged and used as the distillation target across all prototypes. This approach leads to the dilution of information and suboptimal knowledge transfer (refer to Sections 6 and 7 for details). (b) This figure depicts our approach, TAKFL, which treats knowledge transfer from each prototype's ensem- ble separate task and distills them independently. Additionally, KD-based self-regularizatiio technique is introduced to mitigate issues related to the noisy and unsupervised ensemble distilla- tion. Finally, the heterogeneously distilled knowledge is strategically integrated using an adaptive task arithmetic operation, allowing for customized knowledge integration based on each student pro- totype's 's needs.
2 Related Works
Device Heterogeneous FL. Prior works on device heterogeneous FL have considered two distinct approaches with different objectives and settings. The first array of studies focuses on accommo- dating devices with varying compute resources, aiming to train a single global model. Techniques such as static and rolling-hased partial model training allow devices to train a sub-model of the global model tailored to their compute resources [11, 18, 3, 1]. However, this approach does not fully reflect real-world scenarios. In practice, device prototypes such as IoT' and smartphones have unique neural network architectures designed for their specific configurations and underlying tasks, which may not support training varying neural architectures. This highlights a significant limitation in accommodating the full spectrum of device heterogeneity in this approach. The second array of studies addresses more practical scenario where device prototypes with heterogeneous model chitectures participate in FL to enhance their global model performance through mutual knowledge sharing [30, 43, 6]. In this context, KD techniques are used to transfer knowledge among prototypes, where locally updated client models, termed as ensembles, serve as teachers to distill their knowl- edge into each server's student model using unlabeled public dataset. For example, FedDF [30] logit averaging, en-
uses vanilla logit averaging, while Fed-ET [6] applies an uncertainty-weighhee hanced by diversity regularization technique. However, existing works typically focus on settings where prototypes have similar capabilities both model and datasel sizes-and thus neglecting the challenges in more diverse seltings with varying capabilities. This oversight leaves their effective- ness in such settings largely umexplored. In this paper. we aim to study the underexplored diverse heterogeneous device settings. See Appendix A for more detailed discussion on the related works.
Model Editing via Task Arithmetic. Traditional methods for model editing often involve expen- sive joint fine-tuning across multiple tasks, which can limit scalability and democratization [62]. Recently, a promising technique called task arithmetic has emerged as a cost-effective and scal- able method for updating pre-trained models with new information refining undesired behav- ior [53, 37, 32]. The concept of "task vectors" introduced by Wortsman et al. [53] plays a pivotal role in these techniques. For any given task I, a task vector is derived by subtracting the model's pre-trained weights Ppec from its fine-tuned weights 0r on task L i.e. T. = Oft O 0pre- These task vectors act as unique representations for specific tasks. Furthermore, researchers have demonstrated that by summing multiple task vectors (Ta) and integrating them into pre-trained model via
0pee + > Et=1 Ts, one can effectively create model capable of handling multiple tasks [53, 57]. To the best of our knowledge, this work is the first to exlend the notion of task vectors to the federated learning setting, introducing a task arithmetic for knowledge distillation across diverse heteroge- neous device prototypes.
3

---

3 Problem Statement: FL with Heterogeneous Device Prototypes
Consider a cross-device FL setup with a set of M distinct device prototypes M, i.e. M = [ Each device prototype my E M has a distinct neural network architecture (:; (1) parameterized by 0 E Rhy and set of clients C/ with N = |C|| clients in total. Each client ca E C has local private dataset Dr2 ( 34) ) where Rj,A = [ and locally trains the parameters O: of the neural network architecture f on its local dataset. Furthermore, denote TV Ukecy Tyd to be the union of the private datasets for device prototype j. We assume D Di, that a subsample from the population distribution DJ and similarly IDA ~ Di. The union of the private datasets, i.e. D = =UUUMMM DV, is sampled from the entire population D, which is defined as an unknown mixture of the distributions each device prototype sampled its data from, i.e. generically non-i.i.d We formalize this as a mixture of local clients data population, i.e., = Wy. Ej Ex Wi.DDi where 0 S Wj.k S and Ejs Wj.k = 1, and wij.k is unknown.
The ultimate objective is minimize the test error and thus enable accurate inference for each device prototype j. aiming to obtain the optimal parameters for the population dataset:
N
argmin B(a.))>D((( 0²), ))] argminn Ø1
Wi.h Bi (aw)) (:N (2)
Ø/
where E(: .) is the sample-wise loss function (e.g. cross entropy for image classification) and we decompose by total population loss with the linearity of expectation in the mixture. See Fig 4b for visual illustration of heterogeneous device prototype FL.
4 Background: Federated Ensemble Distillation
To address the limitations of standard FL in device heterogeneous settings, Lin et al. [30] proposed ensemble knowledge distillation to transfer knowledge between heterogeneous device prototypes in FL. This procedure consists two stages: (1) local per-prototype FL, and (2) server-side vanilla ensemble distillation. The details of each stage discussed in the following paragraphs.
Local Per-Prototype FL. this context, at each round r subset of clients C/ from each device prototype E M is randomly selected by the server and download their corresponding model ini- tialization 0z. Each client of E C₁,, starting from this model initialization, locally train the model fö on its local private data W.d by taking multiple steps of stochastic gradient descent. Then, they send back their updated parameters (0) (0k)eEs to the server. The server aggregates the received clients parameters, and computes Olug = Ekect Ence P k-
[Pal 0
In classic federated learning formalism,
the parameters 0 ang satisfy,
N Σ E (m (x( v)]
0
0gag E argmin
(2)
01
Vanilla Ensemble Distillation. In this stage, each server model f gets initialized with 0', and undergoes updates using ensemble knowledge distillation. Here, heterogeneous client models from heterogeneous device prototypes, collectively termed as ensembles, serve as teachers, i.e. T :=
(f'(, (((, E M. & E C³). transferring their knowledge to each server student model, i.c. S := O"). For simplicity, we drop the index for each server student model, denoting it as S. The ensemble distillation loss using a mini-batch of data from an unlabeled public dataset, i. æ Doublic can be defined by the following equation:
F(O)) (S(x)) ]. [
CED = KL σ
(AvgLogits)
(3)
FET
where (-) is the softmax function. As illustrated in Eq. 3, vanilla ensemble distillation treats all heterogeneous device prototypes' ensembles equally by uniformly averaging their logits. This way of knowledge integration overlooks the individual strengths and informational value of each pro- totype's 's ensembles. As result, the richer, more informative logits from stronger ensembles are diluted by less informative logits from weaker ensembles, leading to information loss. Furthermore,

---

this averaged logits is used as the distillation target across different-sized student models, irrespec- tive of their intrinsic capacity and the helpfulness of each prototype's 's ensembles. Consequently, this leads to suboptimal knowledge transfer in device heterogeneous FI.. See Section 6 for theoretical analysis and Section for experimental observations.
5 Task Arithmetic Knowledge Transfer and Integration
In this section, we introduce TAKFI,, designed to overcome the fundamental limitations of previous approaches and enhance knowledge transfer across diverse heterogeneous device prototypes, which vary in size-in terms of both model and dataset size. TAKFL consists of two main components: (1) individually transferring knowledge from each device prototype's 's ensembles, and (2) adaptively integrating knowledge via task arithmetic. Detailed descriptions of each component are provided in Section 5.1 and 5.2, respectively. An illustrative overview of TAKFL is presented Figure Ib, and the full algorithm is detailed in Appendix B, Algorithm 1.
5.1 Knowledge Transfer from Individual Device Prototype
We begin by discussing our proposed knowledge transfer framework from each individual device prototype's ensembles. This process consists of two main components: ensemble knowledge trans- fer and self-regularization, each detailed in the subsequent paragraphs.
Ensemble Knowledge Transfer. Vanilla ensemble distillation integrates the knowledge of vary- ing strength ensembles by uniformly averaging their logits. This approach can potentially trans- form or even degrade the overall quality of the knowledge being transferred, leading to suboptimal knowledge transfer. To effectively distill the unique knowledge and contributions of each proto- type's 's ensembles, and to avoid dilution, information loss, and interference from other prototypes' ensembles, we propose transferring the knowledge from each prototype's ensembles separately and independently.
Specifically, let's consider T, (r((,0:]1 C-} denotes the ensembles of device prototype i as teacher and Sj denotes the server student model of the device prototype j. Without loss of generality, we refer to each device prototype's server student model as just student denoted as S. Therefore, the knowledge distillation loss between the teacher ensembles Ti and server student S (T S) is defined below:
CT--- C
KL
{o ( ITL
1
Σ FET:
F(x) EF(T)
a (8(x))
(4)
Scaffolding Student from Noisy Ensemble Distillation. The ensemble distillation process may adversely impact the student, causing to forget its own knowledge acquired through averaged locally updated parameters and be drifted into erroneous directions. This is primarily due to two key factors: (1) The ensemble distillation process introduces noise, mainly because the ensembles' logits inferred on unfamiliar public dataset they have not been trained on. These ensembles are originally trained on local private datasets, which usually differ from the unlabeled public dataset used for distillation. Moreover, other factors such the presence of data heterogeneity within FL. and insufficient training of some ensembles due to limited computational resources can exacerbate this noise, particularly in the initial rounds of federation. (2) The ensemble distillation process lacks supervision from the actual private datasets, which is the ultimate learning objective.
To scaffold the student models from the noisy and unsupervised distillation process, which may cause them to drift into erroneous directions and forget their invaluable self-knowledge, we introduce KD-based self- reglaaauulri technique. Our relllaagultiization technique mitigates these issues by enforcing similarity between the logits of the student and its initial logits (when the student is initialized with averaged parameters) using KL divergence loss defined below:
Caelf = KL σ (S(x; Oavg)) a (S(z))
(5)
Overall Knowledge Transfer Objective. The overall knowledge transfer objective from teacher
ensembles T. of device prototype i to the student S is the combination of the ensemble knowledge distillation loss CT- (Eq. 4) and the relf-regularization loss Cgl (Eq. defined in the following:
CJ TT + + -gg..
(6)
Here, 7 is a hyperparameter controlling the effect of self-regularizzttion term. We associate the knowledge transfer from each device prototype to task T. with the loss C3.
5

---

5.2 Task Arithmetic Knowledge Integration
Centralized Learning
Federated Learning
Herein, we delve into the details of our proposed method for customized integration of the separately distilled knowledge from heterogeneous ensembles. Drawing inspiration from recent advances in model editing via task arithmetic [53], where a pre-trained model's knowledge can be edited via task- specific vectors using arithmetic operation, we propose novel customizable knowledge integration method via task arith- metic. To do centralized learning to federated learning. We conceptualize the averaged locally updated parameters, i.e. Oang as a "pre- trained", similar to those in centralized learning, and the pa- rameters of the distilled model via knowledge transfer objective (Eq. 4), denoted as Ocsstilled "fine-tuned" version of the model (see Fig. 2). Consequently, the task vector т, associated with the knowledge transfer task Cs can be defined by subtracting the distilled parameters from the averaged locally updated parameters as follows:
8t
disttilled
Opro
e avg
r- 8s:
pre
r Oiiatilled
-aay
extend the notion of task vector from
Figure Analogy between task vector in centralized learning and federated learning.
T
0
Ti-S distilled
P Davg-
Essentially, task vectors serve as unique representations for the transferred knowledge from each prototype's 's ensembles the student and encapsulate the distinct contributions of each prototype's ensembles to the student model. To selectively merge the knowledge of each prototype' ensembles into the student, we employ an adaptive task arithmetic operation as follows:
0 meerged = Oang + Σ AAT
(8)
TEM
where X denotes the merging coefficient associated with task vector T;, and they sum to one, i.e. Eiem A, The merging coefficients determine the extent of knowledge integration from each prototype's ' ensembles. Essentially, they enable the student to have customized knowledge inte- gration to achieve maximum performance. The student can determine these merging coefficients based on its own learning capacity and the relative knowledge and helpfulness of other device proto- types' ensembles. This approach provides an effective, low-cost, and scalable knowledge integration strategy in settings with diverse device heterogeneity. In our experiments, we considered this as a hyperparameter and tuned manually or determined it using held-out validation sets which achieves similar results. More details can be found in Appendix F.3.
6 Theoretical Results
We present theoretical understanding on the efficacy of knowledge distillation in device heteroge- neous FL. We argue that vanilla ensemble distillation (VED) diffuses the information from logits, which presents a notable disadvantage for solving (1). This effect is particularly pronounced when the teacher ensembles are from device prototype of small capacity, and the student model is from device prototype of large capacity. By contrast, our proposed method of task arithmetic knowledge integration, mitigates the drawbacks of VED and is able to simultaneously incorporate information from differently sized heterogeneous ensembles, efficiently filling up the capacity of each student with the most informative knowledge, achieving optimal knowledge transfer.
Assumptions and Preliminaries. Standard practice, including the setting in consideration as well as the numerical experiments here, involves overpurumetrized neural networks, that is, the total number of weights far exceeds the training sample size. This implics that the set of weights that minimize the loss is non-unique, and moreover, it has been argued that they form submanifold [8]. This submanifold structure of solution sets will provide the critical source of understanding the subsequent results. In particular, we shall consider knowledge distillation as filling up the capacity of device prototypes' models with basis vectors corresponding to submanifolds that minimize as many device prototypes' data distributions possible.
Since we are interested server-side distillation across heterogeneous device prototypes, we assume optimal conditions at the local per-prototype FL level, meaning that the perfect solution for local per-prototype FL. is achieved. The formal details of the assumptions and statements are presented in Appendix C. In these Propositions we assume that the prototype i and j datasets are disjoint, i.c.

---

D' n DJ = 0. and D = D. We show that the other cases are trivial and uninteresting in the appendix.
Proposition 1. (information loss in VED, informal). Consider the VED procedure in the form of solving (3). Consider two device prototypes with device capacity and solution dimension of Q' Q² and W W2, respectively, and with associated eigenbases ② W Denote WhJ, .i,J = 1, as the capacity allocated by student in order to distill knowledge from teacher j logits.
1. Case 1:
When the capacities are the same, that is Q' = =Z and W = W W1. = W then with VED, there will be some capacity, in the sense of eigenspace, of student prototypes that will be allocated with parameters that do not minimize the student's ss its own data distribution. Case 2: Assume that Q' > Q² and W W1.2 > W2, Then the phenomenon as for Case holds. Moreover, there will be some capacity of student I's model that will be allocated with parameters that do not minimize either of the teacher or student prototype's 's data distribution.
An interesting key mechanism of the proof is that when VED is applied in distilling logits from a small device prototype to a large one, the modeling capacity of W is structurally reduced to that of W2 < W 22, i.e., it is an operation wasteful of the potential model capacity.
Remark 1. This proposition proves that in general, VED is prone to diffuse knowledge already present in students, and leads inefficient and inaccurate under the case that device prototypes have different capacities. VED ends up leading to more er- ronous models entirely as the small information within the small teacher is transferred onto a larger capacity target.
of model capacity. Furthermore,
Proposition? 2. (improve knowledge transfer with task arithmetic, informal). Consider the TAKFL procedure as in the form of computing (8). Consider two device prototypes with device capac- ity and solution dimension of Q- Q² and W1, W2, respectively, and with associated eigenbases ②¹, Wt.
Case 1: In the case that that Q' 2 Q² and W 2 W it holds that the TAKFL with prototype 1 as student preserves the eigenbasis associated to the parameters used to accurately fit the data D
2. Case 2: Assume that Q' = Q2 and W 2 TAKFL yields solution for the student that is at the intersection of the subspaces corresponding to minimizing the two data distributions.
3. Case 3: Assume that Q' Q² and W W In the case of prototype being the student, TAKFL yields a solution that:
(a) retains the approximation accuracy on device I's data distribution.
(b) ensures approximation accuracy to the level of device 2's relative capacity
(c) fills the remaining local capacity device 1 has allocated for device 2's logits with infor- mative new knowledge, unless enforced otherwise.
Remark 2. This proposition proves that general, TAKFL promotes the most efficient allocation of the devices' capacity in order to accurately fit a diverse set of data distributions. With TAKFL, the previously acquired knowledge is entirely preserved. Even under the case that device prototypes have different capacities. TAKFL smartly transfers the most informative knowledge to each proto- type's ' student model based on its own intrinsic capacity. Still, the final statement indicates that in the case that there are many different teachers, while small device prototype serving teacher will not necessurily compromise information, it would still be preferable to allocate that capucity to more informative, larger, teacher model.
Comments. We comment the complexity and convergence aspects of TAKFI. briefly:
1. Computation Time: TAKFL's computation time is O(1) (constant) due to parallclization, as all distillation processes occur simultaneously.
2. Computation Load: The overall computational load scales O(M) (linear) since the distil-
lation tasks are performed independently for each prototype parallel and merged into a singe task arithmetic operation (Eq. 8).
Resource Usage (Memory): The resource usage scales as ((M2) (quadratic) because of the
need to store and process multiple task vectors concurrently.
Optimization Convergence Rate: The convergence rate of the Knowledge Distillation proce- dure is as standard for training nonconvex loss function, e.g. [42].
7

---

7 Experiments
7.1 Main Experimental Setup
Dataset and Architecture. We evaluate our method on computer vision (CV) and natural language processing (NLP) tasks. For CV, we train image classification using CIFARTO/ [24), CINIC- 10 [9], and Tinylmagenet [25]. For NL.,, we fine-tune pre-trained models for text classification on MNLI [52]. SST-2 [45], MARC [22], and AG News [60]. Our architectures include ResNet [17], VGG [44], and ViT [12] for CV, and small BERT variants [47] (-Tiny, -Mini, -Small) for NLP. We simulate a federated non-i.i.d setting using a Dirichlet distribution Dir(at), where a lower a indicates higher heterogeneity [27, 36.. Further details can be found in Appendix F.1 and F.2
Implementation Details. We use the Adam optimizer for both CV and NLP tasks. For CV, local training involves 20 epochs with a learning rate of 0.001, weight decay of 5e-5, and batch size of 64. NLP training is conducted over epoch with a learning rate of 3e-5, no weight decay, and batch size of 32. For distillation, Adam is used with learning rate of 1e-5 and weight decay of 5e-4 for CV, and 3e-5 with no weight decay for NLP. Batch sizes for distillation are 128 for CV and 32 for NLP. The softmax temperature is set at 3 for both tasks, with temperature of 20 for self-regularization. Further details are provided in Appendix F.1 and F.2.
method against standard FL, i.e. Fe- dAvg [33] and SOTA KD-based methods designed for heterogeneous device prototypes FL, includ- ing FedDF [30] and FedET [6]. The evaluation metric is the final top-1 classification accuracy of each device prototype's 's global model on the test dataset, as per the methodology described in [36]. We report the average results and the standard deviation over three independent runs, each with different random seed.
Baselines and Evaluation Metric. We compare
A more detailed version of the experiments, alongside additional experiments and ablation studies, is presented in Appendix D and E.
7.2 Main Experimental Results
In this section, we evaluate the performance of our method, TAKFL, in a federated learning environ- ment that mirrors real-world scenarios with diverse, heterogeneous device prototypes, as illustrated in Fig. 4b. Our experimental setup includes three different device prototype sizes: Small (S) with small model and small dataset, Medium (M) with a medium-sized model and medium-sizee sized dataset, and large (L) with large model and large dataset.
Performance on CV Task. Table presents the performance of TAKFL in the homo-family ar- chitecture setting on the CIFAR-10 and CIFAR-100 [24] datasets (for hetero-family architecture results, see Appendix D.I. Table 4). TAKFL consistently enhances performance across all device prototypes in various scenarios, achieving SOTA results. Notably. in the Dir(0.3) setting on CIFAR- 10, TAKFL improves average performance across all prototypes by 8%, and by 4% on CIFAR-100 100. From Table 1, inconsistent performance improvements are observed with prior KD-based methods, especially for the prototype. While S and M prototypes achieve gains, the L prototype suffers up to a 10% degradation compared to vanilla FedAvg, highlighting the dilution issue where valuable information from larger, more capable device prototypes is diluted by less informative outputs from smaller devices. Moreover, the significant performance improvements TAKFL achieves for each device prototype, particularly for S and M prototypes, illustrate the ineffectiveness of the one-size- fits-all approach used in the existing KD methods. These observations confirm the shortcomings of vanilla ensemble distillation and corroborate our thcoretical findings in Remark and 2. The effectiveness of our self-regularization technique is further supported by these experimental results. For more detailed and insightful analysis see Appendix D.I.I.
Performance on NLP Task. Table 2 presents the results on MNLI [52] and SST-2 [45] datasets (see Appendix D.3 for further experiments). Similar the CV task. TAKFL has consistently improved performance across all device prototypes of varying sizes, achieving SOTA results: 3% average increase on MNLI and 2% on SST-2. The suboptimality of existing KD methods, evident from the results presented here as well. Notably, FedET suffers from a significant performance degradation compared to vanilla FedAvg. This issue stems from FedET's reliance on the confidence scores of neural networks for uncertainty estimates. However, neural networks, especially pretrained language
8

---

Table 1: Performance Results for CV task CIFAR-10 and CIFAR-100. Training data is distributed among S. M. and device prototypes in a 1:3:6 ratio, subdivided among clients using Dirichlet distribution. Public datasets are CIFAR- 100 [24] for CIFAR- 10 [24] and ImageNet- 100 [10] for CIFAR- 100, Client config- urations include 100, 20, and 4 clients for S, M. and L. with sampling rates of 0.1. 0.2, and 0.5, Architectures are ResNet-8, 14, and ResNet- 18 [17] for S, M and I., respectively. All models are trained from scratch for 60 rounds. See Appendix 1 for additional experiments using hetero-family architecture and more details.
Low Darn Hetemgeneity (Dir(0.3)
High Data Heremogeeitty (Dir ())
Daraset
Raseline
S
M
Average
M
Average
FodAve
26.21_2..1
12.11
59.461a.17
17.36
22.01
11.00
51 51.51
52.93
FedDF
49.31=2.15
50.6320.75
49.82+9.us
49.92
34.71a144
36.27
51 08
40.35
CIFAR-10
FedET
19.21
15.0111.41
53.001s.47
52.60
29. 19.
90.96
14.70
15.531e.48
35.36
TAKEL
15.90_1.mm
57 87.93
60.58.2.m
68.14
37 2.
38 28.07
51. 51.19
42.62
TAKFL+Reg
56.37 50.37_n.40 0.46
58.60
65.69
60.22
40. 10.51:110
40.12.1124
23.24
44.22
FedAvg
13:22.0.14
11.39
29.47
21.36
11.86 os
14.6320.65
26.25 64
17.58
FedDF
19.54-0.20
24 24.32=0
29. 29.29
24.38
16 16.00
19.80=0.17
26.59=0
20.83
CIRAR-100
TedET
19.67
20
31.10g1. ca
25.35
11.18 os
18.2220.35
26.40 as
TAKEI
24.48=0.42
27. 27.0020.25
29.8443..0
27.31
22. 90 20.1
23.1322..2
26.98=0.15
24.50
TAKFL+Reg
27.18-0.27
29.1440.20
31. 31.15.000 07
29. 1.5
22.88 87
23.92+0057 B7
28. 28.01+0.84 84
24.94
Table 2: Performance Results for NLP Task on MNLI and SST-2. Training data distribution is similar to the CV task using only Dir(0.5) here. Public datasets are SNLI [2] for MNLI [52] and Sentiment [14] for SST-2 [45]. Client configurations are 8, 4, and 2 clients for S. M. and L, with sample rates of 0.3, 0.5, and 1.0, respectively. Architectures include Bert- Tiny, Bert-Mint and Bert- Small [47] for S, M. and L, initialized from pre-trained parameters and fine- munettned for 20 communication rounds. See Appendix F.2 for more details.
MNLI
SST-2
Baseline
M
L.
Average
S
M
L
Average
FedAvg
16.15
54.47 12.18
57.51
12.79
49.37
54.9811.81
74.71
16.69
72.13
FedDF
14.21
60.4421.0
66.71 09
60.45
74.41
80. 1..71
84.35g1.00
79.82
FodET
18.03+6.32
50.33+7.57
53.80=6.18
50.72
66 66.63+9.14
65.89+16.85
70.03+15.ss
67.52
TAKEL
57.43gn.a1
63.58
68.74
63.25
74.7320.55
82. 82.1720.31
86.93
81.28
TAKFL+Rog
67.61+0.50
63.91=1.05
68. 1.
63.49
74.88+0.4 48
82 $2.10+0.sa sa
87 37.33=0.m os
81.
models (PLMs), tend to be poorly calibrated and overconfident, undermining reliable uncertainty estimates [50, 15, 5, 55].
7.3 Scalability Evaluation
We evaluate the scalability of TAKEL across a spectrum of device prototypes, from extremely small (XXS) to extremely large (XXL), see how well our method adapts from a uniform array of small- size prototypes to a diverse mix of sizes. Each prototype is equipped with appropriately scaled model and dataset sizes, simulating real-world variations in device capabilities.
Figure illustrates TAKFL's ability to effectively scale from to device prototypes. In scenarios where all devices are similarly small, i.e. 3-device setup, TAKFL's performance is slightly better than FedDF. This is because when devices are homogeneously small and similar in capability, they
FedAvg FedET FedDF TAKFL
Avg
3 Devices
Avg
5 Devices
Avg
7 Devices
XXL
2E XXS
XL
SAAS
52.71
25.24 03
33.49 52.30
32,24
24.14 45.93
10.53
XXS
25:03
25.04 24.19
57.05 38.20
66 91
XS
20
XL
32.51
33.82
24.00 b5
52
$0.01
31.77
XS
37
M
00
51
45.70
37:63
40 76
$3.51
5
43.96
43.53
80.16
L
37.85
S
M
Figure 3: Scalability Evaluation of TAKFL. Image classification on CINIC-10 [9] dataset is used to evaluate TAKFL's scalability across device prototypes ranging from XXS XXL.. Training data distributed among prototypes a 1:2:3:4:5:6:7 ratio, further subdivided using (0.(.. Client configurations range from 35 for XXS to 5 for XXL. Architectures span from ResNec1O-XXS for XXS to ResNet50 for XXL prototype, all initialized from scratch and trained over 30 communication rounds. The public dataset is CIFAR- 1000 [24]. See Appendix D.4 for more details.
9

---

do not offer unique contributions that could benefit from more complex distillation strategies. How- ever, as the scenario expands include larger devices like XL and XXL in the 5- and 7-device configurations, TAKFL. significantly outperforms existing KD-based methods. This improvement is driven by the larger devices' ability to offer more signíficant and higher-quality knowledge, which TAKFL effectively distills across all prototypes, contrasting sharply with existing methods that fail to utilize this potential. These experimental observations, corroborated by our theoretical insights in Remark 2, demonstrate TAKFL's superior scalability and effectiveness.
8 Conclusion and Discussion
In this work, we addressed a fundamental issue in standard federated learning: the lack of support for heterogeneous device prototypes. Existing KD-based methods often fall short in real- world scenar- ios, where device capabilities vary widely. To address this, we introduced TAKFL, a novel KD-based method that treats knowledge transfer from each prototype's ensembles as separate tasks and distills them independently. TAKFL susequently integrates the knowledge using an adaptive task arithmetic technique for optimized performance. We also introduced KD-based self-regulation technique to mitigate issues arising from noisy and unsupervised ensemble distillation. The effectiveness of our method is substantiated by both theoretical results and extensive experimentation across CV and NLP tasks, using various datasets and models.
Limitations remain, notably in real-world applicability. While TAKFL's effectiveness in approx- imated woal-world setup has been demonstrated, actual deployment on physical devices and in envi- ronments with extremely large models remains untested due to resource constraints. Experiencing TAKFL in genuine real-world settings could unveil additional challenges or limitations, providing further insights into its scalability and efficiency.
9
Acknowledgment
This work was partially supported by research grant from Cisco Systems, Inc., Project Number 49790. V.K. acknowledges support from the Czech National Science Foundation under Project 24- 11664S. We also gratefully acknowledge the use of the computational infrastructure provided by the OP VVV funded project C7.02. 1. 1.01/0.0/0.0/16 Ww..////16019/00000 01/000000005,, "Research Center for Informatics." which enabled to conduct the experiments presented in this work.
We would like express our sincere gratitude Ang Li, Matias Mendieta, and Guangyu Sun for their invaluable feedback and insightful discussions, which significantly contributed to the develop- ment and refinement of this work. Their thoughtful suggestions and careful review of earlier drafts were instrumental in enhancing the quality of this paper.
10

---

References
[1] Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. Fedrolex: Modee--heteoogeneous federated learning with rolling sub-model extraction. Advances Neural Information Processing Systems, 35:29677-29690,
2022. 3,
[2] Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher Manning. large annotated corpus for learning natural language inference, 2015. 9, 12, 20
[3] Sebastian Caldas, Jakub Konečny, Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXivv1812.072200 2018. 3,
[4] Chuan Chen, Tianchi Liao, Xiaojun Deng, Zihou Wu, Sheng Huang, and Zibin Zheng. Advances in robust federated learning: Heterogeneity considerations. arXiv preprini M.i1:22405009838 2024. 2
[5] Yangyi Chen, Lifan Yuan, Ganqu Cui, Zhiyuan Liu, and Heng Ji. A close look into the calibration of pre-trained language models. arXiv preprint 1eXiv:2211.00155: 2022. 9, 2, 13
[6] Yae Iee Cho, Andre Manoel, Gauri Joshi, Robert Sim, and Dimitrios Dimitriadis. Heterogencous ensem- ble knowledge transfer for training large models in federated learning. arXiv preprint Xix::2204.1270:3 2022. 2,
[7] Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 215-223, Fort Lauderdale, FI., USA, 2011. PMI.R. 12,
[8] Yaim Cooper. Global minima of overparameterized neural networks. SIAM Journal on Mathematics of Data Science, 3(2):67--69) 691, 2021. 6
[9] Luke N. Darlow, Elliot J. Crowley, Antreas Antoniou, and Amos J. Storkey. Cinic-10 not imagenet or cifar-10. 2018. 8, 9, 12, 15, 18
[10] Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: large-scale hier- archical image database. In 2009 JEEE Conference on Computer Vision and Puttern Recognition, pages 248-255, 2009, 9, 10, 18
[11] Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. arXiv preprint 2rXiv:2010.01264, 2020. 3,
[12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenhorn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakoh Uszkoreit, and Neil Houlshy. An image worth 16x 16 words: Transformers for image recognition scale, 2021. 8, 10, 19
[13] Cong Fang, Hanze Dong, and Tong Zhang. Mathematical models of overparameterized neural networks. Proceedings of the IEEF, 109(5):683-703 2021.
[14] Alec Go. Richa Bhayani, and Lei Huang. Twitter sentiment classification using distant supervision. CS224N project report, Stanford, 1(12):2009 2009. 9, 12, 20
[15] Chuan Guo, Geoff Pleiss. Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In International conference on machine learning, pages 1321-1330. PMLR, 2017. 9, 2, 13, 16
[16] Farzin Haddadpour and Mehrdad Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint 1eX.iv:9100.444 2019.
[17] Kaiming He, Xiangyu Zhang. Shanqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015. 8, 99, 10
[18] Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Lcontiadis, Stylianos Venieris, and Nicholas Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout. Advances in Neural Information Processing Systems, 34: 14:112766128899 2021. 3,
[19] Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andrectto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications, 2017. 12, 19
[20] Anil Kag, Durmus Alp Emre Acar, Aditya Gangrade, and Venkatesh Saligrama. Scaffolding student to instill knowledge. In The Eleventh International Conference on Learning Representations, 2022. 19
11

---

[21] Sai Praneeth Karimireddy, Satyen Kale. Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In Inter- national conference on machine learning, pages 5132-5143. PMLR, 2020.
[22] Phillip Keung, Yichao Lu, György Szarvas, and Noah A. Smith. The multilingual amazon reviews corpus. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, 2020. 8,
12, 20
[23] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
[24] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009. 8, 9, 10,
12, 15, 18
[25] Ya Le and Xuan S. Yang. Tiny imagenet visual recognition challenge. 2015. 8, 12, 18
[26] Lin Li, Jianping Gou, Baosheng Yu, Lan Du, and Zhang Yiand Dacheng Tao. Foderated distillation: A survey, arXiv preprint 045.02400408544, 2024. 2
[27] Qinbin Li, Bingsheng He, and Dawn Song. codelacontrastive federated learning. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 10713-10722, 2021. 1,
[28] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Feder- ated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-4:: 450, 2020.
[29] Zhiyuan Li, Tianhao Wang, and Sanjecv Arora. What happens after sgd reaches zero loss?-a mathematical framework. International Conference on Learning Representations, 2021.
[30] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363, 2020.
3, 4, 8, 1,2
[31] Chang Liu, Chenfei Lou, Runzhong Wang, Alan Yuhan Xi, Li Shen, and Junchi Yan. Deep neural network fusion via graph matching with applications to model ensemble and federated learning. In International Conference on Machine Learning, pages 13857- 13869. PMLR, 2022.
[32] Tian Yu Liu and Stefano Soatto. Tangent model composition for ensembling and continual fine-tuning. In Proceedings of the IEEE/CVF International Conference on Compuder Vision, pages 18676-18686, 2023.
[33] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication -effccienneffiitnnt learning of deep networks from decentralized data, In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017. 1,8
[34] Sachin Mehta and Mohammad Rastegari. Mobilevit: Light weight, general-purpose, and mobile-friendly vision transformer, 2022. 12, 19
[35] Alessio Mora, Irene Tenison, Paolo Bellavista, and Irina Rish. Knowledge distillation for federated Iearn- ing: a practical guide. arXiv preprint arXiv:2211.0474, 2022. 2
[36] Mahdi Morafah, Weijia Wang, and Bill Lin. A practical recipe for federated learning under statistical heterogeneity experimental design. IEFF. Tronsactions Arrificial Intelligence, 2023. 8, 18
[37] Guillermo Ortiz- Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard. Task arithmetic in the tangent space: Improved editing of pre-trained models. Advances in Neural Information Processing Systems, 36, 2024.
5,
[38] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019. 18
[39] Kilian Pfeiffer, Martin Rapp, Ramin Khalili, and Jörg Henkel. Federated learning for computationally constrained heterogemeous devices: A survey. ACM Computing Surveys, 55(14)):1-27, 2023. 2
[40] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh. Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748 3748-8763. 8763,.MRR,
2021. 15
12

---

[41] Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and II Brendan McMahan. Adaptive lederated optimization. arXiv preprint arXiv:2003.00295, 2020. 16, 17
[42] Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and heyond. In Interna- tional Conference on Learning Representations, 2018.
[43] Felix Sattler, Tim Korjakow, Roman Rischke, and Wojeiech Samek. Fedaux: Leveraging unlabeled aux- iliary data in federated learning. IEEE Transactions on Neural Networks and Learning Systems, 2021. 3,
[44] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recog- nition, 2015. 8, 10, 18
[45] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings the 2013 Conference on Empirical Methods Natural Language Processing, pages 1631- 1642, Seattle, Washington, USA, 2013, Association for Computational Linguistics. 8, 9, 12, 20
[46] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ily Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15 19:9999-1958, 2014. 1
[47] Tulia Turc, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint
2019. 8, 9, 12, 20
[48] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Feder- ated learning with matched averaging. arXiv preprint arXiv:2002.0644 06400, 2020.
[49] Jianyu Wang, Qinghua Liu, Hao Liang. Gauri Joshi, and H Vincent Poor. Tackling the objective incon- sistency prohlem in heterogencous federated optimization. Advances neural information processing systems, 33:7611 13::761--6222 1-6223, 2020.
[50] Yuxia Wang, Daniel Beck, Timothy Baldwin, and Karin Verspoor. Uncertainty estimation and reduction of pre-trained models for text regression. Transactions of the Association for Computational Linguistics,
10:680-66 696, 2022. 9, 2, 13
[51] Ross Wightman, Hugo Touvron, and Herve Jegou. Resnet strikes back: An improved training procedure in timm. In NeurlPS 2021 Workshop on ImageNet: Pust, Present, and Future. 19
[52] Adina Williams. Nikita Nangia, and Samuel R. Bowman. A broad-coverage challenge corpus for sentence understanding through inference, 2018 8, 9. 12,
[53] Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: aver- aging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pages 23965-23998. PMLR, 2022. 3, 6,4
[54] Yeho Wu, Li Li, Chunlin Tian, and Chengzhong Xu. Breaking the memory wall for heterogeneous federated learning with progressive training. arXiv preprint 2rX4v:2404.133499 2024.
[55] Yuxin Xiao, Paul Pu Liang, Umang Bhatt, Willie Neiswanger, Ruslan Salakhutdinov, and Louis-Philippe Morency. Uncertainty quantification with pre-trained language models: A large-scale empirical analysis. arXiv preprint 4.0.v:2210004114, 2022. 9, 2, 13, 16
[56] Ziyue Xu, Mingfeng Xu, Tianchi Liso, Zibin Zheng, and Chuan Chen. Fedbrb: An effective solution to the small-to-large scenario in devicc--hteroogeeity federated learning. arXiv preprint 172222200.. 2024.1
[57] Prateek Yadav, Derek Tam, Leshem Choshen, Colin A Raffel, and Mohit Bansal. Ties-merging. Resolving interference when merging models. Advances in Neural Information Processing Systems, 36, 2024. 3
[58] Tuo Zhang, Lei Gao, Sunwoo Lec, Mi Zhang, and Salman Avestimchr. Timelyff: Heterogencity-wwan asynchronous federated learning with adaptive partial training. In Proceedingg of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pages 5064-5073, 2023.
[59] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classifica- tion. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2015. 12, 20
13

---

[60] Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for text classi- fication. In NIPS, 2015. 12, 20
[61] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classifica- tion, 2016. 12, 20
[62] Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. comprehensive survey transfer learning. Proceedings the IEEE, 109(1)43376, 2020.
14

---

Appendix
The supplementary materials are organized as follows:
Appendix A: Provides more details related works.
Appendix B: Presents the full algorithm description of TAKFL.
Appendix C: Presents formal thcoretical statements, assumptions, and proofs supporting our method.
Appendix D: Presents detailed experimental results including some additional experiments.
Appendix E: Presents the ablation studies experiments. and implementation details.
Appendix F: Presents
A More Detailed Related Works
(a) Partial Model Training
(b) Heterogeneous Device Prototypes
Figure 4: Overview Two Different Device Heterogeneous FL Settings. (a) the partial model training setting, the objective is to train single global model where heterogeneous devices train a specific sub-model based on their computational resources. This approach necessitates device support for varying neural network architectures, which is impractical as devices typically have spe- cialized architectures designed to match their hardware, software configurations, and underlying machine learning tasks. (b) In the heterogencous device prototypes setting, device prototypes par- ticipate in FL to enhance the performance of their global model by transferring knowledge across prototypes. This setting is more feasible as it accommodates diverse device prototypes with their own specific configurations, including neural network architecture and dataset. However, establish- ing effective knowledge transfer between differently sized prototypes (like IoT' and workstations) and diverse configurations is challenging. In this paper, we address this issue.
Prior works on device heterogeneous FL have considered two distinct approaches with different ob- jectives and settings. The first group of studies focuses on accommodating devices with varying compute resources, aiming train a single global model [11, 3, 58, 54, 56]. Various partial model training techniques have been proposed for this setting, where devices are tasked with training sub- model of a global model according to their compute resources. These include dropout-baad based [3], static [11, 18], and rolling-based sub-model extraction techniques [1]. Federated Dropout builds upon the concept of dropout [46] to extract smaller sub-models. Static sub-model extraction tech- niques like in HeteroFL [11] and FjORD [18] consistently extract designated portions of the global model, whereas FedRolex [1] introduces more flexible rolling method for sub-modes extraction. However, these approaches assume that devices can support various sub-model architectures for training. which does not fully reflect the real-world scenario. In practice, there exist a diverse spec- trum of device prototypes such IoT devices and smartphones each have unique and unhashable neural network architectures tailored to their specific hardware and software configurations and un- derlying machine learning tasks. Consequently, these device prototypes may not support training various neural network architectures, highlighting a significant limitation in accommodating the full spectrum of device heterogeneity in this setting.
The second array studies tackles a more practical scenario where device prototypes with heteroge- neous model architectures participate in FL to enhance their global model performance through mu- tual knowledge sharing. In this context, knowledge distillation techniques are employed transfer knowledge among device prototypes [30, 6, 43]. Here, locally updated client models from various

---

device prototypes, collectively referred to as ensembles, serve teachers to distill their knowl- edge into each server's student model using an unlabeled public dataset. For instance, FedDF [30] utilizes vanilla averaging of all ensemble logits as the distillation target for all server student mod- els. In contrast, FedET [6] employs an ancertainty-weigheee distillation target for all server student models, complemented by a diversity regularization tech- nique. However, methods like FedET rely on the neural networks' confidence scores for uncertainty estimates, overlooking the fact that neural networks are often poorly calibrated and prone over- confidence, which compromises their ability to provide reliable uncertainty estimates [50, 15, 5, 55]. These existing works typically focus settings where device prototypes have similar capabilities, i.e. similar model and dataset sizes, thus neglecting the challenges presented in more diverse set- tings where device prototypes vary significantly in terms of model and dataset size. This oversight limits the effectiveness of these methods in truly diverse and heterogeneous environments. In this pa- per, we introduce TAKFL, which designed to address the limitations of existing methods in these underexplored diverse device heterogeneous settings.
average ensembles' logits as the
Figure illustrates the distinctions between these two different settings studied in the literature. For more information, we refer the reader to recent surveys [35, 26, 39, 4J.

---

B Full Algorithm Description of TAKEL
The full algorithm description of TAKFL is presented in Algorithm
Algorithm TAKFL Algorithm
Require:
number of communication rounds (R), public unlabeled dataset IDpublic server training
iterations I, heterogeneous device prototypes (i E MD) with their associated clients (C') and local datasets ((D) kec:). ). model architecture (f*), local training iterations (Looctt), local learning rate (Mocat)) server distillation iterations (Laistit)), and server distillation learning rate (Atssta))
1: Server Executes:
2: Randomly initialize all device prototype's 's server model {0)) ieM
3: for each round 0, ...... R do
4: C₁ + (randomll select clients from each device prototype) Vi E M
5: for each client C, Vi E M in parallel do
6: 0 + ClientUpdate(k::e 07)
7:
end for
8: O ExEc; Exeej [E 0%
9:
for each device prototype's server student 1,,, .., parallel do for each device prototype's teacher ensembles j 1, M in parallel do
10:
+- Davg
12:
for each server distillation iteration -= 0, ..... Taista do x f- sample mini-batch of data from public dataset Dipublic
14:
0'+1 - 0F Adistill V defined in Eq. 6.
15:
end for
16:
f Ofaear 0 aog
17:
end for
18:
0 + 0 ang + + A3TJ 19: end for
20: 0: f- 0'
21:
end for
22:
function ClientUpdate(N; 0;)
23:
0 <- 0%
24:
for each local update iteration = 0, 1.... Locel do
25:
{x, ) +- sample mini-batch of data from local dataset L3; gr+ + 0 Mocat VE(f'(x 0)),))
26:
27:
end for
28:
0f <- Odiacal
29:
end function
3

---

C Theoretical Results
C.1 Proofs of the Main Propositions
First we present the formal assumptions associated with theoretical derivations.
Assumption 1. Local federated averaging is performed with perfect test accuracy, i.e.,
N
argmin 03
ΣΣ
(2) argmin B(a,y))D (x; B (
Øx
That is, the training error on the datasets {D for the computed Oaey is the same as the test error on
the population distribution Dj. Moreover assume that we can write Ti
Σ s' ((,0))kE
(:00g)))) ). Finally, we assume that the same population distribution E, wi Du is the same that the clients perform their testing on as the server performs distillation on.
These assumptions are made for mathematical practicality while at the same time not starkly unrea- sonable. The local FL the device prototypes perform is generically prone imprecision, especially as the clients' data varies, but this discrepancy is bounded [16]. Similarly the difference in the aver- age of logits and the logit of averages has bounded difference norm [53]. Thus, violations of the Assumption add additional perturbations to quantities derived in the Theoretical analysis without having structural/qualitttvee effects, and thus would only present clutter in the presentation.
Notations. Now we present the notation defining the specific quantities we refer in the derivations below. The set of important quantities is given in Table 3. Note that the formal definitions of the first two quantities are,
+ argmin E B (x:
@.. argmin E [e(p' (z: 0).
0
Table 3: Notation and Definitions
Notation
Definition
Parameters in j's device model that minimize the loss on its population distribution
2
Parameters in j's device model that minimize the loss on I'th population distribution
Q' dim (*)
The total capacity of device prototype j
Q! (e/
Figenbasis for the model device prototype j
WJ dim e7)
Dimension of the solution submamifold e
WT* dim
Dimension of the solution submanifold ()
WI (e()x=1..... wi
Eigenbasis the solution submanifold @
Wilh (e)" was. WSA
Eigenbasis the solution submanifold @
We shall make use the "Choose" combinatorial operator, defined to be Ch(n,p) = pl(n-p) The standard (-)) notation indicates dą: = O(bs) to mean there exists K and C' such that for k: > K,
0x I Cbb.
A recent finding that inspired the methodology in this work is the discovery of the weight disentan- glement phenomenon underlying task arithmetic [37]. Indeed the task arithmetic property provides the ideal circumstance for federated knowledge transfer as we shall see below. Formally, adapting their definition our notation:
(Task Arithmetic Property) holds for a set of vectors {T, if for all j it holds that,
Σ AiTs {{
fx 0 + AsTi) (x;; Pggg)
DE D i E DI \ Di
fi
x;
0 ang
+
(10)
Let us define important property of task arithmetic that shall use in the sequel.

---

(Weight disentanggleentt A parametric function f X x + y is weight disentangled with respect set of task vectors T (TjbjeT and the corresponding supports DT : {Dj)eT
[(; 0o
Σ (.TT))
Σ (5( go());
iET
'ET
where f;(2; 0;7:: O for & D; and i E T, and (o()) = 0 for a E Uier D;.
We now present the formal statements as well the proofs of the main propositions.
Proposition 1. (Information Loss in VED). Let the prototype i and datasets be disjoint, i.e. D' DDD = -., and U D D. Note that this implies the weight disentanglement property. Con- sider the VED procedure in the form of solving (3). Consider two device prototypes with a device ca- pacity and solution dimension of Q² Q² and W W2, respectively, and with associated eigenbases Q¹, W Let the solution set of VED with prototype i student be θ% EVED with dim (an(@vyg)) with eigenbasis W In addition, denote W s, E {1,2} the dimension of the solution set for the student model trained the data from the teacher device's ensembles. We assume that self- distillation is executed appropriately, e.g., W W and W = -.2
w
1. Case 1: Assume that Q² = =Q² and W = W W12 W Then it holds that, in expectation,
dim (ÓL n (2') w')) (e'")n Q'I(WIII(Q' (Q"[)WV(
(Q' - WI + Q'IWL2I(Q' W12y1
W") w7:
This corresponds to the expected capacity of prototype that is tuken up for fitting logits that are not in the span of W', that is, that do not fit the data corresponding prototype 1. Case Assume that Q² > Q² and W W W², Then same quantity as for Case holds. Moreover,
(Q' W')(W')) Way! (W' NWW' ( -0 (Q"W"I(Q')) WI + Q' QWW?!WW1222 w2y!
dim (@VEDn
This corresponds capacity of client 1 that has been allocated but fits, in the model of prototype neither the data of prototype nor of the data of prototype 2.
Proof. Formally.
Σ (F((e,)0)))) (S(x)) 1=1.2
ÓVED argmin LED argmii KL
DEQ1
Since by assumption Ddeg solves the training problem the data associated with device prototype i, the logit is accurate, and thus there is map Oii,j) T, + T C Wii The self distillation, that is, Sj defines a bijective map from W to W and thus does not affect the capacity allocation.
Case 1: In this case, generically (that is, measure zero on some non-atomic sampling on dis- tribution of operators) (i.j) is bijective. Now us compute the expectation of the number of eigenvectors of, e.g. 1/1.2 that are in the complement of the span of W Assuming, for simplic- ity, independence, this would correspond to counting the possible choices within the capacity of Q \ W over the range of possible choices of filling the capacity of Q with vectors in W together with choices of filling it with vectors in W1
Qt-W"
Ch(Q' .)) ChQQ WT) + C,(Q',
Σ
For, e.g., Q² = and W = 112 2 this is J
To derive a scaling rate we can write:
1Q'-w""
Σ
11 - +31
O
(Q' (o BP
G11
Case 2: this case, it must that, at best almost surely, O 1) is injective, but not surjective. This means that distilling from 2 to does not fill the capacity of 121.2 and is thus fundamentally
5

---

wasteful operation, that is 17.1 W < W1 Now let us compute the expectation of the number of eigenvectors of, e.g. W that are in the complement of the span of W Since W1, are being structurally allocated for fitting, the combinatorial expression is the same:
Q Q'-w" W
Ch(Q' - 1) Ch(Q', W + Ch(Q', W1.)
Thus for, e.g., Q' and W W this is, again, 3" The scaling in this case is
(Q' - W') - W. QU!T!(Q' WiJ!+ Q' W
However, we observe that there are vectors the range of W \ ( (WW2) that have been allocated by the VED but lie in neither W nor in W that is, are garbage. We can compute those as the expected number eigenvectors arising from allocating w1l (2, ()W²) that intersect
with Q W (that is, the spare capacity not used for fitting data D This is, using similar principles:
_1 W
Ch(Q' -
Σ
Ch(Q', W1) + Ch(Q2, W1. W2)
This is for, e.g. Q' = 4, W = 2 and W = 1, this would be a To
The scaling here is
(Q' - W'') W QUITTIQ' "')1++"'"""WW w
Proposition 2. (Improve knowledge transfer with task arithmetic).. Consider the TAKFL procedure as in the form of computing (8). Consider two device prototypes with a device capacity and solution dimension of Q² Q2 and W1 W2, respectively, and with associated eigenbases Q W Let the solution set of TAKFL with prototype as student be θ' with dim
W with eigenbasis W In addition, denote Wa,,, s,t E {1,2} dimension of the solution set for the student model trained on the data from the teacher device's 's ensembles. The following statements hold:
In the case that that Q² > Q² and W > W it holds that the TAKFL preserves that the eigenbasis used to model the data D accuracy for device prototype I, that is for student
dim (W" n [o' ( =0
Case I: Assume that Q' Q and W = W Then it holds that,
dim (W" n le' (w" : -0
Moreover, holds that,
θ E Span(W' O w'
Thus, with equal capacity, no information lost in Task Arithmetiic aided knowledge ensemble distillation and capacity is efficiently used to model the data from both prototype I and prototype 2. Case 2: Assume that Q" > Q² and W > W Then it again holds that,
dim (W" n [et (W UW++21) =0
However. while θTA SSaan(W),, it holds that dim (W" n m/ W < W/12
Proof. We consider the case that the prototype and datasets be disjoint, i.c. DÍ n Di = and UM D = D. We shall see how the other cases are easier and the result remain the same.
We can immediately from the weight disentanglement property of Task Arithmetic that,
f (x:; 0 avg + a1T1 a272) 1h (m; + 1:2 02, T2) 1

---

with (2; a17)) for æ & D', g/= a272) for æ & D2 and g g')((z) = 0 for x E D UD? From this, we can immediately conclude the first statement of the Proposition as well the expression
dim (W" n [ (w u2 =
and also, in the case of W = W/2= = W implies
θ E Span 2
For the last statement we observe again in the second Case in the Proposition describing VED, dim (O(2, ))) w1.2 from which we can conclude that, generically
dim (fv : a 1)(22) C W", B(a.g))Dzl(( (a: v). > o}} = -
proving the final statement.
We observe that a key mechanism of the proof is the dimension of the target space the teaching operator (i,j). As an informative model, we can consider coefficients X5 of task vectors as re- stricting the rank, relative to other teachers. For instance, in the previous Proposition, if W 2 and 2 1, then x2 1/2, so as to enforce one vector of W1.2 is a target for the map ( 1), would be appropriately sensible.
The prototype i and j datasets are fully overlapping, i.. D' C Dj, meaning that the two prototypes have the same classes De, Ve E D, then these the changes to the the proposition I and 2 in theory.
Consider that in this case, the logits corresponding to Will and W1, are the same. As such W1l2 E Q¹, and so the propositions change to, trivially:
Proposition 3. (Information Loss in VED). Consider the VED procedure. Consider device prototypes with a device capacity and solution dimension of Q' Q² and W W respectively, and with associated eigenbases Q', W Let the solution set of VED with protorype i as student be W with eigenbasis W In addition, denote W 8, E {1,2)
OVED with dim(O' the dimension of the solution set for the student model trained on the data from the teacher device's ensembles. We assume that self-distillation is executed appropriately, e.g., W1.1 = W and W2.2 2
1. Case I: Assume that Q+ = =Q² and W W = W112 = W2.1 Then it holds that, in expectation,
dim (oLED [et \ (w1) =0
This corresponds the expected capacity of prototype that is taken up for fitting logits that do not fit the data corresponding to prototype 1.
Case 2: Assume that Q- > Q² and W = W1.2 W2, Then the same quantity as for Case holds. Moreover,
dim (θVEDn \ \ww" \WW' uwlz))
This corresponds capacity of client 1 that has been allocated but fits, in the model of prototype I, neither the data of prototype nor of the data of prototype
Now consider that the prototype and j datasets are partially overlapping, i.e. D'n + 0. meaning that the prototypes have the overlap over some of the classes D., then these are the changes to the the proposition and 2 in our theory.
In this case, W112 w/ U W/! that is, the class labels across data prototypes partially overlap (with W/1.2 the overlapping dimensionality), the Proposition for VED changes to:
Proposition (Information Loss in VED). Consider the VED procedure. Consider two device prototypes with a device capacity and solution dimension of Q' Q2 and W W respectively, and with associated eigenbases Q' W.. Let the solution set of VED with prototype i as student be OVED with dim(@vgp)) = W with eigenbasis W7: In addition, denote Wa.t 8, E (1E(1,2) the dimension of the solution set for the student model trained the data from the teacher device's ensembles. assume that self-distillation is executed appropriately, e.g., W1.1 W and W2.2 W2
7

---

0
ay
VED
81.2 any
62
avy
XED
2
02og
11,1
22,1
1.3
Then it holds that, in Arrri, 2122
.1.2
1,2
0
Day
6
Figure 5: Illustration of Geometric Intuition. Each panel presents a different case example. The left and center panels present the geometric intuition of KD for vanilla ensemble distillation (VED) and TAKFL in the case where two different large device prototypes performing knowledge transfer. The planes represent the solution subspaces. The right panel presents circumstance by which two small device prototypes (2, and 3) serve as teacher for transferring knowledge to a larger device prototype 1.
1. Case 1: Assume that Q' Q² and W W22 W1 W expectation,
dimm (OVEDn (w')) o0 (Q"(W"))) W + Q"I W
(Q' - w') -2
This corresponds to the expected capacity of prototype that is taken up for fitting logits that not fit the data corresponding to prototype 1.
Case 2: Assume that Q² > Q² and W W1. Wa, Then the same quantity as for Case
holds. Moreover,
@Q' - WW) Wºy!
dim (ōvEDn \W NoW (w ( WiY! + Q-IW?((W) w2y!
This corresponds to capacity of client that has been allocated but fits, the model of prototype neither the data of prototype I, nor of the data of prototype 2.
Note that the presence of overlapping logits w:1.2 do not appear in the result, it is exactly the same as the original Proposition 1 in the original paper just with W.12 in place of WL.2
The Propositions for TAKFL remain the same.
C.2 Geometric Intuition
In this section aim to provide geometric intuition for the mechanism of VED and Task Arithmetic KD on three different cases. Figure 5 presents the geometry illustration for three different cases. We discuss each case in the following.
Case I: KD between two large prototypes with different data distributions. Consider Figure left panel. This panel corresponds to setting where two large device prototypes with similar total capacity, i.c. Q' Q2 both prototypes to be the same, iee. W 2 2. These would correspond to planes in the ambient space. Therefore, one plane corresponds to the solution subspace of the prototype 1 trained on its own data, i.e. @ subspace in the panel, and the other corresponds to the (theoretical) solution subspace of this prototype trained on prototype 2's data, i.e. (1,2) in the panel. In this case, since the data distributions the prototypes are fairly disparate, this has resulted into near orthogonal subspaces corresponding to these solutions. As which far away from either of the planes corresponding to optimal solution subspaces, and far from the optimal set of parameters, which is their intersection, suggesting loss of knowledge. By contrast, the TAKFL approach, by customizing the merging coefficients and putting each to half, i.e. λ1,1 0.5, can traverse in the tangent space of zero loss surface and get into the intersection subspace which is exactly the optimal solution (0" merged 0b.y + A1,1TI,1 A22
perform knowledge transfer. We consider the solution dimensions of
see from the panel, VED will lead to point

---

Case II: KD between two large prototypes with similar data distributions. Consider Figure 5 center panel. Similar to Case I, this panel corresponds to a setting where two large device prototypes with similar total capacity, i.e. Q' Q² 3 performing knowledge transfer. We consider the solution dimensions of both prototypes to the same, i.e. W W 2. These would correspond to planes in the ambient space. Therefore, one plane corresponds to the solution subspace of the prototype trained on its own data, i.c. @ subspace in the panel, and the other corresponds to the (theoretical) solution subspace of this prototype trained on prototype 2's data, i.e. @ in the panel. In this case, since the data distributions of the prototypes are fairly close. this has resulted into non- orthogonal solution subspaces. As we can see from the panel, while VED could still lead to some information loss, by and large we expect straightforward KD. Our task arithmetic (TA) approach, again by customizing the merging coefficients 11,2 0.5, can traverse in the tangent space of zero loss surface on each plane and get into the intersection subspace, corresponding to the most efficient allocation of device prototype capacity for fitting simultancously the logits corresponding to accurate modeling of device prototype 1 as well as device prototype 2's data distribution.
Case III: KD between two small prototypes and one large prototype. Now consider the right panel in Figure 5. This panel corresponds to a setting where two small prototypes serve as teachers and one large prototype is the student. The flng plane corresponds to the solution subspace of the large prototype on its own data, D The line 01.3 line corresponds to the subspace of solutions in prototype I's parameter space projected into the capacity of the information transferred from device prototype 2. Finally, the line labelled 0 1..3 corresponds to the subspace of solutions in prototype I's parameter space projected into the capacity of the information transferred from device prototype 3. Here, we can see from the relative angle of the lines with respect to the plane that the distribution D1 is closer to the distribution D2 than D3 Comparing this case to the previous cases, 01.9 like case I and 0 0 is like case We can apply the same conclusions here as well regarding the performance of vanilla ensemble distillation and our adaptive task arithmetic approach. We can see from the geometric visualization that knowledge distillation towards 11.3 has more margin of error for prototype 1. Therefore, with the TAKFL approach large prototype 1 can strategically select which prototype to learn more from, and since 022 has closer data distribution to prototype 1, TAK will prioritizes this by putting a larger merging coefficient, i.e 21,2 > \1,3. By contrast, VED lacks this customization and results sub-optimal knowledge distillation.
aey
ang
aig
avy
The geometric intuition discussed here is consistent with our detailed experimental analysis in D.I.I.
C.3 Analytical Properties of Learning Dynamies
Here we provide additional insights from the literature to the nature and properties of learning as it takes place on overparametrized models. Specifically, we comment on literature in the area of Stochastic Differential Equation (SDE) models for SGD training dynamics, and its correspon- dence to the results above. Overparametrrzation has been conjectured to be significant factor in contributing to the (unexpected, by classical variinni tradeoffs) generalization ability of deep neural networks, from number of perspectives [13].
Consider the diffusion model of SGD training for overparametrized NNs provided in [29]. Their analysis relies on the following two assumptions. For our purposes L is shorthand for a client group's 's loss, L(0) Σ EE ((aw)) (((fs((; 0), w)] for some j. which will be identified from the
kEC:
context.
Assumption 2. L: R.R - is C and the solution set I is a -idimensional C2 submanniiol of D for 0 < W < D and rank (Vnk(V = Q - WW
Assumption 3. Assume that U is an open neighborhood of I satisfying that gradient flow starting in U converges to some point in l
From these, [29] derives Theorem 4.6. This Theorem decomposes the random process of the pa- rameter weights driven by SGD after it has reached the solution manifold, e.g., the diffusive random walk of oang in Figure 5 along its respective solution manifold.
Theorem 1. Given L., I and (0) O(0) E U as by Assumptions 2 and 3 the SDE modeling the optimization of by SGD, that is, defining ((X) to be the gradient flow applied to the state random variable X, then for T as long as P\Y (1) E U, v0 S t TT = 1, (LT/1? converges in
9

---

distribution to the stochastic process (T) as 17 - O, with (t) given as.
(t) = =1 - -1
-V) (20 ( (  (, (F2,, di (II)
-
where E = ToT and En,, pp. Σ are given as,
Σ 0 Σ = (In 04 a - -Qb) - = d - aΦ)
(12)
This theorem indicates that the asymptotic flow of SGD on the client training can be decomposed into a covariance-driven random walk in the tangent space, drift to preserve the flow into the tangent plane, the tangent-normal portion of the noise covariance and noise in the normal direction.
This analytical expression provides the probabilistic foundations for the more higher level theoretical results above. In particular, local gradient dynamics, employed by individual device prototypes j using FedAvg on its local clients, yields a flow for the stochastic process defined by the weights. At this point of learning, the weights traversing the solution set, with noise predominantly in the tangent directions. Thus knowledge distillation which preserves this noise structure going to be more effective as far as preserving accuracy across data.
D Detailed Experimental Results
In this section we present a more detailed version of experimental results presented in the main paper Section 7, Additional experimental results are also presented here.
D.1 Main Experimental Results on Computer Vision (CV) Task
The experiments in this section complements the main experimental results in the main paper
Section 7.2.
Experimental Setup. For the evaluation CV task, employ CIFAR-10 and CIFAR-100 [24] datasets. For CI, we use CIFAR-100 as the unlabeled public dataset, while ImageNet- 100, subset of ImageNet [10] with 100 classes (see Appendix F.1.1), is used for CIFAR-100. We distribute the training dataset among the device prototypes in a ratio of 1:3:6 for S. M. and L, respectively. Each device prototype's s data portion is further distributed among its clients using Dirichlet distribution. We apply two levels of data heterogeneity for a comprehensive evaluation: low heterogeneity, i.c. Dir(O.3), and high heterogeneity, i.c. Dir(O.1). Additionally, we configure the number of clients and their sampling rates as follows: 100 clients for S, 20 for M, and 4 for L, with sampling rates set at 0.1, 0.2, and 0.5 respectively. To comprehensively evaluate, we use two distinct architectural settings: the "homo-family" setting, where all device prototypes' architectures are from the same family empllying ResNet8, ResNet14, and ResNet18 [17] for S, M. and L. respectively; and the famitt setting, which diverse architectures are used -iV-- [12] for S, ResNet14 for M. and VGG-16 [44] for L. All models are initialized from scratch, and the communication round is set at 60 rounds. Further details regarding hyper-parameters can be found in Table 13.
Overview of Performance Results. Table 4 presents the performance of TAKFL across diverse architecture settings on the CIFAR-10 and CIFAR-100 datasets. TAKFL consistently improves all device prototypes of different sizes in various cases by significant margin compared the base- lines, achieving SOTA performance. Notably, in the homo-family architecture setting with Dir(0.3) on CIFAR-I0, TAKFL improves average performance across all prototypes by 8%, and by 4% on CIFAR-100. In the vetero-family settings with 1ir(0.1) on CIFAR- 10 and Dir(0.3) on CIFAR- 100, TAKFL enhances performance by ~3% and 1%, respectively. Furthermore, we observe that our self-regularization technique has successfully mitigated issues associated with the noisy and unsu- pervised ensemble distillation process, thereby enhancing performance. Generally, the performance gains from self-regularizzation are more pronounced in low data heterogeneity cases, where proto- types' models perform better and possess higher quality self-knowledge. Thus, self-regularizatioon proves more effective it preserves this higher quality self-knowledge.
10

---

Table 4: Performance Results for CV task CIFAR-10 and CIFAR-I00. Training data is distributed among S. M, and L device prototypes in a 1:3:6 ratio, subdivided among clients using Dirichlet distribution. Public datasets are CIFAR-100 for CIFAR- 10 and ImageNet- 100 for CIFAR- 100. Client configurations include 100, 20, and 4 clients for S. M, and L, with sampling rates of 0.1, 0.2, and 0.5. In homo-family settings, architectures are ResNet8, ResNet14, and ResNet1 8; in hetero-family settings, they are ViT-S, ResNet14, and VGG-16. All models are trained from scratch for 60 rounds. See Appendix F.I for more details.
Homo family Architecture Setting
Daraset
Baseline
Low Data Heterogeneity
High Heterogencity
S
M
Average
M
I.
Average
FedAvg
36.21=2.24
46.41
69.4625.17
47.36
22.01
26. 20..00
51.51es52
32.93
FodDF
19.31.u.35
50. 50.631a73
19.82
19.92
31.71
35. 35.27 14.74
51 51.081400
10.35
CIFAR-10
FedET
49.21
60.01+1111
53 23.00
62.61
29.58
30.96
46. 16.5326
35.36
TAKEL
65.90
57 13.93
600,58
58.11
37. 12.00
38 .0.3
51. 11.49,6.15
12.61
TAKEL+Reg
56.3 -9.47
28.60+o..4 45
1.28
10.51-11.0
10.12:1.24
53. 2
FedAvg
13.22-0.14
21 11.3911.11
29.47
21.36
11 11880
14.6320.64
26. 20.25=1.0
17.58
FedDE
19.54
24.92a0045
29.29
24.38
16.09
19.80=m.17
26.59
20.80
CIFAR-100
FedET
19.67 20.35
26.27
31. 10
25.35
11.18+1.0s
18. 18.2220.55
26. 26.40g0
18.60
TAKFL
24.48.u.4z
27.80
29.84
27.31
22. 02.90..0.1
23. 23.63-0.72
26.98
24.50
TAKEL.RReg
27. 18
29.1 29.144000 20
16.0.
29.
22.8 22.88
92.7
28.01 +0.54
24 94
Hetero-family Architecture Setting
Daraset
Baseline
Low Data Heterogeneity
High Data Heterogeneity
Average
L
Average
FodAve
27.53-0.5
47.30
55.10+a.00
13.31
20.933
25. 62.0.64
36.80
27.78
FeaDE
34.15
54.062100 06
69.07
24.20 74
ga.ce
39.81g5-4 45
32.69
CIFAR-10
FedET
33.24=1.27
58 86.004444
65.5623.40
52.55
24.37
37.77e4.71
43. 25.64
35.26
TAKFL
33.29-..16
57.8420
68.44.0.0
53.12
24.92+1.32
38.07
48.0 48.01.5
37.00
Reg
33.34=3.80
69.01+3.12
4::4
54.19
10
$1.5.
50.26
38 06
FedAvg
8.51guar
22.11aa.s
37.91
22.81
7.0120.47
14.94.u.56
28.51g1.45
16.82
FedDF
10.46
23. 23
36.81
23.58
7.76
18.92 .0.39
29.8 29.81g1
18.83
CIFAR-100
FedET
11.16.u.as
25. 10.30
37.38.0.0
21.65
8. 8,20
20 00.0
28. 28.95
19.27
TAKEL
10.29_n.1
27.14
39. 39.15_nnas
7.88_2..0
21.40
211.20)
TAKFL+Rog
11.25
27.86
38.68
25.93
8. 15
22. 16 p..7
31.95
20.85
D.1.1 Consistency Analysis: Experimental and Theoretical Correlations
In this part, claborate our key experimental observations and their alignment with our theo- retical findings.
Insight 1: From Table 4, it is evident that prior KD-based methods show inconsistent performance across various device prototypes, particularly for the large (L) prototype. For instance, the CIFAR- 10 homo-family setting with Dir(O.3), while small (S) and medium (M) prototypes see performance gains, the prototype experiences up to 10% performance decline compared to vanilla FedAvg. which lacks server-side knowledge distillation. This trend is consistent across other settings, such as CIFAR- 10 Dir(0.1) 1) homo-family and CIFAR- 100 Dir(0.3) homo-family. These outcomes underline the dilution problem inherent in existing methods, where the valuable insights from larger, more pable device prototypes are overshadowed by less informative outputs from smaller devices, thereby degrading the performance of L prototypes. These empirical findings supported by theoret- ical insights as discussed in Remark 1. Specifically, Proposition illustrates that vanilla ensemble distillation (VED) leads to knowledge dilution and inaceuracies due to misaligned device capacity allocations. Moreover, this issue becomes more significant when the smaller device prototype serve teacher.
Insight 2: From Table 4, the suboptimality of existing KD-based methods is evident from the sig- nificant performance improvements of our method, especially for S and M prototypes across various settings. This underscores the ineffectiveness of the mne-size-fits-ale approach these methods em- ploy, where a single averaged logits distillation target is used for all device sizes, proving to be sub- optimal. Our experimental observations regarding the shortcomings of vanilla ensemble distillation methods align with our theoretical findings, as substantiated in Remark and 2. It becomes evident that an efficient knowledge distillation process must allocate capacity in a manner that appropriately corresponds to the information value of the teacher ensemble prototypes.
Our experiments, detailed Table 4, demonstrate TAKFL's adept handling of knowl- edge from various device prototypes under different data heterogeneity conditions. We observed consistent performance gains for small (S) and medium (M) prototypes across both low and high data heterogeneity, compared to vanilla FedAvg. However, in high heterogeneity settings, large (I.)
Insight 3:
11

---

prototypes show less improvement, prompting the question: What can smaller device protorypes offer larger ones?
In low heterogeneity scenarios, large prototypes significantly benefit from the collective knowledge, showing enhanced performance. Conversely, in conditions of extreme heterogeneity, where smaller models contribute less effectively, the performance improvements for larger devices are notably re- duced. This pattern highlights TAKFL's ability to intelligently manage and utilize the available knowledge, selectively distilling information based on the intrinsic capacity and contributions of each prototype, and integrating only the most valuable knowledge from smaller devices when bene- ficial.
By contrast, our analysis of existing KD-based methods shows their failure to effectively discern and utilize the most informative knowledge across prototypes. These methods often overload the capac- ity of larger prototypes with suboptimal or irrelevant information, particularly in high heterogeneity environments, leading to not just stagnation but an accumulation of inefficiencies. These experi- mental observations align with our theoretical insights, outlined in Remark 2, which emphasizes the crucial combinatorial constraint of capacity and diverse information. This further confirms the superiority of TAKFL's adaptive approach knowledge distillation in diverse federated learning environments.
D.2 Additional Experimental Results on CV Task
Experimental Setup. For additional evaluation of the CV task, we conducted experiments on Tiny- ImageNet [25], STL-10 [7], and CINIC-10 [9] datasets using pre-trained models. For TinyIma- geNet [25]. we utilized STL-10 10 [7] the unlabeled public dataset. STL-10 [7] and CINIC-10 [9] both employ CIFAR- 100 [24] as their respective public datasets. We distributed the training datasets among the device prototypes in 2:3:5 ratio for small (S), medium (M), and large (L) prototypes, respectively. The data portion for each prototype was further subdivided among its clients using Dirichlet distribution: Dir(1.0) for TinylmageNet and Dir(0.3) for both STL-10 and CIN Client configurations were set with 4, 3, and 2 clients for S, M, and L, respectively, all with a sampling rate of 1.0. The architectures employed were MobileNetV ViTV2 [34] for M, and ResNet-34 for L, sourced from the TIMM library.' The local training was conducted over 10 epochs using Adam [23] optimizer with a learning rate of le-3 and weight de- cay of le-5. For server-side distillation, the epoch count was 10 for TinylmageNet and for STL-110 and CINIC-10, with a batch size of 128, employing an Adam optimizer with learning rate of le-5 and weight decay of le-5. For TinyImageNet, public dataset images from STL-10 10 were resized to 64, while for STL-10, images were resized to 32x32. No data augmentation was used. The communication rounds is fixed to 40. These experiments were conducted by only I trial. Table 15 details the configurations.
[19] for Mobile-
Table 5 presents the results. The superiority of TAKKL's performance across these challenging datasets using pre-trained models is evident here well.
Performance Results.
D.3 Additional Experimental Results on Natural Language Processing (NLP) Task
The experiments in this section complements the main experimental results in the main paper Section 7.2.
Experimental Setup. For the evaluation of NLP tasks, utilize four datasets: MNLI [52], SST- 2 [45]. MARC [22], and AG-news [60]. The corresponding unlabeled public datasets are SNLI [2] for MNLI, Sentiment [14] for SST-2, Yelp [61] for Amazon, and DBPedia [59] for AG-News. The training data is distributed among the device prototypes in ratio of 1:3:6 for small (S), medium (M), and large (L) categories, respectively, with each portion further subdivided among its clients using a Dirichlet distribution (Dir(0.5)). The client configurations and their sampling rates are set as follows: 8. 4, and 2 clients for S. M, and L categories, respectively, with sampling rates of 0.3, 0.5, and 1.0. The architectures employed for each prototype size are BERT [47] -Tiny, -Small, and -Mini, respectively, each initialized from pre-trained parameters and tested over 20 communication rounds. Additional details regarding hyper-parameters and datasets are presented in Appendix F.2 and Table 17.
12

---

Table 5: Performance Results for CV task TinyImageNet, STL-110, and CINIC-10 using pre-trained models.
Private
Public
Baseline
M
Average
FedAvg
8.97
13.03
15.12
12.37
FedMH
15.08
17.10
17.83
16.67
TinyImageNet STL.-10
FedET
10.60
16.39
17.62
14.87
TAKFL
16.10
17.60
19.03
17.58
TAKFL+Reg
16.55
17.98
19.74
18.09
FedAvg
26.01
34.47
42.88
34.45
FedMH
28.64
34.55
39.25
34.15
STL-10
CIFAR-100
FodET
29.87
33.00
38.26
33.71
TAKFL
29.57
37.57
42.53
36.56
TAKFL+Reg
30.78
37.89
43.38
37.35
FedAvg
44.87
55.49
51.33
50.56
FedMH
45.52
55.75
53.48
51.58
CINIC-10
CIFAR-100
FedET
37.43
52.25
TAKEL
48.21
57.81
52.74
52.92
TAKFL+Reg
47.66
57.54
53.25
52.82
Table 6: Performance Results for NLP Task on 4 Datasets.. Training data is distributed among S, M, and L device prototypes configurations are 8, 4, and 2 clients for S, M. and L, with sample rates of 0.3, 0.5, and 1.0. respec- tively. Architectures include Bert-Ting Teny, Bert-Mini, and Bert-Small for S. M. and L, initialized from pre-trained parameters and fine-tuned for 20 communication rounds. See Appendix F.2 for more details.
1:3:6 ratio, subdivided among clients using Dir(O.5). Client
Private
Public
Baseline
S
M
I.
Average
FedAvg
36.15.0..0 46
14.4742..4
57.51_2.79
49.37
FedDF
54.21
60.4441.91
66.71_1.00
60.45
MNLI
SNLI
FedET
48. 48.0346.32
50.33
18.80_6.18
50.72
TAKFL
57.43
63.58
68.74
63.25
TAKFL+Reg
57.61.0..6
63.91
68.96
63.49
FedAvg
54.98
74.71 2:.22
86.69 86.69_0.06
72.13
FedDF
74.41+2.62
80.71
84.35_11.66
79.82
SST2
Scnt140
FedET
66.6349.14
65.89+16.33
70.05_15.88
67.52
TAKFL
74.73
82.17
86.93_0.4
81.28
TAKFL+Reg
74.88
82.40
87.331..0
81.54
FedAvg
33.7641.18
49.08
59.26_1.44
47.36
FedDF
53.01+1.24
5..87
36.81=0.99
55.06
MARC
Yelp
FedET
52.63
541.28+2.31
56. 56.11_2.44 2261
54.34
TAKFL
55.70
58, 6441..7
59.39_1.m
57.91
TAKFL+Reg
55.96
59. 59.18+1.00
59 61
58 25
FedAvg
83.6449.51
83.47 4:2.55
91.48_2.7 2.22
86.20
FedDF
85.97
89.10+1.85
91.37=1.10
88.81
AG-News
DBPedia
FedET
75.27
81.13
83. 83.19_4.55
79.86
TAKFL
87.37 +1.31
90.11
92.48 92.48_1.12
1.12
89.99
TAKFL+Reg
87.66+1.88
90.3042.05
92 6..
90.19
Performance NLP Task. Table 6 presents the results four different datasets: MNLI, SST-2, MARC, and AG-News. Similar to the CV task, TAKFL. has consistently improved performance across all device prototypes of varying sizes, achieving state-of-the-ard results. On MNLL has en- hanced average performance across all prototypes by 3%, on SST-2 by ~2%, on MARC by 3%, and on AG-News by ~ 1.50%. As observed in the CV task, the suboptimality existing KD-based meth- ods is also evident here. Notably, FedET exhibits very poor performance compared vanilla FedAvg, failing to achieve satisfactory results on all datasets except for the MARC dataset. Particularly, the performance of the L prototype has consistently decreased across all datasets compared to vanilla FedAvg- This behavior can attributed FedET's reliance on neural network confidence scores for uncertainty estimates in its uncertainty-weighted distillation. However, neural networks, espe- cially pretrained language models (PLMs), are often poorly calibrated and prone to overconfidence, which compromises their ability to provide reliable uncertainty estimates [50, 15, 5, 55].
13

---

D.4 Scalability Evaluation
This section complements the experimental results in the main paper Section 7.3.
Experimental Setup. To evaluate the effectiveness and scalability of our method across a broad spectrum of device prototypes, ranging from very small very large sizes, we conduct experi- ments involving 3 to 7 different prototypes. Our objective is to assess how effectively our method adapts from a uniform array of small-size prototypes (3 device prototypes) to diverse mix that includes prototypes ranging from extremely small (XXS) to extremely large (XXL) (7 device proto- types). These experiments involve training image classification models from scratch on the CINIC- 10 dataset, using CIFAR-100 as the unlabeled public dataset. We randomly distribute the dataset among prototypes with dataset ratios set 1:2:3:4:5:6:7 from XXS to XXL. Each dataset portion is further distributed among clients using Dirichlet distribution (Dir(0.5)). The number of clients ranges from 35 from XXS to XXL. respectively. Client sample rates are set at 0.1, 0.1, 0.15, 0.15, 0.2, 0.3, and 0.6 from XXS to XXL. We use a series of ResNet wcchtectures-RRsseeee ReeeeeeeeeeeeeOO-e--- 0-XXXS, ResNet1O-XS, ResNet1O-S, ResNet10-M, ResNet10, R8sNet18, and ResNet50 ately for each prototype. The local training epochs are set at 2, 2, 2, 5, 10, 10, and 20 from XXS to XXL to account for resource constraints, with fewer epochs assigned to smaller devices. We employ the Adam optimizer with learning rate of le-3 and a weight decay of 5e-5 for local training. For XL and XXL, step learning rate scheduler reduces the learning rate by factor of at half epoch. Server-side distillation employs fixed batch size of 128, using the Adam optimizer with learning rate le-3 and weight decay of 5e-5. The softmax temperature is set at for ensemble distillation and 20 for self-regule refu--egulaiiziionn. The number of communication rounds is fixed at 30. These exper- iments are conducted over trials with different random seeds, and the average performance with standard deviation is reported. The entire device prototypes configurations are given in Table 16.
appropri-
The detailed results are presented Tables 7, 8, and 9.
Table Scalability Evaluation. Detailed performance results for 7 device prototypes case.
Baseline
XXS
XS
M
XL.
XXL.
Average
FedAvg
23.17
30.6620.14
32.8140a1
31.772021
37.69
41.7820.06
50.52.a.01
35.49
FedDF
27.98
37 =0.333
40.6140.01
40.26 40.26:0.18
43.83
45.58 45.58-0.19
52, 32.18:0.11
41.13
FedET
25..9
36. 20.1
10.51.0.19
11.00go.16
16. 12.12
18. 18.39
52.714000
11.87
TAKFL
27. 27.30+0.0
36. 26.00
43.31
40.88 80..880.01
48.5240.1
50.9520.04
54.27
43.17
TAKFL+Reg
20. 20.28:0.18
37. 37.1040.45
13.96=1.85
41.83:oora
51.43:0.40 48.77=0a7
34.63-0.84
43.86
Table 8: Scalability Evaluation. Detailed performance results for 5 device prototypes case.
Baseline
XXS
XS
S
M
XL
Average
FedAvg
14
21.0440.76
33. f0.82
38,91 30.74
40.03
32.94
FedDF
28. 20.31
34.6640.00
39.91
38, 24.
46.81
37.59
FedET
26.88
34.
40.27
41.15 41.15
40.81 +0.s7
48. 48.1440.08
38,22
TAKFL
27.91 10.12
37.09 10.0940.11 10.11
40.4611.34
41.0640.02
49.02 10.35
39.11
TAKFL+Reg
28.2440.48
37.30+1.10
40.764a.4
43.09
50.86
40.05
Table Scalability Evaluation. Detailed performance results for 3 device prototypes case.
Baseline
XXS
M
Average
FedAvg
24.19_1.08
33.62
38.911m
32.24
FedDF
27.85_o.10
37.83
37.74
34.47
FedET
26.04_o.67
36.8720.68
37.66
33.52
TAKEL
26.62_o.16
37.32
38.1344..5
34.02
TAKFL+Reg
27.90+0.98
37.63=0.87
38.20
34.58
E
Ablation Studies
E.1
Understanding Merging Coefficient
In this section, we conduct an ablation study to further understand how TAKFL customizes knowl- edge integration and understand how the merging coefficients X; are achieving this. This experiment aims to further understand the trade-offs between customized knowledge integration approach from the one-size-fie aits-all strategy employed in vanilla ensemble distillation and prior works.
14

---

Experimental Setup. Our experimentation focuses on two device prototypes: XXS and XXL, se- lected from the scalability evalulation detailed in Section 7.3, Appendix D.4, and Table 16. We employ the image classification task on the CINIC-10 [9] dataset, starting from scratch. Each proto- type receives a randomly selected, non-overlapping subset of the training dataset-3.5799 for XXS and 25% for
among their clients in a non-i.i.d manner using Dir(0.5). Both pro- totypes have three clients each. The architectures used are ResNet1O-XXS for the XXS prototype and ResNet-50 for the XXL prototype. To focus solely lation process and its evolution with varying X, we pre-train each pro using standard FedAvg for 10 communication rounds, with sample rate of 1.0. Local training involves 20 epochs for XXS and 20 epochs for XXL using an Adam optimizer with a learning rate of 1e-3 and weight decay of 5e-5. The XXL prototype employs step learning rate scheduler that reduces the rate by factor of 0.1 at local epoch 10. For server-side distillation, we utilize a batch size of 128 and an Adam optimizer with a learning rate of 1e-5 and weight decay of 5e-5. CIFAR-100 [24] serves as the un- labeled public dataset. We save the final updated client and server models from both prototypes for further experimentation, focusing in TAKFL. The merging coefficient À is varied linearly from 0 to in increments of 0.05, For simplicity, the XXS prototype is referred to the small (S) prototype and XXL as the large (L) prototype.
the evaulation of the server-side distil-
the impact of merging coefficients without seff-regularization
Discussion. Figure 6 illustrates the significant impact of customized knowledge integration on the performance of both small and large device prototypes compared to the mne-size-fits-bl approach typical of vanilla ensemble distillation in the prior works, at different distillation epochs. Here, TAKEL adeptly manages customization for both small and large prototypes by controlling the merg- ing coefficient >. The merged model for both the small and large student prototypes is obtained using the formula mmerged = -aavg + ((1 - A)Ts + ATZ)- Notably, the performance is benchmarked at x 0.5 in all cases, reflecting similar results to vanilla ensemble distillation (FedDF), where no customization in knowledg transfer occurs. This baseline performance is critical for understanding the effects of further customization.
In small distillation epochs (Laistu < 10), minimal benefit is observed from customized knowledge integration, both small and large prototypes achieve optimal performance at the non-customized ~ 0.5. However, as the distillation process progresses beyond 10 epochs, the influence of À be- comes increasingly pronounced. For À > 0.5, the knowledge from the large prototype's ensembles predominates, enhancing their impact, while for < 0.5, integration is more influenced by the small prototype's ensembles. This pattern suggests that increased distillation epochs enable more effective distillation of each prototype's unique knowledge for extreme cases of extremely small and large prototypes, thereby making the customization benefits evident. In scenarios with small distilla- tion epochs, the absence of significant unique knowledge results in optimal performance at À Rs 0.5. Conversely, as the number of distillation epochs rises (Jdistatt 2 20), the one-size- fits--alllllll strategy proves suboptimal, underscoring the importance of tailored knowledge integration strategies. Opti- mal performance increasingly occurs at À > 0.5, indicating effective leveraging of each prototype's strengths to maximize overall performance. These findings confirm the necessity for customized knowledge integration in environments with significant prototype size variations and support our theoretical insights as detailed in Remark and 2.
E.2 Impact of Public Dataset
In this section, we explore the influence of the public dataset on the performance of TAKEL and existing KD-based methods when the public dataset used for server-side distillation is less simi- lar to the private dataset, which is the actual learning objective. For this analysis we employ the same experimental setup previously outlined in Section 7.2 and Appendix D.I, using the CIFAR-10 10 homo-family architecture. To measure dataset similarity. we compute cosine similarity between the averaged features of datasets, extracted using off-the-she tff-the pre- trairrn CLIP model [40] (CLIP ViT-B/32) available from the official GitHub repository. ?
Discussion. Table 10 presents our results, highlighting a significant observation: the performance of existing methods drastically deteriorates as the similarity between the public dataset and private datasets decreases. In contrast, TAKFL exhibits robustness, suffering much less performance degra- dation under the same conditions. This demonstrates TAKFL's practical utility in real-world sce-
15

---

Distill Epoch
Distill Epoch
Distill Epoch 10
60
60
60
TAKFL Large
Test Accuracy
50
Test Accuracy
50
Test Accuracy
50
TAKFL Small
Pretrain Large
40
40
40
Pretrain Small
FedDF Large
30
30
30
FedDF Small
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Lambda
Lambda
Lambda
Distill Epoch
Distill Epoch 30
Distill Epoch 40
60
60
60
Test Accurocy
50
Test Accuracy
50
Test Accuracy
50
40
40
40
30
0.0
0.5 Lambda
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Lambda
Lambda
Figure 6: Understanding the Impact of Merging Coefficients. This figure showcases server-side knowledge distillation between two device prototypes. XXS and XXL, referred to as small and large, respectively, utilizing CIFAR-100 100 as the unlabeled public dataset. Both prototypes were pre- trained from scratch using standard FedAvg for 10 communication rounds. The CINIC-10 dataset was distributed between the small and large prototypes in ratios of 3.57% and 25% respectively, and further subdivided non-i.i.d. among the clients using Dir(0.5). Each prototype has three clients with sample rate of 1.0. The small prototype utilizes ResNet1O-XXX architecture, while the large prototype employs a ResNet-50.
narios where the server typically lacks knowledge of the private datasets to select closely aligned public dataset for distillation. Notably, FedET underperforms significantly when using less similar public dataset, performing worse than vanilla FedAvg in both low and high data heterogeneity sce- narios. A similar pattern was observed with FedET in the NLP tasks discussed in Section 7.2 and Appendix D.3. This issue is likely due to FedET's dependence ibrated confidence scores from neural networks [15, 55] for uncertainty estimates in its uncertainty- weighted distillation approach.
the overconfident and poorly cal-
Table 10: Impact of Public Dataset on performance results. Same experimental setting described in Section 7.2 and Appendix D. on CIFAR- 10 homo-family setting is used for this experiment. The numbers in parentheses represent the similarity scores between private and public datasets, obtained using a pre-trained CLIP ViT-B/32 model.
Public Dataset
Baseline
Data Heterogeneity (Dir(0.3))
High Heterogencity (DierO.
S
M
Average
S
M
Average
FedAvg
36
46:41.935
44.46
47.36
25.26.0.9
51.51 1a.58
32.93
FedDF
19.31
30.63
29 20.82
49.92
34.71=144
35.27
51.08
$0.35
CIFAR-160 (0:99)
TAKFL
15.90
19.21:0.:
55.01e1.11
57. g7.93
60.58 12.35
52.60+a.47
58. 54.14
52.61
37. 1.04
5.
38. 96 10:1
30.90.2470
11.1916.35
45.5349.4
$5.36
12.62
TAKFL Reg
56.37
58.60 58.60.0.43
65. 69 15.399.1.5
60.22 22
40.51:1.05
40.12:1
23.24:2.1
44.62
FedDF
49.37=1
40.41:4.11
55.0
51.28
31.41-001
30.70
39.8249.16
33.90
TisyImagenet (0.92)
FedET
33 96
37. 26
39, 39.77
56.00 36.99
24 21.12-111
2L.882215
28.91 28.91
25.87
TAKHL
65 200007
36, 3633660.44
60.71
57.42
40 40.08-0011
40 40.2644044
43 43.56/11.9
11.30
TAKFL+Rog
56.28=000
57 $7.14-0..
41. 41.10+1.m 18
46.25=a..0 n.
FedDF
48.99 68.99_n.a
50.06.0.40
55, 15.12
51.39
20 10.80_.0
32 44
44.0
35.36
Coleb-A (0.77)
TAKFL
FedET
28. _3.56_3.0
45.65-278
54.50+1.7
28, 28.50
17.20
13..18
31.52
52.77
31.02:0.s
15.25
36.76g1 se
19.00
23.29 15.01
48.3334n.55
38.70
19.19
46
56.
60. 13:: 38
54.58
30.88 30.88
35 35.96
52 52.68_1 90
39.84
E.3
Impact of FL Optimizer
To further assess the impact of FL optimizer on the performance, we adopt FedOpt (Adam) [41], which is a SOTA FL optimizers as the base per prototype FL optimizer and conduct the same exper- imental setup on CV task discussed in Section xx. The results are presented in Table 11. As we can see, the performance of TAKFL. is robust under different FL optimizers and is consistently achieving SOTA performance.
16

---

Table 11: Impact FL Optimizer. Performance comparison results using FedOpt (Adam) [41] as prototype FL optimizer for the same setting as Table (Section 7.2).
Low Data Heterogeneity (Dir(0.3))
High Data Heterogeneity (Dir(O.1))
Dataset
Baseline
S
M
L
Average
S
M
Average
FedOpt
35.40
41.69
58.29
45.12
23.18
22.03
37 37.82
27.6
FedDF
43.74
43.48
63.36
50.19
29.23
29.46
37.45
32.04
CIFAR-10
FedET
35.28
40.01
63.38
46.22
25.84
28.17
39.26
31.09
TAKFI.
55.95
56.25
65.14
59.11
39.92
43.26
45.64
42.94
FedOpt
10.02
20.91
32.78
21.23
9.22
15.74
21.36
15.44
FedDF
15.34
23.26
30.92
23.17
12.54
19.81
22.58
18.31
CIFAR-100
FedET
9.01
19.59
32.26
20.28
6.85
15.30
23.44
15.19
TAKFL
27.06
28.93
33.76
29.91
22.40
23.38
26.35
24.04
E,4 t-SNE Visualization of Knowledge Transfer
To better understand the effectiveness of knowledge transfer in TAKFL, we conducted a visual- ization study using two device prototypes, XS and XL., on the CIFAR-10 dataset. The prototype configurations follow those used our scalability evaluation in Section 7.3 (more details are avail- able in Appendix D.4). We first pre-trained the two prototypes' global models using FedAvg for 40 communication rounds and then performed knowledge transfer using the TAKFL process for 10 distillation epochs. We have plotted the I-SNE visualizations of both prototypes for FedAvg and TAKFL in Figure 7.
in the t-SNE plots (Figure 7). TAKFL demonstrates substantial enhancement in feature representation for both prototypes compared to FedAvg. The plots show better separation between classes for both XS and XL prototypes. The class clusters more distinct, suggesting that TAKFL, by effectively transferring knowledge between the two prototypes, yielded a clearer representation of the features, even in the smaller XS prototype. Therefore, the t-SNE visualizations confirm that TAKFL is indeed effective in knowledge transfer, further corroborating its superiority and effectiveness.
As
FedAvg
(b) TAKFL
Figure 7: t-SNE Visualization Study, To better understand the effectiveness of TAK'Lss knowledge
transfer, conducted experiment setting involving two device prototypes XS and XL. similar to our scalability evaluation setup in Section 7.3 on CIFAR-10 10 dataset.
17

---

Hyper-parameters and Implementation
In this section we bring the details of the hyper-parametern we used and our implementation. We implement entire code in PyTorch [38] using FedZoo benchmark [36] and release it at https: //github. com/MMorafah/THe experimentation in this paper.
We use two NVIDIA RTX 3090 gpus to conduct the entire
Computer Vision (CV) Task
For comprehensive evaluation of method, consider federated learning image classification training from scratch.
F111 Datasets
Datasets. We experiment with several image classification datasets: CIFAR-10, CIFAR-100, CINIC-10, TinylmageNet, STL-10. The details of each dataset is the following:
CIFAR-10 [24]: CIFAR-10 consists of 60,000 images of size 32x32 RGB across 10 classes, with each class containing 6,000 images.
CIFAR-100 [24]: CIFAR- 100 comprises 60,000 images of size 32x32 RGB distributed across 100 classes, with 500 images per class.
CINIC-10 [9]: CINIC-10 has 270,000 images of size 32x32 RGB across 10 classes, each class containing 27,000 images.
ImageNt [25]: TinylmageNet contains 100,000 images of size 64 RGB across 200 classes.
STL-10 [7]: STL-10 has 100,000 unlabeled images and 13,000 labeled images of size 96x96 RGB across 10 classes.
The ImageNet- as the unlabeled public dataset is constructed by randomly selecting 100 classes from the ImageNet [10] dataset.
F.1.2 Architectures
Experiments in Section 7.2 and Appendix D.I. The architecture that we use are two distinct architectural settings: the "homo-family" setting, where all device prototypes' architectures are from the same family, and the "hetero-family" setting, where architectures do not necessarily belong to the same family. For the homo-family scenario, we employ ResNet-8 for S, ResNet-14 for M. and ResNet- 18 for L. For the hetero-family scenario. we use VIT-S for S, ResNet- 14 for M, and VGG-16 for L. All models are initialized from random initialization.
For the ResNet architecture configuration, we utilize the standard 'BasicBlock' as the building block. This consists of convolutional block, followed by four residual block stages, an adaptive average pooling layer, and a classifier layer. The models within this family differ in terms of the number of repetitions of the residual block and the number of filters in each stage. The configurations for different capacities are detailed below:
R8 is configured with [64, 128, 256, 512] filters and repeats the 'BasicBlock' [2, 2, 2, 2] times.
ResNet-14 14 is configured with [64, 128, 256, 512] filters and repeats the 'BasicBlock' [1, 2, 2, 1]
times.
R is configured with [64, 128, 256] filters, corresponding to the first three stages, and repeats the 'BasicBlock' [1, 1, times.
For VGG-16 [44], max-pooling layers. The configuration of filters for each layer is as follows:
the standard architecture which includes convolutional layers followed by
VGG-16: [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
18

---

The final classification head consists of two linear layers with a hidden size of 512, followed by ReLU activation, and final linear classifier layer.
For VīT-S, we adopt the standard Vision Transformer [12] architecture implementation from Github. We configure ViT-S with 6 attention blocks, each with 16 heads and hidden dimen- sion of 64. The final MLP dimension is set to 256. In our experiments with VīT-S, we set the patch size 4, and the input image size is 32
Experiment in Appendix D.2. The pre-trained MobileNerV3-Large [19], MobileViTV2 [34], and ResNet34 [51] were instantiated using the TIMM library.*
Scalability Experiments in Section 7.3. The architectures in these experiments are inspired by [20]. The details of architectures are as following:
ResNer1O-XXS is configured with [8, 8, 16, 16] filters and repeats the 'BasieBlock' [L, 1, 1]
times.
ResNer1O-XS is configured with [8, 16, 32, 64] filters and repeats the 'BasicBlock' [1, 1, 1]
times.
ResNer1O-S is configured with [16, 32, 64, 128] filters and repeats the 'BasieBlock' [1, 1, 1] times.
ResNerIO-M configured with [8, 16, 32, 64] filters and repeats the 'BasicBlock' [1, 1, 1, 1] times.
RexNer10 configured with [64, 128, 256, 512] filters and repeats the 'BasicBlock' [1, 1, 1] times.
R configured with [64, 128, 256, 512] filters and repeats the 'BasicBlock' [1. 1, 1, I] times.
ResNer50 is configured with [64, 128, 256, 512] filters and repeats the 'BasicBlock' [3,4,6,3] times.
F.1.3 FL configuration and Slyper-paramelee
Base Klyper-parametent The following hyperparameter values apply to all CV experiments unless stated otherwise. We the diversity regularizer coefficient of FedET to 0.1 for our entire experi- mentation per the original paper [6]. We use the Adam optimizer with a learning rate of Ic-5, weight decay value of 5e-5, and batch size of 128 for distillation. The softmax distillation temperature is set to 3, the distillation epoch to 1, and the self-regularizer softmax temperature to 20 for both CV and NLP experiments. Table 13 details the aaper-parameter.
Experiments Section 7.2 and Appendix D.1. For tables 1 and 4, there are 100 clients with the device prototype, 20 clients with the M device prototype, and 4 clients with the I. device prototype. For each round, 10, 4, and clients from each S, M, and L. prototype are randomly sampled respectively for participation. 10% of the data goes the S prototype, 30% M. and 60% to L. The data is distributed to each client among each prototype in a Dirichlet distribution. Table 14 details the FL configuration.
Experiments in Appendix D.2 and Appendix E.2. For tables 5 and 10, there are 4, 3, and clients for S, M, and L device prototypes, respectively. Each round, every client participates in FL. 20% of the data distributed to prototype S, 30% to prototype M. and 50% to prototype L. Table 15 details the FL configuration.
Scalability Experiments in Section 7.3 and D.4. For tables 7, 8, and 9, there are 35, 30, 25, 20, 15, 10, and 5 clients for the prototypes XXS, XS, S, M, L, XL., and XXL, respectively. The sample rate is set to I, 0, 0.15, 0.15, 0.2, 0.3, and 0.6 from XXS to XXL. The data is distributed for each prototype in the ratio 1:2:3:4:5:6:7 from XXS to XXL. Table 16 details the hyper-parameters and configuration.
Validation Set. For TAKFL, the validation set used for the heuristic method (see F.3) is 5% of the training dataset. The validation set and the private dataset does not overlap.
19

---

F.2
Natural Language Processing (NLP) Task
For the NLP task, we fine-tune federated learning text classification task using pretrained models.
F.2.1 Datasets.
All NLP datasets were provided by Hugging Face."
MNLI [52]: MNLI contains 433K sentence pairs, each sentence pair labeled of 'entailment," 'neutral," and 'contradiction."
SNLI [2]: SNLI is similar MNLI, with 570K sentence pairs each labeled one of 3 labels.
SST2 [45]: SST2 consists of 67K phrases, each labeled as sentiment 'positive' or "nega- tive.'
Sentiment140 40 [14]: Sentiment 40 a dataset of 1.6M Twitter messages each labeled with of sentiment values.
MARC [22]: MARC (Multilingual Amazon Reviews Corpus) is dataset with online re- views multiple languages from the Amazon delivery service website. Each review has label which is one of 1-5 stars. We only use the English reviews from this dataset, which results 260,000 English reviews total.
Yelp [61]: The Yelp reviews dataset contains 700K reviews each labeled 1-5 stars from the Yelp service which publishes public reviews of businesses.
AG News [60]] AG News contains 127,600 news article titles. Each article is one of four classifications of news articles.
DBpedia [59]: The DBpedia dataset consists of 630K DBpedia article summaries each labeled one of 14 categorizations.
F222 Architectures
Experiments in Section 7.2 and Appendix D.3. We use three variations of the BERT architec- ture: BERT-Tiny, BERT-Mini, and BERT-Small from [47]. The weights were pre-trained on BookCorpus dataset and extracted text from Wikipedia. Further details regarding each model are described extensively on Github. 6 The tokenizer used for these transformer models are the same ones provided by the authors of [47].
BERT-Tiny contains transformer layers and an embedding size of 128.
BERT-Mini contains 4 transformer layers and an embedding size of 256. BERT-Small contains 4 transformer layers and and embedding size of 512.
F.2.3 FL configuration and hyper-parameter
Base Hyper-parametert For distillation, no weight decay, and batch size 32. The distillation epoch is set to 1, the ensemble distillation softmax temperature 3, and the self-regularizer softmax temperature 20 for all NLP experi- ments. Table 17 details the hyper-parameters.
the Adam optimizer with learning rate of 3e-5,
Experiments in Section 7.2 and Appendix D.3. For tables and 6, we limit the private dataset to 100,000 samples, randomly sampled from the original dataset i.i.d. The public dataset is limited to 30,000 examples sampled i.id as well. There are 8, 4, and clients for the S, M. and L prototypes. The private data split across each prototype in the following proportions: 0.1, 0.3, 0.6. Table 18 details the FL configuration.
Validation Set. The validation dataset used for TAKFL 5,000 samples taken from the original training dataset that does not overlap with the 100,000 private dataset.
20

---

F.3 Hlyper-parametery of TAKEL
Merging Coefficients. We conducted extensive experiments with different merging coefficients on the main 3-device prototype setting of small (S), medium (M), and large (L) discussed in Section 7.2 and Appendix D.1. We empirically observed that the small (S) prototype typically achieves the best performance using uniformly increasing merging coefficient, where the larger the prototype, the larger the merging coefficient, i.e., Às S AM S Az. As we move towards larger prototypes, they benefit more from increasingly skewed merging coefficients towards the larger ones. In the extreme case the large (L) prototype, highly skewed merging coefficients generally led to better performance, i... Äs < AM < Åz. This pattern is intuitive as small prototypes can benefit from everyone while gaining more from the larger, more informative prototypes. However, larger prototypes benefit less from smaller ones, as they typically offer less information, especially in high data heterogeneity cases. Notably, in high data heterogeneity cases, more skewed merging coefficients seemed to be more advantageous the smaller prototypes (S and M) possess lower quality knowledge.
Based on these observations, we designed simple and cost-effective heuristic method that ran- domly instantiates merging coefficients following this intuition. Our heuristic method, presented in leverages these observations by generating candidate merging coefficients that incorporate both uniformly increasing and different degrees skewed merging coefficients. This dual approach enables to explore wide range of merging strategies and identify the most effective configura- tions for different prototypes. The optimal merging coefficient candidate determined using the performance on the held-out validation set.
import numpy np def heuristic (num (arliinnmm___dveeeeee _eviices=,,, c_canddatte==00):
3 candidates [[1/nm_devices for exponent in [1, 5, 10]:
devices for in range [oun_devices]]
4
5
for i in (ange (n_candidatees)
candidate np random. random.betaaa-1, b=100, size-num,devices
candidate candidate exponent
8
candidate np sort (candidate)
9
candidate = candidate mp. up.sum (caandida))e)
10
candidates append (candidaat))))e)
11
return candidates
Listing 1: Implementation of the heuristic method for merging coefficients in Python. The exponent term controls the degree of skewness or peaking the merging coefficients.
Furthermore, we experiment with manually determining the merging coefficients and fixating them throughout the federation. We achieved similar results with this approach compared to adaptively finding the coefficients using the heuristic method and small held-out validation set. We present the merging coefficient candidates that performed reasonably well during our experiments in Table 12.
Table 12:
Details of the experimentally determined merging coefficients for the 3-device prototype
setting discussed in Section 7.2 and Appendix D.I. Coefficients ordered as [Ag, Ass Aul-
Merging Coefficient Candidate
Small Prototype
Medium Prototype
Large Prototype
[0.2, 0.3, 0.5]
[0.1, 0.2, 0.7]
[0.1, 0.2, 0.7]
2
[0.3, 0.3, 0.4]
[0.05, 0.15, 0.8]
[0.01, 0.09, 0.99]
3
[0.2, 0.3, 0.5]
[0.1, 0.2, 0.7]
[0.05, 0.2, 0.75]
[0.05, 0.1, 0.85]
[0.01, 0.19, 0.8]
[0.01, 0.09, 0.90]
5
[0.1, 0.15, 0.75]
[0.05, 0.15, 0.8]
[0.01, 0.09. 0.90]
[0.05, 0.1, 0.85]
[0.05, 0.05, 0.9]
[0.001, 0.009, 0.99]
7
[0.05, 0.15, 0.80]
[0.05, 0.2, 0.75]
[0.001, 0.009, 0.99]
[0:05, 051 0.80]
[0.05, 0.1, 0.85]
[0.001, 0.009, 0.99]
9
[0.3, 0.35, 0.35]
[0.2, 0.3, 0.5]
[0.1, 0.2, 0.71
Self- Suaf-Regulariztiion Coefficient. Extensive experiments were conducted the self-regulation co- efficients for different device prototypes and settings. Although no consistent pattern emerged, we experimentally determined that optimal performance for the small prototype was achieved with self-
21

---

regulation coefficients E 0.1, 0.01, 0.001. For the medium prototype. the coefficients were >M
0.5, 0.1, 0.01, 0.001, 0.0001, and for the large prototype. Yr. 1.0,0 0.5, 0.1, 0.01, 0.001, 0.0001 yielded the best results.
Table 13: Details of hyper-parametert for CV task in Section 7.2, Appendix D.I, and Appendix E.2.
LocaliServer
llyperparameter
Small Projxtype
Medium Protatype
Large Prokotype
Training opochs
100
Batch Size
64
64
64
Optimizer
Local Training
Learning Rate
le:3
le.3
le-3
Weight Decay
5c-5
5e-5
LR scheduler
None
None
St StepLRietep.sie
10,
parma
Optimizer
Adam
Adam
Adam le-5
Learning
Weight Decay
Se-5
Se
Se-5
Training
Batch Size
128
Training Epochs
Fisembll Distillation Softmax Temperature
3
Self Regularizer Softmax Temperature
20
Table 14: Details of Architecture parameters and FL configuration for CV task in Section 7.2 and Appendix D.I.
CIFAR-10
CIFAR-100
Architecture Setting
Device Prototype
Architecture
Parameters
Purameters
Dataset Portion
Clients
Sample Rate
Prototype Prototype Prototype
Res|Net8
1.23M
1.25M
0.1
100
0.1
Homo- Hlomo-Family
ResNet18
1117M
11,22M
0.6
0.5
Prototype Prototype M
VII-S
L.8MM
L79M
0.1
100
0.1
Hetero-Fumily
ResNet14
6.38M
6.43M
Prototype
VGG16
15,25M
15,30M
0.6
0.5
Table 15: Details of Architecture parameters and FL configuration for CV task experiment using pre-trained models in Appendix D.2.
STL-10/CINIC-10
TinyImageNet Parameters
Device Prototype
Architecture
Parameters
Dataset Portion
Clients
Sample Rate
Prototype 5
mobilenetv3-large- mobilevitv2-175 ResNet34
4.21M
4.45M
0.2
4
1.0
Prototype M
13.36M
13.53M
0.3
3
Prototype I.
21.28M
21.38M
0.5
2
1.0
Table 16: Details of Architecture parameters for Scalability Section 7.3, and Appendix D.4.
Deviee Prototype
CINIC-10
Architecture ResNecIO-XXS
Parameters
Dataset Portion
Clicnts
Sample Rate
Local Epochs
Prototype
11K
0.0357
Prototype XS
RosNet10-XS
78K
0.0714
30
0.1
Prototype S
ResNer1O-S
309K
0.1071
25
0.15
Prototype M
ResNer1D.M
1.2M
0.1428
20
0.15
Prototype
ResNetIO
4.9M
0.1785
15
0,2
10
Prototype
Prototype XXL
ResNet50
24M
0.25
5
0.6
20
22

---

Table 17: Details of hyper-parametter for NLP task experiments in Section 7.2 and Appendix D.3.
Local/Server
Hyperparameter
Prototype S
Prototype M
Prototype
training epnchs batch optimizer
32
32
Adam
Adam
Adam
Local Training
Learning rate
3e-5
3e-5
3e-5
Weight decay scheduler
0
None
None
None
optimizer
Adam
Adam
Adam
learning rate
3e-5
3e-5
3e-5
weight decay
3e-5
3e-5
3e-5
Server KD Training
batch size
32
32
32
training epochs
1
1
1
Ensemble distillation softmax temperature
3
3
3
Self-regularizer softmax temperature
20
20
20
Table 18: Details of Architecture parameters and FL configuration for NLP task experiment using
pre-trained models in Section 7.2 and Appendix D.3.
Device Prototype
Architecture
Parameters
Clients
Dataset Portion
Sample Rate
Prototype
BERT-Tim
4.39M
8
0.1
0.4
Prototype M
BERT-Mini
11.17M
4
0.3
0.5
Prototype
BERT-Small
28.77M
2
0.6
1.0
23

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitation of the work is discussed the main paper Section 8.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Ys]]

---

The informal statements are in the main paper Section 6, and the full formal statements with assumptions and proofs can be found in the Appendix C.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
The detailed information can be found in Appendix D, and F.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data also be provided via
often one good way accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models. it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The information can be found in Appendix
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
wwtpcccpppcccc.cc.coooooo /eeddddeeee/ooeeeieeeeeeee ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The information can be found in Appendix D and
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We conduct the main experiments in the main paper Section 7 over 3 inde- pendent trials with different random seeds and average results with the standard deviation is reported in Tables 1, 2, 6, 16, and 10.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors). should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK report I-sigma error bars, but one should state The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The information can be found in Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't t make it into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: : //tp///eens//////ee//ee// httppp..ccpppppeee
Question:
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper focuses advancing technical methodologies in federated learn- ing and does not directly engage with specific societal impacts, as the applications of the proposed method varied and context-dependent, The societal implications would largely depend on the specific use cases implemented by subsequent practitioners. Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.

---

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
This paper primarily focuses technical methodological approach without releasing any models. The potential for misuse depends on specific applications by end users, which are beyond the scope of this study.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The datasets and models used in this work are properly credited in main paper Section 7, Appendix D, and Appendix
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
5

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswith thcode www.coo///aoooo candaataeets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification:
The information can be found in Appendix F.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper deseribe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.


Reviews:
[{"review_id": "6lbARfO6VE", "reviewer": "Reviewer_V7t1", "summary": "This paper introduced a KD-based framework (TAKFL) to address the dilution and diversity issues in heterogeneous FL knowledge transfer learning. The TAKFL distills knowledge from prototypes of varying sizes and incorporates a self-regularization to mitigate noise simultaneously, then integrates these separately distilled knowledge by task arithmetic. Empirical evaluations across various CV and NLP datasets demonstrate the framework's effectiveness.", "questions": "see weakness", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-organized and easy to follow.\n2. The paper novelty introduced a theoretical model to illustrate the efficacy of knowledge distillation in heterogeneous FL.\n3. The paper proposed a new framework, TAKFL, for considering varying sizes of prototypes with different contributed information, and experiments on different CV and NLP datasets show its effectiveness.", "weaknesses": "1. Some baselines are lacking. For example, FedProto [1], which also employs prototypes within device heterogeneous FL, should be included for a more comprehensive comparison. \n2. It seems that the proposed method incurs higher time and storage costs, as it requires the independent learning of multiple student models compared to the vanilla methods. The paper should provide an efficiency analysis that compares the proposed method with existing baselines, highlighting both time and storage metrics.\n3. It would be better to provide a visualization study for a better understanding of the effectiveness of transfer learning from different prototypes.\n\n[1] Tan, Yue, et al. \"Fedproto: Federated prototype learning across heterogeneous clients.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. No. 8. 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your invaluable review and pointing out the **novelty of our framework and theoretical model** and **effectiveness of our framework on both CV and NLP tasks**.\n\n### **Response to 1**\nWe appreciate the reviewer's suggestion and would like to highlight a few differences about FedProto and our work, TAKFL:\n\n**(1) Problem Formulation:** FedProto's problem formulation aims to optimize local models in FL (Eq. 2 in [1]), whereas our problem formulation follows [2], where we aim to optimize device prototypes’ global models (Eq. 1 in our paper).\n\n**(2) Feature Exchange Requirements:** FedProto requires the exchange of local and global average per-label features between devices and the server, necessitating that local architectures produce the same feature matrix dimensions. However, in more general cases where architectures are from different families, this does not hold. Our experimental setup includes the heterogeneous-family case too.\n\n**(3) Evaluation Metrics:** FedProto uses average local test accuracy as the evaluation metric, while our evaluation metric is device prototypes’ global model test accuracies, following [2].\n\nGiven these differences, a direct comparison with FedProto is not entirely feasible or fair. However, considering these factors, we have adopted FedProto’s official implementation code and compared it within our homogeneous-family case experimental setting in Table 1 of our paper. The results, presented in Table 1 of the PDF rebuttal, show that TAKFL consistently outperforms FedProto for both CIFAR-10 and CIFAR-100.\n\nWe will incorporate these comparisons into the final version of the paper to provide a more comprehensive evaluation.\n\n### **Response to 2**\nTo address the reviewer’s comment, we would like to mention that the full algorithm details of TAKFL have been presented in Appendix B. All the distillation processes are being performed in parallel and asynchronously (lines 9-10 in Algorithm [1]), and once done, the distillation task vectors are merged. \n\nTherefore, the computation time remains O(1), and the overall computation cost of the entire server-side distillation process is O(M), as device prototypes do not perform KD one at a time but combine information asynchronously into a singular task arithmetic operation (Eq. 8). Since there are M prototypes, the total computation is O(M). \n\nWe have further conducted an efficiency analysis on a system with two RTX 3090 GPUs and compared our method with prior works for one epoch knowledge distillation process, including time in seconds, GPU memory, and CPU memory in MB for CIFAR-10 experimental setting in Table 1 of our paper. The results are as follows:\n\n| Baseline\t| Time\t| CPU Memory | GPU Memory|\n| ----------- | ------- | ---------- | --------- |\n| FedDF   \t| 55.24   | 1051.89\t| 104.90\t|\n| FedET   \t| 58.81   | 1113.12\t| 107.65\t|\n| TAKFL   \t| 71.55   | 1080.60\t| 266.57\t|\n\nAs we can see, our method achieves similar CPU memory usage compared to the baselines, with a slight increase in GPU memory and time. The additional 16 seconds incurred time can be attributed to our use of the high-level parallelism interface implementation, `concurrent.futures`, for executing callables asynchronously. Furthermore, in FL, servers are typically equipped with substantial computational resources, including multiple GPUs and extensive memory capacities, to manage the orchestration of many clients. These servers can employ highly efficient parallel asynchronous implementations to avoid the incurred time. Given that the entire distillation process occurs at the server, the computation and storage demands are less of a concern since the server has ample resources.\n\nWe will incorporate these results and the efficiency analysis into the final version of the paper.\n\n### **Response to 3**\nWe appreciate the reviewer's suggestion. To better understand the effectiveness of knowledge transfer in TAKFL, we conducted a visualization study using two device prototypes, XS and XL, on the CIFAR-10 dataset. The prototype configurations follow those used in our scalability evaluation in Section 7.3 (more details are available in Appendix D.4). We first pre-trained the two prototypes' global models using FedAvg for 40 communication rounds and then performed knowledge transfer using the TAKFL process for 10 distillation epochs. We have plotted the t-SNE visualizations of both prototypes for FedAvg and TAKFL, included in Figure 1 of the PDF rebuttal file.\n\nAs seen in the t-SNE plots, TAKFL demonstrates a substantial enhancement in feature representation for both prototypes compared to FedAvg. The plots show better separation between classes for both XS and XL prototypes. The class clusters are more distinct, suggesting that TAKFL, by effectively transferring knowledge between the two prototypes, yielded a clearer representation of the features, even in the smaller XS prototype. Therefore, the t-SNE visualizations confirm that TAKFL is indeed effective in knowledge transfer, further corroborating its superiority and effectiveness.\n\nWe will enrich this visualization study by considering knowledge transfer between different size prototypes and incorporate this study into the final version of our paper.\n\n### References\n[1] Tan, Yue, et al. \"Fedproto: Federated prototype learning across heterogeneous clients.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. No. 8. 2022.\n\n[2] Lin et al. \"Ensemble distillation for robust model fusion in federated learning.\" NeurIPS’20", "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": "", "extra_scores": {}}, {"review_id": "v1u8g5oJDF", "reviewer": "Reviewer_RvJG", "summary": "The paper presents a novel framework called TAKFL, which addresses the challenge of transferring knowledge in federated learning across heterogeneous devices, ranging from small IoT devices to large workstations. TAKFL uniquely handles the knowledge distillation by treating the transfer from each device prototype as a separate task, allowing for tailored integration of knowledge to optimize performance. The approach is validated theoretically and through extensive experiments, demonstrating superior results over existing methods across both computer vision and natural language processing tasks.", "questions": "See weaknesses", "limitations": "N.A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "(1) Practical Problem and Innovative Approach\nThe authors address a significant, real-world challenge in federated learning—knowledge transfer across heterogeneous devices. Their novel framework, TAKFL, innovatively treats each device's knowledge transfer as a separate task, allowing for customized integration. This tailored approach is both practical and theoretically sound, making it a substantial contribution to the field.\n\n(2) Clarity and Organization\nThe paper is well-structured, facilitating easy understanding of complex concepts and methods. The clear presentation enhances the accessibility of the content, making it easier for readers to grasp the significance of the proposed solution and its impact on the field.\n\n(3) Strong Experimental and Theoretical Support\nThe authors back their claims with extensive experimental results across multiple tasks and datasets, demonstrating the effectiveness of TAKFL in diverse scenarios. Moreover, the inclusion of theoretical analysis adds depth to the validation, reinforcing the reliability and scalability of their approach. This combination of empirical and theoretical evidence strongly supports the paper's contributions and conclusions.", "weaknesses": "(1) [main concern] Strong Dependence on Hyperparameters\nThe TAKFL framework introduced in the article significantly relies on the setting of hyperparameters, especially during the Task Arithmetic Knowledge Integration process, where the weights for different task vectors are set as hyperparameters and adjusted on a validation set. This might limit the method's generalizability across different real-world applications. To enhance the practicality and robustness of the method, it is recommended that the authors explore more automated hyperparameter optimization strategies to reduce the need for manual tuning and improve the adaptability of the model.\n\n(2) [main concern] Strong Assumptions in the Selection of Public Datasets\nThe experimental design involves the use of public datasets, such as CIFAR-100 and ImageNet-100, for knowledge distillation. This choice seems to be based on two key assumptions: that the public datasets must exhibit high diversity and that the training data distribution can be approximately considered a subset of the public dataset. These assumptions may not always hold in practical applications, so it is advisable for the authors to thoroughly investigate the actual impact of these choices on model performance in future work. Additional experiments could validate the effectiveness of these assumptions, and considerations of these potential limitations should be explicitly stated in the manuscript.\n\n(3) [minor concern] Quantification of Data Heterogeneity and Hyperparameter Selection\nThe authors utilize a Dirichlet distribution to quantify Data Heterogeneity, setting $Dir(\\alpha)$ at 0.3 and 0.1 to simulate varying degrees of data heterogeneity. However, there is insufficient explanation for the choice of these specific values. To enhance the transparency and reproducibility of the research, it is recommended that the authors provide a detailed rationale behind these parameter choices, based on logic and references to previous studies. Moreover, to give readers a more intuitive understanding of the data distribution differences under different $\\alpha$ settings, descriptive statistics or visualizations, such as the distribution of samples across categories, would be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your invaluable review and pointing out that our work makes a **substantial contribution to the field**. We also appreciate the reviewer for the positive remarks on the **practicality, innovation, and clarity of our TAKFL framework**, as well as the **strong experimental and theoretical support** for our approach.\n\n### **Response to 1**\nWe appreciate the reviewer’s feedback regarding the dependence on hyperparameters. We have designed a simple and efficient automated heuristic method to tune the merging coefficients, detailed in Appendix F.3. Based on our extensive experiments, we observed that larger prototypes benefit more from increasingly skewed merging coefficients, while smaller prototypes perform best with uniformly increasing coefficients. This pattern is intuitive as larger prototypes gain less from smaller ones, especially in high data heterogeneity cases.\n\nOur automated heuristic method randomly instantiates merging coefficients following this observed pattern, generating 30 candidate coefficients that include both uniformly increasing and various degrees of skewed coefficients (detailed implementation is presented in Listing 1 of Appendix F.3). The optimal merging coefficient is determined using maximum performance on a small held-out validation set. As detailed in Appendix F.3, we observed similar performance using this automated heuristic method compared to manual tuning (lines 1143-1145). Furthermore, we used this automated heuristic method for our scalability evaluation experiment in Section 7.3, Figure 3 of our paper. This approach reduces the need for manual tuning, enhancing the adaptability and robustness of the TAKFL framework across different real-world applications.\n\n### **Response to 2**\nWe appreciate the reviewer’s feedback regarding the impact of public datasets. Our public dataset selection mainly follows [1] for a fair and consistent comparison with existing works. To address the reviewer’s concern, we have explored the influence of the public dataset on TAKFL's performance and existing KD-based methods when the public dataset used for server-side distillation is less similar to the private dataset, as detailed in Appendix E.2.\n\nIn this experiment, we measured the similarity between different public datasets and the private dataset using CLIP [2]. We fixed the private dataset to CIFAR-10 and used less similar public datasets such as Celeb-A. We observed that the performance of existing methods drastically deteriorates as the similarity between the public and private datasets decreases. In contrast, TAKFL exhibits robustness, suffering much less performance degradation and still achieving SOTA average performance under the same conditions. For instance as presented in Table 10, when using Celeb-A as the public dataset, which is significantly different from CIFAR-10, TAKFL still achieved more than 3% average performance across all prototypes compared to the baselines. This is in contrast to other methods like FedET, which substantially underperformed compared to vanilla FedAvg under the same condition.\n\nThis demonstrates TAKFL's practical utility in real-world scenarios where the server typically lacks knowledge of the private datasets to select a closely aligned public dataset for distillation. Additionally, we limited the size of the public dataset to 60,000 or less across all experiments in throughout the paper.\n\nWe will further highlight these findings in the final version of our paper to emphasize TAKFL's robustness and practical applicability.\n\n### **Response to 3**\nWe appreciate the reviewer’s feedback and would like to further clarify our data heterogeneity setting. We utilized Dirichlet distribution to quantify data heterogeneity, as it is a commonly used method in the FL literature to simulate varying degrees of data heterogeneity [1,3]. The level of data heterogeneity is controlled by $\\alpha$; a lower $\\alpha$ indicates higher data heterogeneity and presents a more challenging setup.\n\nThe chosen values of $\\alpha = 0.3$ for low data heterogeneity and $\\alpha = 0.1$ for high data heterogeneity are standard in the FL literature and follow practical experimental design suggestions from [4] (cited in our paper in reference [36]). To give readers a more intuitive understanding of the data distribution differences under different $\\alpha$ values, we have plotted visualizations showing how clients' data distributions look for CIFAR-10 using these $\\alpha$ values. These visualizations are included in Figure 3 of the PDF rebuttal file. We will include this Figure in the final version of the paper.\n\n### References \n\n[1] Lin et al. \"Ensemble distillation for robust model fusion in federated learning.\" NeurIPS’20\n\n[2] Radford et al. \"Learning transferable visual models from natural language supervision.\" ICML’21\n\n[3] Li et al. \"Federated learning on non-iid data silos: An experimental study.\" IEEE ICDE (2022).\n\n[4] Morafah et al. \"A Practical Recipe for Federated Learning Under Statistical Heterogeneity Experimental Design.\" IEEE TAI (2023).", "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": "", "extra_scores": {}}, {"review_id": "USFRZZJb0q", "reviewer": "Reviewer_F7vu", "summary": "The paper focus on a problem that traditional federated learning methods fail to effectively handle scenarios where devices have widely varying capabilities. It improve existing Knowledge Distillation (KD) methods that are inadequate in these heterogeneous environments. Experimental results show the validity of their proposed method.", "questions": "please refer to Weakness.", "limitations": "please refer to Weakness.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. TAKFL treats knowledge transfer from each device prototype’s ensembles as separate tasks and distills them independently.\n2. The paper is well written with comprehensive experiments", "weaknesses": "1. Assumption on Weight Disentanglement Property: The theorem 2 reply on the assumption of the weight disentanglement property  (line 679-681 in appendix)  is  too strong. In practice, achieving weight disentanglement is challenging. Studies [1,2] demonstrate that there are interferences among task vectors, making disentanglement difficult. [3] achieves disentanglement using Neural Tangent Kernel (NTK) and shows without disentanglement the performance is dropped. Consequently, asserting the disentanglement property is problematic, thereby limiting the theoretical impact. \n\n2. Similar to weight conflicts when doing merge, averge logits also have conflicts, the paper only considers vanilla average logit as the KD loss. However, there are studies to resolve the issue, like [4,5] have proposed better ensemble KD loss designs, more studies online. Therefore, incorporating these methods for comparison is important.\n\n3. The computation cost of the method is high as the number of distillation process is O(M^2) which exponentially increases with the prototype number.\n\n4. The changes compared to [6] are minor. Initially, I believed this method could be simple and efficient. However, upon reviewing the weight disentanglement property, I have concerns about its practical validity, which limits the novelty of the approach.\n\n5. A minor issue is only apply to data with prototype label, thus may limited its impact\n \n[1] TIES-Merging: Resolving Interference When Merging Models\n\n[2]  Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch\n\n[3] Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models\n\n[4] Adaptive Multi-Teacher Multi-level Knowledge Distillation\n\n[5] Agree to Disagree: Adaptive Ensemble Knowledge Distillation in Gradient Space\n\n[6] Ensemble distillation for robust model fusion in federated learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your invaluable review and positive remarks on the strengths of our paper in **improving existing KD methods that are inadequate in diverse device heterogeneous environments** and **comprehensive experiments**.\n\n### **Response to 1**\nStudies [1,2] primarily address parameter interference and redundancy, proposing methods to mitigate this. These studies focus on parameter interference, not weight disentanglement. Parameter interference is not equivalent to weight entanglement.\n\nStudy [3] (NeurIPS’23), our cited reference (reference [37]), provides theoretical insights into task arithmetic, introducing weight disentanglement as a necessary condition. They show that “task arithmetic in non-linear models cannot be explained by NTK,” and weight disentanglement is actually the sole requirement for task arithmetic “where distinct directions in weight space correspond to changes of the network in disjoint regions of the input space ... This allows a model to perform task arithmetic by independently manipulating these weight directions.” They also demonstrate that weight disentanglement is an emergent property of pre-training and propose fine-tuning models in their tangent space to enhance this property. Please note that the quotes are the original text from [3]. Therefore, asserting weight disentanglement property is not problematic per [3].\n\nIn our theoretical framework, we have cited [3] in our paper for weight disentanglement as it is the established theoretical framework to study task arithmetic. Our theoretical framework makes the following contributions: (1) It present the first theoretical framework using the concept of capacity dimension to understand the effectiveness of KD in device heterogeneous FL. (2) Our theoretical analysis of TAKFL is based on the theory in [3] (3) It presents an argument consistent with novel numerical experimental findings, showing that TAKFL performs KD without information loss more effectively.\n\n### **Response to 2**\nWe used vanilla average logits KD in our methodology for knowledge distillation from individual prototypes because it is the most representative and canonical form of KD. This avoids any bias towards the advantages or disadvantages of specific KD implementations, making our comparisons more transparent and informative.\n\nRegarding [4,5], these methods are designed for KD in a centralized learning setting, where teachers are fully trained, and the distillation dataset aligns with actual learning objective, including a cross-entropy supervision loss. However, in FL, the ensembles are noisy, and the distillation dataset differs from the private dataset, lacking a supervision cross-entropy loss (Section 5.1).\n\nWith these considerations, we have compared TAKFL by directly implementing [4,5] in FL as well as incorporating them into TAKFL instead of vanilla KD. The results are presented in Table 2 of the PDF rebuttal. As seen, using [4,5] directly in FL performed poorly, while incorporating [5] into TAKFL achieved better performance in low data heterogeneity compared to TAKFL+Vanilla KD. TAKFL is independent of the specific KD method used. We will incorporate these comparisons into the final version of our paper.\n\n### **Response to 3**\nWe have presented the full algorithm details of TAKFL in Appendix B. All the distillation processes are performed in parallel and asynchronously (lines 9-10 in Algorithm 1). Therefore, the overall computation cost of the entire server-side distillation process is O(M), as device prototypes do not perform KD one at a time but combine information asynchronously into a singe task arithmetic operation (Eq. 8). Additionally, in FL, the server typically has substantial resources to manage the orchestration of many devices, making server-side computation cost not a concern.\n\n### **Response to 4**\nWe highlight the key distinctions and novelties of TAKFL compared to [6]:\n\n(1) [6] overlooks the diversity in device capabilities and applies uniform logit averaging, causing dilution and information loss. TAKFL addresses this by treating knowledge transfer from each prototype's ensembles as separate tasks, distilling them independently to ensure unique contributions are effectively captured.\n\n(2) [6] uses average logits as the distillation target for all student models, failing to provide customized knowledge integration. TAKFL employs a novel task arithmetic approach, allowing each student model to customize knowledge integration based on each prototype's quality and its own capacity, enhancing performance.\n\n(3) [6] neglect challenges associated with noisy and unsupervised ensemble distillation. TAKFL incorporates a novel KD-based self-regularization technique to mitigate these issues.\n\nOur theoretical and experimental findings highlight the deficiencies of [6] in diverse device heterogeneous FL (Remark 1, Proposition 1, Section 7.2, Appendix D.1.1). Our extensive experiments show that TAKFL not only outperforms existing methods but also offers a novel, simple, and efficient solution to a significant practical challenge. In regards to the weight disentanglement property, kindly please see our response to #1 above.\n\n### **Response to 5**\nIt seems that the reviewer is concerned about the impact of TAKFL. We have comprehensively evaluated TAKFL for CV and NLP tasks, including various datasets with differing architectures and data heterogeneity levels. We also introduced a new scalability evaluation showing TAKFL’s adaptability across devices ranging from XXS to XXL, underscoring its broad applicability and consistent performance improvements over existing methods. We believe that the impact of our method extends beyond specific datasets and prototype sizes, providing a versatile solution for diverse device heterogeneous federated learning applications. If this does not address the reviewer’s comment, could the reviewer please clarify and elaborate? The term “data with prototype label” is unclear to us.", "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": "", "extra_scores": {}}, {"review_id": "LQYPvstoHB", "reviewer": "Reviewer_ME5v", "summary": "The authors aim to develop a knowledge distillation method that addresses the challenges posed by heterogeneous device prototypes in federated learning. By capturing the knowledge transfer among device prototypes, the proposed TAKFL tries to preserve each device's unique contribution and prevent knowledge dilution during the learning procedure. The method incorporates a self-regularization technique to address issues of the noisy and unsupervised ensemble distillation. Evaluation of some CV and NLP tasks shows the performance of model accuracy and scalability.", "questions": "Please refer to the concerns and questions mentioned in the weakness part.", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The authors focus on knowledge distillation in heterogeneous settings, which is meaningful to real-world tasks. The overall presentation is clear and easy to understand. The proposed TAKFL shows some primary theoretical analysis of the learning efficiency. The evaluation provides a quantitative analysis to compare model accuracy and scalability with previous methods.", "weaknesses": "While the paper presents a clear presentation and evaluation, there are a few aspects to be strengthened. \n\n1. The concept of \"task arithmetic\" in the context of federated learning is vague. It would be beneficial to explain how this concept enhances the design of the federated learning process.   \n2. While the knowledge distillation process is described, it appears to be largely based on existing methodologies. It would be interesting to explore any novel design elements introduced in TAKFL. Additionally, what is the theoretical convergence order of the proposed method?  \n3. Although TAKFL demonstrates higher model accuracy in performance comparisons, the baselines used, such as FedAvg and FedDF, seem outdated. Incorporating more recent baselines could provide a more rigorous evaluation of TAKFL's effectiveness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your invaluable review and pointing out that our proposed method shows **some primary theoretical analysis of the learning efficiency**.\n\n### **Response to 1**\nTo clarify, we use the concept of task arithmetic to address limitations of previous approaches and enhance knowledge transfer across heterogeneous device prototypes in FL. We extend the notion of “task vector” from centralized learning to FL by considering the averaged locally updated parameters as the “pre-trained” model parameters and the distilled model parameters as the “fine-tuned” parameters (refer to Figure 2 in paper). Task vectors encapsulate the distinct contributions of each prototype's ensembles to the student model, representing transferred knowledge (Section 5.2).\n\nOur task arithmetic approach enhances the design of knowledge transfer across diverse heterogeneous device prototypes in FL by:\n\n**(1) Effective Distillation of Unique Knowledge:** Task arithmetic facilitates the distillation of unique knowledge from each prototype’s ensembles. Prior works neglect individual strengths and average logits, causing dilution and information loss (Sections 4, 5.1, Remark 1).\n\n**(2) Customized Knowledge Integration:** Our adaptive task arithmetic operation (Eq 8) allows the student model to customize knowledge integration based on its capacity and the helpfulness of other prototypes, enhancing performance. Prior works use a one-size-fits-all approach, ignoring the student model’s capacity and prototypes’ helpfulness (Sections 4, 5.2, Remark 2).\n\nWe have further verified the enhancements of our task arithmetic approach through both theoretical analysis and extensive experiments (Sections 6, 7, D.1.1). To the best of our knowledge, our work is the first of its kind to introduce the notion of “task vector” in FL to enhance knowledge transfer across diverse heterogeneous device prototypes.\n\n### **Response to 2**\nWe highlight the key distinctions and novelties of our knowledge distillation methodology compared to existing methods:\n\n**(1) Diversity in Device Capabilities:** Existing methods often overlook device capability diversity and apply logit averaging, causing dilution and interference. TAKFL addresses this by treating knowledge transfer from each prototype as separate tasks, distilling them independently to ensure unique contributions are captured without interference.\n\n**(2) Customized Knowledge Integration:** Existing methods use a single integrated distillation target for all student models, failing to provide customized knowledge integration. TAKFL employs task vectors in FL as a unique representation of transferred knowledge from each prototype. This enables the student model to customize knowledge integration based on each prototype's knowledge quality and its own capacity via adaptive task arithmetic (Eq 8), enhancing performance.\n\n**(3) Handling Noisy and Unsupervised Distillation:** Existing methods often neglect challenges associated with noisy and unsupervised ensemble distillation, which may cause the student to forget its own knowledge and drift into erroneous directions. TAKFL incorporates a novel KD-based self-regularization technique to mitigate these issues.\n\nRegarding the theoretical convergence, TAKFL consists of two processes in each round: (1) per-prototype FL and (2) server-side knowledge distillation and integration via task arithmetic. The convergence order of per-prototype FL is the same as, FedAvg [1], with various per-prototype FL optimizer options, and the server-side KD is similar to centralized neural network training  [2], as the KL loss has the same functional form as cross-entropy. Thus, the complete task arithmetic model merging convergence is simply the sum of these two procedures. We will include this in a Remark in the final revision of our paper.\n\nWe have also obtained the convergence plot for (CIFAR-10, Dir(0.3)) in Table 1 of the paper, included in Figure 2 of the PDF rebuttal. TAKFL shows similar convergence behavior to the baselines while achieving better performance. We will include this plot in the final version of the paper.\n\n### **Response to 3**\nWe appreciate the reviewer highlighting TAKFL's higher performance compared to the baselines. Existing body of works on global device heterogeneous FL consists of two distinct approaches and problem formulations:\n\n**(1) Partial Model Training Methods:** These aim to train a single global model while accommodating heterogeneous sub-models at the devices according to their compute resources, with several SOTA baselines.\n\n**(2) Knowledge Distillation Methods:** These address a more practical scenario where device prototypes with heterogeneous model architectures participate in FL to enhance their global model performance through knowledge distillation. The current SOTA baselines are FedDF and FedET, which we have comprehensively compared in our work. These methods mainly focus on settings where device prototypes have similar capabilities.\n\nOur work follows (2) and introduces TAKFL to address limitations of existing methods in the *underexplored* diverse device heterogeneous settings. We believe our work can significantly contribute to the FL literature and establish a new SOTA benchmark. For a detailed discussion on related works, we kindly refer the reviewer to Appendix A, where we reference recent surveys from 2023 and 2024.\nTo further evaluate TAKFL's effectiveness, we experimented using FedOpt [3], a more advanced FL optimizer for per-prototype FL, and included the results in Table 3 of the PDF rebuttal. As shown, TAKFL achieves SOTA performance results independent of the specific FL optimizer used. We will include this experiment in the final version of our paper.\n\n### References\n[1] Haddadpour, et al. \"Local sgd with periodic averaging: Tighter analysis and adaptive synchronization.\" NeurIPS’19\n\n[2] Aggarwal. “Neural networks and deep learning”. Springer (2018)\n\n[3] Reddi et al. \"Adaptive Federated Optimization.\" ICLR’21", "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": "", "extra_scores": {}}]

Author rebuttal:
We sincerely appreciate all the reviewers' efforts in reviewing and commenting on our work. We are particularly grateful for the positive feedback highlighting the following aspects:

* **The novel theoretical model and framework illustrating the efficacy of knowledge distillation in heterogeneous federated learning (Reviewer V7t1, Reviewer RvJG, Reviewer ME5v).**
* **Effective handling of challenges in transferring knowledge across heterogeneous devices in federated learning and improving existing knowledge distillation methods (Reviewer RvJG, Reviewer F7vu).**
* **Strong experimental validation on CV and NLP tasks with various datasets, demonstrating superior results over existing methods (Reviewer V7t1, Reviewer RvJG, Reviewer F7vu).**
* **Addressing significant real-world challenges in federated learning with meaningful theoretical and practical contributions (Reviewer V7t1, Reviewer RvJG, Reviewer ME5v).**
* **Clear and well-organized presentation, making the paper easy to follow and understand (Reviewer V7t1, Reviewer RvJG, Reviewer F7vu, Reviewer ME5v).**

In summary, we have introduced a novel framework, TAKFL, to address the fundamental limitations of existing knowledge distillation (KD) methods in the *underexplored* diverse device heterogeneous federated learning environments, ranging from small IoT devices to large workstations. Our framework effectively distills the unique knowledge of each device prototype's ensembles by treating them as separate tasks and adaptively integrates them into the student model via a novel task arithmetic approach. Unlike existing methods, TAKFL avoids dilution and information loss, providing customized knowledge integration for each student model based on the knowledge quality and helpfulness of each prototype's ensembles and its intrinsic capacity. Furthermore, in response to the issues associated to noisy and unsupervised ensemble distillation, we have incorporated a novel KD-based self-regularization loss. Our work makes the following key contributions:

* **Framework:** Introducing a new framework to address the fundamental limitations of existing methods in diverse device heterogeneous FL called TAKFL.
* **Theoretical Contribution:** Presenting the first theoretical framework using the concept of capacity dimension to understand the effectiveness of KD in device heterogeneous FL.
* **Experimental Validation:** Conducting comprehensive experimental evaluations for both CV and NLP tasks with various datasets, architectures, and levels of data heterogeneity, demonstrating the effectiveness of our method and achieving state-of-the-art performance by significantly outperforming existing methods. Additionally, we include a new scalability experiment encompassing devices from XXS to XXL to further demonstrate the efficacy of our method.

Our extensive experiments show that TAKFL not only outperforms existing methods but also offers a novel, simple, and efficient solution to a significant practical challenge. Furthermore, to the best of our knowledge, we are the first to introduce the notion of "task vector" in FL to enhance knowledge transfer across diverse heterogeneous device prototypes. We have derived new consistent theoretical and experimental findings detailed in Appendix D.1.1. We believe our work makes a significant contribution to the FL field (as pointed out by Reviewer RvJG), addresses an important practical challenge (as highlighted by Reviewer ME5v), and establishes a new benchmark. We believe that the impact of our method extends beyond specific datasets and prototype sizes, providing a versatile solution for diverse device heterogeneous federated learning applications.

We appreciate the reviewers' constructive comments and have strived to comprehensively address them. To this end, we have provided at least one additional experiment, table, and/or figure for each reviewer:

* In response to Reviewer V7t1's comments #1 and #3, we have provided Table 1 and Figure 1 in the PDF rebuttal.
* In response to Reviewer RvJG's comment #3, we have provided Figure 3 in the PDF rebuttal.
* In response to Reviewer F7vu's comment #2, we have provided Table 2 in the PDF rebuttal.
* In response to Reviewer ME5v's comments #2 and #3, we have provided Figure 2 and Table 3 in the PDF rebuttal.

The PDF rebuttal file containing the associated figures and tables is uploaded with this message.

Once again, we sincerely appreciate the reviewers' time and effort in reviewing our work. If our response comprehensively addresses the reviewers' comments, we kindly ask the reviewers to consider raising their scores. We are happy to address any remaining comments or concerns during the author-reviewer discussion period.

Thank you!

Meta-review: The paper is of high quality, and there is not objection from four reviewers. It requires minor refinements based on the reviewers' comments.
Decision: Accept (poster)