Bridging Generalization Gap of Heterogeneous Federated Clients Using Generative Models
Authors: Ziru Niu
Hai Dong
A. K. Qin
Venue: ICLR 2026 Poster (2026)
Keywords: Federated Learning
Model Heterogeneity
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: This paper presents a model-heterogeneous federated learning framework where clients share feature statistics to generate synthetic data for local fine-tuning, improving generalization while reducing communication and memory costs.

Abstract: Federated Learning (FL) is a privacy-preserving machine learning framework facilitating collaborative training across distributed clients. However, its performance is often compromised by data heterogeneity among participants, which can result in local models with limited generalization capability. Traditional model-homogeneous approaches address this issue primarily by regularizing local training procedures or dynamically adjusting client weights during aggregation. Nevertheless, these methods become unsuitable in scenarios involving clients with heterogeneous model architectures. In this paper, we propose a model-heterogeneous FL framework that enhances clients’ generalization performance on unseen data without relying on parameter aggregation. Instead of model parameters, clients share feature distribution statistics (mean and covariance) with the server. Then each client trains a variational transposed convolutional neural network using Gaussian latent variables sampled from these distributions, and use it to generate synthetic data. By fine-tuning local models with the synthetic data, clients achieve significant improvement of generalization ability. Experimental results demonstrate that our approach not only attains higher generalization accuracy compared to existing model-heterogeneous FL frameworks, but also reduces communication costs and memory consumption.

Published as a conference paper at ICLR 2026
BRIDGING GENERALIZATION GAP OF HETEROGE- NEOUS FEDERATED CLIENTS USING GENERATIVE MODELS
Ziru Niu
Hai Dong
School of Computing Technologies
School of Computing Technologies
RMIT University
RMIT University
Melbourne, VIC, Australia sirr..te wwwu rmit. edu au
Melbourne, VIC, Australia hai donge rew.eeeritttaaaaa. edu.au
A. K. Qin Department of Computing Technologies Swinburne University of Technology Hawthorn, VIC, Australia kqi gggin. .ww..edu..au
ABSTRACT
Federated Learning (FL) is privacy-preserving machine learning framework fa- cilitating collaborative training across distributed clients. However, its perfor- mance is often compromised by data heterogeneity among participants, which can result local models with limited generalization capability. Traditional model- homogeneous approaches address this issue primarily by regularizing local train- ing procedures or dynamically adjusting client weights during aggregation. Never- theless, these methods become unsuitable in scenarios involving clients with het- erogeneous model architectures. this paper, we propose a hoode--hetroogeeouus FL framework that enhances clients' generalization performance on unseen data without relying on parameter aggregation. Instead of model parameters, clients share feature distribution statisties (mean and covariance) with the server. Then each client trains a variational transposed convolutional neural network using Gaussian latent variables sampled from these distributions, and use it generate synthetic data. By fine-tuning local models with the synthetic data, clients achieve significant improvement of generalization ability. Experimental results demon- strate that our approach not only attains higher generalization accuracy compared existing model- heteroo-eeoo FL frameworks, but also reduces communication costs and memory consumption.
INTRODUCTION
The prevalence of the Internet of Things (IoT)) propels federated learning (FL) (McMahan et al., 2017) as widespread technique to process the dispersed data among end clients while ensuring data privacy. Under the regulation of a central server, clients collaboratively train global model without revealing any personal data, and exchange model parameters with the server. Abstaining from data transmission, FL markedly alleviates communication overhead and the risk of data leak- age. However, the inherent data heterogeneity of IoT clients presents significant challenges to FL. Firstly, clients comprise edge devices distributed across various geographical locations, which nat- urally collect usually collect real-time data with frequent distribution shifts in practical scenarios, such as weather forecasting (Mclaughlin & Su 20244 Reddy et al. 2023) and healthcare (Saha et al. 2017).
data (Abeysekara et al. 2023). Secondly, single clients
To overcome both inter-client and intra-client data heterogeneities. improving the generalization ability (i.e., the ability to predict unseen data correctly (Sun et al. 2024)) of clients has become
Corresponding author.

---

Published as a conference paper at ICLR 2026
vital matter in FL (Huang al. 2024). Existing solutions mainly focus on two aspects. including
regularization and weight modification. Regularization methods propose to debias the individual
training process by imposing regularization term on the local objective function of clients (Dinh)
et al. 2020 Li et al. b02 2020b (mmith et al. 2017] Karimireddy et al. 2020 Acar et al. 2021
Gao et al.
2022 Li al. 2019). Weight modification methods aim alleviate the negative impact of biased clients on the global model, including assigning lower weights to biased entities in aggre-
gation (Pillutla et al. 2022; Cao al. 2020: Mu al., 2024, Li al. 2020al Tahmasebian et al.,
2022
or reducing the possibility of selecting these clients (Jee Cho et al., 2022 Li et al. 20222 Niu al. 2024, Zhang et al., 2023 Lai et al., 2021). As result, the global model maintains signif-
icant robustness against the heterogeneous client data. Correspondingly, clients can improve their generalization performance by incorporating the enhanced global model into their local models.
Despite the outstanding progress in addressing data heterogeneity, the aforementioned works are developed based on the assumption that all clients share the same model architecture. In practice, clients are likely to obtain personalized models with diverse and irrelevant architectures (Huang) al. 2024), referred to as model-heteroggeneus FL (see Appendix A for details), making traditional aggregation- based methods prohibitive. In this case, clients cannot learn generalized information through parameter sharing. which necessitates alternative strategies for enhancing the generalization of local models with differing architectures.
Several attempts have been made to improve the generalization of model- honell-hetroogeneuus FL. through
knowledge distillation (Li & Wang 2019 Itahara et al. [003] Cho et al. [022] Nguyee al.. [023]
Cheng et al. 2021] [un & Lyu [022] Sattler et al. 2021I. These works utilize a shared public
dataset and enable clients to exchange knowledge by comparing each other's predictions on the
dataset. However, the prerequisite of a public dataset severely restricts the deployment of these
works, as a public dataset is usually not available in real-world circumstances (Zhu et al., 202Tb.
To address this concern, data-frec model- heddoo-heeroosssouus FL methods have been proposed, which can mainly be divided into two categories. The first category employs generated feature representations for knowledge distillation (Zhu et al. 2021 Luo et al.) (021), eliminating the dependence on public datasets. However, unlike raw samples, knowledge distillation on feature space merely regularizes the head classifier of a local model, while leaving the backbone feature extractor unmodified. The second category involves circumventing knowledge distillation and enabling clients to communicate through alternative messages, such as mean feature representations (Tan et al. 2022 Zhang al. 2024) and a mutual proxy model (Wang et al.. 2024). Although (Tan et al. 2022 Zhang et al. 20224) regularizes local feature extractors by enforcing them to derive unbiased feature representations, the biased head classifiers are neglected in these works, which significantly limits the generalization improvement. For (Wang et al. 2024), training and transmitting an overparameterized proxy model (Wang et al. 2024) will result in substantial communication and memory overhead.
To overcome the disadvantages of prior methods, we expect a comprehensive FL framework that can increase the generalization of model-hhetroggeneuus clients without relying on any public dataset, while maintaining communication and memory efficiency. To this end, we propose VT (Fed- erated Learning with Variational Transposed Convolution), a model-heterogeneed that uses synthetic data to fine-tune local models, thus improving generalization without relying on public dataset. FedVTC achieves this by enabling each client to generate synthetic samples using variational transposed convolutional (VTC) neural network. The VTC model takes low-dimensional Gaussian latent variables input and produces corresponding synthetic samples by upsampling the latent variables. Afterwards, clients fine-tune the local models with the synthetic samples to improve the generalization ability. In communication, clients only exchange the local mean and covariance with the server to mitigate the bias of local feature distributions, leading to substantial communica- tion cost reduction. Compared with representation--asel knowledge distillation (Zhu et al. [021] Luo et al. 202T) that only regularizes the head classifier, fine-tuning with synthetic images can debias both the feature extractor and the classifier in local model, thereby potentially achieving higher generalization performance. In addition, FedVTC enables clients to train the VTC model and local models alternately to avoid extra memory consumption.
FL framework
Similar to variational autoencoder (VAE), a VTC model is initially trained by maximizing the evidence lower bound (ELBO) of local data (Kingma & Welling 2013). Moreover, to strengthen the robustness of VTC against variant input latent variables, we incorporate Distribution Matching (DM) loss (Zhao et al. 2023) into the objective function of VTC for regularization. This strategy
2

---

Published as a conference paper at ICLR 2026
helps FedVTC produce synthetic samples with better training quality, which notably improves the performance of the fine-tuned local models. The contributions of this paper are summarized as follows:
We propose FedVTC, a model-heterogeneous FL. framework based on variational trans- posed convolution. FedVTC uses synthetic data to fine-tune local models for better gener- alization, eradicating the dependence on pre-existing public datasets.
We design novel objective function train the VTC model, which encompasses the stan- dard negative ELBO loss and a DM loss. The DM loss regularizes the training procedure of VTC, enabling VTC to consistently generate high-quality synthetic samples with varying input latent variables.
We conduct extensive experiments to validate the effectiveness of FedVTC. Experiment results show that FedVTC obtains higher generalization accuracy than existing model- heterogeneous FL frameworks over MNIST, CIFARIO, CIFAR100 and Tiny-ImageNet datasets, well as lower equivalent communication costs and memory consumption.
LITERATURE REVIEW
2.1 IMPROVING THE GENERALIZATION OF HOOEE--OMOOSNFOUS FEDERATED LEARNING
The endeavors of the state-of sfateoof-the-ar to enhance the generalization of model-homogeneee FL can mainly be divided into two categories, which are regularization and weight modification. Regular- ization approaches aim to debias the local training procedure by adding regularization term onto the clients' local objective functions. For instance, (Dinh et al. 2020; Li et al. 2021b 2020b) Karim- ireddy al.) 2020 Acarr et al. [201] Gao et al. (2022) prevent client models from overfitting to local optima by restricting the Euclidean distance hetween local and global parameters. Besides, (Smith al. 2017) employs relationship matrix to enforce clients with higher similarities to learn close parameter updates, and (Li al. 2019) alleviates the bias of local training using ridge regression.
Weight modification methods enhance the generalization of the global model by deducting the pro- portion of biased clients participating in global training. Subsequently, clients can also reinforce the generalization ability by incorporating the well-generalized global model into their local models. On one hand, (Pillutla et al.] 2022; Cao et al., 2020 Mu et al. 2024 Li et al., 2020a Tabmasebian al. 2022) replace weighted average with dynamic weighting strategies, and assign lower weights to biased, anomalous and malicious parameter updates in aggregation. On the other hand, (Jee Cho al., 2022 Li et al., 2022, Niu et al. 2024; Zhang et al.) 2023] Lai et al. (02T) design some heuris- tic functions to evaluate the importance of each client, such as training loss and cosine similarity. Moreover, pFedGen (Le et al. 2024) enables unseen clients to generate representation vectors us- ing feature extractor trained by previous clients, and only interacts with clients with the smallest representation distance. Consequently, these works develop a selective client sampling strategy that selects unimportant entities with lower possibilities to mitigate the negative effect of biased clients. Weight modification methods enhance the generalization of the global model by deducting the pro- portion of biased clients participating in global training.
Although these works have achieved remarkable progress in improving FL generalization, they are developed based on the assumption that all clients share a universal model architecture. For model- heterogeneous FL, where clients exhibit diverse model architectures, the aforementioned works be- come infeasible.
2.2 HODEE-HETEERGEENOOUS FEDERATED LEARNING
The most commonly used approach to improve the generalization of model--hterogeenouus FL is knowledge distillation (Li & Wang 2019 Itahara al. 2023 Cho et al. 2022 Nguyen et al.) 2023
Cheng et al. 2021 Sun & Lyu| 2021 Sattler et al.) 2021). With the presence of a public dataset, clients are able to acquire unbiased global knowldege by learning from each other's predictions on the dataset. However, the prerequisite of a public dataset can be a critical bottleneck for these methods, such dataset is not usually available in practice.
3

