Adaptive Latent-Space Constraints in Personalized Federated Learning
Authors: Sana Ayromlou
D. B. Emerson
Venue: NeurIPS 2025 poster ()
Keywords: Federated Learning
Maximum Mean Discrepancy
Data Heterogeneity
Deep Learning
Medical Imaging
Primary area: deep_learning
TL;DR: This paper demonstrates that applying adaptive latent-space constraints in personalized FL algorithms improves performance across a number of challenging benchmark tasks, especially those with significant feature heterogeneity

Abstract: Federated learning (FL) is an effective and widely used approach to training deep learning models on decentralized datasets held by distinct clients. FL also strengthens both security and privacy protections for training data. Common challenges associated with statistical heterogeneity between distributed datasets have spurred significant interest in personalized FL (pFL) methods, where models combine aspects of global learning with local modeling specific to each client’s unique characteristics. This work investigates the efficacy of theoretically supported, adaptive MMD measures in pFL, primarily focusing on the Ditto framework, a state-of-the-art technique for distributed data heterogeneity. The use of such measures significantly improves model performance across a variety of tasks, especially those with pronounced feature heterogeneity. Additional experiments demonstrate that such measures are directly applicable to other pFL techniques and yield similar improvements across a number of datasets. Finally, the results motivate the use of constraints tailored to the various kinds of heterogeneity expected in FL systems.

