Towards Diverse Device Heterogeneous Federated Learning via Task Arithmetic Knowledge Integration

Authors: mahdi  morafah
vyacheslav  kungurtsev
hojin  chang
chen  chen
bill  lin
!  ! " ! ! " " " ! # " # # $%&'%( # )*' #
  ' " " # $%&'%( = # )*' + (! ! " ! + ! " " " + ! # " # )
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 09:28:57

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 .

Main paper: 
 introduction: 
 Federated Learning (FL) has rapidly gained traction as a promising approach to train machine learning models collaboratively across multiple devices, while preserving the privacy of user data. Standard federated learning methods, such as FedAvg [33] , however, are primarily designed for unrealistic 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, 31] . Therefore, standard FL cannot support the participation of heterogeneous devices, all of which could significantly contribute to model training due to their unique and invaluable local datasets. To address this gap, knowledge distillation (KD) techniques have emerged as 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 server student model using an unlabeled public dataset.Despite their success, however, existing KD-based methods for device heterogeneous FL are primarily 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 IoTs 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 heterogeneous settings, primarily due to two reasons: 1 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. 2 Additionally, these methods employ this single integrated distillation target to transfer knowledge across all different size student models. This one-size-fits-all 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 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: 1 The distillation process introduces noise, as the ensembles' logits are inferred on an unfamiliar public dataset, 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 a more effective knowledge transfer framework.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-both model and dataset-and consequently, in strength. TAKFL treats knowledge transfer from each device prototype's ensembles as separate tasks, distilling them independently to ensure that each prototype's unique contributions are accurately distilled without interference. To tackle the challenges associated with noisy and unsupervised ensemble distillation, we incorporate a KD-based self-regularization technique into this individual knowledge transfer process. Subsequently, to selectively 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 ensembles and its intrinsic capacity, aiming to achieve optimal performance. We present theoretical results, grounded on the established theoretical learning properties of overparametrized 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 efficiency, 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 natural language processing (NLP) tasks, utilizing various datasets and architectures, and demonstrate that TAKFL consistently achieves state-of-the-art (SOTA) performance.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 limitations of existing approaches, effectively facilitating knowledge transfer across diverse heterogeneous device prototypes with varying capabilities. 3. We present the first theoretical model for device heterogeneous KD, and demonstrate the effectiveness 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, outperforming existing KD-based methods. Figure 1 : Overview of our approach and its distinction from prior works. (a) This figure illustrates the vanilla ensemble distillation process, where logits from ensembles of various sizes are 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 ensemble as a separate task and distills them independently. Additionally, a KD-based self-regularization technique is introduced to mitigate issues related to the noisy and unsupervised ensemble distillation. Finally, the heterogeneously distilled knowledge is strategically integrated using an adaptive task arithmetic operation, allowing for customized knowledge integration based on each student prototype's needs. 
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 accommodating devices with varying compute resources, aiming to train a single global model. Techniques such as static and rolling-based 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 IoTs 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 a more practical scenario where device prototypes with heterogeneous model architectures 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 knowledge into each server's student model using an unlabeled public dataset. For example, FedDF [30] uses vanilla logit averaging, while Fed-ET [6] applies an uncertainty-weighted logit averaging, enhanced by a diversity regularization technique. However, existing works typically focus on settings where prototypes have similar capabilities-both model and dataset sizes-and thus neglecting the challenges in more diverse settings with varying capabilities. This oversight leaves their effectiveness in such settings largely unexplored. In this paper, we aim to study the underexplored diverse heterogeneous device settings. See Appendix A for a more detailed discussion on the related works.Model Editing via Task Arithmetic. Traditional methods for model editing often involve expensive 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 scalable method for updating pre-trained models with new information or refining undesired behavior [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 t, a task vector is derived by subtracting the model's pre-trained weights θ pre from its fine-tuned weights θ t f t on task t , i.e. τ t = θ f t -θ pre . These task vectors act as unique representations for specific tasks. Furthermore, researchers have demonstrated that by summing multiple task vectors {τ t } T t=1 , and integrating them into a pre-trained model via θ = θ pre +λ T t=1 τ t , one can effectively create a model capable of handling multiple tasks [53, 57] .To the best of our knowledge, this work is the first to extend the notion of task vectors to the federated learning setting, introducing a task arithmetic for knowledge distillation across diverse heterogeneous device prototypes.that Q 1 > Q 2 and W 1 = W 1,2 > W 2 .Then the same quantity as for Case 1 holds.Moreover,dim ΘV ED ∩ Q 1 \ (W 1 ∪ W 1,2 = O (Q 1 -W 1 )(W 1 !)(W 1,2 -W 2 )! Q 1 !W 1 !(Q 1 -W 1 )! + Q 1 !W 2 !(W 1,2 -W 2 )!This corresponds to capacity of client 1 that has been allocated but fits, in the model of prototype 1, neither the data of prototype 1, nor of the data of prototype 2.Proof. Formally,ΘV ED := argmin θ∈Q 1 L ED = argmin θ KL i=1,2 σ f i (x, θ i avg ) , σ (S(x))Since by assumption θ i avg solves the training problem on the data associated with device prototype i, the logit is accurate, and thus there is a map O(i, j) : T i → T j i ⊆ W i,j . The self distillation, that is, S j defines a bijective map from W j to W j and thus does not affect the capacity allocation.Case 1: In this case, generically (that is, measure zero on some non-atomic sampling on a distribution of operators) O(i, j) is bijective. Now let us compute the expectation of the number of eigenvectors of, e.g. W 1,2 that are in the complement of the span of W 1 . Assuming, for simplicity, independence, this would correspond to counting the possible choices within the capacity of Q 1 \ W 1 over the range of possible choices of filling the capacity of Q 1 with vectors in W 1 together with choices of filling it with vectors in W 1,2 :Q 1 -W 1 i=1 i Ch(Q 1 -W 1 , i) Ch(Q 1 , W 1 ) + Ch(Q 1 , W 1,2 )For, e.g., Q 1 = 4 and W 1 = W 1,2 = 2 this is 1 3 . To derive a scaling rate we can write:i i (Q 1 -W 1 )! i!(Q 1 -W 1 -i)! Q 1 ! (W 1 )!(Q 1 -W 1 )! + Q 1 ! W 1,2 !(Q 1 -W 1,2 )! = O (Q 1 -W 1 )(W 1 )!(Q 1 -W 1,2 )! Q 1 !(W 1 )!(Q 1 -W 1 )!+Q 1 !W 1,2 !(Q 1 -W 1,2 )!Case 2: In this case, it must be that, at best almost surely, O(2, 1) is injective, but not surjective. This means that distilling from 2 to 1 does not fill the capacity of W 1,2 , and is thus a fundamentally wasteful operation, that is|T j i | = W 2 < W 1,2. Now let us compute the expectation of the number of eigenvectors of, e.g. W 1,2 that are in the complement of the span of W 1 . Since W 1,2 are being structurally allocated for fitting, the combinatorial expression is the same:Q 1 -W 1 i=1 i Ch(Q 1 -W 1 , i) Ch(Q 1 , W 1 ) + Ch(Q 1 , W 1,2 )Thus for, e.g., Q 1 = 4 and W 1 = W 1,2 = 2 this is, again, 1 3 . The scaling in this case isO (Q 1 -W 1 )(W 1 !)(Q 1 -W 1,2 )! Q 1 !W 1 !(Q 1 -W 1 )! + Q 1 !W 2 !(Q 1 -W 2 )!However, we observe that there are vectors in the range of W 1,2 \ O(2, 1)(W 2 ) that have been allocated by the VED but lie in neither W 1 nor in W 1,2 , that is, are garbage. We can compute those as the expected number of eigenvectors arising from allocating W 1,2 \ O(2, 1)(W 2 ) that intersect with Q 1 \ W 1 (that is, the spare capacity not used for fitting data D 1 ). This is, using similar principles:W 1,2 -W 1 i=1 i Ch(Q 1 -W 1 , i) Ch(Q 1 , W 1 ) + Ch(Q 1 , W 1,2 -W 2 )This is for, e.g., Q 1 = 4, W 1,2 = 2 and W 2 = 1, this would be 3 10 The scaling here isO (Q 1 -W 1 )(W 1 !)(W 1,2 -W 2 )! Q 1 !W 1 !(Q 1 -W 1 )! + Q 1 !W 2 !(W 1,2 -W 2 )! ■ 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 1 , Q 2 and W 1 , W 2 , respectively, and with associated eigenbases Q i , W i . Let the solution set of TAKFL with prototype i as student be Θi T A with dim( Θi T A ) = W v with eigenbasis W v . In addition, denote W s,t , s, t ∈ {1, 2} dimension of the solution set for the student model trained on the data from the teacher device's ensembles. . The following statements hold:In the case that that Q 1 ≥ Q 2 and W 1 ≥ W 2 , it holds that the TAKFL preserves that the eigenbasis used to model the data D 1 's accuracy for device prototype 1, that is for student1 dim W v ∩ Q 1 \ W 1 = 0 Case 1: Assume that Q 1 = Q 2 and W 1 = W 2 . Then it holds that, dim W v ∩ Q 1 \ (W 1 ∪ W 1,2 = 0 Moreover, it holds that, ΘT A ∈ Span(W 1 ∩ W 1,2 )Thus, with equal capacity, no information is lost in Task Arithmetic aided knowledge ensemble distillation and capacity is efficiently used to model the data from both prototype 1 and prototype 2.Case 2: Assume that Q 1 > Q 2 and W 1 > W 2 .Then it again holds that,dim W v ∩ Q 1 \ (W 1 ∪ W 1,2 = 0 However, while ΘT A ∈ Span(W 1 ), it holds that dim W v ∩ W 1,2 = W 2 < W 1,2 .Proof. We consider the case that the prototype i and j datasets be disjoint, i.e. D i ∩ D j = ∅, andM i=1 D ⟩ = D.We shall see how the other cases are easier and the result remain the same. We can see immediately from the weight disentanglement property of Task Arithmetic that,f 1 (x; θ 1 avg + α 1 τ 1 + α 2 τ 2 ) = g 1,1 (x; α 1 τ 1 ) + g 1,2 (x; α 2 , τ 2 ) + g 1,0 (x)with g 1,1 (x; α 1 τ 1 ) for x / ∈ D 1 , g 1,2 (x; α 2 τ 2 ) for x / ∈ D 2 and g 1,0 (x) = 0 for x ∈ D 1 ∪ D 2 . From this, we can immediately conclude the first statement of the Proposition as well as the expressiondim W v ∩ Q 1 \ (W 1 ∪ W 1,2 = 0 and also, in the case of W 1 = W 1,2 = W 2 implies ΘT A ∈ Span(W 1 ∩ W 1,2 )For the last statement we observe again as in the second Case in the Proposition describing VED, dim O(2, 1)(W 2 ) < W 1,2 from which we can conclude that, genericallydim v : v ∈ O(2, 1)(W 2 ) ⊆ W v , E (x,y)∼D 2 l(f 1 (x; v), y) > 0 = W 1,2 -W 2proving the final statement. ■We observe that a key mechanism of the proof is the dimension of the target space of the teaching operator O(i, j). As an informative model, we can consider coefficients λ j of task vectors as restricting the rank, relative to other teachers. For instance, in the previous Proposition, if W 1,2 = 2 and W 2 = 1, then λ 2 = 1/2, so as to enforce one vector of W 1,2 is a target for the map Õ(2, 1), would be appropriately sensible.The prototype i and j datasets are fully overlapping, i.e. D i ⊆ D j , meaning that the two prototypes have the same classes D c , ∀c ∈ D, then these are the changes to the the proposition 1 and 2 in our theory.Consider that in this case, the logits corresponding to W 1,1 and W 1,2 are the same. As such W 1,2 ∈ Q 1 , and so the propositions change to, trivially: Proposition 3. (Information Loss in VED). Consider the VED procedure. Consider two device prototypes with a device capacity and solution dimension of Q 1 , Q 2 and W 1 , W 2 , respectively, and with associated eigenbases Q i , W i . Let the solution set of VED with prototype i as student be Θi V ED with dim( Θi V ED ) = W vi with eigenbasis W vi . In addition, denote W s,t , s, t ∈ {1, 2} 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.,W 1,1 = W 1 and W 2,2 = W 2 . 1. Case 1: Assume that Q 1 = Q 2 and W 1 = W 2 = W 1,2 = W 2,1 . Then it holds that, in expectation, dim Θ1 V ED ∩ Q 1 \ W 1 = 0This corresponds to the expected capacity of prototype 1 that is taken up for fitting logits that do not fit the data corresponding to prototype 1.that Q 1 > Q 2 and W 1 = W 1,2 > W 2 .Then the same quantity as for Case 1 holds. Moreover,dim ΘV ED ∩ Q 1 \ (W 1 ∪ W 1,2 = 0This corresponds to capacity of client 1 that has been allocated but fits, in the model of prototype 1, neither the data of prototype 1, nor of the data of prototype 2. Now consider that the prototype i and j datasets are partially overlapping, i.e. D i ∩D j ̸ = ∅, meaning that the two prototypes have the overlap over some of the classes D c , then these are the changes to the the proposition 1 and 2 in our theory.In this case,W 1,2 = W 1,2 1 ∪ W 1,2 2, that is, the class labels across data prototypes partially overlap (with W 1,2 1 the overlapping dimensionality), the Proposition for VED changes to: Proposition 4. (Information Loss in VED). Consider the VED procedure. Consider two device prototypes with a device capacity and solution dimension of Q 1 , Q 2 and W 1 , W 2 , respectively, and with associated eigenbases Q i , W i . Let the solution set of VED with prototype i as student be ΘiV ED with dim( Θi V ED ) = W vi with eigenbasis W vi . In addition, denote W s,t , s, t ∈ {1, 2} 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., The planes represent the solution subspaces. The right panel presents a circumstance by which two small device prototypes (2, and 3) serve as teacher for transferring knowledge to a larger device prototype 1.W 1,1 = W 1 and W 2,2 = W 2 . θ 1 avg λ 1,1 τ 1,1 + λ 1,2 τ 1.2 θ 1,2 avg VED Θ 1 Θ 1,2 θ 2 avg θ 1,2 avg VED λ 1,1 τ 1,1 + λ 2,1 τ 2,1 Θ 2 Θ 2,1 θ 1 avg θ 1,3 avg θ 1,2 avg Θ 1 Θ 1,2 Θ 1,3EQUATION2 )! This corresponds to the expected capacity of prototype 1 that is taken up for fitting logits that do not fit the data corresponding to prototype 1.that Q 1 > Q 2 and W 1 = W 1,2 > W 2 .Then the same quantity as for Case 1 holds. Moreover,dim ΘV ED ∩ Q 1 \ (W 1 ∪ W 1,2 = O (Q 1 -W 1 )(W 1 !)(W 1,2 2 -W 2 )! Q 1 !W 1 !(Q 1 -W 1 )! + Q 1 !W 2 !(W 1,2 2 -W 2 )!This corresponds to capacity of client 1 that has been allocated but fits, in the model of prototype 1, neither the data of prototype 1, nor of the data of prototype 2.Note that the presence of overlapping logits W 1,2 1 do not appear in the result, it is exactly the same as the original Proposition 1 in the original paper just withW 1,2 2 in place of W 1,2 .The Propositions for TAKFL remain the same.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: The limitation of the work is discussed in 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 are not discussed in the paper. • The authors are encouraged to create a 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 authors should reflect on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
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 = |M|. Each device prototype m j ∈ M has a distinct neural network architecture f j (•; θ j ) parameterized by θ j ∈ R nj and a set of clients C j , withN j = |C j | clients in total. Each client c k ∈ C j has a local private dataset D j k = {(x i , y i )} n j,ki=1 , where n j,k = |D j k |, and locally trains the parameters θ j of the neural network architecture f j on its local dataset. Furthermore, denote D j = ∪ k∈C j D j k to be the union of the private datasets for device prototype j. We assume D j ∼ D j , that is a subsample from the population distribution D j and similarly D j k ∼ D j k . The union of the private datasets, i.e. D = j∈M D j , 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., D = j ω j,: D j = j k ω j,k D j k , where 0 ≤ ω j,k ≤ 1 and jk ω j,k = 1, and ω j,k is unknown.The ultimate objective is to minimize the test error and thus enable accurate inference for each device prototype j, aiming to obtain the optimal parameters for the population dataset:EQUATIONwhere ℓ( The resource usage scales as O(M 2 ) (quadratic) because of the need to store and process multiple task vectors concurrently. 4. Optimization Convergence Rate: The convergence rate of the Knowledge Distillation procedure is as standard for training a nonconvex loss function, e.g. [42] .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: 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.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 crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems. • The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
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 of 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. In this context, at each round r a subset of clients C j r from each device prototype j ∈ M is randomly selected by the server and download their corresponding model initialization θ j r . Each client c j k ∈ C j r , starting from this model initialization, locally train the model f j on its local private data D j k by taking multiple steps of stochastic gradient descent. Then, they send back their updated parameters { θ j k } k∈C j r to the server. The server aggregates the received clients parameters, and computesθ j avg = k∈C j r |D k | k∈C j r |D k | θ j k .In classic federated learning formalism, the parameters θ j avg satisfy,EQUATIONVanilla Ensemble Distillation. In this stage, each server model f j gets initialized with θ j , 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 i (•, θ i k )| i ∈ M, k ∈ C i },transferring their knowledge to each server student model, i.e. S i := f i (•, θ i ). 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.e x ∈ D public , can be defined by the following equation:EQUATIONwhere σ(•) 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 prototype's ensembles. As a 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, irrespective of their intrinsic capacity and the helpfulness of each prototype's ensembles. Consequently, this leads to suboptimal knowledge transfer in device heterogeneous FL. See Section 6 for theoretical analysis and Section 7 for experimental observations.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 conclusions 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. 
task arithmetic knowledge transfer and integration: 
 In this section, we introduce TAKFL, 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 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 in Figure 1b , and the full algorithm is detailed in Appendix B, Algorithm 1.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 transfer and self-regularization, each detailed in the subsequent paragraphs.Ensemble Knowledge Transfer. Vanilla ensemble distillation integrates the knowledge of varying strength ensembles by uniformly averaging their logits. This approach can potentially transform 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 prototype'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 considerT i := {f i (•, θ i k )| k ∈ C i }denotes the ensembles of device prototype i as teacher and S j 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 T i and server student S (T i → S) is defined below:EQUATION)Scaffolding Student from Noisy Ensemble Distillation. The ensemble distillation process may adversely impact the student, causing it 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 are inferred on an 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 as 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 a KD-based self-regularization technique. Our self-regularization 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:EQUATIONOverall Knowledge Transfer Objective. The overall knowledge transfer objective from teacher ensembles T i of device prototype i to the student S is the combination of the ensemble knowledge distillation loss L Ti )S KD (Eq. 4) and the self-regularization loss L self S (Eq. 5) defined in the following:EQUATION) Here, γ is a hyperparameter controlling the effect of self-regularization term. We associate the knowledge transfer from each device prototype i to a task T i with the loss L Ti S .# 0&% # 1- # )*' # (+,-+..%(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 taskspecific vectors using arithmetic operation, we propose a novel customizable knowledge integration method via task arithmetic. To do so we extend the notion of task vector from centralized learning to federated learning. We conceptualize the averaged locally updated parameters, i.e. θ avg , as a "pretrained", similar to those in centralized learning, and the parameters of the distilled model via knowledge transfer objective (Eq. 4), denoted as θ distilled , as a "fine-tuned" version of the model (see Fig. 2 ). Consequently, the task vector τ i associated with the knowledge transfer task L Ti S can be defined by subtracting the distilled parameters from the averaged locally updated parameters as follows:" = # (+,-+..%( -# )*' " = # 1--# 0&%EQUATIONEssentially, task vectors serve as unique representations for the transferred knowledge from each prototype's ensembles to 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:EQUATIONwhere λ i denotes the merging coefficient associated with task vector τ i , and they sum to one, i.e.i∈M λ i = 1. The merging coefficients determine the extent of knowledge integration from each prototype's ensembles. Essentially, they enable the student to have customized knowledge integration 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 prototypes' 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 it manually or determined it using held-out validation sets which achieves similar results. More details can be found in Appendix F.3.• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is 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 are out of range (e.g. negative error rates). Justification: This paper focuses on advancing technical methodologies in federated learning and does not directly engage with specific societal impacts, as the applications of the proposed method are 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 is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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 specific 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 as 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 mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
theoretical results: 
 We present a theoretical understanding on the efficacy of knowledge distillation in device heterogeneous 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 a device prototype of small capacity, and the student model is from a 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 overparametrized neural networks, that is, the total number of weights far exceeds the training sample size. This implies that the set of weights that minimize the loss is non-unique, and moreover, it has been argued that they form a 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 as possible.Since we are interested in 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.e. D i ∩ D j = ∅, and M i=1 D i = 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 a device capacity and solution dimension of Q 1 , Q 2 and W 1 , W 2 , respectively, and with associated eigenbases Q i , W i . Denote W i,j , i, j = 1, 2 as the capacity allocated by student i in order to distill knowledge from teacher j's logits.1. Case 1: When the capacities are the same, that isQ 1 = Q 2 and W 1 = W 2 = W 1,2 = W 2,1 ,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 its own data distribution. 2. Case 2: Assume that Q 1 > Q 2 and W 1 = W 1,2 > W 2 . Then the phenomenon as for Case 1 holds. Moreover, there will be some capacity of student 1's model that will be allocated with parameters that do not minimize either of the teacher or student prototype'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 1,2 is structurally reduced to that of W 2 < W 1,2 , 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 to inefficient and inaccurate use of model capacity. Furthermore, under the case that device prototypes have different capacities, VED ends up leading to more erronous models entirely as the small information within the small teacher is transferred onto a larger capacity target.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 a device capacity and solution dimension of Q 1 , Q 2 and W 1 , W 2 , respectively, and with associated eigenbasesQ i , W i . 1. Case 1: In the case that that Q 1 ≥ Q 2 and W 1 ≥ W 2, it holds that the TAKFL with prototype 1 as student preserves the eigenbasis associated to the parameters used to accurately fit the data D 1 . 2. Case 2: Assume that Q 1 = Q 2 and W 1 = W 2 . TAKFL yields a 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 1 > Q 2 and W 1 > W 2 . In the case of prototype 1 being the student, TAKFL yields a solution that: (a) retains the approximation accuracy on device 1'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 no informative new knowledge, unless enforced otherwise.Remark 2. This proposition proves that in 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 prototype'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 a small device prototype serving as teacher will not be necessarily compromise information, it would still be preferable to allocate that capacity to a more informative, larger, teacher model.Comments. We comment on the complexity and convergence aspects of TAKFL briefly:1. Computation Time: TAKFL's computation time is O(1) (constant) due to parallelization, as all distillation processes occur simultaneously. 2. Computation Load: The overall computational load scales as O(M ) (linear) since the distillation tasks are performed independently for each prototype in parallel and merged into a singe task arithmetic operation (Eq. 8). 
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 CIFAR10/100 [24] , CINIC-10 [9], and TinyImagenet [25] . For NLP, 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(α), where a lower α indicates higher heterogeneity [27, 36] . Further details can be found in Appendix F. Baselines and Evaluation Metric. We compare our method against standard FL, i.e. Fe-dAvg [33] and SOTA KD-based methods designed for heterogeneous device prototypes FL, including FedDF [30] and FedET [6] . The evaluation metric is the final top-1 classification accuracy of each device prototype'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 a different random seed.A more detailed version of the experiments, alongside additional experiments and ablation studies, is presented in Appendix D and E.In this section, we evaluate the performance of our method, TAKFL, in a federated learning environment 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 a small model and small dataset, Medium (M) with a medium-sized model and medium-sized dataset, and large (L) with a large model and large dataset.Performance on CV Task. Table 1 presents the performance of TAKFL in the homo-family architecture setting on the CIFAR-10 and CIFAR-100 [24] datasets (for hetero-family architecture results, see Appendix D.1, 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.From Table 1 , inconsistent performance improvements are observed with prior KD-based methods, especially for the L 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-sizefits-all approach used in the existing KD methods. These observations confirm the shortcomings of vanilla ensemble distillation and corroborate our theoretical findings in Remark 1 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.1.1.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 to the CV task, TAKFL has consistently improved performance across all device prototypes of varying sizes, achieving SOTA results: a 3% average increase on MNLI and 2% on SST-2. The suboptimality of existing KD methods, is 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 Public datasets are CIFAR-100 [24] for CIFAR-10 [24] and ImageNet-100 [10] 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. Architectures are ResNet-8, ResNet-14, and ResNet-18 [17] for S, M and L, respectively. All models are trained from scratch for 60 rounds. See Appendix D.1 for additional experiments using hetero-family architecture and more details. [2] for MNLI [52] and Sentiment140 [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-Mini, and Bert-Small [47] for S, M, and L, initialized from pre-trained parameters and fine-tuned for 20 communication rounds. See Appendix F.2 for more details. We evaluate the scalability of TAKFL across a spectrum of device prototypes, from extremely small (XXS) to extremely large (XXL), to see how well our method adapts from a uniform array of smallsize 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 3 illustrates TAKFL's ability to effectively scale from 3 to 7 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 do not offer unique contributions that could benefit from more complex distillation strategies. However, as the scenario expands to 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 significant 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. 
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 scenarios, 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 a 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 an approximated real-world setup has been demonstrated, actual deployment on physical devices and in environments 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. This setting is more feasible as it accommodates diverse device prototypes with their own specific configurations, including neural network architecture and dataset. However, establishing effective knowledge transfer between differently sized prototypes (like IoTs 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 objectives and settings. The first group of studies focuses on accommodating devices with varying compute resources, aiming to 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 a submodel of a global model according to their compute resources. These include dropout-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 techniques like in HeteroFL [11] and FjORD [18] consistently extract designated portions of the global model, whereas FedRolex [1] introduces a more flexible rolling method for sub-model 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 spectrum of device prototypes such as IoT devices and smartphones each have unique and unhashable neural network architectures tailored to their specific hardware and software configurations and underlying 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 of studies tackles a more practical scenario where device prototypes with heterogeneous model architectures participate in FL to enhance their global model performance through mutual knowledge sharing. In this context, knowledge distillation techniques are employed to transfer knowledge among device prototypes [30, 6, 43] . Here, locally updated client models from various device prototypes, collectively referred to as ensembles, serve as teachers to distill their knowledge 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 models. In contrast, FedET [6] employs an uncertainty-weighted average of ensembles' logits as the distillation target for all server student models, complemented by a diversity regularization technique. 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 to overconfidence, which compromises their ability to provide reliable uncertainty estimates [50, 15, 5, 55] .These existing works typically focus on settings where device prototypes have similar capabilities, i.e. similar model and dataset sizes, thus neglecting the challenges presented in more diverse settings 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 paper, we introduce TAKFL, which is designed to address the limitations of existing methods in these underexplored diverse device heterogeneous settings.Figure 1 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, 4] . 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper primarily focuses on a 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.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: 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?Answer: [Yes] Justification: The datasets and models used in this work are properly credited in main paper Section 7, Appendix D, and Appendix F.Guidelines:• The answer NA means that the paper does not use 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 is used and, if 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation 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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification:Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) 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 to 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. 
 
 Appendix: The full algorithm description of TAKFL is presented in Algorithm 1.Algorithm for each client k ∈ C i r , ∀i ∈ M in parallel do 6: θ i k ← ClientUpdate(k; θ i r ) 7:end for 8:θ i avg = k∈C i r |D i k | k∈C i r |D i k | θ i k 9:for each device prototype's server student i = 1, 2, . . . , M in parallel do 10:for each device prototype's teacher ensembles j = 1, 2, . . . , M in parallel do11: θ ← θ i avg 12:for each server distillation iteration t = 0, 1, 2, . . . , I distill do 13:x ← sample a mini-batch of data from public dataset D public 14:θ t+1 ← θ t -η distill • ∇L Ti S defined in Eq. 6.15:end for 16:τ j ← θ I distill -θ i avg 17:end for 18:θ i r+1 ← θ i avg + M j=1 λ j τ j 19:end for 20:θ i r+1 ← θ i 21: end for θ t+1 ← θ t -η local • ∇ℓ(f i (x; θ t ), y) 27:end for 28:θ i k ← θ I local 29: end function C Theoretical ResultsFirst we present the formal assumptions associated with our theoretical derivations. Assumption 1. Local federated averaging is performed with perfect test accuracy, i.e., argminEQUATIONThat is, the training error on the datasets {D j k } for the computed θ j avg is the same as the test error on the population distribution D j . Moreover assume that we can writeT i = N i k=1 f i (•, θi k )|k ∈ C i = {f i (•, θ i avg )}.Finally, we assume that the same population distribution j ω j D j 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 unreasonable. The local FL the device prototypes perform is generically prone to imprecision, especially as the clients' data varies, but this discrepancy is bounded [16] . Similarly the difference in the average of logits and the logit of averages has a bounded difference norm [53] . Thus, violations of the Assumption add additional perturbations to quantities derived in the Theoretical analysis without having structural/qualitative effects, and thus would only present clutter in the presentation.Notations. Now we present the notation defining the specific quantities we refer to 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, Parameters in j's device model that minimize the loss on its population distributionΘ j := argmin θ E (x,y)∼D j ℓ(f j (x; θ), y , Θ j,k := argmin θ E (x,y)∼D i ℓ(f j (x; θ), yΘ j,kParameters in j's device model that minimize the loss on i'th population distributionQ j = dim(θ j )The total capacity of device prototype j Q j = {e j k } k=1,...,Q j Eigenbasis for the model of device prototype j W j = dim(Θ j ) Dimension of the solution submanifoldΘ j W j,k = dim(Θ j,k ) Dimension of the solution submanifold Θ j,k W j = {e j k } k=1,...,W j Eigenbasis the solution submanifold Θ j W j,k = {e j,k l } l=1,...,W j,k Eigenbasis the solution submanifold Θ j,kWe shall make use of the "Choose" combinatorial operator, defined to be Ch(n, p)= n! p!(n-p)! . The standard O(•) notation indicates a k = O(b k ) to mean there exists K and C such that for k ≥ K, a k ≤ Cb k .A recent finding that inspired the methodology in this work is the discovery of the weight disentanglement 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 to our notation:(Task Arithmetic Property) holds for a set of vectors {τ j } if for all j it holds that,EQUATIONLet us define an important property of task arithmetic that we shall use in the sequel.(Weight disentanglement). [37] A parametric function f : X × Θ → Y is weight disentangled with respect to a set of task vectors T = {τ j } j∈T and the corresponding supportsD T := {D j } j∈T if f (x; θ 0 + T i∈T α i τ i ) = i∈T g j (x; α i τ i ) + g 0 (x),where g i (x; α i τ i ) = 0 for x ̸ ∈ D i and i ∈ T, and g 0 (x) = 0 for x ∈ i∈T D i .We now present the formal statements as well as the proofs of the main propositions. Proposition 1. (Information Loss in VED). Let the prototype i and j datasets be disjoint, i.e. D i ∩ D j = ∅, andM i=1 D ⟩ = D.Note that this implies the weight disentanglement property. Consider the VED procedure in the form of solving (3). Consider two device prototypes with a device capacity and solution dimension of Q 1 , Q 2 and W 1 , W 2 , respectively, and with associated eigenbases Q i , W i . Let the solution set of VED with prototype i as student be ΘiV ED with dim( Θi V ED ) = W vi with eigenbasis W vi . In addition, denote W s,t , s, t ∈ {1, 2} the dimension of the solution set for the student model trained on the data from the teacher device's ensembles. We assume that selfdistillation is executed appropriately, e.g.,W 1,1 = W 1 and W 2,2 = W 2 . 1. Case 1: Assume that Q 1 = Q 2 and W 1 = W 2 = W 1,2 = W 2,1 . Then it holds that, in expectation, dim Θ1 V ED ∩ Q 1 \ W 1 = O (Q 1 -W 1 )(W 1 )!(Q 1 -W 1,2 )! Q 1 !(W 1 )!(Q 1 -W 1 )! + Q 1 !W 1,2 !(Q 1 -W 1,2 )!This corresponds to the expected capacity of prototype 1 that is taken up for fitting logits that are not in the span of W 1 , that is, that do not fit the data corresponding to prototype 1.In this section we 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 5 left panel. This panel corresponds to a setting where two large device prototypes with similar total capacity, i.e. Q 1 = Q 2 = 3 perform knowledge transfer. We consider the solution dimensions of both prototypes to be the same, i.e. W 1 = 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. Θ 1 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 of the prototypes are fairly disparate, this has resulted into near orthogonal subspaces corresponding to these solutions. As we can see from the panel, VED will lead to point which is 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 a loss of knowledge. By contrast, the TAKFL approach, by customizing the merging coefficients and putting each to half, i.e. λ 1,1 λ 1,2 = 0.5, can traverse in the tangent space of zero loss surface and get into the intersection subspace which is exactly the optimal solution (θ* merged = θ 1 avg + λ 1,1 τ 1,1 + λ 1,2 τ 1,2 ).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 1 = Q 2 = 3 performing knowledge transfer. We consider the solution dimensions of both prototypes to be the same, i.e. W 1 = 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. Θ 1 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 of the prototypes are fairly close, this has resulted into nonorthogonal 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 λ 1,1 λ 1,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 simultaneously 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 θ 1 avg plane corresponds to the solution subspace of the large prototype 1 on its own data, D 1 . The line θ 1,2 avg line corresponds to the subspace of solutions in prototype 1's parameter space projected into the capacity of the information transferred from device prototype 2. Finally, the line labelled θ 1,3 avg corresponds to the subspace of solutions in prototype 1'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 D 1 is closer to the distribution D 2 than to D 3 . Comparing this case to the previous cases, θ 1,3 avg is like case I and θ 1,2 avg is like case II. 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 θ 1,3 avg 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 θ 1,2 avg has closer data distribution to prototype 1, TAKFL will prioritizes this by putting a larger merging coefficient, i.e λ 1,2 > λ 1,3 . By contrast, VED lacks this customization and results in sub-optimal knowledge distillation.The geometric intuition discussed here is consistent with our detailed experimental analysis in D.1.1.Here we provide additional insights from the literature as 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 correspondence to the results above. Overparametrization has been conjectured to be a significant factor in contributing to the (unexpected, by classical bias-variance tradeoffs) generalization ability of deep neural networks, from a 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 loss, L(θ) = k∈C j E (x,y)∼D j k ℓ(f j (x; θ), y) for some j, which will be identified from the context.Assumption 2. L : R Q → R is C 3 and the solution set Γ is a W -dimensional C 2 -submanifold of R D for 0 ≤ W ≤ D and rank(∇ 2 L(θ)) = Q -W Assumption 3. Assume that Uis an open neighborhood of Γ satisfying that gradient flow starting in U converges to some point in Γ From these, [29] derives Theorem 4.6. This Theorem decomposes the random process of the parameter weights driven by SGD after it has reached the solution manifold, e.g., the diffusive random walk of θ 1 avg in Figure 5 along its respective solution manifold. Theorem 1. Given L, Γ and θ µ (0) = θ(0) ∈ U as by Assumptions 2 and 3 the SDE modeling the optimization of F 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 (t) ∈ U, ∀0 ≤ t ≤ T ] = 1, θ η (⌊T /η 2 ⌋ converges in distribution to the stochastic process Y (T ) as η → 0, with Y (t) given as,EQUATIONwhere Σ ≡ σσ T and Σ ∥ , Σ ∥ , Σ ∥ are given as,EQUATIONThis 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, as 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 are traversing the solution set, with noise predominantly in the tangent directions. Thus knowledge distillation which preserves this noise structure is going to be more effective as far as preserving accuracy across data.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.The experiments in this section complements the main experimental results in the main paper Section 7.2.Experimental Setup. For the evaluation on CV task, we employ CIFAR-10 and CIFAR-100 [24] datasets. For CIFAR-10, we use CIFAR-100 as the unlabeled public dataset, while ImageNet-100, a 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 data portion is further distributed among its clients using a Dirichlet distribution. We apply two levels of data heterogeneity for a comprehensive evaluation: low heterogeneity, i.e. Dir(0.3), and high heterogeneity, i.e. Dir(0.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-employing ResNet8, ResNet14, and ResNet18 [17] for S, M, and L, respectively; and the "hetero-family" setting, which diverse architectures are used-ViT-S [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 a significant margin compared to the baselines, achieving SOTA performance. Notably, in the homo-family architecture setting with Dir(0.3) on CIFAR-10, TAKFL improves average performance across all prototypes by 8%, and by 4% on CIFAR-100. In the hetero-family settings with Dir(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 unsupervised ensemble distillation process, thereby enhancing performance. Generally, the performance gains from self-regularization are more pronounced in low data heterogeneity cases, where prototypes' models perform better and possess higher quality self-knowledge. Thus, self-regularization proves more effective as it preserves this higher quality self-knowledge. In this part, we elaborate on our key experimental observations and their alignment with our theoretical 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, in the CIFAR-10 homo-family setting with Dir(0.3), while small (S) and medium (M) prototypes see performance gains, the L prototype experiences up to a ∼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) 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 capable device prototypes are overshadowed by less informative outputs from smaller devices, thereby degrading the performance of L prototypes. These empirical findings are supported by our theoretical insights as discussed in Remark 1. Specifically, Proposition 1 illustrates that vanilla ensemble distillation (VED) leads to knowledge dilution and inaccuracies due to misaligned device capacity allocations. Moreover, this issue becomes more significant when the smaller device prototype serve as teacher.Insight 2: From Table 4 , the suboptimality of existing KD-based methods is evident from the significant performance improvements of our method, especially for S and M prototypes across various settings. This underscores the ineffectiveness of the one-size-fits-all approach these methods employ, where a single averaged logits distillation target is used for all device sizes, proving to be suboptimal. Our experimental observations regarding the shortcomings of vanilla ensemble distillation methods align with our theoretical findings, as substantiated in Remark 1 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.Insight 3: Our experiments, detailed in Table 4 , demonstrate TAKFL's adept handling of knowledge 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 (L) prototypes show less improvement, prompting the question: What can smaller device prototypes offer to 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 reduced. 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 beneficial.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 capacity of larger prototypes with suboptimal or irrelevant information, particularly in high heterogeneity environments, leading to not just stagnation but an accumulation of inefficiencies. These experimental observations align with our theoretical insights, as 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 to knowledge distillation in diverse federated learning environments.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 [7] as 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 a 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 a Dirichlet distribution: Dir(1.0) for TinyImageNet and Dir(0.3) for both STL-10 and CINIC-10. 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 MobileNetV3-Large [19] for S, Mobile-ViTV2 [34] for M, and ResNet-34 for L, sourced from the TIMM library. 1 The local training was conducted over 10 epochs using an Adam [23] optimizer with a learning rate of 1e-3 and weight decay of 1e-5. For server-side distillation, the epoch count was 10 for TinyImageNet and 1 for STL-10 and CINIC-10, with a batch size of 128, employing an Adam optimizer with a learning rate of 1e-5 and weight decay of 1e-5. For TinyImageNet, public dataset images from STL-10 were resized to 64×64, while for STL-10, images were resized to 32×32. No data augmentation was used. The communication rounds is fixed to 40. These experiments were conducted by only 1 trial. Table 15 details the configurations.Performance Results. Table 5 presents the results. The superiority of TAKFL's performance across these challenging datasets using pre-trained models is evident here as well.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, we utilize four datasets: MNLI [52] , SST-2 [45] , MARC [22] , and AG-news [60] . The corresponding unlabeled public datasets are SNLI [2] for MNLI, Sentiment140 [14] for SST-2, Yelp [61] for Amazon, and DBPedia [59] for AG-News.The training data is distributed among the device prototypes in a 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 . Performance on NLP Task. Table 6 presents the results on 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-art results. On MNLI, it has enhanced 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 of existing KD-based methods 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 be attributed to FedET's reliance on neural network confidence scores for uncertainty estimates in its uncertainty-weighted distillation. However, neural networks, especially 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] .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 to very large sizes, we conduct experiments 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 a diverse mix that includes prototypes ranging from extremely small (XXS) to extremely large (XXL) (7 device prototypes). 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 to 1:2:3:4:5:6:7 from XXS to XXL. Each dataset portion is further distributed among clients using a Dirichlet distribution (Dir(0.5)). The number of clients ranges from 35 to 5 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 architectures-ResNet10-XXS, ResNet10-XS, ResNet10-S, ResNet10-M, ResNet10, ResNet18, and ResNet50-scaled appropriately 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 a learning rate of 1e-3 and a weight decay of 5e-5 for local training. For XL and XXL, a step learning rate scheduler reduces the learning rate by a factor of 0.1 at half epoch. Server-side distillation employs a fixed batch size of 128, using the Adam optimizer with learning rate of 1e-3 and weight decay of 5e-5. The softmax temperature is set at 3 for ensemble distillation and 20 for self-regularization. The number of communication rounds is fixed at 30. These experiments are conducted over 3 trials with different random seeds, and the average performance with standard deviation is reported. The entire device prototypes configurations are given in Table 16 .The detailed results are presented in Tables 7, 8 , and 9. In this section, we conduct an ablation study to further understand how TAKFL customizes knowledge integration and understand how the merging coefficients λ i are achieving this. This experiment aims to further understand the trade-offs between customized knowledge integration approach from the one-size-fits-all strategy employed in vanilla ensemble distillation and prior works.Experimental Setup. Our experimentation focuses on two device prototypes: XXS and XXL, selected 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 prototype receives a randomly selected, non-overlapping subset of the training dataset-3.57% for XXS and 25% for XXL-distributed among their clients in a non-i.i.d. manner using Dir(0.5). Both prototypes have three clients each. The architectures used are ResNet10-XXS for the XXS prototype and ResNet-50 for the XXL prototype. To focus solely on the evaulation of the server-side distillation process and its evolution with varying λ, we pre-train each prototype using standard FedAvg for 10 communication rounds, with a 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 a step learning rate scheduler that reduces the rate by a 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 unlabeled public dataset. We save the final updated client and server models from both prototypes for further experimentation, focusing on the impact of merging coefficients without self-regularization in TAKFL. The merging coefficient λ is varied linearly from 0 to 1 in increments of 0.05. For simplicity, the XXS prototype is referred to as the small (S) prototype and XXL as the large (L) prototype.Discussion. Figure 6 illustrates the significant impact of customized knowledge integration on the performance of both small and large device prototypes compared to the one-size-fits-all approach typical of vanilla ensemble distillation in the prior works, at different distillation epochs. Here, TAKFL adeptly manages customization for both small and large prototypes by controlling the merging coefficient λ. The merged model for both the small and large student prototypes is obtained using the formula θ merged = θ avg + (1 -λ)τ S + λτ L . Notably, the performance is benchmarked at λ ≈ 0.5 in all cases, reflecting similar results to vanilla ensemble distillation (FedDF), where no customization in knowledge transfer occurs. This baseline performance is critical for understanding the effects of further customization.In small distillation epochs (I distill < 10), minimal benefit is observed from customized knowledge integration, as 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 λ becomes 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 distillation epochs, the absence of significant unique knowledge results in optimal performance at λ ≈ 0.5. Conversely, as the number of distillation epochs rises (I distill ≥ 20), the one-size-fits-all strategy proves suboptimal, underscoring the importance of tailored knowledge integration strategies. Optimal 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 1 and 2.In this section, we explore the influence of the public dataset on the performance of TAKFL and existing KD-based methods when the public dataset used for server-side distillation is less similar 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.1, using the CIFAR-10 homo-family architecture. To measure dataset similarity, we compute cosine similarity between the averaged features of datasets, extracted using an off-the-shelf pre-trained CLIP model [40] (CLIP ViT-B/32) available from the official GitHub repository. 2Discussion. 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 degradation under the same conditions. This demonstrates TAKFL's practical utility in real-world sce-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 as the unlabeled public dataset. Both prototypes were pretrained 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 a sample rate of 1.0. The small prototype utilizes a ResNet10-XXS architecture, while the large prototype employs a ResNet-50.narios where the server typically lacks knowledge of the private datasets to select a closely aligned public dataset for distillation. Notably, FedET underperforms significantly when using a less similar public dataset, performing worse than vanilla FedAvg in both low and high data heterogeneity scenarios. 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 on the overconfident and poorly calibrated confidence scores from neural networks [15, 55] for uncertainty estimates in its uncertaintyweighted distillation approach. 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 experimental 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. 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 in Figure 7 .As we can see in the t-SNE plots (Figure 7 ), 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. In this section we bring the details of the hyper-parameters we used and our implementation. We implement our entire code in PyTorch [38] using FedZoo benchmark [36] and release it at https: //github.com/MMorafah/TAKFL. We use two NVIDIA RTX 3090 gpus to conduct the entire experimentation in this paper.For comprehensive evaluation of our method, we consider federated learning image classification training from scratch.Datasets. We experiment with several image classification datasets: CIFAR-10, CIFAR-100, CINIC-10, TinyImageNet, STL-10. The details of each dataset is the following:• CIFAR-10 [24]: CIFAR-10 consists of 60,000 images of size 32×32 RGB across 10 classes, with each class containing 6,000 images.• CIFAR-100 [24] : CIFAR-100 comprises 60,000 images of size 32×32 RGB distributed across 100 classes, with 500 images per class.• CINIC-10 [9]: CINIC-10 has 270,000 images of size 32×32 RGB across 10 classes, each class containing 27,000 images.• TinyImageNet [25] : TinyImageNet contains 100,000 images of size 64×64 RGB across 200 classes.• STL-10 [7]: STL-10 has 100,000 unlabeled images and 13,000 labeled images of size 96×96 RGB across 10 classes.The ImageNet-100 as the unlabeled public dataset is constructed by randomly selecting 100 classes from the ImageNet [10] dataset.Experiments in Section 7.2 and Appendix D.1. 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 a 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: The final classification head consists of two linear layers with a hidden size of 512, followed by a ReLU activation, and a final linear classifier layer.For ViT-S, we adopt the standard Vision Transformer [12] architecture implementation from Github. 3 We configure ViT-S with 6 attention blocks, each with 16 heads and a hidden dimension of 64. The final MLP dimension is set to 256. In our experiments with ViT-S, we set the patch size to 4, and the input image size is 32×32.Experiment in Appendix D.2. The pre-trained MobileNetV3-Large [19] , MobileViTV2 [34] , and ResNet34 [51] were instantiated using the TIMM library. 4Scalability Experiments in Section 7.3. The architectures in these experiments are inspired by [20] .The details of architectures are as following:• ResNet10-XXS is configured with [8, 8, 16, 16] filters and repeats the 'BasicBlock' [1, 1, 1, 1] times. • ResNet10-XS is configured with [8, 16, 32, 64] filters and repeats the 'BasicBlock' [1, 1, 1, 1] times. • ResNet10-S is configured with [16, 32, 64, 128] filters and repeats the 'BasicBlock' [1, 1, 1, 1] times. • ResNet10-M is configured with [8, 16, 32, 64] [3, 4, 6, 3] times.Base Hyper-parameters. The following hyperparameter values apply to all CV experiments unless stated otherwise. We set the diversity regularizer coefficient of FedET to 0.1 for our entire experimentation per the original paper [6] . We use the Adam optimizer with a learning rate of 1e-5, weight decay value of 5e-5, and a 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 hyper-parameters.Experiments in Section 7.2 and Appendix D.1. For tables 1 and 4, there are 100 clients with the S device prototype, 20 clients with the M device prototype, and 4 clients with the L device prototype. For each round, 10, 4, and 2 clients from each S, M, and L prototype are randomly sampled respectively for participation. 10% of the data goes to the S prototype, 30% to 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 2 clients for S, M, and L device prototypes, respectively. Each round, every client participates in FL. 20% of the data is distributed to prototype S, 30% to prototype M, and 50% to prototype L. 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.For the NLP task, we fine-tune federated learning text classification task using pretrained models.All NLP datasets were provided by Hugging Face.5 • MNLI [52] : MNLI contains 433K sentence pairs, each sentence pair labeled as one of 'entailment,' 'neutral,' and 'contradiction.'• SNLI [2] : SNLI is similar to MNLI, with 570K sentence pairs each labeled one of 3 labels.• SST2 [45] : SST2 consists of 67K phrases, each labeled as sentiment 'positive' or 'negative.'• Sentiment140 [14] : Sentiment140 is a dataset of 1.6M Twitter messages each labeled with one of 2 sentiment values.• MARC [22] : MARC (Multilingual Amazon Reviews Corpus) is a dataset with online reviews in multiple languages from the Amazon delivery service website. Each review has a label which is one of 1-5 stars. We only use the English reviews from this dataset, which results in 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.Experiments in Section 7.2 and Appendix D.3. We use three variations of the BERT architecture: BERT-Tiny, BERT-Mini, and BERT-Small from [47] . The weights were pre-trained on the 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 2 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.Base Hyper-parameters. For distillation, we use the Adam optimizer with a learning rate of 3e-5, no weight decay, and batch size of 32. The distillation epoch is set to 1, the ensemble distillation softmax temperature to 3, and the self-regularizer softmax temperature to 20 for all NLP experiments. 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 be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.• 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 to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via 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 a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , 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 authors 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 be possible for other researchers to have some path to reproducing or verifying the results.https://github.com/huggingface/pytorch-image-modelshttps://github.com/OpenAI/CLIPhttps://github.com/lucidrains/vit-pytorchhttps://github.com/huggingface/pytorch-image-modelshttps://github.com/huggingface/datasetshttps://github.com/google-research/bert 


Decision: The paper is of high quality, and there is not objection from four reviewers. It requires minor refinements based on the reviewers' comments.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 
2. 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.
3. It would be better to provide a visualization study for a better understanding of the effectiveness of transfer learning from different prototypes.

[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.
questions: questions: see weakness
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: (1) [main concern] Strong Dependence on Hyperparameters
The 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.

(2) [main concern] Strong Assumptions in the Selection of Public Datasets
The 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.

(3) [minor concern] Quantification of Data Heterogeneity and Hyperparameter Selection
The 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.
questions: questions: See weaknesses
limitations: limitations: N.A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

2. 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.

3. 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.

4. 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.

5. A minor issue is only apply to data with prototype label, thus may limited its impact
 
[1] TIES-Merging: Resolving Interference When Merging Models

[2]  Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch

[3] Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models

[4] Adaptive Multi-Teacher Multi-level Knowledge Distillation

[5] Agree to Disagree: Adaptive Ensemble Knowledge Distillation in Gradient Space

[6] Ensemble distillation for robust model fusion in federated learning.
questions: questions: please refer to Weakness.
limitations: limitations: please refer to Weakness.
paper_summary: paper_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.
strength_weakness: strength_weakness: While the paper presents a clear presentation and evaluation, there are a few aspects to be strengthened. 

1. 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.   
2. 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?  
3. 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.
questions: questions: Please refer to the concerns and questions mentioned in the weakness part.
limitations: limitations: N/A

Comments: title: Response to Reviewer F7vu
Kind Reminder: End of Author-Reviewer Discussion Period is Tomorrow
Kind Reminder: End of Author-Reviewer Discussion Period is Tomorrow
Kind Reminder: End of Author-Reviewer Discussion Period is Tomorrow
Response to Q1
Kind Reminder: End of Author-Reviewer Discussion Period Approaching
Response to Q2 and Q3
Response to Q1
comment: We sincerely appreciate the reviewer's response and great effort for the continued engagement in the discussion.

Yes, when the labels overlap (case 2 and 3) the task arithmetic property is not held, but that's not a problem. Please note that for case 2 and 3 the statement is trivial and does not require a proof since we already know the logits are in $\mathcal{Q}^1$. Therefore, as a result no need for task arithmetic property for this case as the statement is trivial and no proof is needed.

Please note that Proposition 1 deals with the information loss of vanilla KD (or average logits ensemble distillation) in device heterogeneous FL, while Proposition 2 states the improvement of TAKFL beyond vanilla KD. 

Regarding the improvement of knowledge transfer (TAKFL's improvement) in Proposition 2, under the label overlapping case, the statement becomes trivial. The theorem specifically concerns whether the KD results in parameters that fall within or outside of the space $\mathcal{Q}^1$. When the logits are identical, they are already within that space, meaning there is nothing further to address. As a result, all propositions, which consider the dimensions of the model that do not fit prototype 1's data (i.e., $\mathcal{Q}^1$), remain unaffected by these overlapping logits. This exactly proves our earlier statement that the label overlapping case creates an i.i.d. situation, which can actually make knowledge transfer easier, with the vanilla KD method experiencing minimal information loss. Therefore, as a result TAKFL's guarantee (improvement of knowledge transfer) in proposition 2 becomes trivial.

We would like to also sincerely appreciate the reviewer's great effort in continued engagement in the discussion and providing us feedback. The received feedback indeed helped us to provide a complete theoretical statements for different cases (disjoint and overlapping labels) which we will note this in the final version of the paper.
The problem is when class labels overlap  $\mathcal{D}_{t'} \cap \mathcal{D}_t \neq \emptyset$,  task arithmetic property is not held. As state in lines 679-681 in appendix, the proof heavily rely on task arithmetic property thus "Please note that Proposition 2, that is the guarantees for TAKFL, remains the same." is not true....
**...Continued Response...**

### Case 3 (Partial Overlap)
The prototype $i$ and $j$ datasets are partially overlapping, i.e. $\mathcal{D}^i \cap \mathcal{D}^j \neq \emptyset$, meaning that the two prototypes overlap on at least one of the classes $\mathcal{D}_c$. Then the changes to Proposition 1 and 2 in our theory are detailed in the following.

In this case, $\mathcal{W}^{1,2} = \mathcal{W}^{1,2}_1 \cup \mathcal{W}_2^{1,2}$, that is, the class labels across data prototypes partially overlap (with $W^{1,2}_1$ the overlapping dimensionality). The Proposition for VED changes to:

**Proposition 1** *(Information Loss in VED).* Consider the VED procedure. Consider two device prototypes with a device capacity and solution dimension of $Q^1, Q^2$ and $W^1, W^2$, respectively, and with associated eigenbases $\mathcal{Q}^i, \mathcal{W}^i$. Let the solution set of VED with prototype $i$ as student be $\hat{\Theta}^i_{VED}$ with $\dim(\hat{\Theta}^i_{VED}) = W^{v_i}$ with eigenbasis $\mathcal{W}^{v_i}$. In addition, denote $W^{s,t},\, s,t \in \{1,2\}$ 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., $W^{1,1} = W^1$ and $W^{2,2} = W^2$.

1. **Case 1:** Assume that $Q^1 = Q^2$ and $W^1 = W^2 = W^{1,2} = W^{2,1}$. Then it holds that, in expectation,

   $$
   \dim\left(\hat{\Theta}^1_{VED} \cap \left[\mathcal{Q}^1 \setminus \mathcal{W}^1 \right]\right) = O\left( \frac{(Q^1 - W^1)(W^1)!(Q^1 - W_2^{1,2})!}{Q^1!(W^1)!(Q^1 - W^1)! + Q^1! W_2^{1,2}!(Q^1 - W_2^{1,2})!} \right) 
   $$

   This corresponds to the expected capacity of prototype 1 that is taken up for fitting logits that do not fit the data corresponding to prototype 1.

2. **Case 2:** Assume that $Q^1 > Q^2$ and $W^1 = W^{1,2} > W^2$. Then the same quantity as for Case 1 holds. Moreover,

   $$
   \dim\left(\hat{\Theta}_{VED} \cap \left[\mathcal{Q}^1 \setminus (\mathcal{W}^1 \cup \mathcal{W}^{1,2}) \right]\right) = O\left(\frac{(Q^1 - W^1)(W^1!)(W_2^{1,2} - W^2)!}{Q^1! W^1!(Q^1 - W^1)! + Q^1! W^2!(W_2^{1,2} - W^2)!}\right)
   $$

   This corresponds to the capacity of client 1 that has been allocated but fits, in the model of prototype 1, neither the data of prototype 1 nor the data of prototype 2.

Please note that the presence of overlapping logits $W^{1,2}_1$ does not appear in the result. It is exactly the same as the original Proposition 1 in the original paper, just with $W_2^{1,2}$ in place of $W^{1,2}$.

Please note that Proposition 2, that is the guarantees for TAKFL, remains the same. Indeed, in this case $W^{1,2}$ does not even appear in the original result, and there is no modification needed. Indeed, since the theoretical results consider the complement of $\mathcal{Q}^1$ and all overlapping logits are in $\mathcal{Q}^1$, they do not affect the results.

### **Conclusion Remark**
We will supplement these cases to our theoretical results for completeness in the final version of the paper. We hope that this response comprehensively addresses the reviewer's remaining concerns. 

Once again, we appreciate the reviewer's continued engagement in the discussion and are happy to address any further questions or concerns in the remaining time till end of discussion period in the next few hours.
We sincerely appreciate the reviewer's detailed constructive feedback to enhance or theoretical results. To provide a comprehensive response to the reviewer's point, we present the following cases:

### Case 1 (No Overlap)
The prototype $i$ and $j$ datasets are disjoint, i.e. $\mathcal{D}^i \cap \mathcal{D}^j = \emptyset$, and $\bigcup_{i=1}^{M} \mathcal{D^i} = \mathcal{D}$. Then, our current proposition holds correct with no changes.

### Case 2 (Fully Overlap: Identical Class Labels)
The prototype $i$ and $j$ datasets are fully overlapping, i.e. $\mathcal{D}^i \subseteq \mathcal{D}^j$, meaning that the two prototypes have the same classes $\mathcal{D}_c$, $\forall c \in D$. Then the changes to Proposition 1 and 2 in our theory are detailed in the following.

Consider that in this case, the logits corresponding to $\mathcal{W}^{1,1}$ and $\mathcal{W}^{1,2}$ are the same. As such, $\mathcal{W}^{1,2} \in \mathcal{Q}^1$, and so the propositions change to, trivially:

**Proposition 1** *(Information Loss in VED).* Consider the VED procedure. Consider two device prototypes with a device capacity and solution dimension of $Q^1, Q^2$ and $W^1, W^2$, respectively, and with associated eigenbases $\mathcal{Q}^i, \mathcal{W}^i$. Let the solution set of VED with prototype $i$ as student be $\hat{\Theta}^i_{VED}$ with $\dim(\hat{\Theta}^i_{VED}) = W^{v_i}$ with eigenbasis $\mathcal{W}^{v_i}$. In addition, denote $W^{s,t},\, s,t \in \{1,2\}$ 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., $W^{1,1} = W^1$ and $W^{2,2} = W^2$.

1. **Case 1:** Assume that $Q^1 = Q^2$ and $W^1 = W^2 = W^{1,2} = W^{2,1}$. Then it holds that, in expectation,

   $$
   \dim\left(\hat{\Theta}^1_{VED} \cap \left[\mathcal{Q}^1 \setminus \mathcal{W}^1 \right]\right) = 0
   $$

   This corresponds to the expected capacity of prototype 1 that is taken up for fitting logits that do not fit the data corresponding to prototype 1.

2. **Case 2:** Assume that $Q^1 > Q^2$ and $W^1 = W^{1,2} > W^2$. Then the same quantity as for Case 1 holds. Moreover,

   $$
   \dim\left(\hat{\Theta}_{VED} \cap \left[\mathcal{Q}^1 \setminus (\mathcal{W}^1 \cup \mathcal{W}^{1,2}) \right]\right) = 0
   $$

   This corresponds to the capacity of client 1 that has been allocated but fits, in the model of prototype 1, neither the data of prototype 1 nor the data of prototype 2.

Please note that Proposition 2, that is the guarantees for TAKFL, remains the same.

**...Response continued in the next comment...**
1. As state in lines 679-681 in appendix, the proof  heavily rely on task arithmetic property thus the statement "Relaxing the WD assumption in this context would simply add additional notation and unnecessary complexity to an already lengthy formal property statement without providing additional clarity." is wrong.

2. I didn't see "When class labels overlap, the task vectors associated with training data from different prototypes can include vectors in the tangent space of both prototypes’ loss surfaces." why they are in tangent space?  More extreme, what if there is only one class, the distillation of other prototypes to prototypes $j$ ? Without concrete proof or citation, it is not proper to give such a conclusion. 

3. Combining my question 1,2 this statement is not proper....“As mentioned earlier, this assumption is pedagogically useful, offering clarity and focus without overcomplicating the analysis. The key takeaway remains the "Delta" of information retention/loss between KD and TAKFL, which stays consistent regardless of the baseline error rate. Even if WD does not hold universally, our theoretical results still provide valuable insights into behavior under idealized conditions, which is a common practice in Federated Learning research.”
We appreciate the reviewer’s response and would like to clarify the fundamental distinctions between our methodology and FedProto, particularly in terms of problem formulation and evaluation metrics.

**Problem Formulation:**
While FedProto does address model heterogeneity, its primary focus is on optimizing heterogeneous local models in FL (Eq. 2 in [1]). In contrast, our method follows [2], where the objective is to optimize the global models of heterogeneous device prototypes (Eq. 1 in our paper). Although both methods tackle device heterogeneity, the specific goals and contexts of optimization differ, leading to very different approaches and outcomes.

**Evaluation Metrics:**
FedProto evaluates performance using average local test accuracy, while our evaluation metric is the global model test accuracy of device prototypes, as outlined in [2]. 

Furthermore, we would like to note that our experimental setup significantly differs from the original setting in FedProto’s paper.  As a results, these significant fundamental differences in both problem formulations and evaluation metrics makes direct comparisons between FedProto and our method challenging and unfair and misleading.

However, despite these fundamental differences we have adopted the FedProto’s official implementation code and conducted experiments within our homogeneous-family experimental setting (Table 1 of our paper), and presented the results in Table 1 of rebuttal PDF file. The differences in results, such as FedProto performing lower than FedAvg, could be due to these fundamental differing contexts and setups.

In conclusion, while we agree with the reviewer that FedProto is a model heterogeneous FL method, it is important to note that FedProto is a "personalized FL" model-heterogeneous federated learning per its problem formulation in Eq. 2 of [1]. However, our work is a "global FL" KD-based device heterogeneous federated learning per our problem formulation in Eq. 1 of our paper, following [2]. This distinction is crucial in understanding the differences in performance outcomes.

We sincerely appreciate the reviewer's constructive feedback and suggestions. If there are any further concerns or questions, we would be grateful if the reviewer could kindly elaborate so that we can provide additional clarity.

We appreciate the reviewer's continued engagement and are happy to address any further questions or concerns as the discussion period concludes in the next few hours.

### References

[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.

[2] Lin et al. "Ensemble distillation for robust model fusion in federated learning." NeurIPS’20
Thanks for your response. However, for Q1, FedProto is indeed a method proposed to address the problem of model heterogeneity which utilizes a linear layer to project different feature dimensions into the same for prototype embeddings. Therefore, the targeted problem of FedProto and your method are the same. In addition, the experimental results provided in Table 1 of the PDF rebuttal seem unreasonable: the results of FedProto are even lower than FedAvg, which contradicts the result in FedProto's original paper.
We appreciate the reviewer’s attention to the technical details of weight disentanglement (WD). Formally, the reviewer is correct that the WD theory in [3] assumes exclusively partitioned class labels across device prototypes. In our experiments, there may be overlap in class labels, meaning our setup does not fully satisfy the assumption as formally presented in [3].

However, we contend that this overlap does not undermine the theory’s applicability. When class labels overlap, the task vectors associated with training data from different prototypes can include vectors in the tangent space of both prototypes’ loss surfaces. This overlap creates an i.i.d. situation, which can actually make knowledge distillation (KD) easier, with the Vanilla Average Logits method likely experiencing minimal information loss.

For instance, the task vectors corresponding to training on device prototype 2's data may now contain vectors in the tangent space of the loss surface of device prototype 1's architecture, fit to data exclusively belonging to prototype 2's original training samples. Additionally, vectors could now exist in the tangent space of the loss surface of device prototype 1's architecture, fit to prototype 1's data, due to the overlapping classes. This overlap creates an i.i.d. condition, simplifying KD. In this case, Vanilla Average Logits should also suffer minimal information loss.

Relaxing the WD assumption in this context would simply add additional notation and unnecessary complexity to an already lengthy formal property statement without providing additional clarity. However, we are open to including a remark summarizing this point for completeness.

While we appreciate the rigorous technical discussion, we respectfully disagree that this makes the theory inadequate. As mentioned earlier, this assumption is pedagogically useful, offering clarity and focus without overcomplicating the analysis. The key takeaway remains the "Delta" of information retention/loss between KD and TAKFL, which stays consistent regardless of the baseline error rate. Even if WD does not hold universally, our theoretical results still provide valuable insights into behavior under idealized conditions, which is a common practice in Federated Learning research. Introducing additional complexity to account for WD variations would make the results more opaque without significantly altering the conclusions.

Finally, we believe that theoretical details should enhance the pedagogical and explanatory value of the theory. Simplifying assumptions, while not always factually accurate, are essential for clarity and understanding. For example, convergence theory often assumes continuously differentiable functions, despite the widespread use of ReLUs and max-pooling in modern architectures. Such theory remains valuable in providing insights into an algorithm’s properties and relative advantages, and we believe the same applies to our theoretical framework.

Once again, we appreciate the reviewer's continued engagement in the discussion and are happy to address any further questions or concerns in the remaining time till end of discussion period by tomorrow on August 13th.
Your response is not true because even though you don't have overlapping samples, the classes themselves overlap.

As mentioned in the problem statement in [3]:  Consider $T $ tasks, with every task $ t $ consisting of a triplet $\mathcal{D}_t, \mu_t, f_t^*) $ where $ f_t^* : \mathcal{D}_t \to \mathcal{Y} $ a target function (e.g., labels).  

Additionally, as illustrated in Fig. 6 in [3], the eigenfunction localization clearly demonstrates disentanglement in the case of $ x \in RESISC45 $ and $ x \in cars $. The former refers to Remote Sensing Image Scene Classification, while the latter pertains to car classification. The labeling functions $f_t^*$ are significantly different.

However, in your data splitting, if we denote the data distribution of class $ c $ as $ D_c$, for any prototype group $ i $, there are samples belonging to the class $ x^i_c \in D_c $. Consequently, the intersection among different prototypes is not a null set $\mathcal{D}_{t'} \cap \mathcal{D}_t \neq \emptyset$.

As a result,  the Theory is inadequate.
Dear Reviewer ME5v,

We sincerely appreciate your extensive efforts in reviewing and commenting on our work. We hope that our rebuttal has comprehensively addressed your comments and concerns.

As the author-reviewer discussion period is approaching to its end by tomorrow on August 13th, we kindly ask you to reach out if you have any further questions or points of discussion. Your feedback is highly valued, and we are eager to address any additional remaining concerns you may have.

Thank you very much for your attention to this matter.
Dear Reviewer RvJG,

We sincerely appreciate your extensive efforts in reviewing and commenting on our work. We hope that our rebuttal has comprehensively addressed your comments and concerns.

As the author-reviewer discussion period is approaching to its end by tomorrow on August 13th, we kindly ask you to reach out if you have any further questions or points of discussion. Your feedback is highly valued, and we are eager to address any additional remaining concerns you may have.

Thank you very much for your attention to this matter.
Dear Reviewer V7t1,

We sincerely appreciate your extensive efforts in reviewing and commenting on our work. We hope that our rebuttal has comprehensively addressed your comments and concerns.

As the author-reviewer discussion period is approaching to its end by tomorrow on August 13th, we kindly ask you to reach out if you have any further questions or points of discussion. Your feedback is highly valued, and we are eager to address any additional remaining concerns you may have.

Thank you very much for your attention to this matter.
We would like to extend our sincere gratitude for the reviewer's great effort in reviewing our work and continued engagement in the discussion. We also appreciate the reviewer for raising his score to 5.

### **Response to Q1**
To answer Q1, we would like to note that the goal of device prototypes $i$ and $j$ participating in FL and doing server-side knowledge transfer with each other is benefiting and learning from their underlying datasets $\mathcal{D^i}$ and $\mathcal{D^j}$ (as noted in problem formulation Eq. 1), where indeed $\mathcal{D^i} \cap \mathcal{D^j} = \emptyset$ and $\bigcup_{i=1}^{M} \mathcal{D^i} = \mathcal{D}$ (lines 618-619, 679-680).

In all of our experimentation, the data partitions of prototypes, i.e. $\mathcal{D^i}$'s, are all indeed disjoint with no sample overlap. Therefore,  $\mathcal{D^i} \cap \mathcal{D^j} = \emptyset$ does hold in our work and theoretical framework. 

We hope that this answers the reviewer's question.

Once again, we appreciate the reviewer's continued engagement in the discussion and are happy to address any further questions or concerns in the remaining time till end of discussion period by tomorrow on August 13th.
Although the novelty is limited, I will raise my score from 4 to 5 based on the authors' experiments. If there are any additional points I may have overlooked, I'm open to discussing them.
Thank you response.

Q1: First, I agree with the author's conclusion that "while perfect WD across all merging coefficient ranges is challenging to achieve, strong WD within the effective range (0, 1] is possible." Then, as stated in [3], the WD refers to different tasks (classes) as shown in the problem statement in Section 2, Property 3, and Fig. 6. In Property 1 (Task Arithmetic) of [3], it is noted that this occurs with non-intersecting task supports $\mathcal{D}_{t'} \cap \mathcal{D}_t = \emptyset$. 

However, in the problem setting (lines 131-135), this paper focuses on various network architectures, and in Table 13 of the appendix, it is shown that heterogeneity occurring within each device setting.  As a result  $\mathcal{D}_{t'} \cap \mathcal{D}_t \neq \emptyset$, thus might be too strong to use this assumption.

Q2: the answer is clear to me.
Q3: the answer is clear to me.

As a result,  the inadequate Theory, the application of task arithmetic to FedDF, limited the novelty of this paper.
Dear Reviewers, 

We sincerely appreciate your extensive efforts in reviewing and commenting on our work. We hope that our rebuttal has comprehensively addressed your comments and concerns.

As the author-reviewer discussion period is approaching to its end on August 13th, we kindly encourage you to reach out if you have any further questions or points of discussion. Your feedback is highly valued, and we are eager to address any additional remaining concerns you may have.

Thank you very much for your attention to this matter.

Anonymous Authors
We appreciate the reviewer’s detailed feedback and engagement during the discussion period. 

### **Response to Q2**

Yes, the second row in Table 2 in the rebuttal PDF file refers to CIFAR-100. We apologize for any confusion, as we were working diligently to prepare the results. To address the reviewer's question, we directly replaced the KD loss in FedDF with AMTML-KD [4] and AE [5] in Table 2, resulting in the baselines labeled AMTML-KD [4] and AE [5]. Additionally, we replaced the vanilla KD loss in our TAKFL framework with [4,5], resulting in baselines TAKFL+[4] and TAKFL+[5].

The reason why [4,5] are performing worse in FL could be that these methods are designed for KD in a centralized learning setting, where teachers or ensembles are fully trained and robust, and the distillation dataset aligns with the 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. This discrepancy likely contributes to the reduced effectiveness of these methods in the device heterogeneous FL setting.

### **Response to Q3**

We appreciate the reviewer's comment and would like to clarify the distinctions between overall computation time, computational load, and resource usage.

* **Computation Time:** TAKFL's computation time is O(1) (constant) due to parallelization, as all distillation processes occur simultaneously.
* **Computation Load:** The overall computational load scales as O(M) (linear) since the distillation tasks are performed independently for each prototype in parallel and merged into a singe task arithmetic operation (Eq. 8). 
* **Resource Usage (Memory):** The resource usage scales as O(M^2) (quadratic) because of the need to store and process multiple task vectors concurrently.

While it's true that resource usage grows with O(M^2), it’s important to note that the entire KD process occurs on the server side. In FL, servers typically possess substantial resources. Given these resources, the resource usage incurred by TAKFL is generally less of a concern in practical applications. Moreover, in real-world scenarios, the number of device prototypes (M) tends to be limited, often encompassing categories such as IoT devices, mobile devices, and workstations. Therefore, while the theoretical resource usage grows with O(M^2), the practical impact is mitigated by the server's capabilities and the limited number of prototypes involved.

We appreciate the reviewer's continued engagement in the discussion and are happy to address any further questions or concerns.
We appreciate the reviewer’s detailed feedback and engagement during the discussion period. 

### **Response to Q1**
We would like to clarify our position on the use of weight disentanglement (WD) in our work.

**Response to Weight Disentanglement Concerns**

Task Arithmetic (TA) was first introduced in reference [51] of our paper and has demonstrated practical effectiveness. [3] provides the theoretical foundation for TA, identifying WD as the only necessary condition to perform TA effectively. Importantly, [3] demonstrates that WD is an emergent property during pre-training, absent in randomly initialized models (Section 6.2 of [3]). Fig 3 of [3] shows that WD remains sufficiently strong within the merging coefficients range (0, 1] across tasks, even without linearization. Linearization primarily enhances WD outside this range, suggesting that models are weight-disentangled enough in practice to enable task arithmetic and achieve strong performance within the effective range.
This effective range is crucial, as both the original TA paper [51] and [2,3] set task vector coefficients within (0, 1], ensuring that WD is maintained. Additionally, [7] uses `torch.clamp` to constraint task vector coefficients within this range, further validating the effectiveness of WD in this context. Our work, following [51,2,3,7], also ensures that the merging coefficients remain within this range, supporting robust WD and effective TA (Eq. 8, Section 5.2, Appendix F.3).

**Response to Studies [1,2,7]**

Studies [1,2] primarily address parameter interference and redundancy rather than WD. [1] discusses interference between task vectors due to redundant parameters and sign disagreement, proposing TIES-MERGING to mitigate this. [2] identifies redundant delta parameters in SFT LMs and introduces DARE to enhance model merging. [7] suggests that merging coefficients play a pivotal role in the average accuracy of the final MTL model and proposes AdaMerging to optimize these coefficients.

We couldn’t find any argument in studies [1,2,7] stating that achieving WD is difficult, causing interference within different task vectors. It does not appear to us that there is any formal correspondence between parameter interference and redundancy as these are considered in [1,2,7] and the notion of  WD as defined in Eq. 4 of [3]. We would appreciate it if the reviewer could please clarify and elaborate, which would help us provide a more comprehensive response regarding [1,2,7]. Moreover, if the reviewer suggests that [7] handles the difficulty better through adaptive TA, it is important to note that our approach also employs adaptive TA, ensuring that we address this challenge effectively.

**Conclusion**

In conclusion, while perfect WD across all merging coefficient ranges is challenging to achieve, the strong WD within the effective range (0, 1] is sufficient for TA to perform well. Our work, following [51,2,3,7], ensures that the merging coefficients remain within this range, maintaining the robustness of WD and enabling effective TA as it is evident from the improvements in our extensive experiments across both CV and NLP tasks.

Finally we believe that this concern is largely orthogonal to the specific theoretical results we claim. Regardless, we believe that this assumption is pedagogically useful, providing clarity and focus in our theoretical results without overcomplicating the analysis. The key takeaway is the "Delta" of information retention/loss between KD and TAKFL, which remains consistent regardless of the baseline error rate. Even if WD does not hold generically, our theoretical results are still informative as they indicate behavior in certain idealized, limited behavior circumstances, in particular cases with significant degrees of freedom of capacity. This is common practice in studying Federated Learning. Introducing additional complexity to account for WD variations would make the results more opaque without significantly altering the conclusions.

We appreciate the reviewer's continued engagement in the discussion and are happy to address any further questions or concerns.
**Question 1:**

Thank you for your response. I still insist that using disentanglement theory as proof is not fair because the model cannot achieve perfect disentanglement. In equation 6 of [3], they define disentanglement error to measure the weight disentanglement of a model. In Fig. 3 of [3], they found that non-linear models are not weight-disentangled enough, but linearized models are more weight-disentangled. Thus, they enhance task arithmetic performance via linearization, fine-tuning in the tangent space to their pre-trained initialization, which is the same as training a kernel predictor with kernel kNTK (as shown on pages 6-7 in [3]).

Normally, due to the difficulty in achieving weight disentanglement, there is interference within different task vectors (which could be introduced by non-perfect disjoint regions of the input space). As a result, recent studies focus on resolving the interference when combining different task vectors [1,2,7].

Therefore, assuming perfect weight disentanglement is unfair when comparing it with logit distillation. I feel that one way to consider a more concrete theory might be by showing that the degree of influence on the parameter space is smaller than in the logit space.

**Question 2:**

Thank you for your efforts in conducting extra experiments. In Table 2 in  rebuttal pdf the second row should be CIFAR-100?
May I know how the new knowledge distillation losses were added? Were they used to replace the KD loss in FedDF? Because it is a little surprising that they perform much worse than the original FedDF.

**Question 3:**

It is a little tricky to say that parallelization can reduce the computational cost, since the amount of resource usage still grows in O(M^2).

If you can resolve these remaining issues, I will increase my score.

[1] TIES-Merging: Resolving Interference When Merging Models

[2]   Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch

[3]  Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models

[7] AdaMerging: Adaptive Model Merging for Multi-Task Learning