---

Published as a conference paper at ICLR 2026
9h
2
h4
fk
y
Forward propagation
x
Reparameteriatton
(a)
(b)
Upsampling
Figure I: A comparison between: (a): a standard FL client and (b): client in FedVTC. (b), "Ø"
and "O" respectively standard for element-wise addition and product.
To overcome the dependence on a public dataset, (Zhu al., 2021 Luo al., 2021 propose to use generative models, such as generative adversarial network (GAN) and a multivariate Gaussian sampler, to generate feauture representations for knowledge distillation. However, unlike raw sam- ples that progress throuth all layers, feature representations merely go through h the head classifier of a neural network. Consequently, the distillation scheme in (Zhu et al. 2021: Luo et al. 202I) only debiases the classifier of local models, which significantly limits the generalization improvement. Additionally, instead of knowledge distillation, (Tan et al. 2022 Zhang et al. 2024 Wang et al, 2024) allows clients exchange gooolllarchitectur--nneeee
messages for generalization, such
as the mean feature representations (Tan et al. 2022] Wang et al.) (024) and a common proxy model (Wang al.] (024).. For one thing, representation sharing (Tan al. 2022] Wang et al.) 2024) regularizes the backbone feature extractors in local models by enforcing clients to derive un- biased feature representations, while neglecting the head classifiers. For another thing, training and
transmitting proxy model (Wang et al. (024)) might cause undesired communication and memory costs.
2.3 FEDERATED TRAINING OF GENERATIVE MODELS
Most existing FL schemes for training a generative model have several practical limitations. For instance, (Stanley Jothiraj & Mashhadi 2024) and (Peng et al. 2025) propose train global dif- fusion model in FL networks. However, they require a public dataset to facilitate knowledge distil- lation among client models, which is not usually accessible. (Zhang et al. (022) trains global gen- erative adversarial network (GAN) by maximizing the discrepancy between client models. Clients the downloaded GAN produce training samples from unseen classes, thereby enhancing the performance of zero-shot prediction. However, the transmission of client models might cause over- whelming communication overhead, (Shi al.) (224)) utilizes contrastive language-image pre- training (CLIP) model to guide clients' local training for better few/zero-shot performance, while pre- pre-existing CLIP model is not usually available.
METHODOLOGY
3.1 PRELIMINARIES
Suppose a network consisting of a central server and a set of clients K. = {1, ..., K} with non-iid local datasets (D1, Dx). Each client (1 I < K) obtains a local model ft. IRd -> TR, which can be decomposed into backbone feature extractor 9e IRd + RR and a head classifier hu IRP + R. formula, fú h& gk- Let {(x;, 91) vi)ssssn denote the collection of data in the local dataset D; with total of пœ = Dal samples. As shown in Figure for each sample x Rd with the corresponding label y E (1.2....... C)2 fa takes a as input and makes a prediction
J. where C is the total number of all possible classes. Firstly, is transformed latent variable
2 E Iº' using the feature extractor 9t (i.c., z = 9a (x)). Secondly, forwarded to the classifier
hy to produce prediction y hy (z).
4

---