Adaptive Latent-Space Constraints in Personalized Federated Learning
Sana Ayromlou" Vector Institute & Google Toronto, Ontario, CA ayromlous@gmail. com
D. B. Emerson Vector Institute
Toronto, Ontario, CA
david. hww.dd..merrooeettttiiisssitteee
Abstract
Federated learning (FL) is an effective and widely used approach to training deep learning models decentralized datasets held by distinct clients. FL also strength- ens both security and privacy protections for training data. Common challenges associated with statistical heterogeneity between distributed datasets have spurred significant interest in personalized FL (pFL) methods, where models combine aspects of global learning with local modeling specific to each client's unique char- acteristics. This work investigates the efficacy of theoretically supported, adaptive MMD measures in pFL. primarily focusing on the Ditto framework, state-of- the-art technique for distributed data heterogeneity. The use of such measures significantly improves model performance across variety of tasks, especially those with pronounced feature heterogeneity. Additional experiments demonstrate that such measures are directly applicable other pFL. techniques and yield similar improvements across number of datasets. Finally, the results motivate the use of constraints tailored the various kinds of heterogeneity expected in FL systems.
Introduction
Federated learning (FL.) has become indispensable tool for training deep learning models distributed datasets. Such a setting arises naturally in many scenarios where, for various reasons including privacy, security, and resource constraints, training data should or must reside in disparate locations. Federally trained models receive updates from larger collection of training data than models trained on single data silos and, as such, often demonstrate better performance, generalizability. and robustness [34] In typical FL system, clients each hold a distinct dataset on which a model is be collaboratively trained. A server is used perform model aggregation of some kind, requiring transfer of model weights but never raw data. In each round, clients perform model training using their local dataset. After a period of local training, the server aggregates the individual models, or subsets thereof, and sends the aggregated information back to the clients for another round of training.
While data heterogeneity is challenging in standard model training, is particularly pernicious in FL settings due to the increased prevalence of heterogeneity in distributed datasets [19 40). Many works have aimed address the effects of data heterogeneity in FL since the original mechanism of FedAvg was introduced [30]. One branch of study considers techniques for robust global optimization, seeking to restrict or correct local divergence [4 28 20 45). However, such approaches, referred to here as global FL methods, typically train single model for all clients. In many instances, system and statistical heterogeneity limit the existence of a shared model that performs well across all FL participants [6 3 This has led to the rise of personalized FI. (pFL.) methods, which federally train distinct models each client. Specifically, this work focuses on improving cross-silo pFL. methods, where clients represent small number of reliable institutions with sufficient training resources [3]
*Correspondigg author. Work done while at the Vector Institute.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

A common technique for addressing heterogeneity during federated training is the augmentation of each client's local loss with penalty discouraging large deviations in model weights or representa- tions from a reference model during training. In this work, we propose integrating optimizable MMD measures [9] into penalty-based pFL. frameworks, modifying the Ditto and MR-MTL. algorithms as representative examples [21 24]. Specifically, Multi-Kernel MMD (MK-MMD) [10] and MMD-I
[23] are applied to constrain divergence between models' latent representations. Ditto remains one of the best performing pFL methods available, especially in challenging heterogeneous settings. Several
recent benchmarks show that it delivers state-of-the-art performance in many settings E 29 As
such, it is used here as the primary pFL approach for experimentation.
The contributions of this work are three-fold. First, it proposes using theoretically supported statistical distance measures as penalty constraints for dealing with heterogeneity in pFL, moving beyond previously considered paired feature-based constraints. Second, it demonstrates that leveraging the adaptability of MK-MMD or MMD-D measures through iterative re-optimization provides notable performance improvements settings with high feature heterogeneity where existing methods, such as Ditto, under-perform. Finally, experiments with either natural or controllable levels of label and feature heterogeneity highlight the strengths and weaknesses of the types of drift penalties investigated. While Ditto is the main focus, additional results show that the proposed methodology is directly applicable to other penalty-based optimization techniques in FL and provides similar benefits.
2 Related Work
The most common formulation of penalty-based constraints in FL discourages model weight drift from set of reference weights in the g² norm. This has been applied successfully in global FL
methods like FedProx (20) and pFL approaches such as Ditto [2] and MR-MTL 24). Recent work, has modified this constraint to consider feature representations. Both MOON [8 and PerFCL [46] use contrastive losses to encourage feature representations to be close, or far, from some reference. In certain settings, MOON demonstrates marked improvements over FedProx as a global FL approach, partially motivating the investigation here. However, the contrastive losses used in MOON and PerFCL are not adaptive and require careful hyperparameter tuning. Furthermore, in (39. PerFCL underperformed contemporary pFL approaches, including Ditto.
As part of the work proposing Ditto, the authors experiment with two alternatives to the standard (2 norm to measure weight divergence, symmetrized KL. divergence [12] and Fisher-weighted squared distances 44. Therein, no performance improvements are achieved. In addition, these alternative regularizers are still applied to quantify differences in model weights, rather than latent spaces. Finally, the measures are not trainable as part of the learning process.
Prototype-based PFL methods, such as FedProto [8 share some similarities with this work. They aim to constrain local representations tied class labels. or prototypes, from drifting too far from their global counterpart. However, these approaches are generally only applicable to classification tasks by design. Prototypes are also most useful near the final layers of model, where class separation becomes sharp. Finally, drift is measured with e² or tl norms rather than the adaptive measures proposed here. There is, however, an opportunity for future work to explore the utility of integrating adaptive MMD measures into these techniques. Appendix Be provides additional discussion of other, less related, pFL methodologies and how the modifications proposed here might be applied.
A few works have considered integrating MMD into FL systems. FedMMD incorporates a fixed MMD measure on local features as means of modifying server-side aggregation in FedAvg [IT]. This is strictly a global FL. approach and applies MMD to modify model aggregation rather than local learning. In [42) a static MMD measure constructed and used constrain a model's latent space during client-side training. However, as with FedMMD, the proposed algorithm is not pFL approach and does not consider kernel optimization. also lacks robust experimentation. Non-adaptive MMD measures have also been used successfully to overcome cnature-distribution shifts when federally training generative models for local data augmentation 4. Finally, various forms of MMD have been used in other areas, most widely in hypothesis testing and domain adaptation [IO] 24 25 41 [7]. To our knowledge, this is the first study considering the utility of substituting or augmenting the loss regularizers of Ditto and MR-MTL. with an adaptive, feature-based regularizer in the forms of MK-MMD and MMD-D. Because weight-based, e² penalties are common in other global and pFI. methods, the results have broader implications beyond improving any single pFL technique.
2

---

3 Methodology
At its core, FL enables multiple clients to train models collaboratively on distributed data while sharing only model parameters with server, rather than transferring raw data. Throughout, assume there are N clients with datasets, Di...., ,, D such that D; {I y)) Let -1 +... + n Each client incorporates a local loss function & parameterized by some of weights wc In the sections to follow, the Ditto algorithm is used as an exemplary approach for experimentation. However, the methodology is transferrable other pFL methods, with additional experiments showing similar benefits for MR-MTL through the proposed techniques.
Ditto maintains two models, sharing the same architecture, on each client i, global model with weights w and a local model with weights w/.". Algorithm summarizes Ditto training, where T denotes the number of FL rounds, denotes the number of local training steps performed by each elient, and Φ is the initial set of weights for both the global and local models. In practice, any optimizers may be used to train the global and local models, with potentially different learning rates. Likewise, the number of local training steps may differ. Here, however, as in the original work, these are coupled for simplicity. For standard Ditto, the measure d(w (6) w) ||
(1) wl12-
Algorithm 1: Ditto algorithm with FedAvg aggregation and batch SGD for local optimization.
Input: : N, T, 8, A, 1), W.
Set L1 (6) = i for each client i. do
for t T
for each client in parallel do
Set u (x)
for s iterations, draw batch do
- (1) (2)
ww v/V (b wg (0)
w1 = w v (e. (1:w(" + dww m))
(0)
(1)
end
Send G to server for aggregation.
end
N ny W
end
3.1 Beyond Weights: Feature-Drift Constraints
The e--wiight constraint imposed by Ditto has been effective in many settings. However, the relationship between the type and degree of data heterogeneity to the most effective kind of drift- constraint remains under-explored. In this work, we experiment with both weight- and feature-drift penalties, analyzing their impact on learning personalized models in various settings. To define the adaptive constraints investigated below, consider splitting model into one or several stages, with intermediate outputs representing latent features. For simplicity of presentation, the setting of single latent space is discussed here, where the extension to additional layers is straightforward. The global and local models, along with their weights, wG = [0c, &G) and wy.- = [0r,, Φr.), are decomposed into feature extractor [(; 0) and classifier g(;;)) For an input local predictions are generated as i = gf (z; 0L); LL). Global predictions are produced analogously. In Algorithm I the reference weights, are also decomposed as w = [0, o).
Rather than penalizing divergence of the local model from the global one via the 6 distance between weights, measure of drift between latent features is defined. Such constraints may be constructed in two ways: paired or unpaired. In paired approaches, which are widely used [I8 46), for each input I, the aim is to minimize dif (x; OL), fz; 8c)). This enforces wise alignment, but disregards statistical interpretations of the feature space. In contrast, we propose an unpaired approach which reduces the distance between the probability distributions of the global and local latent spaces, allowing greater flexibility for models to adapt based on data heterogeneity. Here, cosine similarity used bascline to represent paired feature alignment by maximizing similarity.
3

---

3.2 Maximum Mean Discrepancy
Let X be some topologcal space and P and Q Borel probability measures on X. Any symmetric and positive-definite kernel k X x X + IR induces unique Reproducing Kernel Hilbert Space (RKHS), Hg. into which P and Q are embedded and their distance measured [3]. This is written
MMD (P, Q; Ha) ((P,)")<PPPP)) k(x, x'), + ((yx'')~(Q.)) k(y.v/) 2 (P.))))((P.)) k(r, y). In this work, X C R where m is model's latent-space dimension, and the feature extraction modules f(x; 0.) and f(x; 0) produce distributions on X.
MMD measures are powerful tool for efficiently measuring the distance between two distributions. However, their effectiveness depends heavily on the kernel, k. As such, kernel optimization techniques are crucial. Herein, the primary goal of kernel selection is to maximize test power, i.e. the ability of the measure to accurately distinguish two distributions. Two kernel optimization techniques are considered to produce strong measures of feature-space discrepancies: MK-MMD and MMD-D. The former optimizes a linearly weighted combination radial basis functions (RBFs) of varying length scale. The latter leverages a learnable deep kernel to construct strong distance metric. RBFs are not the only kernel type used in MMD, but they are a common and nigh-performing choice 35
3.3
Multi-Kerne MMD
In
10 a set of possible kernels defined as
d
d
K = k
k Σ
Byki
33 33=1,,0
(. (l..,.a)
where {ky) is a set of symmetric and positive-definite functions, kn : X x X - R. Note that any k uniquely defines an RKHS and associated measure
MK-MMD"(P,Q;3H5)
Σ B M.M j=1
Intuitively, any k E K. produces measure of discrepancy between two distributions. The novel idea of [10] is to optimize the coefficients, Sy. to minimize the Type-Il error in testing if two distributions, P and Q, are the same for fixed Type-I error. That is, one engineers an RKHS that has the best chance of properly detecting when two distributions differ through specification of an optimal kernel.
Denoting by 3 E R2d the coefficients of the kernels ki, Type-Il error minimization expressed as
MK-MMD² (P, Q; H2)
3. arg max
O,
320
where 2, is the variance of MK-MMD Hc). This variance has the form dP Q H₂) = Qu.8 where Qk is the respect to P and Q.
covariance matrix between kernels, kj, with
Let X and Y be sets of samples drawn from P and Q, respectively. To approximately solve Equation ①. empirical estimates of each MMD2 (P, Q; Hs,). gathered into vector n [mi. ma) and
Qu. denoted Q.., are computed. Following 10. the objective is approximated and rewritten as the quadratic program
B. = arg min ST(22 + EL)
(2)
g m=1
820
where 0 is stabilizing shift and I is the d x d identity matrix.
In the experiments, c= 1c-3. The kernels selected to form IC are RBFs of the form ky y) --,,ll=-vl for a set of {_ Note that this optimization is only valid if m contains at least one positive entry. If this is the case, [10] suggests selecting the kernel, ky. that max- imizes the ratio MMD Q: Hx, (P,,) Q; Hr,). Finally, linearly scaling estimates for both MK-MMD Q; Ho) and Q2 are provided in [10] Unfortunately, these estimates performed poorly empirically. As such, they are not used in the experiments to follow.

---

3.4 MMD-D
Some studies have shown that simple RBF kernels, varying only in length-scale, 7 may underperform in more complex spaces where distributions vary significantly across regions [23 [5]. To address this, the use of deep kemels,, which can adapt the diverse structures needed in different regions, has been proposed. Here, the structure of [23] is used. Given inputs x: and y, a deep kernel is constructed by applying kernel function, to the output of a featurization network, ye, parameterized by w as k(p((z; w) (p(;; w)). Alone, such kernel may inadvertently learn to treat distant inputs as overly similar. To mitigate this, safeguard is introduced defining the full deep kernel as
(x, y)
()p((())))) ((r:w)) + eq(s, y);
where q is separate kernel acting on the unmodified input.
As in Section 3.3 RBFs are used such that k<x,, 3) - - and q(2, y) = e The variables of e, 7k. Tg. and w are optimized monolithically. Again, the goal of such optimization is to minimize the Type-II error of the induced measure which, when using the kernel in Equation (3) is denoted MMD-D" (P,C Q: Hx)) Similar to MK-MMD, this error minimization takes the form
(M,
max
(4)
where o² (P, Hy.) now the variance MMD-D?(P,Q:HH.)
For samples X and y drawn from P and Q, respectively, the estimate for MMD-D Q; H.))
derived [23] Equation 2] is used. To estimate (P, Q; H₂). the regularized estimator of [23
Equation 5] is applied with 1e-8. fixed number of steps using an Adam optimizer [26 with
a learning rate of 1e-3 is applied to approximately solve Equation (4) and learn kernel (x, y).
W
Hidden Layer
Hidden Layer
Hidden Layor
Output Layers
Model Base
[xd(fi(x; 11, f1(x; 0.))
pad( (xx (s())ss)
+ Allly 1112
(2d[f2(x;0 01,)) fe(x;;0))
8(y, wL)
Hidden Layer
Hidden Layer
Hiddon Laper
Output Layers
Model
X
WL
Feature-drift constraints applied to three latent spaces of a model, where (x.y y) represent a batch of labeled data. In Ditto, the frozen global model (top) is used to constrain the local model (bottom) during client-side training.
Figure 1:
3.5 Adaptive MMD Measures in FL
Provided feature-extraction maps for the global and local models in the Ditto framework, fa(;;; 0:) and f((;; L.,)) the MMD measures defined in the previous section are used to augment the local model's loss function during training with weights μ4. This process is exhibited in Figure [ for three intermediate latent spaces. Note that (x, y) in the figure represents batches of data rather than individual data points. These measures may be used in isolation combined with the standard Ditto penalty of AlIwo wIl.. Both MK-MMD and MMD-D aim to adapt non-linear kernels measure the distance between two distributions more accurately.

---

In a static setting, where P and Q are fixed, the optimization processes for these approaches yield strong measures to distinguish these distributions. Here, these techniques are used to measure and constrain the differences in distributions of model latent spaces which evolve significantly throughout the FL. training process. As the model feature maps evolve, previously optimal kernels become sub-optimal. To address this, we propose taking advantage of the optimization procedures discussed above to periodically re-optimize the MMD measures using samples of training data. In this way, strong measures of latent-space divergence are maintained throughout training.
4 Datasets
To quantify the utility of the proposed measures, four datasets are used. Each poses unique challenges in the FL setting. The first set of experiments considers several variants of CIFAR- 10 with differing levels of label heterogeneity. This heterogeneity is synthetically induced using Dirichlet allocation, a common strategy in previous work 77 [ (3. Smaller values of the allocation parameter, a, result in more heterogeneous clients. Values of (5.0, 0.5, } are used to split data between five clients. a generated dataset with controllable levels of feature
The second dataset, referred to as Synthetic, heterogeneity across eight clients. It is an extension of the approach used in [20]. Heterogeneity depends on parameters a and B. with larger values corresponding to more heterogeneity. Two datasets are generated with = B = 0.0 and = 3 = 0.5. Inputs have 60 numerical features mapped to one of 10 possible classes. Appendix | provides details on the generation process.
The final two datasets focus on real clinical tasks with natural client splits and heterogeneity. Fed- ISIC2019 is drawn from the FLamby benchmark [8] and consists of 2D dermatological images to be classified into one of eight melanoma categories. The data is split across six clients. The RxRx1 dataset (36] is composed of 6-channel fluorescent microscopy images, each highlighting different cellular organelles. The task is to classify which genetic treatment, if any, the cells received based on the image. The original dataset contains 125K+ images across 1,108 classes. Following [4, only samples corresponding to the 50 most common classes are selected and three of the image channels are used. The data is partitioned into four clients based on the hospitals where the images were collected. This dataset exhibits significant feature heterogeneity. Both Fed-ISIC2019 and RxRxI are quite challenging for global FL methods like FedAvg.
For all tasks, performance is measured through standard accuracy. with the exception of Fed-ISIC2019. Because the labels in Fed-ISIC2019 are highly imbalanced, balanced accuracy is used better quantify model performance and matches the metric used in the FLamby benchmark.
5 Experimental Setting
In the experiments, hyperparameter sweeps are conducted to calibrate items such as learning rate for all methods. A full list of the hyperparameters considered and their optimal values appears in Appendix B Other parameters, such as batch size, are detailed in Appendix C The metrics reported in the results the average of values across three training runs. Accompanying standard deviations are gathered in Appendix G Within each FL training run, the final metric is the uniform average of each client's performance on their respective test sets
Each task uses a different model architecture. For the CIFAR- 10 datasets, a CNN is used with two convolutional layers, each of which are followed by batch normalization (BN) and max pooling. The final two stages are fully connected (FC) layers ending in classification, matching the benchmark in 3. Experiments are conducted with up to three latent spaces constrained with the MMD-based drift penalties. These are the output of the first FC layer, the output of the second BN layer, and the output of the first BN layer. When referring to the constraint "depth,"" 1 corresponds to constraining only the FC layer, 2 corresponds to adding the second BN layer, and refers to adding the first BN layer (e.g. Figure T]. During experimentation with CIFAR-10, 10, the optimal "depth" for the MMD constraints was uniformly 1. Because adding additional constraints did not improve performance, the remaining experiments simply applied MMD-based constraints at single latent layer. For more detailed discussion and results, see Section 6.2
2 All code is found at: http:///ttthhoooeeerrtttt/ttt/tt

---

The model for the Synthetic datasets is simple two-layer DNN followed by a softmax layer. The first layer serves as the feature extractor for constraint purposes. The second layer produces classification. For Fed-ISIC2019, EfficientNet-IR BO [37, pretrained on ImageNet, is fine-tuned. The final linear layer serves as the classifier, and the rest of the model functions as the feature extractor. This architecture is also used in the Flamby benchmark [8]. Finally, pretrained ResNet-18 18 model is fine-tuned for the RxRx1 task, as done in [4]. The MMD penalties are applied to the output of the model prior to the final classification layer.
During local training, two ways of updating the MMD measures are considered. In the first approach, the measures are updated every 20 steps of model training using 2 batches of sampled training data. Alternatively, the measures are updated after every training step using the same batch of data. In either setting, the MMD-D kernel is optimized using AdamW. For MK-MMD both approaches are reported and denoted with subscript 20 or -1, respectively. For the CIFAR-10 10 and Synthetic datasets, MMD-D is trained using the periodic update approach with 5 optimization steps, while on Fed-ISIC2019 and RxRx1, it updated after every training step with 25 optimization steps.
5.1 Communication, Privacy, and Computation Implications
With the proposed modifications to Algorithm [ communication costs between clients and the server remain the same as standard Ditto. Only the global model is exchanged for aggregation and feature representations remain local. As such, the privacy properties are also nearly identical to Ditto, except that local training data periodically used to optimize the MMD kernels on each client. There also a small uptick in memory overhead with the need to store a modest number of feature activations and to train a set of kernel weights for MMD-D. The largest impact associated with the proposed approach is the computational cost of adapting the MMD measures during training. Optimizing the kernels for batch of training data adds (m33)) and (kkq) arithmetic operations to each local training step for MK-MMD and MMD-D, respectively, where m is the latent- spent--pace dimension, & is the number of kernel optimization steps, and g is the arithmetic operations associated with an SGD step for yp For more details on computational costs and runtime complexity, see Appendix D
6 Results
In the first set of experiments, performance of the MMD-based feature constraints is considered in isolation. That is, the traditional Ditto constraint in Algorithm [ is replaced with one of the proposed adaptive measures. As a baseline, training with cosine-similarity measure to penalize feature dissimilarity is also explored. The results are reported in Table D
Across all datasets, and heterogeneity levels. at least one of the MMD-based penalties outperforms the cosine-similarity baseline. In the case of the Synthetic, Fed-ISIC2019, and RxRx1 datasets, the performance gaps are fairly large. These results validate the effectiveness of the proposed measures in guiding federated training by adaptively limiting the brasere-based penalties within the Ditto framework significantly outperforms FedAvg. The only setting where this is not the case is CIFAR-10 with 5.0. This is the most homogeneous setting, and the convergence of FedAvg is minimally impacted. Below, combining the standard Ditto penalty with MMD- based measures boosts performance past FedAvg, even for a 5.0.
drift. In most cases, using
For the second set of experiments, the MMD-based measures are combined with the original Ditto penalty to consider the complementary utility of applying both constraints. The resulting local model loss function for client i is written
& 6; (/
+ 2
-11l+ + pd
(5)
where represents a batch of training data and d(.,,) -) denotes one of the MMD-based measures. For comparison, standard Ditto is also used to federally train models for each task. This is equivalent to setting μ = 0 in Equation (5). The results for each dataset are summarized in Table 2
Comparing the results of Tables I and 2 vanilla Ditto outperforms both MMD-D and MK-MMD when used in isolation for all variants of CIFAR- 10 and Fed-ISIC2019. However, for the Synthetic and RxRx1 I datasets, which demonstrate notable feature heterogeneity, using either MMD-D or MK- MMD yields large performance improvements compared to standard Ditto. For RxRx 1, including the Ditto penalty produces small, additional improvement with the MMD-D measure, but can
7

---

Table 1: Average performance when replacing the standard Ditto constraint with various feature-drift penalties. Bold and underline indicate the best and second best value across pFL. methods. Subscripts for CIFAR-10 10 and Synthetic indicate values for a and = 3. respectively.
Feature Drift Constraint
Dataset
FedAvg
Ditto
Cos. Sim.
MMD-D MK-MMD-1
MK-MMD20
CIFAR-IO0.
71.220
84.930
84.212
83.789
84.136
84.439
CIFAR-100.5
75.575
702
75.167
75.094
75.678
76.564
CIFAR-1O5,0
77.284
77.658
67.298
67.729
68.718
68.832
Synthetic
84.733
89. 19.129
89.975
90.237
91.418
90.066
Synthericu.s
85.458
85.533
90.199
91.270
91.137
90.262
Fed-ISIC2019
64.057
71.350
61.269
64.302
60.168
62.677
RxRx1 35.207
65.629
64.985
67.478
65.861
67.078
Table 2: Average performance when augmenting the standard Ditto constraint with various fea- ture drift constraints. Bold and underline indicate the best and second best value across methods. Subscripts for CIFAR-10 and Synthetic indicate values for a and a = B, respectively.
+ Feature Drift Constraint
Dataset
FedAvg
Ditto
Cos. Sim.
MMD-D
MK-MMD_1
MK-MMD200
(IARR-10.3
71.220
84.930
84.924
85.214
84.723
84.900
CIFAR-100.5
75.575
80.702
80.669
80.696
80.936
80.976
CIFAR-IOs.
77.284
77.658
78.052
77.739
77.578
77.739
Synthetic
84.733
89.129
89.187
89.458
89.183
89.258
Synthetic
85.458
85.533
87.783
89.695
88.154
88.10
Fed-ISIC2019
64.057
71.350
70.970
72.226
71.169
71.267
RxRx1
35.207
65.629
67.027
67.755
65.984
66.892
degrade accuracy with MK-MMD. Similarly, adding the Ditto penalty for the Synthetic datasets hurts performance across the board. By construction, the Synthetic datasets have aspects that are adversarial to weight-based constraints, as the feature and label generation process incorporates normally distributed local weight variations. As such, these constraints may hinder learning at certain point. As noted by the dataset creators, RxRx1 exhibits deep feature heterogeneity, which is especially challenging for FL algorithms. The results suggests that the proposed MMD-based measures are particularly well-suited for the kind of underlying heterogeneity in these datasets, while standard Ditto constraints are less effective.
The results Table 2 also show that combining the traditional Ditto constraint with an MMD-based one produces the best results for Fed-ISIC2019 and CIFAR-10. in all but one setting. In the case of CIFAR-10, the improvements are modest but persistent across different levels of heterogeneity. Interestingly, despite yielding less favorable results as standalone constraint, pairing the Ditto and cosine similarity penalties does provide best performance for the most homogeneous version of CIFAR-10. In this setting, the MMDD-based constraints remain second best. When using MMD-D in combination with Ditto for Fed-ISIC2019, the improvement is more prominent, especially given how challenging the task is 8 39.. Finally, the combined results imply that the flexibility offered by the deep kernel of MMD-D produces better results than MK-MMD, especially for complex datasets.
The results above demonstrate that replacing or augmenting the traditional Ditto penalty with the adaptive MK-MMD or MMD-D constraints yields performance benefits. Notably, the proposed approach is transferrable to other pFL methods that use penalties to guide local training. To exhibit the utility of the adaptive constraints other settings, experiments augmenting MR-MTL, another high- performing pFL. method, with MK-MMD or MMD-D latent-space penalties are conducted. The results are reported in Appendix[F Integrating the adaptive measures produces accuracy improvements over the standard MR-MTL algorithm in all settings, with the exception of FedISIC-2019.

---

6.1 Interaction of the Ditto and MMD Constraints
For the Synthetic and RxRx1 tasks, the traditional Ditto penalty under-performs compared to applying either MK-MMD or MMD-D in isolation. However, for the other datasets, a natural question arising from the results is whether the benefits of adding the MMD-based penalties are captured by simply increasing the Ditto constraint weight. To consider this, performance is measured for variety of near the optimal value for the CIFAR- and Fed-ISIC2019 datasets. Ditto is both augmented with the MMD-based measures and considered alone. The results are shown in Figure
Accuracy for CIFAR-10 (a 5.0)
Accuracy for CIFAR-10 (a 0.5)
Average Accuracy
77.7
Average Accuracy
77.6
80.5
77.5
80.0
77.4
79.5
Ditt
Ditic
77.2
MMD-D
79.0
MMD-D
0.1
0.5
1.0
2.0
0.1
0.5
1.0
2.0
Ditto Loss Weight (A)
Ditto Loss Weight (A)
Accuracy for CIFAR-10 (a 0.1)
Accuracy for Fed-ISIC2019
Average Accuracy
Accuracy
73.0
84.5
71.0
83.5
70.0
833.0
69.0
82.5
Ditto
MK-MMD
MK MMD
82.0
MMD-O
67.0
MMD-D
0.1
0.5
1.0
2.0
0.1
1.0
2.0
5.0
10.0
Ditto Loss Weight (A)
Ditto Loss Weight (A)
Figure 2: Results for CIFAR-10 and Fed-ISIC2019 when varying around the optimal value. The best result is reported for MMD-based penalty weights, μ. drawn from {0.01, 0.1,1.03.
For a = 5.0, adding MMD-D outperforms Ditto-only for all values of X. MK-MMD also improves performance for Ả at or above the best value. In the 0.5 setting. augmentation with MK-MMD is better or equivalent to standard Ditto with increasing x. In this setting, MMD-D under-performs for 0.5 but is equivalent or better than Ditto for other values. In the high-heterogeneity regime, the results are more mixed. For the optimum X, adding MMD-D provides a performance boost, but decays more rapidly than Ditto as increases. The addition of MK-MMD simply yields comparable performance. At the highest levels of heterogeneity, weaker constraints, which permit more substantial local model divergence, can be advantageous. As such, the stronger penalties provided by augmentation with the MMD-based measures may be less helpful. For Fed-ISIC2019, both MMD-D and MK-MMD maintain an advantage over Ditto for nearly all values of MMD-D achieves the highest peak value by good margin.
Overall, the results demonstrate that the benefits of the new constraints are not simply replicated by increasing A. Rather, these constraints capture beneficial aspects of drift not considered by the purely weight-based penalty. Intuitively, the weight- and feature-based constraints impact training dynamics differently. Weight-drif penalties weakly constrain model to a neighborhood of reference weights. However, with model depth and nonlinearity, small weight perturbations can still produce large deviations in representations and outputs. We conjecture that combining such penalties with strong feature constraints can produce a more uniform drift penalty, providing complementary benefits in certain settings. Furthermore, in many contexts, single form of heterogeneity will not characterize all drift between FI. clients. Thus, penalizing drift in distinct ways is likely to be advantageous.
6.2 Ablation Studies
In the hyperparameter sweeps for CIFAR-10, the best results are obtained with a constraint depth of 1. This is illustrated on the left of Figure 3 For fixed À 0.1. accuracy degrades with the

---

introduction of additional penalties. We conjecture that constraining too many layers in a relatively small model, as used in the CIFAR-10 experiments, produces heavy constraint and hinders learning. This relationship informed the decision to apply MMDD-based constraints to the latent space of single layer for the remaining datasets. Further performance improvements may indeed exist for other datasets with different constraint configurations, as this simply limited the configuration search space.
Accuracy for CIFAR-10 (a 0.1)
Synthetic Dataset Accuracy for Fixed Kemels
Average Accuracy
85.0
Average Accuracy
91.5
84.0
83.0
91.0
82.0
61.0
90.5
MK-
80.0
MMD-D
79.0
Ditto NKG-MD
90.0
Fixed (0.0)
MK MMD (0.0)
MMD - (0.9)
Ditto MMD
Fixed (0.5)
MK-MDD (0.5) MMD-D (0.5)
78.0
89.5
1
3
5
9
13
18
Kernel Depth
Number Kernels
Figure 3: Evolution of accuracy on the CIFAR-10 (a: = 0.1) dataset (left) with an increasing number of constraints at various model depths for the MMD measures with and without the standard Ditto drift penalty. Accuracy for Synthetic datasets (right) using varying number of fixed kernels with uniform weight compared to using the kernel optimization procedures of MK-MMD or MMD-D.
Improvements to test power when optimizing MMD kernels has been established in previous work 23 35.. However, validate that the kernel optimization procedures translate to task improvements, the performance of these approaches is compared to using a fixed set of RBF kernels linearly combined with uniform weights. In the experiment, the first r' kernels are selected from the original 3.5 to 2 with exponents evenly spaced in increments of 0.25, For example, when r = 5, kernels corresponding to ?/j = 2 2-3.5 2-3.25 2- 2- 2-2.5 are used. Each receives uniform weight equal 1/r such that the MMD measure is the average value of the individual measures induced by the fixed kernels. Results for the two Synthetic datasets are shown on the right in Figure 3 Kernel optimization with MK-MMD and MMD-D both outperform all fixed kernel settings. Notably, the 18-kernel case is equivalent to the MK-MMP setting in previous experiments, but with the kernel weight optimization procedure disabled. The optimization processes clearly generate substantive improvement in accuracy in this setting.
18 generated by "/ values from 2 23355
7 Conclusions, Limitations, and Future Work
In this work, two adaptive MMD-based measures are proposed to constrain feature representation divergence between models in FL. The Ditto and MR-MTL frameworks for pFL are expanded to incorporate such measures, penalizing local-model drift from a model aggregated with FedAvg on the server. We find that using these measures of feature-space distributions alone, or in tandem with traditional weight- based distances, produces marked performance improvements across several tasks, including important real- world, clinical benchmarks. The main results and ablation studies demon- strate that, rather than simply providing a parallel penalty to the standard weight-drift constraints, the MMD-based measures differ kind. These measures appear especially effective settings with significant feature heterogeneity between clients, such as in the Synthetic and RxRx datasets. Both measures are effective, but MMD-D empirically provides more consistent improvements.
While the results of this work are promising, some limitations are worth noting. The experiments focus on the utility of the MMD measures in the context of pFL frameworks with weight-based penalties. However, other methods, such as FedProto, may also benefit from introducing such penalties. Moreover, while the experiments indicate settings where the proposed measures offer notable benefits, this work does not theoretically quantify the conditions under which such measures are most effective, Some preliminary thoughts on such theory are provided in Appendix H but more work is required. Finally, as noted in Section [ the kernel optimization procedures introduce non-trivial computational costs. Avenues for reducing such costs are discussed in Appendix D The above questions and limitations will be the subject of future work.
10

---

Acknowledgments
The authors would like extend heartfelt thank you to Fatemeh Tavakoli of the Vector Institute for her help during the rebuttal period. Her efforts made this publication possible and we are truly grateful. We would also like to thank the anonymous referees for their helpful comments and suggestions that improved the quality of this work.
References
[1] M. G. Arivazhagan, V. Aggarwal, A. K. Singh, and S. Choudhary. Federated learning with personalization layers, 2019.
[2] S. Caldas, S. K. Duddu, P. Wu, T. Li,J. Konečný, H. B. McMahan, V. Smith, and A. Talwalkar. LEAF; A benchmark for federated settings. In Workshop on Federated Learning for Data Privacy and Confidentiality, 2019.
[3] D. Chen, D. Gao, W. Kuang, Y. Li, and B. Ding. pFI.-hench: a comprehensive benchmark for personalized federated learning. In Proceedings of the 36th International Conference Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2022. Curran Associates Inc.
[4] H. Chen, A. Frikha, D. Krompass, J. Gu, and V. Tresp. FRAug: Tackling federated learning with non-IID features via representation augmentation. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4849-4859 2023.
[5] L. Collins, H. Hassani, A. Mokhtari, and S. Shakkottai. Exploiting shared representations for personalized federated learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pages 2089-2099, PMLR, 18-24 Jul 2021. URL httphtttt///pceeediii
mlr. press/ v/////13////// cols/oollnsnlle htm.
[6] Y. Deng, M. M. Kamani, and M. Mahdavi. Adaptive personalized federated learning. 2020.
[7] E. Diao, Ding, and V. Tarokh. HeteroFL: Computation and communication efficient federated learning for heterogeneous clients. In International Conference on Learning Representations, 2021.
[8] O. du Terrail, S.-S. Ayed, E. Cyffers, F. Grimberg, C. He, R. Loeb, P. Mangold, T. Marchand, O. Marfoq, E. Mushtaq, B. Muzellec, C. Philippenko, S. Silva, M. Teleñczzk, S. Albarqouni, S. Avestimehr, A. Bellet, A. Dieuleveut, M. Jaggi, S. P. Karimireddy, M. Lorenzi, G. Neglia, M. Tommasi, and M. Andreux. FLamby: Datasets and Benchmarks for Cross-Silo Federated Learning in Realistic Healthcare Settings. In NeurlPS 2022 Thirty-sixth Conference on Neural Information Processing Systems, Proceedings of NeurlPS, New Orleans, United States, November 2022.
[9] A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf. and A. Smola. A kernel two- sample test. Journal of Machine Learning Research, 13(25):723-773, 2012. URL http 7/jmlr www.oo/g.///eeeee.. papers/ v 13/ gretton12a. /atperrs////grrtton1aa
html
[10] A. Gretton, D. Sejdinovic, H. Strathmann, S. Balakrishnan, M. Pontil, K. Fukumizu, and B. K. Sriperumbudur. Optimal kemel choice for large-scale two-sample tests. Advances in neural information processing systems, 25, 2012.
[11] K. Hu, Y. Li. S. Zhang, J. Wu, S. Gong, Jiang, and L. Weng. FedMMD: federated weighting algorithm considering non-IID and local model deviation. Expert Systems with Applications, 237:121463, 2024. ISSN 0957-4174. doi: http:://ddbo.ooo//o.0011//////// 121463. URL
[12] H. Jiang, P. He, W. Chen, X. Liu, J. Gao, and T. Zhao. SMART: Robust and efficient fine-tuning for pre-trained natural language models through principled regularized optimization. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault, editors, Proceedings of the 58th Annual Meeting of the Association for Computational I Linguistics, pages 2177-2190, Online, July 2020. Association for Computational Linguistics.

---

[13] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings, R. G. L. D' D'Oliveira, H. Eichner, S. El Rouayheb, D. Evans, J. Gardner, 7. Garrett, A. Gascón, B. Ghazi, P. B. Gibbons, M. Gruteser, Z. Harchaoui, C. He, I.. He, Z.. Huo, B. Hutchinson, J. Hsu, M. Jaggi, T. Javidi, G. Joshi, M. Khodak, J. Konecný, A. Korolova, F. Koushanfar, S. Koyejo, T. Lepoint, Y. Liu, P. Mittal, M. Mohri, R. Nock, A. Özgür, R. Pagh, H. Qi. D. Ramage, R. Raskar, M. Raykova, D. Song, W. Song, S. U. Stich, Z. Sun, A. T. Suresh, F. Tramèr, P. Vepakomma, J. Wang, L. Xiong, Z. Xu, Q. Yang, F. X. Yu, H. Yu, and S. Zhao. Advances and open problems in federated learning. Foundations and Trends@ in Machine Learning, 14( 2): 1-200 2021. ISSN 1935-8237. doi: 10.1561/2200000083. URL http:ttppp// h/tp:///xx.o doi org 10. 10w1115112220 1111111122220211333
[14] S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5132-5143. PMLR, 13-18 Jul 2020.
[15] N. Kotelevskii, M. Vono, A. Durmus, and E. Moulines. FedPop: A bayesian ap- proach for personalised federated learning. In S. Koyejo, S. Mohamed, A. Agar- wal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Informa- tion Processing Systems, volume 35, pages 8687-8701. Curran Associates. Inc.. caturpp..c/papeee httppp/rrrieeeepppe/eeeee Conf 2022. URL htttps: Ponfrr- pdd 33540067727001111111100066666663
[16] V. Kulkarni, M. Kulkarni, and A. Pant. Survey of personalization techniques for federated learn- ing. In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4), pages 794-797, 2020.
[17] K. Li, J. Lu, H. Zuo, and G. Zhang. Multi-source contribution learning for domain adaptation. IEEE Transactions on Neural Networks and Leurning Systems, 33(10):5293-3307 2022.
[18] Q. Li, B. He, and D. Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021.
[19] T. Li, A. K. Sahu, A. Talwalkar, and V. Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-66)) 2020.
[20] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith. Federated optimization in heterogeneous networks. In I. Dhillon, D. Papailiopoulos, and V. Sze, editors, Proceedings of Machine Learning and Systems, volume 2, pages 429-450, 2020.
[21] T. Li, S. Hu, A. Beirami, and V. Smith. Ditto: Fair and robust federated learning through personalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 6357-6368. PMLR, 18-24 Jul 2021.
[22] X. Li, M. Jiang, X. Zhang, M. Kamp, and Q. Dou. FedBN: Federated learning on non-III features via local batch normalization, 2021.
[23] Liu, W. Xu, J. Lu, G. Zhang, A. Gretton, and D. J. Sutherland. Learning deep kernels for non-parametric two-sample tests. In International conference on machine learning, pages 6316-6326. PMLR, 2020.
[24] Z. Liu, S. Hu, S. Wu, and V. Smith. On privacy and personalization in cross-silo federated learning. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS $22, Red Hook, NY, USA, 2022. Curran Associates Inc.
[25] M. Long, Y. Cao, J. Wang, and M. Jordan. Learning transferable features with deep adaptation networks. In International conference on machine learning, pages 97-105. PMLR, 2015.
[26] I. Loshchilov and F. Hutter. Fixing weight decay regularization in ADAM, 2018.
[27] Y. Mansour, M. Mohri, J. Ro, and A. T. Suresh. Three approaches for personalization with applications to federated learning, 2020.
12

---

[28] K. Matsuda, Y. Sasaki, C. Xiao, and M. Onizuka. FedMe: Federated learning via model exchange. In Proceedings of the 2022 SIAM International Conference on Data Mining (SDM). pages 457,467, 2022.
[29] K. Matsuda, Y. Sasaki, C. Xiao, and M. Onizuka. Benchmark for personalized federated learning. IEEE Open Journal of the Computer Society, 5(1):2-13, Jan 2024. ISSN 6644-1266
[30] H. B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. y. Arcas. Communication- efficient learning of deep networks from decentralized data. Proceedings the tional Conference on Artificial Intelligence and Statistics (AISTATS), 2017.
the 20th Interna-
[31] P. Passban, T. Roosta, R. Gupta, A. Chadha, and C. Chung. Training mixed-domain translation models via federated learning. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2576-2586, Seattle, United States, July 2022. Association for Computational Linguistics.
[32] S. J. Reddi, Z. Charles, M. Zaheer, Z. Garrett, K. Rush, J. Konečný, S. Kumar, and H. B. McMahan. Adaptive federated optimization. In ICLR 2021, 2021.
[33] B. Schölkopf and J. Smola. Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond. The MIT Press, 2018.
[34] C. Sun, A. Shrivastava, S. Singh, and A. Gupta. Revisiting unreasonable effeetiveness of data in deep learning era. In 20/7 IEEE International Conference on Computer Vision (ICCV), pages 843-852, 2017.
[35] D. J. Sutherland, H.-Y. Tung, H. Strathmann, S. De, A. Ramdas, A. J. Smola, and A. Gretton. Generative models and model criticism via optimized maximum mean discrepancy. In Sth International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings, 2017.
[36] M. Sypetkowski, M. Rezanejad, S. Saberian, O. Kraus, J. Urbanik, J. Taylor, B. Mabey, M. Victors, J. Yosinski, A. R. Sereshkeh, I. Haque, and B. Earnshaw. RxRxI: dataset for evaluating experimental batch correction methods. In 2023 IEEE/CVF Conference Computer Vision and Pattern Recognition Workshops (CVPRW), pages 4285 4294, 2023.
[37] M. Tan and Q Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference on machine learning, pages 6105-61 14, PMLR, 2019.
[38] Y. Tan, G. Long, L. Liu, T. Zhou, Q. Lu, J. Jiang, and C. Zhang. FedProto: Federated prototype learning across heterogeneous clients. In AAAI Conference on Artificial Intelligence, 2022.
[39] F. Tavakoli, D. B. Emerson, S. Ayromlou, J. T. Jewell, A. Krishnan, Y. Zhang, A. Verma, and Razak, comprehensive view of personalized federated learning on heterogeneous clinical datasets. In Kaivalya Deshpande, Madalina Fiterau, Shalmali Joshi, Zachary Lipton, Rajesh Ranganath, and fñigo Urteaga, editors, Proceedings of the 9th Machine Learning for Healthcare Conference, volume 252 of Proceedings of Machine Learning Research. PMLR, 16-17 Aug 2024.
[40] H. Wang, K. Sreenivasan, S. Rajput, H. Vishwakarma, S. Agarwal, J. Sohn, K. Lee, and D. Papailiopoulos. Attack of the tails: Yes, you really can backdoor federated learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc.
[41] Z. Wang, C. Zhou, B. Du, and F. He. Self-paced supervision for multi-source domain adaptation. In Lud De Raedt, editor, Proceedings of the Thirty-Firsl International Joint Conference on Artificial Intelligence, IJCAI-22, pages 3551-3557. International Joint Conferences on Artificial Intelligence Organization, 2022. Main Track.
[42] X. Yao, T. Huang, C. Wu, R.-X. Zhang, and L. Sun. Federated learning with additional mechanisms on clients to reduce communication costs, 2019. URL, pittps://arxin http//orv.vo//////////// abs/ abs/1908.0 10t899999809991
13

---

[43] Y. Ye. Interior Point Algorithms: Theory and Analysis. John Wiley & Sons, Inc., USA, 1997. ISBN 0471174203.
[44] Tao Yu, Eugene Bagdasaryan, and Vitaly Shmatikov. Salvaging federated learning by local adaptation, 2020.
[45] R. Zhang. Q. Xu, J. Yao, Y. Zhang, Q. Tian, and Y. Wang. Federated domain generalization with generalization adjustment. In 2023 IEEE/CVF Conference on Computer Vision and Pattern 2023.
Recognition (CVPR), pages
[46] Y. Zhang, Y. Xu, S. Wei, Y. Wang. Y. Li, and X. Shang. Personalized federated contrastive learning. In 2022 IEEE International Conference on Big Data (Big Data), pages 4218-4222 2022.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly state the paper's contributions, which include proposing the use of adaptive statistical distance measures, MK-MMD and D-MMD, constrain the divergence of local models from the server-side aggregated model for addressing various types of heterogeneity in pFL. These claims are supported by empirical results on four different datasets with varying levels and types of heterogeneity, which highlight the strengths and weaknesses of the different drift penalty methods.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
The paper includes a final section on Conclusion, Limitations, and Future Work, which also discusses limitations such as computational overhead of iteratively optimizing the kernels to better measure statistical distances and the need for additional investigation into thcoretical framework quantifying the conditions under which the proposed measures are most effective.
Justification:
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 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 few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
15

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
This paper proposes the use of theoretically supported statistical measures in FL setting. The theoretical underpinning of the measures used appear in works cited throughout. The primary contribution of this work, however, is proposing framework for their use in Fl. and extensive numerical experiments showing the effectiveness thereof.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any 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 short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
The paper provides all the details about the datasets, models, training con- figurations, and hyperparameter selection procedure, allowing for the reproduction of all experimental results.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not 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 dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, 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 instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
16

---

(c) If the contribution is new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, 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 have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes]
The paper does provide open access to code to reproduce results. All the datasets used in the experiments are public and should be downloadded/eqqueteed separately,
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
ottpe//ooooooeoassiooooooooooiee mublcc/guuideesoodettuiiiiii.. public/uuides/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.oo ccc wacc/publc//ae htaliccgguies//ee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
The paper includes detailed descriptions of the training and test setups, in- cluding hyperparameters, data generation and splits, and optimizers, in the Experimental Settings section, with additional details provided in several Appendices.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
17

---

Justification: All results are reported as the average of three runs with different, fixed random seeds, ensuring reproducibility. However, due the large number of experiments and the extensive hyperparameter search involved, experiments were limited to three runs. Additional computational resources would be required to produce tight confidence intervals in most settings. The standard deviations associated with the main experimental results are reported in Appendices and G. Nevertheless, the mean performance metrics are consistently higher across all datasets and heterogeneity types when incorporating MMD-based penalties.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
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).
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 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 range (c.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide all information about the type of GPUs used, as well as the required execution time and memory for each dataset's experiments, in Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: This research does not violates any aspects of the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
18

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
This work proposes pFL. solution designed for training in environments with data-sharing constraints, which can have broad impact in domains with strict privacy regulations, such as healthcare. To support our claims, we go beyond using only synthetic data and CIFAR-I0 10 by including two challenging medical imaging datasets, demonstrating the effectiveness of our method in real-world healthcare scenarios. Improvements in model performance, even small ones, have the potential to meaningfully improve patient care and healthcare outcomes.
Justification:
Furthermore, federated learning is widely recognized as solution for addressing data constraints. Adding to the growing set of literature showing real-world benefits with this methodology is important building towards wider adoption of the methodology in fields where it can break down barriers to model development.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA 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 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: To the best of our knowledge, this work does not introduce any aspects with high risk for misuse, including data models.
Guidelines:
The answer NA means that the paper poses 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.
19

---

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The paper cites the code, datasets and models used, acknowledging the original creators and their contributions.
Guidelines:
The answer NA means that the paper does 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 used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com /dtassts has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- 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 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code to reproduce the experiments of this paper publicly available. These assets are strictly in the form of code and library elements. Model assets and dataset artifacts will be made available request, due their size.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl 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 anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. 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 details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects.
Guidelines:
20

---

The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included 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.
15. 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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve research with human subjects any aspects that required IRB approvals.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending 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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
The work does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve L.LMs as any important, original, non ntansaandard components.
Please refer to our LLM for what should or should not be described.
21

---

A Synthetic Dataset Generation Details
Label heterogeneity of varying levels of intensity can be imposed on public datasets using Dirich- let allocation. Inducing controllable levels of feature heterogeneity is more challenging. To simulate feature-level heterogeneity, we extend the approach outlined 20]. Therein, gen- erate data for k clients, samples (Xxx Y&) are synthesized using simple one-layer network as y arg max(softmax] (Wz + b)). To produce more complicated and multi-layer relationship, the generative function is extended to use two-layer network as
y arg max (softmax (W2 (T- (Wiæ + b1)) + b2))
where x IR W1 E R20x60 b E R200 W22 E IR 10x20 by E R and T is the temperature for smoothing the output of first layer. For data generation, T =
For layer i {1, we model W ~ VV 1). : AV, 1). ut ~ NO(0, a), and va ( Σ).
where Σ is diagonal covariance matrix with E(i,)) = j- Each element in the mean vector Uy is drawn from (B,, 1), with B, ~ (0, B). Therefore, controls the variation among local models, while 3 controls the differences between local data across clients. We vary a, B to generate two levels of heterogeneous datasets, with (a, () set to (0,0) and (0.5, 0.5).
B Hyperparameter Details
The hyperparameters studied in each of the experiments, along with their optimal values, are outlined in this section. For each dataset, the only hyperparameter tuned when using FedAvg the learning rate (LRR) for their respective local optimizers. The LR is selected from the set (0.00001,0.00) 1,..000000.000 0.0100000000000000 are displayed in Table For any fixed set of parameters, three runs are performed with random seeds of (2021, 2022, 2023
The optimal hyperparameters found for each method and dataset
B.1 CIFAR-10
Following the implementation in pFL-bench BI, we use an SGD optimizer with momentum 0.9, For all levels of heterogeneity, the same parameter ranges are considered. For Ditto, the client-side LR are (0.0001,0.00) 0.001,0.01,
and À E (0.1, 0.5, 10).
For cosine similarity, when applied alone, the LR values are (0.001, 0.01). and μ has values {0.01, 0.1, 1.0]. When combined with Ditto, LR {0.001, P (0.01, 1.0,, and E (0.1, 2, 10).
In experiments applying only MMD-D and MK-MMD, the LR selected from {0.001, 0.01). The value of is drawn from (0.01,0.1, 1.03. Constraint depths of 1, 2, and 3 are considered with the optimal depth always being 1. As such, when combining the standard Ditto penalty and an MMD-based measure, depth of 1 is always used. In such settings, the LR values are {0.001, 0.01 ju drawn from {0.01, 0.1,1 1.0), and À is selected from (0.1, 1, 0.5, 1, 10).
B.2 Synthetic
Following the setup in 20). an SGD optimizer with momentum 0.9 and weight decay 0.001 is used. For both levels of heterogeneity, the same parameter ranges are considered. For Ditto, the LR is selected from {0.0001, 0.001, 0.01, 0.1) and E (0.01, 1.0.
In the remaining experiments, based on the Ditto results, the LR is fixed to 0.001 for 3) (0,0) and to 0.01 for data with (a. B) = (0.5, 0.5). The hyperparameter μ for cosine similarity, MMD-D, and MK-MMD is selected from the set (0.01, 0.1, 1.0). When these methods combined with Ditto, we additionally select from {0.01, 0.1} while continuing to tune fu within the same range.
B.3 Fed-ISIC2019
In 33 using an AdamW optimizer for this dataset led to improved performance in the FL setup, compared to using SGD, as was done in [ Based on these findings, Adam W is used in the experiments. Similar to the previous two datasets, for Ditto search for the optimal LR in the
22

---

Table 3: Best hyperparameters for each of the methods across all datasets.
CIFAR-10
Synthetic
Fed-ISIC2019
RxRxI
0.1
0.5
5.0
(0.0, 0.0)
(0.5, 0.5)
FedAvg
LR
0.01
0.01
0.01
0.001
0.01
0.0001
0.0001
LR
0.01
0.01
0.01
0.001
0.01
0.001
0.0001
Ditto
0.1
0.1
0.5
0.01
0.1
1.0
0.01
LR
0.01
0.01
0.01
0.001
0.01
0.001
0.0001
MMD-D
0.01
0.01
0.01
1.0
1.0
1.0
0.01
LR
0.01
0.01
0.01
0.001
0.01
0.001
0.0001
MK-MMD
μ
1.0
1.1
1.1
1.0
0.1
1.0
0.01
LR
0.01
0.01
0.01
0.001
0.01
0.001
0.0001
Cos. Sim.
0.1
0.1
0.1
1.0
1.0
1.0
1.0
LR
0.01
0.01
0.01
0.001
0.01
0.001
0.0001
Ditto
0.1
0.1
0.5
0.01
0.01
1.0
0.01
MMD-D
0.01
0.01
0.01
1.0
1.0
1.0
0.01
LR
0.01
0.01
0.01
0.001
0.01
0.001
0.0001
Ditto
X
0.1
0.1
1.0
1.00
0.01
1.0
0.01
MK-MMD
μ
0.1
0.1
0.1
0.1
1.0
0.1
0.1
LR
0.01
0.00
1000
0.001
0.01
0.001
0.0001
Ditto
0.1
T00
0.5
1.00
0.01
0.1
0.01
Cos. Sim.
0.1
0.1
1.0
1.0
1.0
0.1
0.01
set {0.0001, 0.001, 0.01, 0.1}. and tune À over (0.1, 1}. Among these, an LR of 0.001 consistently outperformed other values.
Fixing the LR to 0.001, the hyperparameter j is tuned for cosine similarity, MMD-D, and MK-MMD from the set {0.01, 0.1, 1.0). When these methods are integrated with Ditto, the same range for is searched, along with À E 10.
B.4 RxRx1
Following 4 we use an AdamW optimizer to fine-tune pretrained ResNet- 18 model. For Ditto, the LR is tuned within the set {0.00001, 0.0001, 0.001, 3. and the regularization parameter is taken from {0.01, 0.1, 1}. An LR of 0.0001 consistently yielded the best performance across various values.
With the learning rate fixed at 0.0001, we optimize the hyperparameter μ for cosine similarity, MMD-D, and MK-MMD from the set {0.01, 0.1, When these methods are combined with Ditto, À = 0.01 and the same range of search for an optimal is maintained.
C FL Settings and Checkpointing
In this section, dataset specific FL settings for the experiments are outlined and the checkpointing strategy is described. Regardless of dataset, each client participates in every round of FL training. That is, there is no client subsampling. It should noted that, in principle, the Ditto approach allows for different types of optimizers, and learning rates, to be applied to the global and local models during client-side optimization. However, in the experiments, the same LR and optimizer type are applied to both models, which is also done in the original work.
For all MK-MMD experiments, in constructing the kernel space K., 18 different RBF kernels are used of the form k; (x., :) = - 1 The values {7,}4=1 ranged from 3.5 2 with exponents evenly spaced in increments of 0.25. Further, when periodically optimizing MK-MMD or MMD-D every 20 steps, the number of batches of data used for optimization, z, varied. For
Fed-ISIC2019, 64, and for the remainder of datasets = 50.
23

---

For all variants of CIFAR-I0, there are five clients and FL proceeds for 10 server rounds. Within each round, clients perform five epochs of local training using a standard SGD optimizer with momentum set to 0.9 and batch size of 32. For the Synthetic data experiments, we generate eight clients, each with 5000 samples. Federated training is conducted over 15 server rounds, with each round consisting of 5 local training epochs. Following [20], we use a batch size of 10 and optimize the models using SGD with a momentum of 0.9 and a weight decay of 0.001. Following the settings of the FLamby benchmark (8) there are 15 rounds of federated training with six participating clients for Fed-ISIC2019. During an FL round, each client trains for 100 steps using a batch size of 64, An AdamW optimizer is used with default parameters. Finally, for RxRxI, four clients are present and 10 rounds of federated training are run. Each client performs five epochs of local training using a batch size of 32. Again an AdamW optimizer is applied using default parameters.
In all experiments, the final model for each client is checkpointed and evaluated on held-out test set to quantify performance. This is commonly applied approach in FL. 20 14 31. For FedAvg, the final global (aggregated) model is saved. For all other approaches, the local models are persisted. During experimentation, we also considered checkpointing each client's model after each round of local training based on local validation performance. Because results were similar or worse for all approaches with such checkpointing, only those using the latest checkpoint are presented.
D Compute Resources and Runtime Complexity: Ditto Experimentation
All experiments were performed on a high-performange computing cluster. For the CIFAR- 10 and Synthetic datasets, an NVIDIA T4V2 GPU with 32GB of CPU memory was used. Each FedAvg and Ditto experiment CIFAR-10 took approximately 10 and 20 minutes, respectively. The addition of optimization-based losses increased training time, with MMD-D, MK-MMD and MK-MMD20 experiments requiring approximately 2, I, and 1.5 hours, respectively. Training for the Synthetic data took under 5 minutes on the same GPU for all methods.
For the Fed-ISIC2019 and RxRx1 datasets, we used an NVIDIA A100 GPU with 64GB and 100GB of CPU memory, respectively. With red-ISIC2019, FedAvg and Ditto required and 3.5 hours of training time, while incorporating MMD-D, MK-MMD -1. and MK-MMD20 increased the training time to upproximately 4, 4, and 7 hours, respectively. For RxRx1, FedAvg took 2 hours, Ditto took 2.5 hours, and the addition of MMD-D, MK-MMD 1. and MK-MMD22 extended training times to 3,5, and hours, respectively.
The objective of this work is demonstrating that integrating the proposed MMD measures improves performance in heterogeneous data settings. Thus, no concerted effort made to reduce the overhead associated with optimization of the MK-MMD or MMD-D kernels. However, number of possibilities for cost reduction exist. For MK-MMD, general. python-based, quadratic program (QP) solver called qpth is used. Experimenting with more heavily optimized libraries or QP solvers specifically tailored to the MK-MMD system may yield large reductions. It also possible that the QP systems need not be solved to significant precision, decreasing the number of íterations used by the solvers. Finally, pre-conditioning techniques, such as using previous solutions as iteration starting points could markedly compress computation time. For MMD-D, future work could include studying utility trade-offs associated with various configurations, including kernel update frequencies, training data sizes, and optimizers. In the experiments, the featurization networks, şo, are also of moderate size, consisting of four linear layers and activations with a hidden dimension of 50. Contracting or changing this architecture would likely speed up optimization.
D.1 Runtime Complexity Analysis
The use of the proposed, adaptive MMD measures comes with an increase in runtime complexity due to the arithmetic operations associated with either solving the quadratic program for MK-MMD or performing SGD steps with MMD-D. In this section, we quantify these increases using Ditto as the reference algorithm. Borrowing the notation of Algorithm let p denote the number of arithmetic operations for single step of SGD with batch size b to update the model weights, W', or wa- For Ditto, the total cost across all N clients is O(N . T p). For simplicity, assume we perform kernel optimization for either MK-MMD or MMD-D after every iteration of local model training and have a single latent-space penalty.
[
24

---

For MMD-D, let q be the arithmetic operations associated with a single step of SGD for the deep kernel, i,, with batch size b. Denote by k the number of kernel optimization steps taken. Then the total cost across all clients becomes (N. . s (p + kg)). In most cases, p will be significantly greater than kg, as the primary models are expected to be much larger than 4p.
With MK-MMD, a convex quadratic program needs to be constructed according to Equation (2) However, the cost of this construction is generally smaller than the cost of solving the system for small data sizes. A QP solver based on interior-point methods requires O mm)) iterations, each costing (mm)) arithmetic operations, where m is the size of the latent space 43 Hence, the total cost across all clients is O(N T s (p + m²-9)). As in the case of MMD-D. the quantity m will often be much smaller than p. as d simply represents the dimension of the constrained latent space.
E Extended Related Work: Other pFL Approaches
This section discusses less similar, but still related, approaches in pFL and how the measures proposed in this work might be incorporated, if applicable. Sequentially split models are common in pFL. Quintessential examples include FedPer and FedRep [].. These techniques sequentially decompose models into feature maps and classifier layers and only exchange model sub-components, most often the feature maps, with the server for aggregation. As result, only certain modules incorporate global information through aggregation. Parallel-split PFL methods employ a similar trick, but decompose the models differently. In these methods, two or more modules process input in parallel. At least one of the modules is meant to be aggregated, incorporating global information, while others remain strictly local. Parallel modules could constitute complete classifiers, as in APFL [6] or lower-level feature maps, as in PerFCL [46] or FENDA-FL (39]. Combining partial model exchange, whether sequential or parallel, with penalty based pFL approaches like Ditto can take many forms, including straightforward modifications. For example, during local training, the feature maps in FedRep may be constrained to not drift too far from those of a FedAvg model, as in Ditto, using the MK-MMD or MMD-D measures proposed in this work.
Other PFL methods, fitting into less broad categories, exist. FedBN 22 normalization layers from global exchange, as these layers can struggle to handle heterogeneity well. HypCluster [27] partitions clients into clusters and trains isolated models for each cluster. When the number of clients is small, as in this work, such partitioning is impractical and forgoes global information. Other methods, such FedPop [15], are explicitly designed for cross-device settings rather than cross-silo FL. They take advantage of large-scale client populations with relatively small datasets, rather than focusing on smaller client pools. Each of these techniques incorporate ideas that narrow their applicability with respect to the setting considered in this work.
proposes excluding batch
F MR-MTL Experiments
The integration of the proposed adaptive MMD measures is not limited to the Ditto algorithm. This appendix demonstrates the extensibility of the main results to other existing approaches. In the experiments below, MR-MTL is augmented to include either MK-MMD or MMD-D as latent-space penalties. The training process of MR-MTL is similar to Ditto except that, instead of training separate global model to constrain local model training, the local models are simply penalized when drifting too far from the average of all local models [24. Thus, the algorithm resembles FedProx, but without replacing local models with the averaged one prior to client-side training.
While MR-MTI. is competitive with Ditto for certain tasks, it tends to be a slightly less effective approach. However, the removal of the separate global-model training process has significant differeetial-privace benefits and also markedly reduces the computational cost that comes with training two models simultaneously (24] Thus, MR-MTL is useful in settings where privacy is of the utmost importance or compute constraints intervene.
Table summarizes the experimental results. As observed with Ditto, combining the beight-based e² constraint of MR-MTL with the adaptive, latent-space constraints proposed in this work has visible benefits. For all variants of CIFAR-10 10 and the RxRx1 dataset, using either MK-MMD or MMD-D markedly boosts performance. Similarly, the Synthetic datasets also see improvements with MMD-D for both heterogeneity levels, and MK-MMD lifts performance when a = 3 = 0.5. Fed-ISIC2019 the lone datasct that does not see improvements when augmenting with these constraints.
25

---

Table 4: Average performance (and standard deviation) when augmenting the MR-MTL constraint with the proposed MMD penalties. Bold and underline indicate the best and second best value. Subscripts for CIFAR-10 and Synthetic indicate values for and a = B, respectively.
Dataset
MR-MTL
+ MM-MMD20
+ MMD-D
CIFAR-100.
79.516 (1.674)
81.269 (0.400)
80.307 (1.705)
CIFAR-1O0.
73.361 (1.027)
74.333 (0.742)
74.446 (0.529)
CIFAR- 10s.0
68.224 (1.223)
69.487 (1.180)
70.353 (0.149)
Synthetico.o
90.879 (0.267)
90.708 (0.257)
91.142 (0.031)
Synthetico.5
86.750 (3. (89)
90.958 (0.981)
90.337 (0.499)
Fed-ISIC2019
70.628 (0.911)
68.180 (2.735)
70.366 (0.970)
RxRx1
64.065 (1.663)
65.791 (0.363)
66.673 (0.480)
The setup for the MR-MTL experiments mirrors that of Ditto. The same FL settings and checkpointing described in Appendix C are used, and the same models are trained. The values reported in Table 4 the average of three distinct runs. The optimal hyperparameters for each method and dataset displayed in Table The hyperparameter ranges swept are reported thereafter. Note that the ranges subsets of those covered in the Ditto experiments, only neighborhoods of the optimal values there are considered.
Table Best hyperparameters for methods in the MR-MTL experiments across all datasets.
CIFAR-10
Synthetic
Fed-ISIC2019
RxRx1
0.1
0.5
5.0
(0.0, 0.0)
(0.5, 0.5)
LR
0.001
0.001
0.001
0.001
0.001
0.0001
0.0001
MR-MTL
0.1
0.1
0.1
0.01
0.01
0.1
0.01
LR
0.001
0.001
0.001
0.001
0.001
0.0001
0.0001
+ MMD-D
A
0.1
0.1
0.1
0.01
0.01
0.1
0.01
1.0
1.0
1.0
1.0
1.0
0.01
0.01
LR
0.001
0.001
0.001
0.001
0.01
0.0001
0.0001
+ MK-MMMD2
0.1
0.1
0.1
0.01
0.01
0.1
0.01
0.01
0.01
0.01
1.0
1.0
0.01
0.1
CIFAR-10 For all levels of heterogeneity, the same parameter ranges are used. For MR-MTL, the client-side LRs are f0.001, 0.01, and X (0.1, 1.0). In experiments applying either MMD-D and MK-MMD alongside the traditional MR-MTL constraint, the LR is fixed at 0.001, A = 0.1, and y is drawn from f0.01, 0.1, 1.0).
Synthetic For both levels heterogeneity, the same parameter ranges are searched. For MR-MTL, the LR is (0.001, 0.01) and À E (0.01, 0.1, 1.0]. When MK-MMD or MMD-D are combined with MR-MTL, we fix 0.01 while tuning j E {0.1, 1.03 and LR (0.001, 0.013.
Fed-ISIC2019 For MR-MTL, we search for the optimal LR in the set {0.0001, 0.0011,, and tune over (0.01, 0.1, 13. When MK-MMD or MMD-D are used with MR-MTL, the LR constant at 0.0001 and À = 0.1. The value of is tuned in the range (0.01, 0.1, 1.0].
RxRx1 For MR-MTL,, the LR tuned within the {0.0001, 0.001), and À is taken from (0.01, 0.1, 1.0). For fixed learning rate of 0.0001 and = 0.01, the hyperparameter l is op- timized, when combining MR-MTL and MMD-D or MK-MMD, within 001
G Standard Deviations of Main Results
Due to space constraints, and for clarity of presentation, the standard deviations of the main results of Section 6 are deferred to this appendix. Table 6 reports condensed version of Tables D and
26

---

Table 6: Average performance (and standard deviations) for various constraint configurations. Bold indicates the best values. Subseripts for CIFAR-10 and Synthetic indicate values for a and = 3 respectively. Finally + indicates that the MMD constraint is applied in tandem with the standard Ditto constraint, rather than as a replacement.
Dataset
Ditto
MMD-D
MK-MMD->
MK-MMD20
CIFAR- 10
84.930 (0.142)
83.789 (0.184)
84.136 (0.346)
84.439 (0.304)
CIFAR- 1 10g.5
80.702 (0.121)
75.094 (0.532)
75.678 (0.000)
76,564 (0.125)
CIFAR- 1IFAR-10s.0
77.658 (0.268)
67.729 (0.391)
68.718 (0.262)
68.832 (0.174)
Synthetic
89.129 (2.200)
90.237 (0.405)
91.418 (0.339)
90.066 (2.622 628)
Synthetico.5
85.533 (0.250)
91.270 (0.309)
91.137 (0.694)
90.262 (0.144)
Fed-ISIC2019
71.350 (1.191)
64.302 (1.546)
60.168 (2.979)
62.677 (0.604)
RxRxI
65.629 (1.936)
67.478 (1.200)
65.861 (1.151)
67.078 (0.394)
Dataset
Ditto
MMD-D
MM-MMMD_1
+ MK-MMD20
CIFAR- 10
84.930 (0.142)
85.214 (0.079)
84.723 (0.501)
84.900 (0.318)
CIFAR-I00: 100.5
80.702 (0. 121)
80.696 (0.148)
80.936 (0.224)
80.976 (0.365)
CIFAR- 10,.0
77.658 (0.268)
77.739 (0.376)
77.578 (0.139)
77.739 (0.271)
Synthetic
89.129 (2.200)
89.458 (0.348)
89.183 (2.249)
89.258 (2.269)
Synthetic
85.533 (0.250)
89.695 (0.186)
88.154 (1.817)
88.104 (0.721)
Fed-ISIC2019
71.350 (1.191)
72.226 (1.241)
71.169 (1.369)
71.267 (1.037)
RxRx1
65.629 (1.936) 67.755 (0.672)
65.984 (0.774)
66.892 (0.851)
with the corresponding standard deviations over three runs. It should be noted that in pFL., especially under heterogeneous settings, certain degree of performance variation expected and has also been reported in previous work [29]
H Initial Thoughts on Optimal Conditions for Adaptive MMD Measures
This work empirically demonstrates that the proposed adaptive, feature-based MMD penalties beneficial in the context of pFL algorithms imposing local drift constraints. The experiments are specifically designed to provide insight into settings where such penalties provide the most benefit. The results suggest that the presence of strong feature heterogeneity is one such setting. In contrast, when label heterogeneity is induced via Dirichlet allocation and training data is otherwise independent and identically distributed (IID), as in the CIFAR-10 10 datasets, the benefits are less pronounced.
In most settings, theoretical support for FL algorithms centers on demonstrating convergence through bounding the magnitude of model updates or local gradients by some contracting value with respect to the server rounds (20, [8, 21, 32). Proving various aspects of optimality compared with other algorithms is more rare, especially in non-IID settings. Some exceptions of interest are the conver- gence rates analysis of SCAFFOLD 14 and the optimal regularization parameter theory of Ditto and MR-MTL 24]
In (21) it is shown that there exists a regularization parameter value, X", yielding optimal average test performance across heterogeneous clients for the Ditto framework in the context of federated linear regression. Beyond the simplified model setting, the setup assumes that clients' heterogeneity expressed as perturbations in their labeling functions from a common mapping. While this heterogeneity is different in kind from the feature heterogeneity for which the approach proposed herein empirically improves performance, the theoretical model may be useful building similar results for latent-space penalties.
possible avenue is considering a shared latent-space distribution and labeling function across clients but varying linear feature maps. As such, the input feature spaces of each client would have differing statistical properties induced by the inverse feature maps. Nonetheless, with the correct feature mappings, labels are consistently and accurately applied. Provided that the target latent-space
27

---

has structured statistical properties, such as being Gaussian mixture, and the feature maps are sufficiently conditioned or simple, deriving optimal MMD kernels may be possible. From there, showing that there exists optimal weight, μ4, for such kernel would yield the desired result.
The setting above is specifically manufactured to be difficult for standard Ditto-hased approach. By construction, feature maps are not within neighborhood of one another, implying that penalizing drift from a central model is likely detrimental. While the setting artificial, is related real problems, where, for example, different medical imaging devices produce unique artifacts but fundamentally similar decision criterion when these artifacts are accounted for during feature processing.
Characterizing the settings under which the MMD constraints studied in this work are optimal remains open and interesting question. While the suggested theoretical tack may require modifications, the experimental results this paper provide insights into the design such theoretical machinery.
28


Reviews:
[{"review_id": "ovOOQ6Qr7U", "reviewer": "Reviewer_zAAU", "summary": "This paper proposes and investigates the application of adaptive Maximum Mean Discrepancy (MMD) measures, specifically Multi-Kernel MMD (MK-MMD) and MMD-D, as regularization terms within the Ditto framework for personalized FL. The primary contribution lies in exploring these theoretically supported statistical distance measures to constrain feature representation divergence between local and global models. The paper conducts extensive experiments across various datasets, including CIFAR-10, a synthetic dataset, Fed-ISIC2019, and RxRx1, demonstrating that these MMD-based penalties, either alone or combined with the traditional Ditto penalty, yield improved model performance, particularly in settings with pronounced feature heterogeneity.", "questions": "1. Could you elaborate on the criteria used for selecting the optimal depth for MMD constraints, particularly for datasets other than CIFAR-10?  This is mainly because the paper states that for CIFAR-10, the optimal depth was uniformly 1. A more detailed analysis or discussion on why deeper constraints did not yield better results, and whether this observation generalizes to all datasets, would be valuable.   \n\n2. The paper mentions that for some datasets (Synthetic and RxRx1), including the Ditto penalty with MMD-based measures, degrades accuracy or hurts performance across the board. Could you provide a more in-depth explanation or theoretical intuition for this observation? \n\n3. Regarding the computational overhead of kernel optimization procedures, the paper acknowledges this as a limitation. Could the authors suggest specific directions for future work to reduce this overhead? \n\n4. Could the authors consider including standard deviations for the reported average performance metrics in tables, especially for key results, as this would give a better sense of the consistency of the improvements?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper presents a technically sound exploration of adaptive MMD measures in personalized FL. The methods are clearly articulated, which provides mathematical formulations for MMD, MK-MMD, and MMD-D, and details their integration into the Ditto algorithm. \n\n2. The experimental setup is comprehensive. It utilizes both synthetic and real-world clinical datasets with varying levels of heterogeneity, which adequately tests the proposed methods. \n\n3. The results are systematically presented and analyzed. This highlights the conditions under which the MMD-based constraints are most effective. \n\n4. Ablation studies support the claims regarding kernel optimization and the benefits of a single constraint layer.   \n\nWeaknesses:\n\n1- The computational overhead associated with the iterative optimization of kernels for MMD measures. \n\n2- The paper lacks a theoretical framework that quantifies the specific conditions under which the proposed MMD measures are most effective. \n\n3. The paper primarily focuses on the Ditto framework, and the general applicability of these MMD measures to other personalized FL methods, such as FedProx and MR-MTL, is suggested but not empirically validated in detail.   \n\n4. The experimental results are reported as the average of only three runs. This would limit the ability to fully assess the statistical significance and variability of the reported performance improvements.", "was_revised": true, "final_justification": "I appreciate the response, and I still believe that some of my concerns are still valid. I will not change my score.", "rebuttal": "Thank you for your efforts with our manuscript and thorough review. We appreciate your comments on the technical contributions of our proposed method and comprehensive experiment setup. Below is an itemized list of your questions and remarks. Each is followed by a discussion and description of how they were addressed in revising the paper. Note that related questions or remarks may be grouped together. We look forward to additional conversations with you.\n\n**The paper lacks a theoretical framework that quantifies the specific conditions under which the proposed MMD measures are most effective.** \n\nThis is accurate and a limitation that we recognize in Section 7 of the paper. While there is some theory supporting the weight-based penalty used by Ditto, theoretical characterization of the settings where MMD-based latent-space penalties become optimal is likely challenging enough to warrant a paper of its own.\n\n**The paper primarily focuses on the Ditto framework, and the general applicability of these MMD measures to other personalized FL methods, such as FedProx and MR-MTL, is suggested but not empirically validated in detail.**\n\nWe have worked to expand our experimental results to better substantiate the claim that our proposed method is applicable and beneficial in other personalized federated learning algorithms. Specifically, we augmented the MR-MTL algorithm to include the MK-MMD and MMD-D constraints and evaluated such integration across all seven datasets. The results are presented below and will be added to the paper upon revision.\n\n| Dataset        | MR-MTL | MR-MTL+MKMMD | MR-MTL+MMD-D |\n|----------------|--------|---------------|------------------|\n| **CIFAR (0.1)**      | 79.515 | **81.268**        | 80.307           |\n| **CIFAR (0.5)**     | 73.360 | 74.332        | **74.446**           |\n| **CIFAR (5.0)**      | 68.224 | 69.487        | **70.352**           |\n| **Synth (0.0)**  | 90.879 | 90.708        | **91.141**          |\n| **Synth (0.5)**  | 86.750 | **90.958**        | 90.337           |\n| **Fed-ISIC2019**    | **70.627** | 68.180        | 70.366           |\n| **RxRx1**          | 64.064 | 65.790        | **66.673**           |\n\nThese results demonstrate the extensibility of the proposed MMD measures to other methods. Both MK-MMD and MMD-D variants largely improve performance compared to the baseline, MR-MTL, with the exception of the Fed-ISIC2019 dataset.\n\n**Could you elaborate on the criteria used for selecting the optimal depth for MMD constraints, particularly for datasets other than CIFAR-10? This is mainly because the paper states that for CIFAR-10, the optimal depth was uniformly 1. A more detailed analysis or discussion on why deeper constraints did not yield better results, and whether this observation generalizes to all datasets, would be valuable.**\n\nThank you for this suggestion, a deeper discussion on this has been added to Sections 5 and 6.1. CIFAR-10 served as the foundation for our first set of experiments. In observing that additional constraints at various depths did not improve performance for either MMD approach, we restricted our experiments to single layer constraints. Additional performance improvements may indeed exist for other datasets with different constraint configurations, but this simplified our search space. Our conjecture is that constraining too many layers in a relatively small model, as used in the CIFAR-10 experiments, produced too strong a constraint, hindering learning.\n\n**The paper mentions that for some datasets (Synthetic and RxRx1), including the Ditto penalty with MMD-based measures, degrades accuracy or hurts performance across the board. Could you provide a more in-depth explanation or theoretical intuition for this observation?**\n\nThe Synthetic dataset is constructed to be adversarial to weight constraints, in some sense. The data points and labels are generated by transformations drawn from the same distribution, but with normally distributed variations. As such, weight-based constraints are likely to hinder learning at a certain point. The RxRx1 dataset was specifically chosen for experimentation due to the rich feature heterogeneity noted by the dataset creators, which is especially challenging for FL algorithms. We believe the proposed adaptive latent-space constraints are well suited for this heterogeneity. We have added a bit of this discussion to the revised manuscript.\n\n**The computational overhead associated with the iterative optimization of kernels for MMD measures.**\n\n**Regarding the computational overhead of kernel optimization procedures, the paper acknowledges this as a limitation. Could the authors suggest specific directions for future work to reduce this overhead?**\n\nWe believe that there are several potential avenues for each method to help reduce the overhead associated with the kernel optimization procedures. The directions discussed below have been added to the revised manuscript under Appendix D.\n\nFor MK-MMD, a general, python-based, quadratic program (QP) solver was used (qpth). Experimenting with more heavily optimized libraries or QP solvers specifically tailored to our system is a straightforward step. It is also possible that the systems need not be solved with significant precision, reducing the number of iterations required by the solvers. Finally, pre-conditioning techniques, such as using previous solutions as iterative starting points could reduce computation time.\n\nFor MMD-D, future work could include studying utility trade-offs associated with different configurations including varying kernel update frequencies, training data sizes, and optimizers. The featurization networks are also somewhat large, consisting of four linear layers and activations with a hidden dimension of 50. Compressing or changing this architecture would likely speed up optimization.\n\n**The experimental results are reported as the average of only three runs. This would limit the ability to fully assess the statistical significance and variability of the reported performance improvements.**\n\n**Could the authors consider including standard deviations for the reported average performance metrics in tables, especially for key results, as this would give a better sense of the consistency of the improvements?**\n\nAs noted in the paper, all experiments were run with three random seeds, and the reported values reflect the averages. Per your request, we now include the corresponding standard deviations for the key results presented in Table 1 and Table 2:\n\n| Dataset    | Ditto        | MMD-D        | MK-MMD -1      | MK-MMD 20      | +MMD-D        | +MK-MMD -1    | +MK-MMD 20     |\n|------------|------------------|------------------|-------------------|-------------------|-------------------|------------------|-------------------|\n| **CIFAR (0.1)**  | 84.930 ± 0.142   | 83.789 ± 0.184   | 84.136 ± 0.346   | 84.439 ± 0.304  | **85.214 ± 0.079**   | 84.723 ± 0.500  | 84.900 ± 0.317     |\n| **CIFAR (0.5)**  | 80.702 ± 0.120  | 75.094 ± 0.532  | 75.678 ± 0.000   | 76.564 ± 0.125    | 80.696 ± 0.148   | 80.936 ± 0.224  | **80.976 ± 0.365**    |\n| **CIFAR (5.0)**  | 77.658 ± 0.267  | 67.729 ± 0.591   | 68.718 ± 0.261   | 68.832 ± 0.173   | **77.739 ± 0.375**   | 77.578 ± 0.139  | **77.739 ± 0.270**   |\n| **Synth (0.0)**  | 89.129 ± 2.200  | 90.237 ± 0.405  | **91.418 ± 0.339**   | 90.066 ± 2.628   | 89.458 ± 0.347   | 89.183 ± 2.249  | 89.258 ± 2.268   |\n| **Synth (0.5)**  | 85.533 ± 0.250  | **91.270 ± 0.308**   | 91.137 ± 0.694   | 90.262 ± 0.144   | 89.695 ± 0.185   | 88.154 ± 1.816  | 88.104 ± 0.721   |\n| **Fed-ISIC2019**    | 71.350 ± 1.191   | 64.302 ± 1.546  | 60.168 ± 2.978   | 62.677 ± 0.603   | **72.260 ± 1.240**    | 71.169 ± 1.369  | 71.267 ± 1.036   |\n| **RxRx1**      | 65.629 ± 1.935  | 67.478 ± 1.199  | 65.861 ± 1.151   | 67.078 ± 0.394   | **67.755 ± 0.671**   | 65.984 ± 0.774  | 66.892 ± 0.851  |\n\nIn personalized federated learning, especially under heterogeneous settings, a certain degree of performance variation is expected and has also been reported in previous work [1].\n\nIn our experiments, we specifically designed the datasets to reflect different types of heterogeneity. CIFAR-10 was partitioned to emphasize label heterogeneity, while the Synthetic dataset was generated to showcase feature heterogeneity among clients. Among the real-world datasets, Fed-ISIC2019 and RxRx1 can contain both forms, though Fed-ISIC2019 primarily exhibits label heterogeneity, whereas RxRx1 is more dominated by feature heterogeneity.\n\nOur results support the hypothesis that Ditto, which applies a weight drift constraint, is particularly effective in handling label heterogeneity. However, it tends to underperform when feature heterogeneity is the dominant factor. This is evident in the observed accuracy improvements. Additionally, although less prominent, we also observe that incorporating the MMD measures on top of Ditto provides further gains in label heterogeneous settings existing in the rest of the datasets.\n\n[1] Matsuda, K., Sasaki, Y., Xiao, C., & Onizuka, M. (2023). Benchmark for personalized federated learning."}, {"review_id": "txa9RZjHh7", "reviewer": "Reviewer_T2e1", "summary": "This paper extends the Ditto framework for personalized Federated Learning (pFL), by replacing or augmenting the way of penalizing divergence of the local model from the global model, with adaptive latent-space regularization using Maximum Mean Discrepancies (MK-MMD and MMD-D). The paper shows that using these alternative measures is beneficial in scenarios with high feature heterogeneity between clients.", "questions": "### Questions\n\n1) In the paper, it is stated that the effectiveness of MMD measures is highly sensitive to the choice of kernel; however, only the RBF kernel is evaluated in the experimental section. Would it be possible to provide more insights on the choice of kernels other than RBF?\n\n2)  ⁠Why have Ditto results not been included in Table 1? A comparison with Ditto may be beneficial to interpret results.\n\n3) Could the authors please elaborate more on what they mean by the increasing advantages of using MMD measures in an iterative training process in the sentence in rows 178-179?\n\n### Minor suggestions\n\n1) The introduction should give more relevance to MMD-D and MK-MMD. Additionally, references to other algorithms in rows 39-45 may be included and better characterized in the related work section, while symbols should be introduced in the methodology section.\n\n2) The methodology section provides an overview of the Ditto algorithm and the MMD measures, which is acceptable. However, paragraph 3.5 should provide a more thorough description of the main innovative contributions of the proposed approach, including how it could be adapted or integrated into pFL methods, as well as its advantages.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 1, "strengths_and_weaknesses": "### Strengths\n\n1) Building upon existing literature, the paper formally and clearly characterizes the Ditto and MMD measures used to discourage divergent model weights or feature maps in pFL.\n\n2) The proposed framework is simple and generalizable to pFL algorithms other than Ditto.\n\n### Weaknesses\n\n1) The paper proposes the application of existing measures, taking inspiration from Ditto. However, the use of MMD measures in FL applied to latent feature maps is not entirely new in the literature [1, 2], and the relationship between the proposed method and existing literature has not been sufficiently discussed.\n\n2) The related works section could be improved. Specifically, an overview of techniques for pFL based on methods other than regularization techniques, and how the proposed framework positions with respect to these methods, is absent. Due to page limitations, you can also consider expanding the related works in the appendix.\n\n3) The computational and communication costs of the proposed combinations of MMD penalty-based pFL methods have not been discussed.\n\n4) The experiment section is not convincing. Specifically, the settings considered are overly simplistic, using too few clients and vanilla models. It would be beneficial to include experiments with other widely used networks in FL, such as ResNet20 or MobileNetV2. These networks are lightweight yet well-established in the existing literature. Additionally, increasing the number of clients, ideally to at least hundreds, would provide a more realistic representation of real-world scenarios.\n\n[1] Yao, Xin, et al. \"Federated learning with additional mechanisms on clients to reduce communication costs.\"\n\n[2] Hu, Kai, et al. \"FedMMD: a federated weighting algorithm considering non-IID and local model deviation.\"", "was_revised": true, "final_justification": "After revising all the reviews and the rebuttal answers provided by the authors, I am very satisfied with the answer provided and now lean towards acceptance.", "rebuttal": "Thank you for your thoughtful comments about our work. We appreciate the opportunity to more deeply discuss the results and make improvements to the manuscript. Below is an itemized list of your questions and remarks. Each is followed by a discussion and description of how they were addressed in revising the paper. Thank you again for your time.\n\n**The paper proposes the application of existing measures, taking inspiration from Ditto. However, the use of MMD measures in FL applied to latent feature maps is not entirely new in the literature [1, 2], and the relationship between the proposed method and existing literature has not been sufficiently discussed.**\n\nThank you for identifying this gap in our related work discussion. The cited work does indeed consider incorporating MMD measures within an FL framework and should have been discussed. We have added a thorough discussion to our Related Work section to better differentiate our approach. In [2], a fixed MMD measure is used to modify server-side aggregation weights, rather than influence client-side training. The approach is a global FL method and does not optimize the MMD kernel. In [1], a fixed multi-kernel MMD measure is used to constrain the feature space of client-side models during training. The objective is to learn a single global model, rather than personalized ones, and no kernel optimization is applied. As demonstrated in our ablation studies (Figure 3, right), we compare fixed MMD kernels with various bandwidths against adaptive kernel optimization. The results indicate that none of the fixed kernels match the performance of MK-MMD or MMD-D.\n\n**The related works section could be improved. Specifically, an overview of techniques for pFL based on methods other than regularization techniques, and how the proposed framework positions with respect to these methods, is absent. Due to page limitations, you can also consider expanding the related works in the appendix.**\n\nWe have significantly restructured the Introduction and Related Works section to include a broader overview of existing pFL methods, including the references mentioned in the previous remark. We have also added an additional appendix section, where we discuss other pFL frameworks and how they relate to our work. Specifically, we discuss sequential and parallel-split pFL methods such as FedPer, FedRep, APFL, and FENDA-FL, prototype based methods in FedProto, and other methods such as FedBN, FedPop, and HypCluster.\n\n**The computational and communication costs of the proposed combinations of MMD penalty-based pFL methods have not been discussed.**\n\nThank you for pointing this out. In our original submission, we did discuss computational overhead in Appendix D, but we failed to refer to that section in the main body. So it was not well-represented to readers. We now include a reference to these details in Section 7. We have also expanded the appendix to include a more thorough discussion of ways to reduce computational cost.\n\nThe communication costs of the proposed approach remain identical to that of Ditto. Regardless of constraint applied, each client exchanges only the global model with the server for aggregation. In Section 3.1, we have added a brief discussion of this, as it was not clear in our original submission.\n\n**The experiment section is not convincing. Specifically, the settings considered are overly simplistic, using too few clients and vanilla models. It would be beneficial to include experiments with other widely used networks in FL, such as ResNet20 or MobileNetV2. These networks are lightweight yet well-established in the existing literature. Additionally, increasing the number of clients, ideally to at least hundreds, would provide a more realistic representation of real-world scenarios.**\n\nIn this work, we have focused on cross-silo FL, which considers a small or moderate number of clients with larger training data sets. However, we did not make this clear in our introduction and have revised it. The setting of cross-silo FL is common, especially in medical or financial data settings, where clients represent a modest collection of institutions. The datasets of the FLamby benchmark are a nice example.\n\nIn terms of models, we have tried to stay consistent with benchmarks and datasets where possible. The model for CIFAR-10 is drawn from the benchmark in [3], and the model for Fed-ISIC2019 matches the one used in the FLamby benchmark [4]. For RxRx1, we use ResNet-18, which is a smaller version of the ResNet-50 model used in the work that introduced the dataset [5] and is one of the models trained in [6]. We have added notes in the paper to better motivate our choice of models.\n\nWe hope this discussion helps provide a better background for the experimental setup considered.\n\n**In the paper, it is stated that the effectiveness of MMD measures is highly sensitive to the choice of kernel; however, only the RBF kernel is evaluated in the experimental section. Would it be possible to provide more insights on the choice of kernels other than RBF?**\n\nThis is a great question and something that we did not communicate clearly in the original text. Other kernels do exist such as Laplace or inverse multiquadric kernels [7]. Using other kinds of kernels may well improve our results. The comment around MMD sensitivity to kernel choice was meant to imply that fixing a kernel of any type can be quite sub-optimal, motivating the optimization procedures of MK-MMD and MMD-D and the periodic re-optimization procedure we propose.\n\nWe have edited the paper to make this more clear and also note that other choices of kernel, beyond RBFs, are possible.\n\n**⁠Why have Ditto results not been included in Table 1? A comparison with Ditto may be beneficial to interpret results.**\n\nThank you for the suggestion. We now include the Ditto results in Table 1.\n\n**Could the authors please elaborate more on what they mean by the increasing advantages of using MMD measures in an iterative training process in the sentence in rows 178-179?**\n\nAs the training process proceeds, the distributions over the latent space induced by the global and local models evolve. While they may stabilize later in training, the optimal choice of MMD construction at one stage of training is unlikely to be optimal at another. Generally, this is true of any fixed distance measure. However, the MK-MMD and MMD-D measures can be re-optimized. We take advantage of this fact to ensure that we maintain strong measures of latent-space divergence by periodically re-optimizing the kernels based on data samples throughout the training trajectory. We have added additional discussion to make this more clear in the manuscript.\n\n**The introduction should give more relevance to MMD-D and MK-MMD. Additionally, references to other algorithms in rows 39-45 may be included and better characterized in the related work section, while symbols should be introduced in the methodology section.**\n\nThank you for the suggestions. We have significantly reworked the Introduction and Related Work sections. We now give more weight to MMD-D and MK-MMD in the Introduction and have moved the bulk of FL algorithm discussions to the Related Work section. This section has also been expanded to incorporate additional methods. Finally, we have removed the introduction of technical symbols in these sections, as recommended.\n\n**The methodology section provides an overview of the Ditto algorithm and the MMD measures, which is acceptable. However, paragraph 3.5 should provide a more thorough description of the main innovative contributions of the proposed approach, including how it could be adapted or integrated into pFL methods, as well as its advantages.**\n\nWe agree that, throughout the paper, more explicit discussion of our contributions would be helpful. We have added more clarity around our contributions to the Introduction, Related Works, and Adaptive MMD Measures in FL sections. Specifically, our contributions are:\n\n1. Proposing the use of theoretically supported statistical distance measures as penalty constraints for dealing with heterogeneity in pFL, moving beyond previously considered paired feature-based constraints. \n\n2. Showing that adapting the MK-MMD or MMD-D measures through periodic re-optimization provides notable performance improvements in settings with high feature heterogeneity, where existing methods, such as Ditto, under-perform. \n\n3. Finally, experiments with either natural or controllable levels of label and feature heterogeneity highlight the strengths and weaknesses of the types of drift penalties investigated. These experiments provide a starting point for understanding when the proposed measures perform most effectively.\n\n[1] Yao, Xin, et al. Federated learning with additional mechanisms on clients to reduce communication costs.\n\n[2] Hu, Kai, et al. FedMMD: a federated weighting algorithm considering non-IID and local model deviation.\n\n[3] Daoyuan Chen,  et al.  pFL-bench: a comprehensive benchmark for personalized federated learning.\n\n[4] J. O. du Terrail,  et al. FLamby: Datasets and Benchmarks for Cross-Silo Federated Learning in Realistic Healthcare Settings.\n\n[5] Maciej Sypetkowski, et al. RxRx1: A dataset for evaluating experimental batch correction methods.\n\n[6] Haokun Chen, et al. FRAug: Tackling federated learning with non-iid features via representation augmentation.\n\n[7] Sutherland, et al.  Generative models and model criticism via optimized maximum mean discrepancy."}, {"review_id": "Z5X3NjRkw9", "reviewer": "Reviewer_AkTQ", "summary": "This paper introduces adaptive statistical constraints - specially Multi Kernel Maximum-mean Discrepeancies (MK-MMD) and deep MMD (MMD-D) - to the panelty-based (e.g., Ditto) personalized federated learning (PFL) framework. It empirically demonstrates that the proposed methods improve performance in tasks with feature and label heterogeneity compared to existing feature-space and weight-based constraints. \n\nHowever, this paper’s contribution seems incremental. It essentially applies well-established statistical distance measures (MK-MMD, MMD-D) to the latent-space distributions within Ditto, without clear theoretical innovation. The justification to choose MMD is not established, and empirical results shows that performance gain from using MMD is not significant.", "questions": "- 1. What is the difference of the proposed methods from the prototype based PFL? As the prototype-based PFL also apply constraints in the latent space, discussion about this should be included in the related works. \n- 2. Why should MMD be combined to the penalty-based pFL (e.g. Ditto). It seems that applying MMD is orthogonal to the penalty-based constraints. \n- 3.  In section 3.1, the authors claim that this paper proposes an unpaired approach. In MMD-D, however, feature-drift constraints depicted in figure 1 seems to be paired. Please correct me if I missed something. \n- 4. What is commputational and communication costs of the proposed schemes to additionally apply MMD on top of Ditto?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 1, "strengths_and_weaknesses": "- Strenghts\n  - 1. This paper considers important problems in PFL, feature-drift. \n  - 2. This paper is well written and hence is easy to follow. \n\n- Weaknesses\n  - 1. The major concern of the reviewer is the novelty of the paper. MMD constraints are not new, nor is their application in deep learning contexts unusual. Merely applying these known measures to Ditto feels straightforward rather than innovative. The authors must articulate precisely what novel insight or substantial methodological development distinguishes this work from straightforward integration.\n  - 2. The authors rely heavily on empirical validation without providing theoretical conditions for why MMD constraints are the best or good to be combined with the penalty-based PFL such as Ditto. \n  - 3. Even in the emprifical validation, performance gain from using MMD is not significant. Especially, in table 2, comparing to Ditto and the proposed schemes, the gain is less than 1% in most cases (only synthetic_0.5 has the large gain).\n  - 4. While the proposed schemes has extra communications and computation (to exchange  $f_i(x;\\bar{\\theta_i})$ for all $i$ compared to the baseline (e.g., Ditto), this paper lacks computational and communication efficiency analysis, which weakens claims of practical utility.", "was_revised": true, "final_justification": "My previous concerns have been addressed by the author's detailed explanation and additional experiments, and hence I increased my score. I still believe that the paper would be strengthened by including an intuitive discussion of why combining latent-space constraints with weight-based constraints creates synergy, especially in feature-heterogeneous FL.", "rebuttal": "We would like to thank you for your time and effort with our manuscript. Your feedback thus far has motivated a number of improvements to our work. Below is an itemized list of your questions and remarks, followed by a discussion and description of how they were addressed in revising the paper. We look forward to further discussing our work with you and any additional questions you may have.\n\n**The major concern of the reviewer is the novelty of the paper. MMD constraints are not new, nor is their application in deep learning contexts unusual. Merely applying these known measures to Ditto feels straightforward rather than innovative. The authors must articulate precisely what novel insight or substantial methodological development distinguishes this work from straightforward integration.**\n\nWe appreciate the chance to make the novelty of our contributions more clear.\n\nWhile MMD measures are well-established in other contexts, they are an under-explored tool in pFL. In this work, we propose the first use of adaptive MMD measures as latent-space penalties for dealing with heterogeneity in pFL, moving beyond previously considered paired feature-based penalties. Ditto is investigated as a representative example, but the use of such measures readily generalizes to other pFL settings.\n\nA small collection of existing work incorporates MMD into FL algorithms, outside of pFL, but all of them use fixed kernels. We demonstrate that taking advantage of the inherent adaptability of MK-MMD and MMD-D measures through periodic re-optimization during training provides notable performance improvements in settings with high feature heterogeneity, where existing methods under-perform. \n\nFinally, experiments with either natural or controllable levels of label and feature heterogeneity highlight the strengths and weaknesses of the types of drift penalties investigated, providing a starting point for understanding when the proposed measures perform most effectively.\n\nIn the edits to our paper, we have striven to state these contributions clearly and to distinguish our work from the existing literature. We hope the above discussion provides a concise and direct summary of this work’s novelty.\n\n**The authors rely heavily on empirical validation without providing theoretical conditions for why MMD constraints are the best or good to be combined with the penalty-based PFL such as Ditto.**\n\nThis is indeed a limitation that we recognize in the conclusion of our work. There is some existing theory around the weight-based penalty of Ditto that may serve as a good starting point. However, a complete characterization of settings where MMD-based latent-space penalties become optimal is, unfortunately, beyond the scope of this work. However, we have endeavored to construct our experiments in such a way as to suggest settings where our proposed latent-space penalties are most useful. Those being settings where datasets are anticipated to incorporate feature heterogeneity, as in the Synthetic, RxRx1, and Fed-ISIC2019 datasets.\n\n**Even in the empirical validation, performance gain from using MMD is not significant. Especially, in table 2, compared to Ditto and the proposed schemes, the gain is less than 1% in most cases (only synthetic_0.5 has the large gain).**\n\nThe gains with the MMD measures are modest for the CIFAR-10 dataset, which considers synthetically induced label heterogeneity rather than feature heterogeneity. The accuracy gains for both Synthetic datasets and RxRx1 exceed 2%. The improvement for Fed-ISIC2019 is almost 0.9%, which is somewhat smaller, but still a notable improvement for a clinical task. In many benchmarks, the margins between the top performing methods are often less than 1%, and it is rare for an approach to uniformly improve performance, which we see here [1,2].\n\nThe table below summarizes the improvements over Ditto with the best performing MMD approach across the datasets. We hope that this discussion provides more confidence in the significance of our observed improvements.\n\n| Dataset        | CIFAR (5.0) | CIFAR (0.5) | CIFAR (0.1) | Synth (0.0) | Synth (0.5) | Fed-ISIC2019  | RxRx1 |\n|----------------|-------------|-------------|-------------|--------------|--------------|--------|--------------|\n| **Acc. Gain**      | +0.283      | +0.274      | +0.081      | +2.289       | +5.885       |  +0.876  |   +2.126   |\n\n**While the proposed schemes has extra communications and computation (to exchange \n f_i(x;\\hat{\\theta}_i) for all i) compared to the baseline (e.g., Ditto), this paper lacks computational and communication efficiency analysis, which weakens claims of practical utility.**\n\nA nice benefit of our proposed modification to Ditto is that no additional communication costs are introduced. As in the original Ditto formulation, only the global models are exchanged with the server for aggregation after client-side training. Each client maintains its own MMD kernel and optimizes it on local data. We realize that this could have been made more clear in the paper and have included remarks on communication and memory implications in Section 3.1.\n\nIn our original submission, computational costs compared to FedAvg and Ditto were discussed in Appendix D. However, we failed to refer to that section in the main body, which reduced its visibility. A reference to these details now appears in Section 7. We have also expanded the appendix to include a more thorough discussion of potential ways to reduce computational cost.\n\n**What is the difference of the proposed methods from the prototype based PFL? As the prototype-based PFL also apply constraints in the latent space, discussion about this should be included in the related works.**\n\nWe have refactored and expanded the Introduction and Related Works section to capture a wider swathe of pFL methods and better position our work within the existing literature. A primary addition to this is a discussion of prototype-based pFL methods, of which FedProto [3] is a canonical example. Prototype-based pFL methods aim to learn uniform feature representations for each data label. In FedProto, for example, only the prototypes are aggregated between clients, rather than model weights. In the approach here, we consider the distribution of all points in the feature space together. Prototype methods are also generally applicable to classification-based tasks only. \n\nSome prototype methods for pFL also have some notion of constraint. In FedProto, the local prototypes are constrained, via the $l^2$-norm, to not diverge too far from the aggregated prototypes. The method proposed here could be used to replace this norm as well.\n\n**Why should MMD be combined to the penalty-based pFL (e.g. Ditto). It seems that applying MMD is orthogonal to the penalty-based constraints.**\n\nThis is a great question. We agree with the characterization of the two constraints as “orthogonal,” in some sense. It is perhaps why combining them can have a positive impact. In many natural datasets, it is unlikely that a single form of heterogeneity characterizes all drift between clients in an FL system. As such, we believe that incorporating losses that penalize drift in distinct ways can be beneficial and may be the source of improvements for the Fed-ISIC2019 and CIFAR-10 datasets when combining the two constraints. We have added additional discussion on this to the end of the Results section.\n\n**In section 3.1, the authors claim that this paper proposes an unpaired approach. In MMD-D, however, feature-drift constraints depicted in figure 1 seems to be paired. Please correct me if I missed something.**\n\nThank you for pointing this out. We were ambiguous in our notation of $(\\mathbf{x}, \\mathbf{y})$ in Figure 1. This is meant to symbolize a batch of labeled data rather than a single input-label pair. We have updated the figure caption and the reference to figure in the text to make this clear.\n\n**What is computational and communication costs of the proposed schemes to additionally apply MMD on top of Ditto?**\n\nAbove, we briefly discuss the communication cost associated with our proposed approach. In short, there is no additional communication overhead. The computational overhead associated with applying periodic optimization as part of the Ditto training framework is either, forming and solving the QP problem on a subset of training data for MK-MMD or performing a certain number of SGD steps on a modest, but not insignificant, network and set of variables with MMD-D. The focus of this work is demonstrating the utility of using these measures rather than optimizing the computational overhead. As such, there is a fair increase in training time when using these measures, but there is a lot of opportunity to reduce that overhead. The training time increases are discussed in Appendix D, Opportunities to reduce the cost have been added to this appendix as well.\n\n[1] K. Matsuda, Y. Sasaki, C. Xiao, and M. Onizuka. Benchmark for personalized federated\nlearning.\n\n[2] Daoyuan Chen, Dawei Gao, Weirui Kuang, Yaliang Li, and Bolin Ding. pFL-bench: a comprehensive benchmark for personalized federated learning.\n\n[3] Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across heterogeneous clients."}, {"review_id": "GteJK4CA9b", "reviewer": "Reviewer_Vrrq", "summary": "This paper proposes an improvement to the Ditto framework for personalized federated learning (pFL) by incorporating adaptive latent space constraints using statistical distance measures — specifically MK-MMD and MMD-D — instead of the traditional ℓ2 norm on weights. These constraints aim to better handle feature heterogeneity by minimizing distributional drift in the latent representation space rather than simple weight differences. The authors conduct experiments on diverse datasets (CIFAR-10, Synthetic, Fed-ISIC2019, RxRx1) and show that the proposed methods offer consistent improvements, particularly under significant heterogeneity. The approach is theoretically grounded and broadly applicable to other pFL frameworks beyond Ditto.", "questions": "1. More Experiments: Although the paper claims that the method has wide applicability, the current validation is limited to the Ditto framework. Experimental validation on other mainstream or different types of pFL frameworks would be more convincing.\n2. Privacy Considerations: The reliance on feature embeddings for regularization raises potential privacy concerns. Specifically, intermediate activations can retain semantic information about raw inputs, which—if combined with model inversion attacks or a compromised global model—may lead to input reconstruction or data leakage. \n3. Complexity Concern: How do the local computation cost and storage usage compare with Ditto and FedAvg? Is there a trade-off here?\n4. Novelty Concern: As a mature analysis method, MMD is very likely to improve PFL. Although the authors have made some improvements, I am still concerned about the novelty of the insight and technical contribution. I hope the authors can explain their novelty more clearly.\n\nI will consider increasing my score if my questions are answered.", "limitations": "See Questions above.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths：\n 1. Clear Problem Focus and Innovation: The paper precisely identifies feature heterogeneity as a crucial challenge in Personalized Federated Learning (pFL). It proposes to address this by introducing adaptive, theoretically supported MMD constraints within an existing advanced framework (Ditto). This approach, building upon a strong baseline, is more convincing.\n2. Constraints in Latent Space: Compared to imposing constraints in the raw data space, applying constraints in the latent space is a more advanced and flexible way to handle feature heterogeneity. It operates on abstract representations of data, promising to capture deeper distributional relationships.\n3. Modularity and Extensibility: Although the research focuses on the Ditto framework, the paper suggests that these adaptive MMD constraints can be directly applied to other pFL settings. This indicates a modular design and good generality, facilitating broader adoption.\n\nWeaknesses:\nSee Questions below.", "was_revised": true, "final_justification": "Questions regarding the validation of the experimental framework and the novelty of the method have been addressed. The authors also reported a relatively honest comparison of their methods in terms of storage, computational overhead, and computational time. While the proposed method still has room for improvement in computational time, this could be a promising direction for future research. Overall, the proposed method is interesting, and the evaluation criteria used can effectively support innovation across different approaches. I improved my score after discussing this with the authors.", "rebuttal": "Thank you for your thoughtful review and providing us the opportunity to improve our work. Below is an itemized list of your questions, followed by a discussion and description of how they were addressed in revising the paper. We look forward to continuing to answer your questions about our work and further improving the presentation of our ideas.\n\n**More Experiments: Although the paper claims that the method has wide applicability, the current validation is limited to the Ditto framework. Experimental validation on other mainstream or different types of pFL frameworks would be more convincing.**\n\nWe appreciate this feedback. To better support the conclusion that our proposed method is generalizable to other pFL algorithms, we have expanded our experimental results. Specifically, we augmented the MR-MTL algorithm to include the MK-MMD and MMD-D constraints and evaluated such integration across all seven datasets. The table below summarizes the results.\n\n| Dataset        | MR-MTL | MR-MTL+MK-MMD | MR-MTL+MMD-D |\n|----------------|--------|---------------|------------------|\n| **CIFAR (0.1)**      | 79.515 | **81.268**        | 80.307           |\n| **CIFAR (0.5)**     | 73.360 | 74.332        | **74.446**           |\n| **CIFAR (5.0)**      | 68.224 | 69.487        | **70.352**           |\n| **Synth (0.0)**  | 90.879 | 90.708        | **91.141**          |\n| **Synth (0.5)**  | 86.750 | **90.958**        | 90.337           |\n| **Fed-ISIC2019**    | **70.627** | 68.180        | 70.366           |\n| **RxRx1**          | 64.064 | 65.790        | **66.673**           |\n\nBoth MK-MMD and MMD-D variants largely improve performance compared to the baseline, MR-MTL, with the exception of the Fed-ISIC2019 dataset. We believe these results better support our claim of transferability and also provide more evidence to the effectiveness of our approach in general.\n\n**Privacy Considerations: The reliance on feature embeddings for regularization raises potential privacy concerns. Specifically, intermediate activations can retain semantic information about raw inputs, which—if combined with model inversion attacks or a compromised global model—may lead to input reconstruction or data leakage.**\n\nPrivacy is an extremely important consideration in FL, and there are downsides to the Ditto approach in terms of privacy guarantees due to the additional data iterations associated with simultaneously training two models, as noted in [1]. However, the method proposed here does not modify the information exchanged between the clients and server nor does it expose the feature embeddings of individual clients to others. As such, the privacy implications are nearly identical to that of Ditto, with the exception that training data is periodically used to optimize the MMD kernels. We have noted this in Section 3.1 as an important point.\n\n**Complexity Concern: How do the local computation cost and storage usage compare with Ditto and FedAvg? Is there a trade-off here?**\n\nThis is a good question. Throughout the manuscript, we have endeavored to make the communication,  computation, and storage costs associated with our work more clear.  In our original submission, we did discuss computational overhead compared to FedAvg and Ditto in Appendix D, but we failed to refer to that section in the main body. So it was not well-represented to readers. We now include a reference to these details in Section 7. We have also expanded the Appendix to include a more thorough discussion of ways to reduce computational cost, as the present overhead is fairly high.\n\nThe communication costs associated with the approach are identical to Ditto, as no changes are made to the information exchanged. We now explicitly note this in Section 3.1. The memory overhead of the proposed method is also nearly identical to Ditto, with the exception of storing a modest number of feature activations and training a set of kernel weights. We have added a remark on this to Appendix D as well.\n\n**Novelty Concern: As a mature analysis method, MMD is very likely to improve PFL. Although the authors have made some improvements, I am still concerned about the novelty of the insight and technical contribution. I hope the authors can explain their novelty more clearly.**\n\nThank you for the opportunity to better describe our contributions. The technical contributions in this paper can be summarized as follows.\n\n1. We propose the first use of adaptive and theoretically supported measures as penalty constraints for dealing with heterogeneity in pFL, moving beyond previously considered paired feature-based constraints, using Ditto, and now MR-MTL, as representative examples.\n\n2. While a small collection of existing work applies MMD in different ways outside of pFL, all of them used fixed kernels. We demonstrate that taking advantage of the adaptability of MK-MMD and MMD-D measures through periodic re-optimization during training provides notable performance improvements in settings with high feature heterogeneity, where existing methods, such as Ditto, under-perform. \n\n3. Finally, experiments with either natural or controllable levels of label and feature heterogeneity highlight the strengths and weaknesses of the types of drift penalties investigated.\n\nWe have endeavored to make these novel contributions more clear throughout the manuscript and hope that the presentation has been significantly improved.\n\n[1] Ziyu Liu, Shengyuan Hu, Zhiwei Steven Wu, and Virginia Smith. On privacy and personalization in cross-silo federated learning."}]

Author rebuttal:

Meta-review: This paper explores adaptive MMD-based constraints within the Ditto framework for personalized federated learning. Reviewers agree the paper is well written, technically sound, and supported by systematic experiments and ablations. The major concerns regard the limited novelty of applying MMD in this setting, modest empirical improvements, lack of large-scale or diverse experimental validation, and incomplete discussion of costs and prior related work. The rebuttal and discussions sufficiently address the raised concerns, and all the reviewers lean toward the acceptance rating. The AC stands with the reviewers for the decision.
Decision: Accept (poster)