Published as a conference paper at ICLR 2026
Client &
Lyee
=
B))(())
D
Server
(#)=-1
D
{sr hm
V2
Class
{z)y ci Bza((z)
DE
Dx
(a/)gee
Exe
Figure 2: VTC trained with loss function L. Lre + DKL. Lre (blue) is the reconstruction loss between the original samples and the generated samples, and DeL (red) is the KL-divergence between the distributions of the local and global latent variables.
The objective of Modee-hheterogeneore FL is to find the optimal set of local models (fi,....., fi minimizing the local empirical loss, that is:
f Jk arg min f1....KK
Σ
Cly, sx(r))
(m.D))D,
C(y. [x())) the classification loss (e.g., cross-entropy) of model fx on the sample (ær, y). In addition, to evaluate generalization performance, consisting of unseen samples (Sun et al. 2024), that is, Vk, Do O Dx Ø.
each fo on Do. which is a test dataset
3.2 VARIATIONAL TRANSPOSED CONVOLUTION
As Figure Th) shows, aside from the fundamental local model fa. FedVTC introduces a variational transposed convolution (VTC) model the IR" - Rdd each client k. VTC is upsampling network that takes a latent variable z E Rº as input to produce an enlarged data sample æ' E R (Dumoulin & Visin 2016). Similar to a VAE (Kingma
Welling 2013). for each sample with the corresponding 2, client k samples random variable v from the distribution N(u|2, Σx), and forwards to d'y to derive x' (blue arrow in Figure b)). The covariance matrix Σg RPxp can be learned via the gradient method. To make Eg differentiable, FedVTC applies the well-known reparametrization trick as per VAE and lets v z + 0 ④ 6 (red arrow in Figure 6 TRP is random Gaussian noise with distribution reparameterized standard deviation (SD), and "o" represents element-wise product. In this case, for each entry in Σx. we have (Ex), ); of and (Ex)s; )g for any i + j. and each (1 sis p) can be learned through gradient descent.
[Ne(0, Tx [01, (2,,02:.... ................................ 0pl indicates the
Following the design of VAE (Kingma & Welling 2013), VTC is trained by maximizing the evi- dence lower bound (ELBO), which calculated as:
ELBO = log Pv, (z'|z, σR) - Dru (x))),.((/
(2)
log Per (x' ok) represents how the can generate real samples based on the latent variable
0k C According to (Kingma & Welling, 2013), this term can be well approximated by minimizing the empirical reconstruction loss Cre ( negative Kullback-Leibler (KL))-ivergencee -D (4 (=(=))[p(()) measures how qg. (2|æ² (the
x x as Figure 2 shows. The
distribution of z learned by the feature extractor 9&) approaches the real distribution pl p(z). Follow-
ing (Kingma Welling 2013), we can assume p(z) N(, with μ denoting the unbiased mean of z which is usually O by default. In FL settings, we follow (Tan et al.) 2022) and esti- mate ya using c", which the global prototype of class y. Figure showss global prototype cl = = ΣKEK, c₁ is calculated as the expectation of all local prototype cg for all k E Ky,, with K., C K. indicating the set of all clients containing data of class y,. A local prototype of client is the mean of the local feature representations, i.e. c% D Exr,)) 9k. (x)) where DV Dę in- dicates the collection all samples of class y in local dataset D₁. In consequence, for two Gaussian distributions gg. (xr) 2z, .kk and p(z), the KI KL-divergence be expressed analytically:
5

---

Published as a conference paper at ICLR 2026
1 1 -=== + - log [
(ag (x)l=,,))llp(()) Therefore, we derive a negative ELBO loss L. optimize gk, the and mz:
(3)
DEl
Le
Σ
Σ
Fre (2 xzi) +))+ + (g (z/1==))lp(z))
D
yEYr
im
(4)
1 [(2i- c() - c²) 2:ll2 122
Σ
|a
tr(a) - p- p-llgg Ex|
D
vE
3// = {y: # 20, S SySC) C) indicates all possible classes from client kss data.
As Figure 2 shows, mitigate the bias of the local distribution ((z, 24),, each client will update the local class-wise prototypes {ck DE>E, and the SD σ to the server. In response, the server obtains the global prototypes (c") and SD through aggregation, and sends {c*), back clients.
3.3 REGULARIZING VTC FOR BETTER GENERALIZATION
Owing to the randomness of the VTC process, the is likely to generate variant synthetic samples {æ'} with unknown training quality. To address this issue, we propose to regularize the training procedure of VTC with a distribution matching (DM) loss Lam
D
ID Σ
Lam = Σ
Σ
llxx
c²113 113
Σ DE
llxx (vx:())) - c²
(5)
NEVA D
D
This loss enforces 2hx to generate high-quality samples in which gt can acquire feature distribution close the unbiased global distribution (Zhao al., (2023). Subsequently. the local model fx will learn unbiased parameter updates from the synthetic data, resulting in better generalization performance.
Overall, the loss function Lte to train a VTC model is defined as:
Lie = La + ALam
(6)
À 0 is the coefficient of regularization.
3.4 FEDVTC OVERVIEW
The comprehensive FedVTC framework is presented in Algorithm Specifically, for each iteration every active client k regularly trains the local model as in traditional FL. In addition, trains t/ and d using loss function Lte (Equation 6). Afterwards, k uploads the local prototypes and SD to the server for aggregation. The server averages the received results and returns the global prototype and SD to clients.
For memory efficiency, FedVTC applies alternating strategy in VTC training. For line 12 in [ we first freeze thy and a and incorporate the step +- (e--) V 92 Lie into the standard training of Jk- In this case, line in Algorithm I can be written as fi fk. 1/ fa (C Lre) (with yA included in ft)). Then we freeze f' and optimize t/ and with Lic- As a result, FedVTC will not extend the maximum memory usage, as ths and σ; usually require less memory space than fü for training. In contrast, if we optimize fū. 4/k and σr; simultancously, the maximum memory usage will be enlarged to their accumulated memory usage.
Once FL is complete, every client l uploads the local VTC 4/ to the server, then the server derives a global VTC & through aggregation, and broadcasts &b along with {c'..., )), a to all clients. In order aggregate VTC parameters, FedVTC lets all v/a 's have the same architecture. For each c (y E {1...., C)), every client & samples latent variables from distribution. M(c",, E), and generates the corresponding synthetic samples by forwarding these variables to th. E RRx is diagonal

---

One noteworthy point in FedVTC is that, clients only need to upload the local VTC models once Published as a conference paper at ICLR 2026
Algorithm FedVTC
Require: FL iterations T, fine-tuning rounds T, local epochs E, clients K (l. K) local
models (fi,.... fk): datasets Di...., Dr J. local VTCs (e.....,.. &K). local SDs (oI.... o )
initial SI) a, initial prototypes {cl y.,, e learning rate 7, number of generated samples S.
1: for = 1, 2.., do
2:
server does:
3: randomly sample a set of participating clients K C K.
4: for y {I,... C.:
5:
Kiy + K broadcast e to clients in JCG
6:
7:
broadcast clients in K.. every client 1: E K. does:
8:
9:
Ok
10:
for epochs 1,... E:
11:
fk JK 11 SnC.
Update fa with SGD.
12:
9: + 9k: - 1/ on Lrcs V/ t Uhk myVeCie ok + oh V Lre.
for y Yu:
14:
ck + [D EeEDy (z(
b Compute local prototype.
15:
upload c to the server.
16:
upload a the server.
17:
server does:
18: for y {1, mr. CJ:
19:
c + ICC DKEK; ck
> Update global prototype.
20:
- Te EREK Oke
Update global SD.
21: end for
22: every client E IC does:
23:
upload thx the server.
24: server does:
25:
1
EA=1 r'k-
obtain global VTC.
K
broadcast w.
and a to every client & E
27: for every K. do
28: create synthetic dataset Dí with S samples using wi, o and .c.,)...c.).
29: for rounds 1, T:
30: L' Di (x()).
31:
fk < Jk v.C.
fine-tuning with the synthetic dataset.
end for
33:
return w
matrix with Ei equal to the squared i-th element in a. Afterwards, each client fine-tunes its local model by running classification tasks on the total C' synthetic samples. this stage, local fine-tuning processes run in isolation, and no communication occurs between clients and the server.
instead of every FL. iteration (see line 23 in Algorithm []. In the experiment, we compare the performance of FedVTC in cases of singular VTC transmission and per-round VTC transmission. The results show that these two modules exhibit marginal differences in terms of accuracy (see Table 3}. Therefore, in the ultimate VedVTC framework, local VTCs are only communicated once for communication efficiency.
4 EXPERIMENTS
4.1 EXPERIMENT SETUP
Datasets. We evaluate FedVTC on the following datasets: MNIST (LeCun 1998) contains a train- ing set with 60,000 samples and validation set with 10,000 samples from 10 classes. CIFAR10

---

Published as a conference paper at ICLR 2026
Dir
FedProto
FedTGG
FodGen
CCVR
FedType
PFedAFM
FedVTC
0.1
83,4+0.4
85,4+-0.5
81.1+1.0
05.5+0.4
83.9+0.3
85.6 + 0.5
88.7++0.7
MNIST
1.0
86.1
86.8
81.910.
86.210.3
86.7:0..5
87.2 0.4
90.1=0.11
0.1
39.3++0.7
40.9+0.2
39.1+-0.2
39.3+0.2
39.7+00.3
38.9
46,9++0.6
CIFARIO
1.0
41.1+0.2
42.1+0.77
39.7
41.0++0.1
40.2+0.6
40.5 0.5
51.7+00.7
0.1
26.2+H0.2
29.2++0.3
27.
27.6+0.9
31.3+1..5
36.3+0088
CIFARI00
1.0
30.4+0.2
01.8+0.5
00.8
29.3+0.4
34.440.6
34.9
40.4+0..3
Tiny-
0.1
23.7+t0.3
23.6+0.11 24.9+
24.1+0.3 26.6
2. 24.8+0.4
30.2+0.3
ImageNet
1.0
26.4=0.2
27.2+00.3
26.1
31.6+00
27.3 1.3+00 0.2
05.8+00.
Table 1: The average generalization accuracy (in %) of each client's local model on the global validation dataset (with mean + SD).
and CIFARI00 (Krizhessky) 2009) contain a training set with 50,000 samples and validation set of 10,000 samples from 10 and 100 classes. Tiny-ImageNet (Chrabaszez et al. 2017) contains a training set with 100,000 samples and a validation set with 10,000 samples from 200 classes.
Comparative methods. We compare FedVTC with five state-of- state-of-the-arr ahe moodll-eterogeeeeou FL frameworks. 1. FedProto (Tan al. 2022) enables clients regularize the local feature extractors using class-wise global prototypes. 2. FedTGP (Zhang et al. 2024) extends the design of Fed- Proto by optimizing the decision boundary between inter-class prototypes using gradient descent. 3. FedGen (Zhu et al., (221)) and 4. CCVR (Luo et al.) (021) utilizes virtual feature representations to fine-tune the classifiers of clients, with FedGen using a generative adversarial network (GAN) and CCVR using Gaussian distribution to generate the representations. 5. FedType (Wang et al. 2024) allows clients to collaboratively train a proxy model using conformal prediction and enhance generalization by distilling knowledge through the proxy model. 6. FFedAFM (Yi et al. 2025) clients employ a proxy feature embedding model for knowledge distillation.
System implementation. We simulate virtual network consisting of K 100 clients with participation rate of 0.1 (i.e. |AC*| = 10). Both the server and all clients operate on one NVIDIA Geforce RTX 4070 GPU with 12 GB of RAM space. For heterogeneous client models, we follow the settings in (Zhang et al. 2024; Wang et al. 2024), where clients are divided into five uniform clus- ters with five model architectures: ResNet- 18, ResNet-34, ResNet-55, ResNet-101 and Res (He et al. 2016). For non-iid local data distribution, we follow (Luo et al. (021) and distribute the training sets of MNIST, CIFAR1O, CIFARI00, and Tiny- ImageNet unevenly among clients using two Dirichlet (Dir) distributions with parameters 0. and 1.0. The experiment platform is imple- mented with Py PyTorch 2.0 (Paszke et al. 2019). The VTC model has simple architecture with four transposed convolution layers (see Appendix B.2).
Parameter settings. Following (Tan et al., 2022 Wang et al., 2024), the total number of iterations is set to T 100. For FedVTC, the number of additional fine-tuning rounds is set to T = 5. For fairness, we also let the baselines train extra T rounds with full client participation so that all methods have the same training rounds. The learning rate 77 is set to le-4 for MNIST (Luo et al.) 202T) and 0.01 for others (Tan et al. (002)). The local epoch E is set to five, and the batch size is set to 16 universally (Wang al. 2024). For fairness, the number of synthetic samples S is set to 500 for MNIST and CIFARTO, 10, and 1000 for CIFAR100 and Tiny-ImageNet, which is identical to the amount of generated representations in (Luo al.) [02II. We set the regularization parameter
as X = 0.1 (see Appendix B.4). The settings of and d can be found in Appendix B.I
4.2 EXPERIMENT RESULTS
Generalization accuracy. We evaluate each client model with the global validation sets in the aforementioned four datasets for generalization accuracy. Each method runs for three independent trials, and the average results are recorded. As shown in Table FedVTC consistently outperforms the comparative methods across all datasets and distributions, demonstrating a remarkable general- ization ability over non-iid data, As supplementary, the graphs depicting accuracy per iteration are included in Appendix B.3
8

---

Published as a conference paper at ICLR 2026
FedProto
FedTGP
FedGen
CCVR
FedTTyp
pFedAFM
FedVTC
MNIST
0.94
0.94
13.15
46.61
12.37
1.94
0.70
CIFARIO
3.93
3.93
54.58
807.46
37.45
5.88
2.90
CIFAR 100
39.32
39.32
91.07
825.16
72.85
41.27
26.49
Tiny-ImageNet
78.64
78.64
844.82
923.46
179.61
147.10
52.71
Table 2: The total volume of data transmission (in GB).
MNIST
CIFARIO
CIFARI00
Tiny-ImageNet
Dir(0.1)
Dir(10)
Dir(0.1)
Dir(10)
Dir(0.1)
Dir(10)
Dir(0.1)
Dir(1.)) 1.0)
88.7(0.7)
90.1(0.1)
46.9(0.6)
$1.7(0.7)
36.3(0.8)
40.4(0.3)
30.2(0.3)
35.8(0.5)
89.7(0.2)
92.3(0.3)
45.4(0.3)
54.3(0.5)
37.7(0.4)
39.6(0.6)
31.8(0.7)
38.0(L.)
I
0.70
2.90
26.49
52.71
2
5.63
8.77
32.36
123.22
Table 3: The (*)generalization accuracy (%) and (** joomunicaation cost (GB) of Fed with singular VTC transmission and FedVTC with regular VTC transmission per round.
MNIST
CIFAR10
CIFAR100
Tiny-ImageNer
Dir(0.1)
Dir(1.)
Dir(0.1)
Dir(1.)
Dir(0.1)
Dir(1.)
Dir(0.1) Dir(1.))
L.
87.2(0.4)
89.9(0.3)
39.7(0.5)
45,4(0.4)
28.9(0.5)
35.2(0.2)
28.1(0.1)
33.4(0.3)
Cie
88.7(0.7)
90.1(0.1)
46.9(0.6)
51.7(0.7)
36.3(0.8)
40.4(0.3)
30.2(0.3)
35.8(0.5)
Table 4: The generalization accuracy of two training modules FedVTC, The first module uses the standard Le loss, and the second module uses the regularized Lie Le XLdm loss.
Communication efficiency. As shown in Table FedVTC obtains the least overall communication cost among all methods. Specifically, for CCVR, clients share the local prototypes and covariates with server. Although transmitting
p-dimensional prototypes introduces marginal com- munication cost, this cost grows exponentially for transmitting the covariance matrices of p dimension. For FedGen and FedType, the server respectively broadcasts a generator (2-layer per- ception (Zhu et al. 1021 and proxy model (ResNet-18 (Wang et al.) (2024) to clients, while these overparametrized models usually consume massive communication resources transmission. In comparison, FedProto and FedTPG only transmit prototypes with significantly less communication costs. Although FedVTC transmits additional messages other than prototypes, including SDs and VTC models, it still reduces the total communication costs compared with FedProto and FedTGP, as no information is transmitted in the phase of fine-tuning.
Memory requirement. We calculate the required memory space of each method by accumulating the magnitudes of parameter weights, gradients and activations involved in training (Pfeiffer et al. 2023). As depicted in Figure 3 compared with other methods, FedVTC accounts an equivalent or less memory space across all model architectures. We attribute this to the strategy of alternately training local models and VTC models in FedVTC. In contrast, FedType, which simultaneously trains local models and the proxy model, demands much higher memory space for training.
(a) MNIST.
(b) CIFAR10.
(c) CIFAR100
(d) Tiny-ImageNet.
Figure 3: FedVTC (green) has equivalent or less memory requirement compared with others.

---

Published as a conference paper at ICLR 2026
Ablation study. We compare the performance of FedVTC with two modules of VTC transmission, including singular (TC models are communicated only once) and regular (TC models are commu- nicated every round) transmissions. As shown in Table 3 the generalization accuracy of singular VTC transmission is slightly lower (sometimes even higher) than regular VTC transmission, while the latter causes much higher communication cost than the former. Therefore, apply the sin- gular VTC transmission strategy in the ultimate FedVTC framework for communication efficiency. Furthermore, we compare two training modules of FedVTC, where the VTC is trained with Ce and Lie Le + dm respectively. As shown in Table 4 the DM-based regularization effectively increases the generalization performance of FedVTC all scenarios.
1
2
3
4
5
7
8
9 *
7
MNIST
CIFAR10
Figure 4: A fraction of raw images (top) and synthetic images (bottom) MNIST and CIFARIO.
Privacy preservation. FedVTC transmits prototypes, SDs and VTC models, which brings no ad- ditional privacy concerns according to (Luo al. 2021 Tan al. 2022). The reason is that, is almost impossible to recover the raw data from prototypes and SDs without accessing the raw rep- resentations. As Figure shows, the VTC-generated images diverge drastically from real images, making it extremely difficult to capture any sensitive information from these images.
5
CONCLUSION
This paper proposes FedVTC, model-heteerggeneuus FL framework based on variational trans- posed convolution (VTC). By fine-tuning local models with VTC-generated samples, FedVTC ef- fectively overcomes the www-generaiization bottleneck of FI. clients in hooeee-hetroogeneous settings without relying on any public datasets. the future, we aim to consolidate FedVTC with the early- stopping technique (Niu al. (024) to mitigate the additional computation cost of training VTC. We also plan to amplify the adaptability of FedVTCC to data from unseen classes with techniques like zero-shot learning (Chen et al.. 2024) and model diffusion (Croitoru al. (023).
ACKNOWLEDCEMENTS
This work is funded by the Australian Research Council under Grant No. DP220101823. DP200102611, and LP180100114.
REFERENCES
Prabath Abeysekara, Hai Dong, and A. K. Qin. Data-driven trust prediction in mobile edge computing-based iot systems. IEEE Transactions Services Computing, 16(1):246-260, 2023. doi: 10.1109/TS..201.3112889.
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based dynamic regularization. In Proceedings of the 2021 International Conference on Learning Representations (ICLR 2021), 2021.
Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. Fedrolex: model-hnttroggnncuu federated learning with rolling sub-model extraction. In Proceedings of 31st Conference Neural Infor- mation Processing Systems (NeurIPS 2022), 2022. ISBN 9781713871088.
Sebastian Caldas, Jakub Konečny, H. Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. In NeurlPS Workshop Federated Learning for Data Privacy and Confidentiality, 2018.
Xiaoyu Cao, Minghong Fang. Jia Liu, and Neil Zhenqiang Gong. Fltrust: Byzantine- robusst feder- ated learning via trust bootstrapping. arXiv preprint vrXiv:20/2.33995 2020.
10

---

Published as a conference paper at ICLR 2026
Xin Chen, Xiaoling Deng, Yubin Lan, Yongbing Long, Jian Weng. Zhiquan Liu, and Qi Tian. Ex- planatory object part aggregation for zero-shot learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(2):851-863 2024. doi: 10.1 10.11099PAMML2223333335
Sijie Cheng, Jingwen Wu, Yanghua Xiao, Yang Liu, and Yang Liu. Fedgems: Federated learning of larger server models via selective knowledge fusion. arXiv preprint arXiv:2110.11027 2021.
Yae Jee Cho, Andre Manoel, Gauri Joshi, Robert Sim, and Dimitrios Dimitriadis. Heterogeneous ensemble knowledge transfer for training large models in federated learning. In Proceedings of the Thirty-First International Joint Conference Artificial Intelligence (IJCAI-22), pp. 2881-2887, 2022. doi: 2020)/ Main Track.
Patryk Chrabaszez, Hya Loshchilov, and Frank Hutter. A downsampled variant of imagenet as alternative to the cifar datasets. arXiv preprint arXiv:1707.08877 2017.
Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor Ionescu, and Mubarak Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(9): 10850-10869, 2023. doi: 10. 20.1000/
Enmao Diao, Jie Ding, and Vahid Tarokh. (fl: Computation and communication efficient federated learning for heterogeneous clients. In Proceedings of the 2021 International Conference on Learning Representations (ICLR 2021), 2021.
Canh T. Dinh, Nguyen H. Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. In Proceedings of the 34th Conference on Neural Information Processing Systems (NeurIPS 2020), 2020. ISBN 9781713829546.
Vincent Dumoulin and Francesco Visin. guide convolution arithmetic for deep learning. arXiv preprint arXiv:1603.07285, 2016.
Liang Gao, Huazhu Fu, Li Li, Yingwen Chen, Ming Xu, and Cheng-Zhong Xu. Feddc: Federated learning with non-iid data via local drift decoupling and correction. In Proceedings of the 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR 2022), pp. 10112- 10121, June 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR 2016), pp. 770-778, 2016.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two-time-scale update rule converge local nash equilibrium. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp- 6640, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860999
Samuel Horváth. Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout. In Proceedings of the 35th Conference on Neural Information Processing Sys- tems (NeurIPS 2021), volume 34, pp. 12876-12889, 2021
Chun-Yin Yin Huang, Kartik Srinivas, Xin Zhang, and Xiaoxiao Li. Overcoming data and model heterogeneities in decentralized federated learning via synthetic anchors. arXiv preprint arXiv:2405.11525. 2024.
Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto. Distillation-based semi-supervised federated learning for communication-efiiienn collaborative training with non-iid private data. IEEE Transactions Mobile Computing, 22(1):191-205. 2023.
Yae Jee Cho, Jianyu Wang, and Gauri Joshi. Towards understanding biased client selection in federated learning. In Proceedings of the 25th International Conference on Artificial Intelligence and Statistics (AISTATS 2022), volume 151, pp. 10351-10375, 2022.

---

Published as a conference paper at ICLR 2026
Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin K Leung, and Leandros Tassiulas. Model pruning enables efficient federated learning on edge devices. IEEE Transactions Neural Newworks and Learning Systems, 2022.
Zhida Jiang, Yang Xu, Hongli Xu, Zhiyuan Wang. Jianchun Liu, Qian Chen, and Chumming Qiao. Computation and communication efficient federated learning with adaptive model pruning. IEEE Transactions Mobile Computing, pp. 1-18, 2023.
Sai Pranceth Karimireddy, Satyen Kale, Mchryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning (ICML 2020), volume 119, pp. 5132-5143, 2020.
Minjae Kim, Sangyoon Yu, Suhyun Kim, and Soo-Mook Moon. DepthFL: Depthwise federated learning for heterogeneous clients. In Proceedings of the Eleventh International Conference on Learning Representations (ICLR 2023), 2023.
Diederík Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114, 2013.
arXiv preprint
Alex Krizhevsky. Learning multiple layers of features from tiny images. University of Toronto, 2009.
Fan Lai, Xiangfeng Zhu, Harsha V. Madhyastha, and Mosharaf Chowdhury. Oort: Efficient feder- ated learning via guided participant selection. In Proceedings of the 15th USENIX Symposium Operating Systems Design and Implementation (OSDI 2021), pp. 19-35, 2021. ISBN 978-1- 939133-22-9.
Van An Le, Nam Duong Tran, Phuong Nam Nguyen, Thanh Hung Nguyen, Phi Le Nguyen, Truong Thao Nguyen, and Yusheng Ji. Enhancing the generalization of personalized federated learning with multi-head model and ensemble voting. In 2024 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pp. 205 216, 2024.
Yann LeCun, 1998. URL http:://yamet lecun. wtt.:/o///eeeee http:/cc///xx/mmmitsste.o
Ang Li, Jingwei Sun, Pengcheng Li, Yu Pu, Hai Li, and Yiran Chen. Hermes: An efficient federated learning framework for heterogeneous mobile clients. In Proceedings of the 27th Annual Interna- tional Conference on Mobile Computing and Networking (MobiCom 2021), pp. 420-437, 2021a. ISBN 9781450383424.
Chenning Li, Xiao Zeng, Mi Zhang, and Zhichao Cao. Pyramidft: fine-grained client selec- tion framework for efficient federated learning. In Proceedings of the 28th Annual International Conference on Mobile Computing and Networking (MobiCom 2022), pp. 158-171, 2022. ISBN 9781450391818.
Daliang Li and Junpu Wang. Fedmd: Heterogenous federated learning via model distillation. arXiv preprint
2019.
Suyi Li, Yong Cheng, Wei Wang, Yang Liu, and Tianjian Chen. Learning to detect malicious clients for robust federated learning. arXiv preprint arXiv:2002.002// 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of the 2020 Conference on Ma- chine Learning and Systems (MLSys 2020), pp. 429-450, 2020b.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In Proceedings of the 2021 International Conference on Ma- chine Learning (ICML 2021), pp. 6357-6368, 2021b.
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of fedavg non-iid data. arXiv preprint 1907.02 2019.
12

---

Published as a conference paper at ICLR 2026
Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, and Jiashi Feng. No fear of heterogene- ity: Classifier calibration for federated learning with non-lID data. Proceedings of the 35th Conference on Neural Information Processing Systems (NeurIPS 2021). 2021.
Connor Mclaughlin and Lili Su. Personalized federated learning via feature distribution adaptation. Proceedings of the 38th Conference on Neural Information Processing Systems (NeurIPS 2024), 37:77038-77059, 2024.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communicatoon-fficcenttt learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282, 2017.
Xutong Mu, Ke Cheng. Yulong Shen, Xiaoxiao Li, Zhao Chang, Tao Zhang, and Xindi Ma. Feddme: Efficient and robust federated learning via detecting malicious clients. IEEE Transactions on Dependable and Secure Computing, 21(6):5259-527)) 2024. doi: 1 1.1119/TSC.22223337726444
Duy Phuong Nguyen, Sixing Yu, J. Pablo Muñoz, and Ali Jannesari. Enhancing heterogeneous fed- erated learning with knowledge extraction and multi-model fusion. In Proceedings of the SC '23 Workshops of the International Conference on High Performance Computing, Network, Storage, and Analysis (SC-W 2023), pp. 36 43, New York, NY, USA, 2023. ISBN 9798400707858, doi: 10. 10.1145//324406.36622222
Ziru Niu, Hai Dong, A. K. Qin, and Tao Gu. Firce: Resource-efficien federated learning with early- stopping strategy. IEEE Transactions Mobile Computing, 23(12)14514414422, 2024. doi:
N0.1009/TMC.22244444700000
Adam Paszke, Sam Gross, Francisco Massa, Adam Lever, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Yang. Zach De Vito. Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: an imperative style, high-performance deep learning library, volume 32, pp. 8024 8035. Curran Associates Inc., 2019.
Zihao Peng, Xijun Wang, Shengbo Chen, Hong Rao, Cong Shen, and Jinpeng Jiang. Federated learning for diffusion models. IEEE Transactions on Cognitive Communications and Networking, pp. 1-1, 2025. doi: 10. R0.1099C...222255550003
Kilian Pfeiffer, Ramin Khalili, and Joerg Henkel. Aggregating capacity in FL through successive layer training for compuuatioonally-contttaaai devices. In Proceedings of the 37th Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
Krishna Pillutla, Sham M. Kakade, and Zaid Harchaoui. Robust aggregation for federated learning. IEEE Transactions on Signal Processing, 70:1 1154,114, 2022. doi: 10.1 10.11099...2222.315553
Nalavala Partheswar Reddy, Shaik Khaja Hassain, P Baraneedharan, and Shaik Abdul Rahman. In- ternet of bhied remote weather forecasting system using legacy sensors. In Proceedings of the 9th International Conference on Advanced Computing and Communication Systems (ICACCS 2023), volume 1, pp. 23095-200, 2023. doi: 10.1 1LL119/ICACCC55222220000. 10115522..
Himadri Nath Saha, Supratim Auddy, Subrata Pal, Shubham Kumar, Shivesh Pandey, Rocky Singh, Amrendra Kumar Singh, Priyanshu Sharan, Debmalya Ghosh, and Sanhita Saha. Health mon- itoring using internet of things (iot). In Proceedings of the 8th Annual Industrial Automa- tion and Electromechanical Engineering Conference (IEMECON 2017), pp. 69-73, 2017. doi: 10. 10.1109/// 100119//EEECON201...7777777
Felix Sattler, Tim Korjakow, Roman Rischke, and Wojeiech Samek. Fedaux: Leveraging unla- beled auxiliary data in federated learning. IEEE Transactions Neural Networks and Learning Systems, 34()):533155433 2021.
Jiangming Shi, Shanshan Zheng, Xiangbo Yin, Yang Lu, Yuan Xie, and Yanyun Qu. Clip-guided federated learning on heterogeneous and long-tailed data. AAAI'24, AAAI Press, 2024. ISBN 978-1-57735-887-9,
13

---

Published as a conference paper at ICLR 2026
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. Proceedings of the 31st Conference on Neural Information Processing Systems (NIPS 2017), 30, 2017.
Fiona Victoria Stanley Jothiraj and Afra Mashhadi. Phoenix: A federated generative diffusion model. Companion Proceedings the ACM Web Conference 2024, WWW '24, pp. 1568-1577, New York, NY, USA, 2024, Association for Computing Machinery.
Lichao Sun and Lingjuan Lyu. Federated model distillation with noise-free differential privacy. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence (IICAI-2I), pp. 1563-1569, 2021.
Zhenyu Sun, Xiaochun Niu, and Ermin Wei. Understanding generalization of federated learning via stability: Heterogeneity matters. In Proceedings of The 27th International Conference Artificial Intelligence and Statistics (AISTATS), volume 238, pp. 676-684, 2024.
Farnaz Tahmasebian, Jian Lou, and Li Xiong. Robustfed: A truth inference approach for robust federated learning. In Proceedings of the 2022 International Conference on Information and Knowledge Management (CIKM 2022), pp. 1868-1877, 2022. ISBN 9781 9781450392365.
Yue Tan, Guodong Long. Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang, Fed- proto: Federated prototype learning across heterogencous clients. In Proceedings of the 2022 AAAI Conference on Artificial Intelligence (AAAI 2022), volume 36, pp. 8432 8440, 2022.
Jiaqi Wang, Chenxu Zhao, Lingjuan Lyu, Quanzeng You, Mengdi Huai, and Fenglong Ma. Bridging model heterogeneity in federated learning via uncertainty-ban based asymmetrical reciprocity learn- ing. In Proceedings of the 41st International Conference Machine Learning (ICMI. 2024), 2024.
Mang Ye, Xiuwen Fang, Bo Du, Pong C. Yuen, and Dacheng Tao. Heterogeneous federated learning: State-of- the-art and research challenges. ACM Computing Surveys, 56(3), 2023. ISSN 0360-0300.
Liping Yi, Han Yu, Gang Wang, Xiaoguang Liu, and Xiaoxiao Li. pfedafim: Adaptive feature mix- ture for data-level personalization in heterogeneous federated learning on mobile edge devices. In 2025 IEEE s1 International Conference on Data Engineering (ICDE), pp. 1981-1994, 2025.
Jianqing Zhang, Yang Liu, Yang Hua, and Jian Cao. Fedtgp: Trainable global prototypes with adapiiv--mariin-eenaaedd contrastive learning for data and model heterogeneity in federated learning. In Proceedings of the 2024 AAAI Conference on Artificial Intelligence (AAAI 2024), volume 38, pp. 16768-1666 16776, 2024.
Jianyi Zhang, Ang Li, Minxue Tang, Jingwei Sun, Xiang Chen, Fan Zhang, Changyou Chen, Yiran Chen, and Hai Li. Fed-CBS: A heterogeneity-awee awware client sampling mechanism for federated learning via class-imbalance reduction. In Proceedings of the 40th International Conference on Machine Learning (ICMI. 2023), volume 202, pp. 41354-4138 2023.
Lan Zhang, Dapeng Wu, and Xiaoyong Yuan. Fedzkt: Zero-shot knowledge transfer towards resource-constrained federated learning with heterogeneous on-device models. In 2022 IEEE 42nd International Conference on Distributed Computing Systems (ICDCS), pp. 928 928-938. IEEE, 2022.
Ganlong Zhao, Guanbin Li, Yipeng Qin, and Yizhou Yu. Improved distribution matching for dataset condensation. In Proceedings of the 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR 2023), pp. 7856-7865, 2023.
Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogeneous federated learning. In Proceedings of the 2021 International Conference on Machine Learning (ICML 2021), pp. 12878-12889, 2021.
A DEFINITION OF MODEL HETEROGENEITY
To our best knowledge, the definitions of handel-hetrrogeneous federated learning" in the state-of- the-art can be divided into the following two categories.
14

---

Published as a conference paper at ICLR 2026
A.1 TAB-MODEL--BSSER MODEL HETEROGENEITY
Sub-model-baae model heterogeneity, also named partial model heterogeneity (Ye 2023), is developed based on the assumption that all local models are a subset of the global model. Specifi- cally, the server preserves global model and employs the typical dropout technique (Caldas al, 2018) to extract a fraction of the global model's parameters, i.e., sub-model and allocates the sub- models to clients. For example, (Caldas et al. 2018) generates sub-models by randomly pruning neurons in the global model. (Horváth et al. 2021) and (Diao et al.] 2021 consistently extract the left-most neurons of the global model to form sub-models. (Alam et (002) dynamically extracts sub-models using sliding window across the neurons in the global model. (Jiang et al. 2023 Li ct al. 202 la Jiang et al., 2022) selectively prune the unimportant neurons in the global model with the least heuristic scores (like parameter norms). (Kim et al. 2023) vertically extracts sub-models by pruning the top layers the global model.
In these works, although clients have different model architectures due to different dropout strate- gies, the parameters between clients are usually co-dependent, and can be aggregated using sub- model aggregation shown in Figure 5
WA(0.40,
0.61,
0.40
0.40,
0.61)
-0.10,
0.61
0.15
-0.73
WA(0.15,
0.15
-0.73
-0.10
-0.73
-0.1)
Client
Server
Client j
Figure 5: The sub-model aggregation scheme for sub-model-bases model-heteroogneeous FL ("WA"
stands for weighted average).
A.2 COMPLETE MODEL HETEROGENEITY
In completely heeereehetroosssuuus FL (shown in Figure the longer maintains global model, and clients obtain isolated local models with diverse, independent, and irrelevant architec- tures (Ye et 2023). In this scenario, clients can no longer acquire generalized parameter updates through model aggregation. As a replacement, clients apply alternative strategies to learn gener- alization information, such as sharing knowledge over public dataset transmitting parameter- independent messages, as discussed in Section 2.2
Without ambiguity, in this paper, the definition of "model heterogeneity" refers to complete model heterogeneity (Appendix A.2) where all aggregation (e.g., standard aggregation and sub-model aggregation) methods become prohibitive. Correspondingly, the goal FedVTC to improve the generalization performance of clients in the settings of complete model hetero- geneity without parameter aggregation involved.
B EXPERIMENT DETAILS
B. FOUNDAMENTAL SETTINGS
The fundamental experimental settings are summarized in Table
B.2 STRUCTURE OF THE TRANSPOSED CONVOLUTION MODEL
The detailed information the structure of the VTC model is listed in Table
15

---

Published as a conference paper at ICLR 2026
Server
PIM
PIM
PIM
Client i
Client
Client k
Figure 6: In complete mooee--heeroggeneuu FL, model architectures between clients are inde- pendent, irrelevant, and cannot be aggregated altogether. Alternatively, clients communicate some Parameter- Inndeeede-Iddeeedeenn Messages (PIM) with the server for knowledge sharing, such as prototypes and feature covariates.
Channel
Input dimension
Latent dimension p
Class C
Initial SD o
MNIST CIFAR10
1
1 x 28 x 28
980
10
3
3 x x
4096
10
1 E R
CIFAR100
3
3 x 32 x 32
4096
100
Tiny-ImageNet
3
x x
12768
200
Table 5: Fundamental experimental settings.
B.3 ACCURACY GRAPHS
The graphs reflecting the generalization accuracy against the FL iterations are shown in Figure
B.4 HYPERPARAMETER TUNING
As stated previously, the number of synthetic samples per class S is set 500 for MNIST and 1000 for CIFAR10, CIFARI00 and -mag-ema following (Luo et al.] 2021 for faimess.. Despite this, additionally tune S the MNIST dataset to study the impact of different S values. As shown in Figure Sal when S is too small, the generalization accuracy is low as the effect of fine-tuning weak. When S too large, many redundant synthetic samples will be generated, which may degrade the effectiveness of fine-tuning.
In addition, we evaluated the coffecient À of the DM loss (Eq. (5) As shown in Figure 8b turns out that x= = 0.1 is the optimal setting based on the empirical results.
16

---

Published as a conference paper at ICLR 2026
MNIST
Index
Components
Input shape
Output shape
TranposCCom/2(eemmee-,,, strddde
1
BatchNorm220
20 x x x
x
LeakyRel.U(0.01)
2
BatchNorm22d(
16 x x
32 14x 14
LeakyReLUU(0.0))
TaunpsseCon/2kkeeee-4a.ddddddd stridee BatchNorm220
3
32 14 14
32 14x14
LeakyReLUU(0001)
TrunspsseCon/22keeeeee44.. str
4
BatchNorm22()
32 x 14 14
1 1 28 x
Sigmoid(
CIFARIO & CIFAR100
Index
Components
Input shape
Output shape
TTansposConn22deeeeee3,,, padding=1, stridee)
1
BatchNorm2d0
64 x 8 x 8
32 8 x 8
LeakyReLU(0000))
Sumnsporcoom//kmmmmm444p.dnnnnge stridee)))
2
BatchNorm220)
32 x 8 X
64 16 16
LeakyReL.U(0.01)
TransposCCon/22Kkeeee33333 pttp://eeeeeaa.aiinnggln..eeeeee stride)
3
BatchNorm220
64 16 16
64 16 16
LeakyReLU((0001)
TranppssCCom/2KKeeeee444nppdipg 2tid))
4
BatchNorm220)
64 X 16 16
3 X 32 x 32
Sigmoid()
Imay-ImageNet
Index
Components
Input shape
Output shape
TrangpossCCnn/22Keeeeee13,, padding=1, stride-) BatchNorm22()
1
128 x 16 x 16
64 x 16 x 16
LeakyReLU(0.01)
2
BatchNorm220)
64 x 16 16
64 32 x 32
LeakyReLU(0.01)
3
BatchNorm220
64 32 32
64 32 32
LeakyReLUU(0.0))
TrangossCon22ddeeeeeea.n.. padding=l,, stride=2) BatchNorm22()
32 32
3 x x
Sigmoid()
Table 6: Structural information on the transposed convolutional neural network.
17

---

Published as a conference paper at ICLR 2026
HHIST
CIPAR110
0.25
0.70
0.6
0.5
FedPrate
Fedhrats
0.00
FedPrato
0:15
02
CCWR
0.00
CCVR
Pebliic
Podhicc
0.5
Fodric
0.0
FLSS Fine-tnne soit
0.
EL The- eath
0.00
EL Tuee col
20
50
20
120
20
60
120
20
60
20
100
(a) MNIST Dir (0.1).
(b) CIFARIO Dir (0.1).
(c) CIFARI00 Dir (0.1).
0.8
Accureyy
Fodhiss de
Poll de
Fealbate
TedTGP
Fodsen
Federe
0.0
0.0
0.0
20
80
130
20
60
150
130
(d) MNIST Dir (1.0).
(c) CIFARIO Dir (1.0).
(f) CIFARI00 Dir (1.0).
Tiny Imagstett
Thy Imagetist
0.50
0.32
0.30
0.22
0.25
FedFricto
0.20
FedFFroto
FedGinn
FodTOP
2015
PesTOP
FedGinn
0.18
FadTypa
0.18
0.02
0.33
0.00
0.32
100
25
200
Round
Round
(g) Tiny- Dir (0.1).
(h) Tiny-ImageNet Dir (1.0).
Figure 7: The per-round average generalization accuracy (i.e., the mean of each local model's s aceu- racy on the global validation set). The vertical dashed line splits the standard FI. iterations and the local fine-tuning procedure.
18

---

Published as a conference paper at ICLR 2026
90.0
89.8
Accuracy (%)
(%
89.6
89.4
88-
89.2
89.0-
88.0
88.6
1000
2000
3000
4000
0.0001
0.001
0.01
0.1
$
Regularization coffecient
(a) Generalization accuracy MNIST with Din(1.0) (b) Generalization accuracy MNIST with Dir(1.0)
vs. number of synthetic samples per class.
vs. the regularization coffecient
Figure 8: The generalization accuracy with different values of X.
Before fine-tuning
After Fine-tuning
Improvement
MNIST Dir(O.1)
84.9%
88.7%
3.8% t 3.3%
MNIST Dir(1.0)
86.8%
90.1 1%
CIFARIO Dir(O.1)
40.5%
46.9%
6.4%
CIFAR1O Dir(1.0)
41.6%
51.7%
10.1%
CIFARI00 Dir(O.1)
27.6%
36.3%
8.7% 7.6%
CIFARIO0 Dir(1.))
32.8%
40.4%
Tiny-ImageNet 1ir(0.1)
25.1%
30.2% 35.8%
1 5.1%%
Imag-lm Dir(1.0)
29.8%
6.0%
Table 7: The generalization accuracies of FedVTC before and after fine-tuning with the synthetic samples.
B.5 QUALITY OF SYNTHETIC SAMPLES
In this paper, the objective of generating synthetic samples is to produce training data for fine-tuning instead of generating realistic images. Therefore, traditional metrics measuring how generated im- ages resemble real images (like FID score Heusel et al. 2017) do not accurately reflect the quality of the synthetic samples. As a replacement, we compare the accuracy of FedV'TC before and after fine-tuning examine the quality of the synthetic samples. As depicted Figure ( and Table the accuracy of FedVTC significantly increases after fine-tuning, demonstrating good quality of the synthetic samples.
C USE OF LARGE LANGUAGE MODELS
Large Language Models are only used to correct spelling and grammatical mistakes in this paper.
19


Reviews:
[{"review_id": "CXgfhBBqWW", "reviewer": "Reviewer_Ltp9", "summary": "This paper proposes a method to improve the generalization of personalized Federated Learning, specifically enhancing accuracy on unseen data. The authors suggest generating synthetic data at the client side to fine-tune the global model after the federated training phase. This approach improves the performance of local models without requiring access to a public global dataset. To achieve this, the paper introduces a variational transposed convolutional neural network for synthetic data generation. Experimental results demonstrate significant improvements in the average client accuracy across multiple settings, including four datasets and two non-IID scenarios.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1. The paper is clear to understand and follow\nS2. The approach of fine-tuning the local models of clients using synthetic data is not new but the method using to generate the synthetic data is novel. \nS3. Significant improvement in average accuracy among clients are shown.", "weaknesses": "W1. Computation and Communication Overhead: The proposed method requires the server to train and distribute the VTC model to clients. In addition, clients must generate synthetic data locally. These steps introduce non-trivial computational overhead compared to standard federated learning. Please discuss the communication and computation requirements of the proposed method and compare them with those of other existing approaches.\n\nW2. Privacy Preservation: Although the authors discuss privacy preservation in the paper, the examples shown in Figure 4 do not convincingly demonstrate that data privacy is maintained. It would be beneficial to quantitatively assess the similarity between the original samples and the corresponding generated data. For instance, the authors could use the method proposed in Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli, “Image quality assessment: from error visibility to structural similarity,” IEEE Transactions on Image Processing, vol. 13, no. 4, pp. 600–612, 2004, which introduces the Structural Similarity Index (SSIM).\n\nW3. The authors compares FedVTC with 5 other state of the art. However, the authors should discuss the approach in this paper \"Enhancing the Generalization of Personalized Federated Learning with Multi-head Model and Ensemble Voting\" (https://ieeexplore.ieee.org/abstract/document/10579121) which also try to improve the generalization of Personalized Federated Learning.  It will be better if the author could show that proposed method outperform such kind of approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tKLgmUJMED", "reviewer": "Reviewer_xMVi", "summary": "This paper proposes FedVTC, a model-heterogeneous federated learning framework based on Variational Transposed Convolution. The framework generates synthetic data to fine-tune local models, thereby improving the generalization of heterogeneous clients without relying on public datasets. The paper compares FedVTC with several existing methods, but it lacks comparisons against strong works from 2025.", "questions": "1. I would like to know the meaning of the symbol $V_i$ in Formula 5.\n\n2. Could you add experiments comparing with the outstanding works from 2025 to provide a clearer comparison of experimental results?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper proposes using FedVTC to generate synthetic data, which can be used to fine-tune local models and improve their generalization ability, thereby eliminating the reliance on public datasets.\n\n2. FedVTC avoids exposure of raw data by transmitting prototypes, covariances, and VTC models, thus preventing additional privacy risks.\n\n3. A new objective function is designed for training the VTC model, which includes the standard negative ELBO loss and a distribution matching (DM) loss, regularizing the training process and improving the quality of the generated synthetic data.", "weaknesses": "1. Some symbols in the formulas are not defined, such as the symbol $V_i$ in Formula 5, which lacks a definition.\n\n2.  The proposed FedVTC framework does not introduce an entirely new solution or framework, but rather builds upon existing methods, which results in a somewhat limited level of innovation.\n\n3. The comparative experiments seem to only compare with works published before 2025, without including comparisons with excellent works from 2025.\n\n4. The flowchart in the paper is somewhat simplistic. For example, in Figure 1, the elements, colors, and arrows are used in a rather basic way, and the relationships and operational details between the steps are not fully demonstrated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8ILWUeJIIe", "reviewer": "Reviewer_MpGA", "summary": "This paper addresses the challenge of data and model heterogeneity in federated learning (FL), where clients possess non-IID data and diverse model architectures, leading to limited generalization. The authors propose FedVTC, a model-heterogeneous FL framework that leverages variational transposed convolution (VTC) networks to generate synthetic data for fine-tuning local models. Instead of sharing model parameters, clients communicate feature distribution statistics (prototypes and standard deviations) with the server. The VTC model is trained using a composite loss combining evidence lower bound (ELBO) and distribution matching (DM) terms. Experiments on MNIST, CIFAR10, CIFAR100, and Tiny-ImageNet demonstrate improved generalization accuracy, reduced communication costs, and efficient memory usage compared to existing methods.", "questions": "Please respond to the Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The use of variational transposed convolution to produce synthetic samples for fine-tuning local models is a distinctive approach, avoiding dependency on public datasets or parameter aggregation.  \n\nThe method is evaluated on four benchmark datasets under varying non-IID settings (Dir(0.1) and Dir(1.0)), demonstrating superior generalization accuracy over five state-of-the-art baselines.  \n\nFedVTC achieves lower communication costs than FedProto, FedTGP, FedGen, CCVR, and FedType and maintains comparable or reduced memory usage through alternating training of local and VTC models.  \n\nThe authors argue that transmitting prototypes and standard deviations does not expose raw data, supported by visual evidence of synthetic samples diverging from real images.", "weaknesses": "Although FedVTC is compared to five baselines, the discussion overlooks potential overlaps or distinctions with contemporary approaches such as those using diffusion models or zero-shot learning.  \n\n Generalization performance is measured solely by accuracy; additional metrics like per-class precision/recall or robustness under extreme non-IID conditions would strengthen the claims.  \n\nThe requirement for homogeneous VTC models for aggregation may limit applicability in fully heterogeneous settings where clients cannot support the same VTC structure.  \n\nThe paper does not analyze cases where VTC fails to generate useful synthetic data (e.g., under high data skew or low client participation), nor does it explore the impact of varying latent dimensions on performance.  \n\nWhile memory usage is addressed, the additional computational cost of training VTC networks is not quantified or compared to baseline methods.  \n\nNo formal convergence guarantees or analysis are provided for the alternating optimization of local models and VTC networks, which is critical for FL methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Kqs1iczfWO", "reviewer": "Reviewer_fKa4", "summary": "This paper addresses the generalization challenge in model-heterogeneous federated learning where clients have different model architectures, making standard parameter aggregation infeasible. The proposed framework, FedVTC, avoids parameter sharing and instead has clients exchange feature distribution statistics (mean and covariance) with the server. Each client then trains a local variational transposed convolutional network, using latent variables sampled based on the aggregated statistics, to generate synthetic data. By fine-tuning their local models on this generated data, clients improve generalization without needing a public dataset.", "questions": "Please refer to above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The motivation of the method is described accurately and in detail, the algorithm process is also explained clearly, and the specific ablation experiments are very comprehensive.", "weaknesses": "1. Since each client needs to generate high-quality samples during training, please provide a more detailed analysis of the associated computational cost. Is this process computationally expensive, and if so, how does it affect training efficiency?\n\n2. The paper claims to address model heterogeneity. Please clarify whether the proposed method can handle heterogeneous architectures (e.g., ResNet, MobileNet, CNNs) rather than merely different-sized variants of the same backbone (e.g., ResNet-18 vs. ResNet-50). \n\n3. Since the primary goal of the proposed method is to improve generalization under heterogeneous model settings, the experiments should explicitly evaluate cross-domain generalization. In particular, results should be presented for each client model tested across all global test datasets. Currently, Figure 7 lacks details about the experimental setup, and Table 1 ambiguously refers to “all datasets and distributions.” Please describe these configurations more precisely to ensure reproducibility and clarity.\n\n4. The comparative analysis omits several VAE-based federated learning methods, which are highly relevant to the proposed approach. Including these baselines would provide a more comprehensive evaluation of the method’s relative advantages.\n\n5. In heterogeneous model scenarios, achieving generalization capability is likely related to class-specific feature representations. Even if data distributions differ across clients, a visualization (e.g., feature embeddings or attention maps) illustrating shared generalizable features would greatly enhance the interpretability of the results.\n\n6. In the privacy validation experiments, the visual differences between reconstructed and original samples are evident to the human eye. However, for a machine learning system, such differences should be quantitatively evaluated (e.g., via PSNR, SSIM, or feature similarity metrics). A quantitative privacy comparison would make the evaluation more rigorous.\n\n7. Several detail-related issues should be addressed. The abbreviation “TC” is undefined and should likely be “VTC”; please ensure consistency throughout the manuscript. Additionally, in line 195, “e.g.” should be written as “e.g.,” following standard academic style.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)