Toward Enhancing Representation Learning in Federated Multi-Task Settings
Authors: Mehdi Setayesh
Mahdi Beitollahi
Yasser H. Khalil
Hongliang Li
Venue: ICLR 2026 Poster (2026)
Keywords: Contrastive learning
federated learning
knowledge transfer
multi-task learning
mutual information
representation learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: This paper proposes a new contrastive learning loss and, leveraging it, designs a novel federated multi-task learning algorithm that addresses both task and model heterogeneity among users.

Abstract: Federated multi-task learning (FMTL) seeks to collaboratively train customized models for users with different tasks while preserving data privacy. Most existing approaches assume model congruity (i.e., the use of fully or partially homogeneous models) across users, which limits their applicability in realistic settings. To overcome this limitation, we aim to learn a shared representation space across tasks rather than shared model parameters. To this end, we propose *Muscle loss*, a novel contrastive learning objective that simultaneously aligns representations from all participating models. Unlike existing multi-view or multi-model contrastive methods, which typically align models pairwise, Muscle loss can effectively capture dependencies across tasks because its minimization is equivalent to the maximization of mutual information among all the models' representations. Building on this principle, we develop *FedMuscle*, a practical and communication-efficient FMTL algorithm that naturally handles both model and task heterogeneity. Experiments on diverse image and language tasks demonstrate that FedMuscle consistently outperforms state-of-the-art baselines, delivering substantial improvements and robust performance across heterogeneous settings.

Published as a conference paper at ICLR 2026
TOWARD ENHANCING REPRESENTATION LEARNING IN FEDERATED MULTI-TASK SETTINGS
Mehdi Setayesh, Mahdi Beitollahi, Yasser H. Khalil, and Hongliang Li
Huawei Noah's 's Ark Lab, Montreal, Canada mehdi... set ayesalysslle ayeshh@huaweeiii com yasser. wha.hhaliii i 1111111 Bhuawei.... com
mahdi wwwtteiiieeeeee www.ahhipaatters.oomm
hongliang. wwwwwwa.ee....mmmm.. com
ABSTRACT
Federated multi-task learning (FMTL) seeks to collaboratively train customized models for users with different tasks while preserving data privacy. Most existing approaches assume model congruity (i.e., the use of fully or partially homoge- neous models) across users, which limits their applicability in realistic settings. To overcome this limitation, we aim to learn a shared representation space across tasks rather than shared model parameters. To this end, we propose Muscle loss, a novel contrastive learning objective that simultaneously aligns representations from all participating models. Unlike existing multi-view or multi- con- trastive methods, which typically align models pairwise, Muscle loss can effec- tively capture dependencies across tasks because its minimization equivalent to the maximization of mutual information among all the models" representations. Building on this principle, we develop FedMuscle, a practical and communication- efficient FMTL algorithm that naturally handles both model and task heterogene- ity. Experiments on diverse image and language tasks demonstrate that FedMus- cle consistently outperforms oftte-of-the-arr baselines, delivering substantial im- provements and robust performance across heterogeneous settings.
INTRODUCTION
Federated learning (FL) enables a group of users to collaboratively train models while preserving
privacy by not sharing their local data (Konečný et al., [016 AbdulRahman et al. 2021 1)
Most
conventional FI. algorithms assume that users share the same model architecture (Smith et al. 2017
Ll 2019 Dinh et al. 2020 Li al. 2021b) or perform the same task (Diao et al. 2021 Alam
et al. 2022 Zhu et al. 2022 Setayesh et al. (023). which limits their applicability in diverse real- world settings (Cai et al. 2024 Wang et al. (024). As illustrated in Figure practical FL scenarios may involve users with heterogeneous model architectures and tasks, each using task-specific local dataset. For example, with recent advances foundation models (FMs), users can select a pre-trained FM based their resource constraints and task requirements, and then fine-tune it using a local dataset (Bommasani et al.) 2022 Shao et al.) 2024 Zheng et al.] 2025). Task heterogeneity in FL scenarios has motivated the development of various federated multi-task learning (FMTL)
algorithms Park et al. 2021 He et al. 2024 Chen et al. 2023 Lu et al. 2024 Jia et al. 2024).
Early works in FMTL focused on personalization tasks (i.e., non-IID data across users) and aimed to learn a customized model with the same architecture for each user (Smith et al. 2017 Marfoq et al., 2021 Recent studies have addressed a wider range of tasks by dividing users' models into shared encoder and task-specific predictors (Jia et al. 2024, Lu et al. 2024). As a result, most existing FMTL algorithms assume that users employ fully or partially homogeneous model architectures. From broader perspective, our central idea is that sharing model parameters, whether fully or partially, ultimately aims to establish shared representation space for effective knowledge transfer among users' tasks. Therefore, the objective of an FMTL algorithm can be reframed as learning shared representation space across tasks.
Contrastive learning (CL) is a widely used technique for learning a shared representation space (He et al. 2020 Misra & Maaten 2020 Chen et al. 2020). The core idea of CL is to bring similar instances, known as positives, closer together in the representation space while pushing dissimilar ones, known as negatives, farther apart (Radford et al. 202I). One popular CL loss function is In- foNCE (Oord et al. 2019), which aligns the representations of two models. For scenarios with more

---

Published as a conference paper at ICLR 2026
DLDOEE Deaad
Sana
leago eiis
Ditact
Mille 2!
#1
Toeis srop our
honerced dee
HPRT
NO Newa
Sipple
Reproaaatt
Aggeggied Hara
Appoopate
Weightngg Vickar o
Coaticiin Verie
Vicse of
Sialt
The server compara du wiighting andfiic lor vacien appoppted rproee merrecs
Somantic 1550
UXKC
Roperon al
poo nod Macco las binsss
nouthvel of haad
Figure FedMuscle overview. Users may have heterogeneous models, tasks, local datasets. shared public dataset is used by all users to align the representation spaces of their models via the Muscle loss function.
than two models, prior works typically adopt pairwise alignment approaches, applying the InfoNCE loss each pair of models (Yang al. 2022] Girdhar al. 2023 Xue al. (024a)). However, pair- wise alignment cannot effectively capture the dependencies among representations obtained from all Citraation,Cccchetti
models (Kromilaas et al. 2025). To address this limitation, [iichhtti al. (2025) recently proposed a Gramian- contrastive loss that aligns all models' representations simultaneously. Neverthe- less, their approach lacks theoretical justification.
In this paper, we propose Muscle (Multi (Multi-task// task/modal Systematic Contrastive Learning), novel CL loss function that effectively captures dependencies among the representations of all models multi- model scenarios. This loss incorporates theoretically grounded weighting coefficients to emphasize more dissimilar negatives, an important mechanism absent in prior works. We show that minimizing the proposed Muscle loss is equivalent maximizing lower bound on the mutual information (MI) between the models' outputs, thereby enhancing knowledge transfer across tasks. Furthermore, design novel FMTL algorithm called FedMuscle by leveraging the Muscle loss. As shown in Figure [ in FedMuscle, users transmit their private knowledge on a shared public dataset to the server rather than sharing their local datasets model parameters. Not revealing local model parameters provides an additional layer of privacy protection, especially when users employ pre- trained FMs (Du et al., 2025). Upon receiving low-dimensional representations of public data from users, the server computes the weighting coefficients and aggregated representations that each user needs align its representation space with those of others by minimizing the Muscle loss.
In summary, (1) we propose Muscle loss, a CL objective that applies to representations from any number of models and systematically captures dependencies among them. (2) We develop FedMuscle, a novel FMTL algorithm that addresses both model and task heterogeneity across users by aligning their representation spaces using Muscle loss. (3) We empirically validate the effective- ness of FedMuscle on various computer vision (CV) and natural language processing (NIP) tasks, demonstrating its superior performance over state-of-the-an baselines under settings with model and task heterogeneity. (4) We show that our proposed systematic CL framework can be seamlessly inte- grated into multi-modal FL algorithms, such as CreamFL (Yu et al. 2023), to enhance performance by replacing the heuristic knowledge distillation approach with Muscle loss.
RELATED WORK
Federated Multi-Task Learning Various FMTL algorithms have been proposed to address task heterogeneity in FL. For example, FeSTA (Park al. 2021), FedBone (Chen et al. 2023), FedHCA² (Lu et al.) 2024). FedLPS (Jia et al. 2024), and FedDEA (Wei et al.) 2025) have been introduced to enhance the generalization performance of all tasks by leveraging knowledge from
2

---

Published as a conference paper at ICLR 2026
lated tasks. However, all of these FMTL algorithms assume model congruity, meaning that all users employ fully or partially homogeneous model architectures. Another line of research is federated semi-supervised learning (FSSL) (eong et al. 202I). with algorithms such as CoFED (Cao et al. 2023) and SAGE (Liu et al. 2025), where users collahoratively generate pseudo-labels for a public dataset. However, most FSSL. approaches can only handle personalization tasks.
Model Heterogeneity in FL Model heterogeneity in FL. may arise from users' diverse and lim- ited storage, communication, and computational capabilities (Setayesh et al., 2023). Some existing works assume global model architecture. where each user's 's local model parameters are subset of the global model's '' parameters (Diao et al. 2021] Hong et al. 2022). However, such approaches limit the flexibility for users choose arbitrary model architectures. To overcome this limitation, one potential solution is to adopt knowledge distillation (KD))based FL algorithms (Li & Wang 2019 Park al. 2024). For example, FCCL (Huang et al. 2022), KT-pFL. (Zhang et al. 2021). and FedDF (Lin et al. (0200) leverage KD and an unlabeled public dataset to address model heterogene- ity. However, in these approaches, knowledge can only be transferred among models with the same logit size (i.e., models associated with the same task). FedHeNN (Makhija et al., (022) addresses this issue by introducing a proximal term into users' local loss functions. Centered kernel alignment (CKA) is used in FedHeNN to measure similarity among representation matrices. Nevertheless, the reliability of CKA as a similarity metric is still under investigation (Davari et al. 2023
Representation Learning in FL CL is a widely used technique for representation learning. Re- cently, incorporating CL into FL has shown promise in improving the performance of users' models. Federated SimCLR (Louizos et al., 2024) employs the SimCLR loss (Chen et al., 2020) as the users' local loss function. In MOON (Li et al. 2021a), each user applies CL to bring the representations obtained from its local model closer to those of the global model and farther from those of its pre- vious local model. However, both Federated SimCLR and MOON assume homogeneous models across users. FedPCL (Tan et al., 2022) and FedTGP (Zhang et al., 2024) utilize CL with a fo- cus on class-wise prototypes, which makes them inapplicable to multi-task scenarios. multi-moda federated learning (MFL) (Yu et al. 2023 Qi & Li 2024; Che et al. 2024; Chen et al., 2024) is another line of research that relies representation learning. CreamFL (Yu et al.) (023) is of the most effective MFL algorithms, in which local contrastive regularization (LCR) is applied at the users and global-local contrastive aggregation (GCA) is performed the server to learn global model. Although the focus of our work is not on obtaining a global model, our experiments show that replacing LCR and GCA in CreamFL with our proposed systematic CL using the Muscle loss improves the performance of the global model on multi-modal retrieval task.
3 PROBLEM FORMULATION
Notations In this paper, we represent vectors by boldface lowercase letters (e.g., 2) and sets by calligraphic letters (e.g.. 2). The cardinality of set 2 and the number of elements in vector z are de- noted by |z| and |z|, respectively. We define [N] {1, 2, ..... _N3 and {z, Ji= {z1, #2,.... 2 The cosine similarity between two normalized vectors z; and z; is given by z; 2j 2. Zn-
Problem Setting We consider an FL setting with one server and N users, where each user has a local, task-specific labeled dataset. Let D" denote the local dataset of user n [N]. Users may employ heterogeneous models based on their tasks and available resources (e.g., storage capacity and computational capability). Let 0" denote the parameters of user n's local model. The optimization problem ming EE[""(0" (Ø")) ] can be solved locally by each user n to find a customized model for its task, where L is the loss function corresponding
user n's task, and the expectation is taken over its labeled dataset D However, the key goal of FMTL is to enable users to train their local models under the orchestration of the server, thereby capturing common agnostic information. This coordination can improve the generalization capability of the users' models, leading to better performance compared to training solely on their local datasets. Unlike existing FMTL algorithms that employ a parameter-sharing mechanism (e.g.. shared encoder) to capture common information across tasks, we use joint representation learning among the users' models.
For each user [.V. we decouple its local model 0" into representation model w" and task-specific prediction head d", such that 0= = { (w", 6²). We assume that the output features of all representation models have the same dimension d We also consider that users have ac-
'This requirement does not preclude model heterogeneity across users and can he relaxed by appending lightweight, learnable projection head that maps each user's latent representation to common dimension.
3

---

Published as a conference paper at ICLR 2026
cess to shared unlabeled public dataset. This public dataset may be uni-modal depending on the users' tasks. Moreover, it can be easily obtained by collecting data from rele- vant domains (Huang et al.] (022), using publicly available datasets (Yu et al.] (023), generating synthetic data samples (Huang et al. 2024). Let 'D denote the shared public dataset. In each com- munication round of FMTL., each user first trains its local model 0 on its local dataset 'D' for E local epochs. Then, users extract representation vectors corresponding to samples from the public dataset and transmit them the server. In particular, for each data sample æi D. user feeds it into w" to obtain a normalized representation vector denoted by 2 E Rd. Considering a batch of size B from the public dataset, each user n transmits a representation matrix Z" E IR Bxd to the server, where each row of Z" corresponds to sample in the batch. As illustrated in Figure I based on our proposed Muscle loss function (to be presented in Section the server obtains an aggregated matrix and a weighting coefficient vector for each user by using the representation matrices received from the other users. The complete proposed FMTL algorithm is presented in Section 5
modal,
THE CONTRASTIVE LEARNING FRAMEWORK
In this section, we first explain why pairwise alignment approaches may not fully capture the depen- dencies among the representations of more than two models. Next, we propose Muscle as new CL loss and show that it can systematically capture such dependencies. Finally, we theoretically demon- strate the relationship between the Muscle loss and the MI among the models' representations.
4.1 PRELIMINARIES
InfoNCE is a popular CL loss function, and its variants have been widely used in the literature to improve the quality of learned representations between two views of the same data (Chen et al.) 2020 or across two modalities (Radford al. 202I). The idea behind the InfoNCE loss function is to learn a representation space that contrasts samples from two different distributions. Given and {2.T 1 as the representations corresponding to a batch of B samples, =² and z," are considered positive pairs if i = j (i.e., the representation vectors originate from the same data sample), and negative pairs if i + j. Without loss of generality, we treat z as the anchor. That is, we identify its positive and negative counterparts and define the loss function accordingly. The InfoNCE loss is given as follows (Oord et al. 2019):
(= i==1, )
exp(22 2 2 /TT,m)
mn InfONCE (2,) = log
(I)
ΣjE[B exp( 22," 23"²) /T,m)
where Tn,m is temperature parameter that moderates the effect of similarity. To extend InfoNCE to more than two distributions (e.g., representations obtained from multiple modalities, views, or mod- els), pairwise alignment approaches apply the InfoNCE loss to contrast samples from each pair of distributions (Tian et al. 22200] Yang et al., [222]] Wang & Sun 2022). In particular, given x, as the anchor, we have Pairwise = 2m [x) LaisNnE )² However, since Pairiise (2,") accounts
only for pairwise dependencies, it cannot effectively capture dependencies among all distributions
or learn a representation space that contrasts their samples jointly.
4.2 OUR PROPOSED MUSCLE LOss FUNCTION
We introduce the Muscle loss, which goes beyond pairwise alignment by jointly considering repre- sentations from multiple models. To this end, we focus on the N-tuple of representation vectors in learning representations across N models. In particular, all representation vectors {2(" J.V mm mztn considered positives for the anchor representation vector 2 Also, any combination of represen- tation vectors from all models in [N]] \ {ml} (i.c., one from each model except model n) is considered negative for the anchor vector 2," if at least one of the vectors in the combination corresponds to data sample j F i. We define the Muscle loss function as follows:
f(z {2)TT) NN 1m=1, mgnin
Mussele (2,)) log
Ejern f(2 {2mt I.N
m= m
where f is a function that assigns high values to positive N N-tuples and low values to negative N- tuples. For the sake of brevity, we define J" = {j- (Ji..... ........)) jm E [B], m E [N] \ {m})-
In Appendix A we show that the optimal value of [(z {=] jM m+t))) proportional to the
4

---

Published as a conference paper at ICLR 2026
probability density ratio p(zī" {zj, h Im=1, mm= )/222))))))))))))))))))))))))) )p( 12 JM 1m=1, m,tt)) and that, the Muscle loss function can be expressed as follows:
Mussle (2,") = - log ((......)) exp(z" EmE(NTAAn 21th //m//m,m)) (N)
(3)
Ejern Oij exp(2"" EmEE[NRiMM 2 /T.m) (
where 7/.m is a temperature parameter that moderates the effect of similarity between the rep- resentations of models n and m, given the correlations among representations of all N mod- els. Additionally, we have jj- = exp(-2 EmE(NN) Emve[M (N() Tmm."."" Jm zmd ), where mm.. 1/T - 1/T As shown in equation (3), Oj. where j J serves as a weighting coefficient in the Muscle loss function. These coefficients depend on the similarity among the rep- resentations of non-anchor models. In Appendix [ we show that the proposed Muscle loss captures dependencies among representations more effectively than pairwise alignment, owing the use of these weighting coefficients.
Remark In Appendix C we show that Tu,m > -(N-1) which implies that mm in Cjj is always
positive. Consequently, greater dissimilarity among the representation vectors { 3 mn sults in a larger weighting coefficient 0z. This leads the Muscle loss function to place greater emphasis on increasing the dissimilarity between the anchor vector x; and negatives that exhibit higher dissimilarity among themselves. In Appendix D we provide simple example to illustrate the impact of the weighting coefficients oj in the Muscle loss function.
The following theorem establishes the relationship between proposed Muscle loss function and the MI among the models' representations.
Theorem Given MMaacte (2+) in equation the mutual information I(2;;; (2.T " m==, ment) lower-bounded as follows:
[(2;; M mm=l, (N 1) log Mu'Mssete (z?),
(4)
where E denotes the expectation over random batch of data samples.
Proof. See Appendix
minimizing the Muscle loss is equivalent to maximizing a lower bound on T(2P; {>]" 1 m//n Thus, the Muscle loss facilitates knowledge transfer from models m E [N]) {n} to model n by aligning the representations of model n with those of the other models. In Appendix
Remark 2 Given inequality
provide discussion on why the Muscle loss is more effective for facilitating knowledge transfer among models compared to the pairwise alignment from the MI perspective.
Our proposed Muscle loss function enables a systematic CL framework by incorporating weighting coefficients aj. In the next section, we demonstrate how this systematic CL framework can be leveraged to develop novel FMTL algorithm.
FEDMUSCLE METHODOLOGY
In this section, we propose FedMuscle, a novel FMTL algorithm designed to address both model and task heterogeneity users' models using the Muscle loss function defined in (3). As discussed in Section 3 given a batch of public data samples, each user n E [N] sends a representation matrix Z" to the server. For each user n, the server computes an aggregated matrix and weighting coefficient vector, which are then sent back to the user to compute the CL loss locally as follows:
users. To achieve this, FedMuscle aligns the representation spaces of
a...... exp (xF 57
C
Σ log tE[l
(5)
Ejern 0z exp 2 87)
(N) m where, based on 3). have aß exp(-} E Eme(N(M) Em [n.m) mm jm = and
8 Eme[N](( (x) 3m TT,m Note that, for each user the server can compute az and 8I for all
(N)
j E J" using the representation matrices nh received from the other users [N] \ (n).
5

---

Published as a conference paper at ICLR 2026
Algorithm Training Procedure of FedMuscle
1:
Input:
Public dataset D: local dataset D" x E [N] number of local epochs E; number of communication rounds R: number of CL epochs T: batch size B; number of sclected representation matrices M. Output: customized local model 0 for each user n E [V] on task.
3: Randomly initialize 0 fw"", Ø for each user E [.]]
4: for each communication round E [R] do
5: for each user E [N] in parallel do
6: {u=", + LocalUpdatte((" DD" FF)
7: for each CL epoch 1 E [2'] do
8:
for each batch of data samples from 2 do
9:
Obtain the representation matrix z" and send it to the server.
10:
Receive S and a from the server, where S and a" are computed based on the M repre- sentation matrices randomly selected by the server for user
11:
Update w" by minimizing the CL loss function (5).
12: end for
13: end for
14: end for
15:
end for
16: function LocalUpdate( D", E)
17: for cach local epoch e [B] do
18: Update 0" by minimizing the loss function L corresponding to user n's task on D".
19: end for
20:
Return [w", ""}
Let S and a denote the aggregated matrix and the weighting coefficient vector computed by the server for user n, respectively. The rows of 5 and the elements of a² correspond to 8) and az for all j J", respectively. Algorithm T summarizes the training procedure of FedMuscle. The users collaboratively train their local models over R communication rounds. In each communication round r E [R). each user n [N] performs local epochs to update its model O using its local dataset DD The users then perform 7 CL epochs to update their representation models by aligning representation spaces of their models. Specifically, in each communication round r E [R] and each CL epoch t e [T), each user sends the representation matrix Z", obtained from a batch of the public dataset D, to the server and receives the aggregated matrix S and the weighting coefficient vector ax from the server. The user then minimizes the CL loss function L in equation (
For each user n E [N], the communication cost of the FedMuscle algorithm in the uplink direction depends the size of the representation matrix Z". which is B x d. The communication cost in the downlink direction, however, depends on the sizes of S and a.. If the server uses the repre- sentation matrices from all the other users, i.e., [N] [n), to compute S and az" their dimensions are N x d and BN-1, respectively. Note that the large sizes of S and a² may incur high communication costs in the downlink direction, as more parameters need be transmitted from the server to user particularly when N is large. To reduce communication costs in FedMuscle, we randomly select only M of the available representation matrices from the N users to compute S" and a for user n in each communication round and CL epoch. Let A C [N] (n) denote the set of users whose representation matrices are selected to compute the elements of S and ax
for user n, where [W'| = M. Thus, we have J [j (jm (mm)m Jm E [B], m E EV* to specify the index tuples used computing the terms of the CL loss function La. in equation (5)
Consequently, in FedMuscle, we have 5 E R|a"|xdd and a E RRIJ"), where | = BM.
6 EXPERIMENTS
6.1
EXPERIMENTAL SETUP
Datasets and Tasks We conduct our experiments using datasets such as CIFAR- 10, CIFAR- 100
(Kiizhessyy) 2009), common objects in context (COCO) (Lin et al. 2014). Yahoo! An- swers Zhang et al. 2015). Pascal VOC (Everingham et al. 2015) and Flickr30K (Plummer et al, 2015). Inspired by the experimental setup in (Chen al. 2020), consider the following three CV tasks: image classification using CIFAR- 10 (IC10), image classification using CIFAR- 100 (IC100), and multi-label classification using COCO (MLC). We refer to this uni-modal setup Setup1,

---

Published as a conference paper at ICLR 2026
which enables fair comparison with existing baseline algorithms. To evaluate the performance of FedMuscle in a more challenging, multi-moda setting, we further expand Setup1 by adding two additional tasks: semantic segmentation using COCO (SS) and text classification using Yahoo! An- swers (TC). We refer to this multi-modal setup as Setup2. For the uni-modal Setup1, we consider that the samples in the public dataset are entirely drawn from one of the following datasets: CIFAR- 100, COCO, or Pascal VOC. For the multi-modal Setup2, we consider that the samples in the public dataset are entirely drawn from Flickr30K. In Setup2, users with CV tasks obtain representations from the images in the public dataset, while users with the TC task obtain representations from the image captions. The size of the shared public dataset set to 5000. For additional details on the number of training and test samples for each task, please refer to Appendix G
Implementation Details We use heterogeneous FMs as the users' local models. Specifically, our experiments involve heterogeneous ViT models (Dosovitskiy et al. [202]] pre-trained on ImageNet- 21K (Deng et al.] 2009), heterogeneous SegFormer models (Xie et al. (02T) pre-trained on ImageNet- IK and fine-tuned on ADE20K (hoou et al. 2017), and heterogencous BERT- based mod- els (Devlin et al. 2019 Sanh et al. 2019) pre-trained on a large corpus of English tex Details of the selected models and assigned tasks for users in Setup1 and Setup2 are provided in Appendix HH We set d (i.e.. the output dimension of the representation models) to 256. To reduce the com- putational cost for users, we fine-tune their pre-trained FMs using parameter-efficient fine-tuning
(PEFT) (Zhang et al. 2023b). Specifically, we use low-rank adaptation (LoRA) (Hu et al. 2022
with rank of 16. The models are trained using the Adam optimizer (Loshchilov & Hutter 2019)
with a learning rate of le-3. We set the batch size B to 32 and the value of M in Algorithm to
3. The temperature parameters (4) and (3) are set to 0.2 and 0.15, respectively. Unless otherwise stated, the number of local epochs E, communication rounds R. and CL epochs T' are set 1, 150,
and 1, respectively. We implement FedMuscle and the baseline algorithms using PyTorch (Paszke al. 2019) and run all experiments on NVIDIA Tesla V100 GPUs.
Baselines We compare the performance of our proposed algorithm, FedMusele, with Fe- dRCL (Seo et al., 2024) as a CL-based FL algorithm; SAGE (Liu et al.) 2025) and CoFED (Cao et al. 2023) as FSSL algorithms; FedDF (Lin et al.) 2020) as a KD-based FL algorithm; and Fed- 2,2.. a prodel-agnostic FL algorithm. Additionally, we compare FedMuscle with SimCLR (Chen al. 2020), pseudo-labeling (Cascante-Boni) BoniiloBnnlaaaa al. 202I), and local train- ing, where each user's 's model is trained individually. Specifically, SimCLR and pseudo-labeling utilize the public dataset for self-supervised and semi-super learning, respectively, while local training relies solely on each user's 's labeled local dataset.
HeNN (Makhija et al. 2022).
Performance Metrics For the IC10, IC100, and TC tasks, accuracy on the test samples is used as the performance metric. For the MLC and SS tasks, micro-average FI-score (micro-FI) (Wu & Zhou 2017) and mean intersection over union (mloU) (Everingham et al.) 2010) are used as the performance metrics, respectively. All the evaluation metrics are expressed as percentages (%), and higher values indicate better performance. Furthermore, to provide an overall evaluation metric for an algorithm's 's performance, the average per-user performance improvement relative to the local training baseline, denoted as A, can be employed (Lu et al. 2024). Specifically, we have A = N Ene[N] (MR ME MM Miear where M and M represent the performance metrics of the considered algorithm and the local training baseline, respectively, for the task of user [M].
6.2 BENCHMARK EXPERIMENTS
Comparison with Baselines Table I presents the performance of FedMusele compared to the baseline algorithms in the uni-modal Setup1. Based on the results in Table our observations are as follows: (1) The proposed algorithm, FedMuscle, consistently achieves better overall performance in terms of A compared to the baseline algorithms. (2) Using public datasets with detailed images, such as Pascal VOC and COCO, leads to improved performance for FedMuscle
(3) Even when the public dataset is derived from CIFAR-100, which contains less detailed images, FedMuscle still enhances users' performance on their respective tasks. The performance of FedMuscle using a synthetic public dataset is presented in Appendix D These results confirm that FedMuscle remains
3 Prre-trained FMs are downloaded from Hugging Face at awwww. face. co/mooeee////eoooooooe In Appendix || we present a training curve illustrating how decrease in the Muscle loss corresponds to an increase in A communication rounds.
7

---

Published as a conference paper at ICLR 2026
Table 1: Performance of FedMuscle compared to the considered baseline algorithms in Setupl.
Algurithm
Public Dumase
User
Medel
Tosk
Eval. Metric
FedMuscle (Ouss
SAGE
PelifeNN
CaFFD 2023
FedDF 2012
FaIRCI 2024
SimCL R
Pscudo labeling 2021
Local
Training
VIT Base
MLC
Ft
46.35
41
41.27
47.47.8
42.43.0.1
44.77.64
40.80.44
45.93
42.17.60
Paccll VOC
Wi.lmall
33
4830s
43
40
43.7
44.10
42.57mm
42533-c
VIE-Ros
IC100
Acamcy
24
23.0
Wi--mmlll
IC100
Acuunyy
19
35
23.70s
24
27.23
22.77a1
24
Tim
Accuncy
66
033331.0
43.00
03.00.0
4080
1.
43.77113
43:77
-26.70
+2.07
+3.97
-1.84
Bose
micro
41,97 kM
4200. x
50.81 kM
43.17.0
41.77.ex
42.50 cm
47.25
42.17.00
Smmll
m1
0O00
Wir.agge
5000
5530 km
42
45.3
50335
Bise
ICT00
Accorayy
37
25
km
23331
24:331 km
24100.0
25
25.87
23:20
24.77
WT--oull
IC100
Actunyy
25
34734 mm
25
VIT-Tay
ICT0
Accinsy
63
43.03
4500751
40090s PM
03
44832
47.37530
44:13.19
43.77mm
193
+28.63
-2.31
-2.57
+9.85
-025
+2.07
+4.49
+4.85
0.00
ViT-Bise
MLC
micr-FF
4n
42. mm
41.49a1M
43.73 am
42.40.mm
41.77
43
42.17
CHAMR-10
VIT-Smull
MTC
MLC
micn-Fi
46.00.
15
44.97.mm
45.0. x 47.10s w4
43.03..M
44.07.m0
43.90
km
43.87cc
VIT Bise
1C100
Accuracy
03.45
25.07
am
25231 in 27.77 km
25.27 om
19:23 14
2693 km 24.77 oc
WT Small
KCIM
Accuracy
29
25004 km 2530 xM 27227 us
2013500
27.25
22331
1s 2880 us 20.00
Vii-Tig
ICTM
Accuracy
38.376
41336-0 xm
24.45 m 43.30l am
43.20.00
50003
1 45.77
43.77-m
& (15)
+16.88
+0.1
2285
+5.99
-1.42
42.37
-4.94
+6.20
0.00
Computaiion Cost. (TeraFLOPS)
525
1171
34
586
322
78
549
152
29
effective regardless of the chosen public dataset. However, it can further improve users' overall performance when the public dataset contains feature-rich samples.
Performance of FedMusele in Multi-Modal Setup2 cle across broader range of tasks, including both
Table 2 the performance of FedMus-
CV and NLP tasks. The results in Table 2
demonstrate that even with the addition of two users performing the SS task and two users performing the TC task, FedMuscle still improves the overall performance of users their respective tasks in terms of A. Despite the differences in modalities, the of CL. in FedMuscle aligns representations from various modalities into a shared representa- tion space, enabling each user to capture common task-agnostic information from the other users.
Table 2: Performance of FedMuscle in Setup2.
User
Model
Task
Foal Matric
FedMascle (Ohrs)
Algorittas
Local Training
ViT-Base
MI.C
micro-FI
47.80geaa
42. 17
VIT-Small
MI.C
mirr-Fi
51 $1.05geas
43.67gm0
VIT Large
MLC
micro TT
48,00 400
VIT Base
1C100
Accunecy
16.15.46
24.77100
Vil-Small
IC10
Ancwasy
28.70-423
24.70-mm
Ti
IC10
Aucunsy
01,00-
43.77.00
SegForme-H8i R0
88
mbal)
11
31.75gmm
SegFormes-B
88
mltt)
33,40,19
32.45gna
BERT Base
TC
Accuracy
45.95.08
Impact of Muscle Loss in FedMuscle The re-
10
DivilBERTT-Baa
TC
Accurscy
54,10
56 56.03
(99
+14.39
0,00
sults in Figure 2 show that, compared to the Soumian-based contrastive loss recently proposed by Cicchetti et al. (2025), the Muscle loss func- tion improves A by 11.2%, 28.4%, and 11.1% when Pascal VOC, COCO, and CIFAR-100 are used as the public datasets, respectively These improvements stem from the fact that the Muscle loss has solid theoretical justifications and can bet- ter capture dependencies among the represen- tations obtained from multiple models. In ad- dition, the computation of the Musele loss is modular and partially offloaded to the server. By contrast, the Gramian- based loss requires computing the determinant of Gramian matri- ces, which imposes (M + +1) times higher com- putational cost on the users. In Appendix D we provide more details on the features offered by the Muscle loss compared to the Gramian- based loss. Furthermore, a comparison between the A values in Figure 2 and Table D reveals that even pairwise alignment in FedMuscle yields better performance than baseline algorithms. This highlights the efffectiveness of the FedMuscle algorithm under model and task heterogeneity.
30
Pairwiso Alignment
-267
Ceaziinn-bssed Loss
-23.3
Muscle Loss (Ouss
20
A (9%)
+15.3
-15.2
+20.
Pascal VOC
COCO
CIFAR-100
Figure 2: Performance of FedMuscle in Sctup1 using the proposed Muscle loss function, compared to the Gramian-based contrastive loss and pairwise alignment.
Integration of Muscle Loss into a Multi-Modal FL Algorithm To evaluate the effectiveness of our proposed CL framework, we integrate it into an MFL algorithm. We also aim to demon-
strate its effectiveness in setting with a larger number of users, where the models are trained from scratch under a non-IID partition of local data. To this end, we adopt the CreamFL (Yu et al. 2023)
Detailedd results be found in Appendix K
8

---

Published as a conference paper at ICLR 2026
Table 3: Global model performance the image-text retrieval task in CreamFL
Yu et al.
2023 setup.
Eval.
12:
12
12.
12
Local Traningg
48.36
79,06
89924
37.46
74,05
36.59
0,00
1K Text Images
CreamFi
49.20
80.76
89.92
57.83
74.82
86.53
+0.93
Cream Muscle Loss (Ours)
49.32
9002
38.11
74.39
86.37
+1.17
Local Training
34.78
52.44
90
13.72
43.28
57.47
0.00
5K Test Images
24.45
53.48
66.92
17.96
43.74
58.35
40.35
(Ours)
18.20
+1.94
30
23.9--19.22-00
30
-27.1
2.5
20
4 (16)
20
132.8
10.2
15
R=150, 1=1, 17- (100, 500, 1000, 3000, 5000)
15.7
25), T-3,
10
R [1, 10, 25). 120-500
10
25
0
50
100
150
200
250
500
350
Number of Communication Kounds
Communication Cost
Figure Impact number of local epochs E, communication rounds H. and CL epochs I' on Fed- number of communication rounds It on the communi- Musele performance. COCO used as the public cation cost of FedMuscle. COCO is used as the public dataset.
Figure Impact of the public dataset size [D|
dataset.
setup, which consists of 10 uni-modal image users with the CIFAR- 100 dataset, 10 uni-modal text users with the AG News dataset (Zhang et al. 2015), and 15 multi-modal users with the Flickr30K dataset. Dirichlet distribution with a = 0.1 is used for non-IlI data partition. CreamFL aims to learn global model at the server using representations derived from public dataset shared among the users and the server. COCO used as the public dataset. We integrate the Muscle loss into the CreamFL setup by replacing its local contrastive regularization and global-local contrastive aggre- gation methods with our proposed Muscle loss. More details are provided in Appendix M Table 3 shows that integrating the Muscle loss into the CreamFL setup improves the global model's perfor- mance in terms of image-to-ter recall (i2t_R) and text-to-image recall (12i_R) metrics.
6.3 ABLATION STUDIES
Effeet of the Number of Communication Rounds, Local Epochs, and CL Epochs Figure illustrates the impact the number communication rounds R, local epochs E, and CL epochs T on the performance of FedMuscle. To ensure fair comparison, set E R 150 for all cases. Note that higher values E and R correspond to increased computation and communication costs, respectively. The results in Figure 3 show that greater number of communication rounds leads to improved performance for FedMuscle Additionally, when the number of communication rounds low, FedMuscle can still achieve improved performance by increasing the number of CL epochs.
Effect of the Public Dataset Size Figure illustrates the impact of the public dataset size |D] on FedMuscle's 's communication cost and its performance. To provide comprehensive comparison, we have also included two cases from Figure in Figure ( T 3 and T 5). A larger [D] leads higher communication costs as more representation vectors are transmitted between users and the server. However, as shown in Figure 4 FedMuscle is flexible enough to achieve high performance in terms of A while maintaining low communication cost by appropriately adjusting the number of communication rounds Rt the number of CI. epochs T, and the public dataset size (D].
Effeet Non-lII) Data Distribution and Number Users We consider a setup where
users perform one of the following four tasks: IC100, IC10, image classification using Food- 101 (Bossard et al., 2014) (ICF), and image classification using Caltech-101 (Fei-Fei et al.) 2004)) (ICC). Specifically, users 10-12 perform ICC. For those performing the same task, the local data samples are divided across them using Dirichlet distribution with a 0.1. All users have 100 training samples and 2000 test samples. In Appendix O
1-3 perform IC100, users 4 perform IC10, users 7-9 perform ICF, and
we have illustrated the non-lID data partitioning across users.
Detailed results be found in Appendix

---

Published as a conference paper at ICLR 2026
Table 4: Performance of FedMuscle in the non-llI setting, where assigned to the same task have hetero- geneous local data distributions. The public dataset derived from Pascal VOC.
Algoritim
(9)
11
Local Training
45.2
54.2
3088
94.3
94.8
79.2
40.0
36.2
13.6
76.7
758
62.6
0.00
FallMsekk (Ours)
58,4
96.8
49,9
83.4
+15.14
58.0
554
94.2
93.7
50.3
42.5
824
790
410.18
12 57.7 54.
414
98.1
928 86.0
47.8
40.3
25.6
81.8
780
71.0
+17.80
Furthermore, we consider heterogeneous models: users 1, 4, 7, and use ViT-Base, users 5, 8, and 11 use ViT-Small, and users 3, 6, 9, and 12 use ViT-Tiny.
Table presents the performance of FedMusele under the aforementioned non-Ill> data partitioning across users for three scenarios, where the number of users is set to 4, 8, and 12. The results in Table 4 show that FedMuscle improves the performance of each user on its respective task. Thus, the overall performance, measured in terms of A, is improved by FedMuscle regardless of the data, task, and model heterogeneity among users, and the number of users in the system.
Effeet of Number of Selected Representation Matrices We present an ablation study on the number of selected representation matrices M in Table [ The results are obtained under Setup1, using Pascal VOC as the public dataset. As shown the table, increasing M improves overall performance in terms of A, but also increases the communication cost per round. Notably, the performance gains from increasing M beyond 3 (e.g., from to or 5) are marginal. Therefore, adopt M = 3 with random between performance and communication overhead. We also find that random selection provides sufficient diversity across rounds to facilitate knowledge transfer among all tasks.
selection in our experiments, this offers balanced trade-off
Table 5: Performance of FedMuscle under different number of selected representation matrices M.
3
+17.90 +24.17
+25.70
+27.53
+27.74
Communication Cost per (GB)
361.565
Effeet of Batch Size We investigate the impact of batch size B on FedMuscle's performance in Setupl when COCO used as the public dataset. Note that by increasing the batch size, the approximation used for deriving inequality (23) in Ap- pendix becomes more accurate. Thus, based on The-
30
orem minimizing the Muscle loss provides tighter
lower bound on maximizing the MI between the models' outputs. The results in Figure aasso show that by increas- ing the batch size, the performance of FedMuscle can be improved in terms of the overall performance metrie A. However, based on the provided discussion in Section know that larger batch sizes can incur higher com- munication costs in the downlink direction. To achieve a balanced trade-off between performance and commu- nication overhead, we adopt B = 32 in our experiments. Additional ablation studies on the output dimension of the representation models and temperature parameters (4) and Tr,m are presented Appendix P
20
10
0
8
16
64
128
Batch Size (B)
Figure 5: Impact batch size on FedMus- cle's 's performance in Sctup1. COCO is used the public dataset.
7 CONCLUSION
In this work, we proposed a new CL loss function, called Muscle, which systematically and effec- tively captures dependencies among representations obtained from multiple models. These depen- dencies are modeled through theoretically grounded weighting coefficients in the Muscle loss. Fur- thermore, we showed that minimizing the Muscle loss equivalent to maximizing lower bound on the MI between the models' representations. Building on the Muscle loss, we designed Fed- Muscle, a novel and effective FMTL algorithm that enables collaborative training of heterogeneous models in the presence of task heterogeneity across users. To this end, FedMuscle aligns the rep- resentation spaces of users' models using shared public dataset to capture common, task-agnostic information. Through extensive experiments, users' performance on their respective tasks compared to wrate-ofthe art baseline algorithms under both model and task heterogeneity. A discussion of possible future work is provided in Appendix Q
showed that FedMuscle consistently improves
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILIT STATEMENT
For theoretical results, provide the complete proofs along with their underlying assumptions in and S All data used in this study are publicly available. The code for running our algorithm also provided in the supplementary material. Furthermore, we provide all experimental details in Section 6 and Appendices G H and M We have cited all original sources for the models, baselines, and datasets used this work. Please refer to Section for further details.
Section 4 and Appendices A B
REFERENCES
Sawsan AbdulRahman, Hanine Tout, Hakima Duld-Slimane, Azzam Mourad, Chamseddine Talhi, and Mohsen Guizani. A survey on federated learning: The journey from centralized to distributed on-site learning and beyond. IEEE Internet Things J. 8(7):5476-5497, Apr. 2021.
Samiul Alam, Luyang Liu, Ming Yan, and Mi Zhang. FedRolex: Modee-hhetssggeneuuu federated learning with rolling sub-model extraction. In Proc. Conf. Neural Inf. Process. Syst. (NeurIPS), New Orleans, LA, Dec. 2022.
Rishi Bommasani, Drew Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.072 07258, Jul. 2022.
Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool. Food-101 101-Minn discriminative components with random forests. In Proc. Eur. Conf. Comput. Vis. (ECCV), Zurich, Switzerland, Sept. 2014.
Ruisi Cai, Xiaohan Chen, Shiwei Liu, Jayanth Srinivasa, Myungjin Lee, Ramana Kompella, and Zhangyang Wang. Many-tash task federated learning: new problem setting and simple bascline. Proc. JEEÉ/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Seattle WA, Jun. 2024.
Xingjian Cao, Zonghang Li, Gang Sun, Hongfang Yu, and Mohsen Guizani. Cross-Silo heterogeneous model federated multitask learning. Knowledge- Based Systems, 265:1 10347, Apr. 2023.
Paola Cascante-Bonin Fuwen Tan, Yanjun Qi, and Vicente Ordonez. Curriculum labeling: Revisiting pscudo- labeling for cemi-supervised learning. In Proc. AAAI Conf. Artif. Intell., Feb. 2021.
Liwei Che, Jiaqi Wang. Xinyue Liu, and Fenglong Ma. Leveraging foundation models for multi-modal feder- ated learning with incomplete modality. In Joint Eur. Conf. Mach. Learn, and Knowl. Discovery Databases (ECMI. PKDD), Vilnius, Lithuania, Sept. 2024.
Ting Chen, Simon Komblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proc. Int. Conf. Mach. Learn. (ICML), Jul. 2020.
Xinyue Chen, Yazhou Ren, Jie Xu, Fangfei Lin, Xiaorong Pu, and Yang Yang. Bridging gaps: Federated multi- view clustering in heterogeneous hybrid views. In Proc. Advances Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2024.
Yiqiang Chen, Teng Zhang, Xinlong Jiang, Qian Chen, Chenlong Gao, and Wuliang Huang. Fedbone: Towards large-scale federated multi-task learning. Proc. Inf. Workshop on Trustworthy Federated Learning in Conjunction with IJCA (FL-IJCAI). Macau, Aug. 2023.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated re- current neural networks on sequence modeling, arXiv preprint arXiv: 1472.3555, Dec. 2014.
Giordano Ciechetti, Elconora Grassucei, Luigi Sigillo, and Danilo Comminiello. Gramian multimodal repre- sentation learning and alignment. Proc. Int. Conf. Learn. Represent. (ICLR), Singapore, Apr. 2025.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In Proc. Int. Conf. Mach. Learn. (ICMI.), Jul. 2021.
MohammadReza Davari, Stefan Horoi, Amine Natik, Guillaume Lajoie, Guy Wolf, and Eugene Belilovsky. Reliability of CKA as a similarity measure in deep learning. In Proc. Int. Conf. Learn. Represent. (ICLR), Kigali, Rwanda, May 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. Proc. IEEF/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Miami, Fl., Jun. 2009.

---

Published as a conference paper at ICLR 2026
Jacob Devlin. Ming-War Bing-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. Proc. Conf. North Amer. chapter Assoc. Comput. Linguistics: Human Lang, Technologies (NAACL-HLT), Minneapolis, MN, Jun 2019.
Enmao Diao, Jie Ding, and Vahid Tarokh. HeteroFL: Computation and communication efficient federated learning for heterogeneous clients. In Proc. Int. Conf. Learn. Represent. (ICLR), May 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikow, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakoh Uszkoreit, and Neil Houlsby, An image is worth 16x 16 words: Transformers for image recognition at seale. Proc. Int. Conf. Learn. Represent. (ICLR), May 2021.
Hao Du, Shang Liu, Lele Zheng, Yang Cao, Atsuyoshi Nakamura, and Lei Chen. Privacy in fine-tuning large language models: Attacks, defenses, and future directions. arXiv preprint arXiv:2412. 16504, Apr. 2025.
Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman. The pascal visual object classes (voc) challenge. Int. J. Comput. Vis., 88(2):303-338, Jun. 2010.
Mark Everingham, SM Ali Eslami, Luc Van Gool, Christopher KI Williams. John Winn, and Andrew Zisser- man. The pascal visual object classes challenge: A retrospective. Int. J. Comput. Vis., 111(1);98-138) 136,, Jan. 2015.
Li Fei-Fei, Rob Fergus, and Pietro Perona. Learning generative visual models from few training examples: An incremental bayesian approach tested on 101 object categories. In Conf. Comput. Vis. Pattern Recognit. workshop, Washington, DC, Jun. 2004.
Rohit Girdhar, Alaaeldin El-Nouby, Zhuang Liu, Mannat Singh, Kalyan Vasudev Alwala, Armand Joulin, and Ishan Misra. ImageBind: One embedding space bind them all. In Proc. IEEE/CVF Conf. Comput. Vis. Puttern Recognit. (CVPR), Vancouver, Canada, Jun. 2023.
Chaoyang He, Emir Ceyani, Keshav Balasubramanian, Murali Annavaram, and Salman Avestimehr. Spread- GNN: Decentralized multi-task federated learning for graph neural networks molecular data. In Proc. AAAI Conf. Artif. Intell., Vancouver, Canada, Feb, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Las Vegas, NV, Jun. 2016
Kaiming Hc, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proc. IEEE Conf. Comput. Vis. Pattern Recognit., Jun. 2020.
Junyuan Hong, Haotao Wang, Zhangyang Wang, and Jiayu Zhou. Efficient Split-Mix federated learning for on on-demand and in-situ customization. In Proc. Int. Conf. Learn. Represent. (ICLR), Apr. 2022.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen, LoRA: Low-rank adaptation of large language models. In Proc. Int. Conf. Learn. Represent. (ICLR), Apr. 2022.
Chun- Yin Huang, Kartik Srinivas, Xin Zhang, and Xiaoxiao Li. Overcoming data and model heterogeneities in decentralized federated learning via synthetic anchors. In Proc. Int. Conf. Mach. Learn. (ICML), Vienna, Austria, Jul. 2024.
Wenke Huang, Mang Ye, and Bo Du. Learn from others and be yourself in heterogencous federated learning. In Proc. IEEE/CVF Conf. Comput. Vis. Patten Recognit. (CVPR), New Orleans, LA, Jun. 2022.
Chaouki Ben Issaid, Praneeth Vepakomma, and Mehdi Bennis. Tackling feature and sample heterogeneity in decentralized multi-task learning: sheaf-theoretic approach. arXiv preprint arXiv:2502.00 145, Jun. 2025.
Wonyong Jeong, Jachong Yoon, Eunho Yang. and Sung Ju Hwang. Federated semi- seni-supervised learning with inter-client consistency disjoint learning. In Proc. Int Conf. Learn. Represent. (ICLR), May 2021.
Yongzhe Jia, Xuyun Zhang, Amin Beheshti, and Wanchun Dou. FedLPS; Heterogeneous federated learning for multiple tasks with local parameter sharing. In Proc. AAAI Conf. Artif. Intell., Vancouver, Canada, Feb. 2024.
Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Dis- tributed machine learning for on-device intelligence. arXiv preprint 1rx102111. Oct. 2016.
Panagiotis Koromilas, Efthymios Georgiou, Giorgos Bouritsas, Theodoros Giannakopoulos, Mihalis A Nico- laou, and Yannis Panagakis. A principled framework for multi-view contrastive learning. arXiv preprint 7077095507.0697, Jul. 2025.
12

---

Published as a conference paper at ICLR 2026
Alex Krizhevsky, Learning multiple layers of features from tiny images. Tech Report, Apr. 2009.
Daliang Li and Junpu Wang- FedMD: Heterogenous federated learning via model distillation. In Int. Work- shop on Federated Learning for User Privacy and Data Confidentiality in Conjunction with NeurIPS (FL- NeurlPS), Vancouver, Canada, Dec. 2019.
Qinbin Li, Bingsheng Hc, and Dawn Song. Model-contrastive federated learning. In Proc. JEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Vancouver, Canada, Jun. 2021a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization heterogeneous networks. In Proc. Adaptive & Multitask Learning Workshop, Long Beach, CA, Jun. 2019.
Tian Li, Shengyuan Hu, Ahmad Beirami. and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In Proc. Int. Conf. Mach. Learn. (ICML), Jul. 2021b.
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In Proc. Advances Neural Inf. Process. Syst. (NeurIPS), Dec. 2020.
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and Lawrence Zitnick. Microsoft COCO: Common objects context. Proc. Fur. Conf. Comput. Vis. (ECCV), Zurich, Switzerland, Sept. 2014.
Yijie Liu, Xinyi Shang, Yiqun Zhang, Yang Lu, Chen Gong, Jing-Hao Xue, and Hanzi Wang. Mind the gap: Confidenee discrepancy guide federated semi-superrised lcarning across pseudo-mismatch. In Proc. IEEE/CVF Conf. Comput. Vis. Puttern Recognit. (CVPR), Nashville, TN, Jun. 2025.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proc. Int. Conf. Learn. Represent. (ICLR), New Orleans, LA, May 2019.
Christos Louizos, Matthias Reisser, and Denis Korzhenkov. A mutual information perspective on federated contrastive learning. In Proc. Int. Conf. Learn. Represent. (ICLR), Vienna Austria, May 2024.
Yuxiang La, Suizhi Huang, Yuwen Yang, Shalayiding Sirejiding, Yue Ding, and Hongtao Lu. FedHCA?: To- wards hetero-client federated multi-task learning. In Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Seattle WA, Jun. 2024.
Disha Makhija, Xing Han, Nhat Ho, and Joydeep Ghosh. Architecture agnostic federated learning for neural networks. In Proc. Int. Conf. Mach. Learn. (ICML), Baltimore MD, Jul. 2022.
Kevis-Kokiin Maninis, Ilija Radosavovic, and lasonas Kokkinos. Attentive single-tasking of multiple tasks. In Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR). Long Beach, CA. Jun. 2019.
Othmane Marfoq. Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under a mixture of distributions. Proc. Advances in Neural Inf. Prucess. Syst. (NeurIPS), Dec. 2021.
Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invariant representations. In Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Jun. 2020.
Jishnu Mukhoti, Tsung- Yu Lin, Omid Poursaeed, Rui Wang. Ashish Shah, Philip HS Torr, and Ser-Nam Lim. Open vocabulary semantic segmentation with patch aligned contrastive learning. In Proc. IREE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Vancouver, Canada, Jun. 2023.
Ryumei Nakada, Halil Ibrahim Gulluk, Zhun Deng, Wenlong Ji, James Zou, and Linjun Zhang. Understanding multimodal contrastive learning and incorporating unpaired data. In Proc. Int. Conf. on Artif. Intell, and Stat. (AISTATS), Valnniia, Spain, Apr. 2021.
Aaron den Oord, Yazhe Li. and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748. Jan. 2019.
Sangjoon Park, Gwanghyun Kim, Jeongsol Kim, Boah Kim, and Jong Chul Ye. Federated split task-agnostic vision transformer for COVID-19 CXR diagnosis. In Proc. Advonces in Neural Inf. Process. Syst. (NeurIPS),
Dec. 2021.
Sejun Park, Kihun Hong, and Ganguk Hwang. A kernel perspective on distillation-basel collaborative learning. In Proc. Advances in Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2024.
13

---

Published as a conference paper at ICLR 2026
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- perforerrrmacc deep learning library. In Proc. Advances in Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2019.
Bryan A Plummer, Liwei Wang. Chris M Cervantes, Juan C Caicedo, Julia Hockenmaier, and Svetlana Lazeb- nik. Flickr30k entities: Collecting region- to- to-phrase correspondences for richer mmage-to-sentene models. In Proc of IEEE Int. Conf. Comput. Vis. (ICCV), Santiago, Chile, Dec. 2015.
Fan Qi and Shuai Li. Adaptive hyper-graph aggregation for aggnstcc federated learning. In Proc. IEEE/CVF Conf. Comput. Vis. Puttern Recognit. (CVPR), Seattle WA, Jun. 2024.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sas- try, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In Proc. Inr. Conf. Mach. Learn. (ICML., Jul. 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In Proc. IEEE/CVF Conf. Compua. Vis. Pattern Recognit. (CVPR), New Orleans, LA, Jun. 2022.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. In Proc. Int. Workshop Energy Efficient Mach. Learn. and Cogn. Comput. (EMC2) Co-located with NeurlPS, Vancouver, Canada, Dec. 2019.
Seonguk Seo, Jinkyu Kim, Geeho Kim, and Bohyung Han. Relaxed contrastive learning for federated learning. In Proc. IEEE/CVF Conf. Comput. Vis. Puttern Recognit. (CVPR), Seattle WA, Jun. 2024.
Mehdi Setayesh, Xiaoxiao Li, and Vincent WS Wong. PerFedMask: Personalized federated learning with optimized masking vectors. In Proc. Int. Conf. Learn. Represent. (ICLR), Kigali, Rwanda, May 2023.
Wenqi Shao, Chonghe Jiang, Kaipeng Zhang, Yu Qiao, Ping Luo, et al. Foundation model is efficient mul- timodal multitask model selector. In Proc. Advances in Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2024.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi pulti-task learning. In Proc. Annu. Conf. Neural Inf. Process. Syst. (NIPS), Long Beach, CA, Dec. 2017.
Emese Sükei, Elisabeth Rumetshofer, Niklas Schmidinger, Andreas Mayr, Ursula Schmidt- Erffurth, Günter Klambauer, and Hrvoje Bogunović. Multi-modal representation learning in retinal imaging using self- supervised learning for enhanced clinical predictions. Scientific Reports, 14(1):26802, Nov. 2024.
Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with Moreau envelopes. In
Proc. Advances Neural Inf. Process. Syst. (NeurIPS), Dec. 2020.
Yue Tan, Guodong Long, Jie Ma, Liu, Tianyi Zhou, and Jing Jiang. Federated learning from pre-trained models: A contrastive learning approach. In Proc. Advances in Neural Inf. Process. Syst. (NeurIPS), New Orleans, LAA, Nov. 2022.
MTCAJ Thomas and A Thomas Joy, Elements of information theory. Iiley-Interscience, 2006.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding. In Proc. Eur. Conf. Comput. Vis. (ECCV), Aug. 2020.
Feng Wang and Huaping Liu. Understanding the behaviour contrastive loss. In Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Vancouver, Canada, Jun. 2021.
Jiaqi Wang, Qi Li, Lingjuan Lyu, and Fenglong Ma. pFedClub: Controllable heterogeneous model aggregation for personalized federated learning. In Proc. Advances in Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2024.
Zifeng Wang and Jimeng Sun. Transtab: Learning transferable tabular transformers across tables. In Proc. of Conf. Neural Inf. Process. Syst. (NeurIPS), New Orleans, LA, Dec. 2022.
Yipan Wei, Yuchen Zou, Yapeng Li, and Bo Du. Towards unified modeling in federated multi-task learning via subspace decoupling, arXiv preprint ar arXiv:2505.2411 May 2025.
Xi-Zhu Wu and Hui-Hua Zhou. A unified view of multi-lahel performance measures. In Proc. Int. Conf. Mach. Learn (ICML), Sydney, Australia, Aug. 2017.
14

---

Published as a conference paper at ICLR 2026
Enze Xie, Wenhai Wang. Zhiding Yu, Anima Anandkumar, Jose M Alvarez. and Ping Luo. SegFormer: Simple and efficient design for semantic segmentation with transformers. Proc. Advances Neural Inf. Process. Syst. (NewrPPS), Dec. 2021.
Le Xue. Ning Yu. Shu Zhang, Attemis Panagopoulou. Junnan Li, Roberto Martin-Martin, Jiajun Wu, Caim- ing Xiong, Ran Xu, Juan Carlos Nichles, al. Ulip-2: Towards scalable multimodal pre-training for 3D understanding. In Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognir. (CVPR), Scattle WA, Jun. 2024a.
Yihao Xue, Siddharth Joshi, Dang Nguyen, and Baharan Mirzasoleiman. Understanding the robustness of multi-modal contrastive learning distribution shift. In Proc. Int. Conf. Learn. Represent, (ICLR), Vienna, Austria, May 2024b.
Chuanguang Yang, Zhulin An, Linhang Cai, and Yongjun Xu, Mutual contrastive learning for visual represen- tation learning. In Proc. AAAI Conf. Artif. Intell., Feb. 2022.
Qiying Yu. Yang Liu, Yimu Wang, Ke Xu, and Jingjing Liu. Multimodal federated learning via contrastive representation ensemble. In Proc. Int. Conf. Learn. Represent. (ICLR), Kigali, Rwanda, May 2023.
Jianqing Zhang, Yang Liu, Yang Hua, and Jian Cao. FedTGP: Trainable global prototypes with adaptive- margin-enhance contrastive learning for data and model heterogeneity in federated learning. In Proc. AAAI Conf. Artif. Intell., Vancouver, Canada, Feb. 2024.
Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang, Wenchao Xu, and Feijic Wu. Parameterized knowledge transfer for personalized federated learning. In Proc. Advances in Neural Inf. Process. Syst. (NeurIPS), Dec. 2021.
Lymin Zhang, Anyi Rao, and Mancesh Agrawala. Adding conditional control to text-to-image diffusion mod- els. In Proc. IEEF/CVF Int. Conf. Comput. Vis. (ICCV), Paris, France, Oct. 2023a.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks text classification. In Proc. Annu. Conf. Neural Inf. Process. Syst. (NIPS), Montréal Canada, Dec. 2015.
Zhuo Zhang, Yuanhang Yang, Yong Dai, Qifan Wang, Yue Yu, Lizhen Qu, and Zenglin Xu. FodPETuning: When fedcrated learning meets the parameter-effiiient tuning methods pre-trained language models. In Proc. Annu. Meeting Assoc. Comput. Linguistics (ACL), Toronto, Canada, Jul. 2023b.
Hongling Zheng. Li Shen, Anke Tang, Yong Luo, Han Hu, Bo Du Yonggang Wen, and Dacheng Tao. Learning from models beyond fine-tuning, Not. Mach. Intell., 7:6-17, 1,,JJaa Jan. 2025.
Bolci Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ADE20K datasct. In Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recognit. (CVPR), Honolulu, Hawaii, Jul. 2017.
Zhuangdi Zhu, Junyuan Hong, Steve Drew, and Jiayu Zhou. Resilient and communication efficient Icarning for heterogencous federated systems. In Proc. Int. Conf. Mach. Learn. (ICML), Baltimore MD, Jul. 2022.
15

---

Published as a conference paper at ICLR 2026
We provide more details and results about work in the appendices. Here are the contents:
Appendix
A
Derivation of proposed Muscle loss function.
Appendix
Effect of aj in capturing dependencies.
The relation between (()
and (N-1)
Appendix
TR
Tk,
Appendix
Effect of the weighting coefficient aj in toy example.
Appendix
Proof of Theorem T
Appendix
Discussion on the MI related to Musete (z,) and L Pairwise (2.))
Appendix
Details of the local training and test datasets experiments.
Appendix
Details of the selected models and assigned tasks for users.
Appendix Impact decreasing Muscle loss on A.
Appendix
Performance of FedMusele using synthetic public dataset.
Appendix
Impact of Muscle loss function in FedMuscle.
Appendix
Features offered by the Muscle loss.
Appendix
M Integration of Muscle loss into CreamFL.
Appendix
Effect local epochs, communication rounds, and CL epochs.
Appendix
Non-IID data partition among
Appendix
Impact of the output dimension and temperature.
Appendix
Discussion on possible future work directions.
Appendix R
Limitations.
Appendix
Discussion on convergence of FedMuscle.
Appendix
The use of large language models (LLMs).
A
DERIVATION OF OUR PROPOSED MUSCLE Loss FUNCTION
Considering equation
we have
[( {27t N Jm=1, mytin
E E. Musseee (2/²)
-E
Σjer f(z,, {2/m J.
mm=
where E denotes the expectation over random batch of data samples. Equation (6) can be interpreted as the categorical cross-entropy for correctly classifying the positives, where the fraction inside the log(.) function represents the predicted probability. The optimal expected categorical cross-entropy obtained -E p(correctly classifying positives), where
pccorrectyy classifying positives)
{2.m)) 1amiN N
m xx²"
TyEFmr(en (( I m
IN EjEgn p (z1m 1= msi 2,7
Nyrer (hhnnns) (2.7 1m=1, mon
1mmt, pl (z 12
x?²
p((x mmm. m
1.N
(8)
p( (1/g) 1mm
Σj
p( (2m 1.
Immi, matess
16

---

Published as a conference paper at ICLR 2026
By comparing equations (6) and is obtained follows:
we observe that the optimal value of f (x) {x
1.
m== mmn
{2 1.W mm metie 23²
p
f" (2
M.
Jm= mott) x
i.N {8)m Jmm mattn
I.N p a Jm Fmmm mytln
(10)
p (z.) p
m 3m Jmml,
IN
where denotes proportionality, implying equivalence up multiplicative constant.
In the InfoNCE loss, f is associated with the cosine similarity between two vectors (i.e., positive pairs). However, as inferred from equation (TO). in our proposed Muscle loss, I* should be related to the similarity more than two vectors. To this end, we define a cosine similarity matrix, each element of which is the cosine similarity between two vectors.
Let C gNNN denote the cosine similarity matrix corresponding to the set of vectors {ux E Rd k E [N]] where d is the dimension of each vector in the set. We have Ck. us ui k, E [N]. Using matrix C, we consider that the following expression holds:
((us) IKEINY p p(us)
exp
Σ
Σ
C
Ti (N) AN)
(11)
2 NEINT
exp
Σ Σ
Σ
uh u//Tkd (N) Thll
(12)
2
KEIN]
(N)
((N)
where TRll = Tie is a temperature parameter that moderates the effect similarity between
vectors us and us given the correlations among N vectors. Note that we set = 2 in
expression 12 it simplifies p(uc)p(us)
p(us. exp [u: mn/T(?) Since there
additional
vectors influencing the similarity between uy and uy when N = 2, we can set Th. Th. where
Thtt is the temperature parameter used in the InfoNCE loss. Therefore, for N = 2, pf us.. w.)
pl 201
exp (u uy i. which aligns with the formulation used in the InfoNCE loss (Oord et al. 2019
Tian et al. 2020). Furthermore, for a set with more than two vectors, expression (12) accounts for interactions among all vector pairs by incorporating cosine similarity across all pairs, along with distinct temperature parameters.
Based on expression
12 the optimal function f* in (T) obtained as follows:
8t (2)), {2] mm mpn)
JNN
exp (1 (2 EmEENNA 2.
TH (N) Tn.m
+ Σ EMEINTANN Em'EINTY(N,. 27th 2/m T mmme
(N)
exp 1 EmEINIM EmPEINIV 2mm 2 / (oN
(13)
(EmE(NMa) 2 2m / Tuym 27t
exp
(14)
exp
2 2)m 2 EmE(NN)) Em (/--"-1) - 1/ - )
(N)
mm
= aj exp
3mm
(NN)
(15)
mE (m)
where aj = exp
(N) E [m][N] EmEENN mm 2m. 231h and m
1/T
I/-(N) tion
Combining
15) and ② results in the Muscle loss function Ln Li (2!) defined in equa-
17

---

Published as a conference paper at ICLR 2026
B EFFECT OF aj IN CAPTURING DEPENDENCIES
In this section, we show that, compared to pairwise alignment TTin et al. 2020 Wang & Sun
2022 Yang et al.) 2022 Xue al. 2024a), our proposed Muscle loss more effectively captures dependencies among representations obtained from more than two models. This improvement arises from the use of weighting coefficients aj. j E J
We consider special case in which there is no dependency among the representation vectors. In this case, aj for all J Thus, we have
(N) Muscle (2 log ((......)) exp(2?" EmEINIIAU 2/" Tn,m
Ejer oj exp(2," [mE(NN](n)
(n) 3jm 3 TT..e))
(N)
exp(2% EmELNTMAn 23" 2/t/ /T(
log
ΣjeJn exp(z) Eme(N]Ma) 2)m /T.m) (N)
exp(2?' 22"h (TT.))
(b)
Σ log
EjmE(B) [[B] exp(2,) 2 (n.m)
(c)
CR.m InioNCT(22)
Pairwise (2.),
(16)
where equality (a) follows from the assumption that a₁ = j E J"; equality (b) follows from the logarithmic product rule; and equalities (c) and (d) follow from the InfoNCE and pairwise alignment formulations provided in Section 4.1
Equality (16) shows that, when there no dependency among the representation vectors, our pro- posed Muscle loss and pairwise alignment equivalent. Nevertheless, the assumption of no de- pendency among representation vectors is generally incorrect, as some representation vectors may originate from the same data sample. While pairwise alignment approaches inherently rely this assumption, our proposed Muscle loss can effectively capture dependencies among representation vectors through the of weighting coefficients aj. j J
18

---

Published as a conference paper at ICLR 2026
C THE RELATION BETWEEN (N) Tel
AND (N-1)
As discussed in Appendix A the temperature parameter Tl moderates the effect of similarity between vectors ug and uy when modeling the probability density ratio among N vectors in (22 In this section, we show that the temperature parameter T,,
(N)
(N-1) which serves the same purpose for
N vectors, should be smaller than Thet
(N) From (2), for N vectors, have
P (fux)km)) exp [R (ux))
(N) us 2AL WA/TTR
(17)
Similarly, for N vectors, we have
p ((uu) 1-11 [ P/ (us)
N-1
(N-1)
exp
ux ww w TR.
(18)
2
k=
LEIN (1
Now, we aim derive an expression similar
(18) using
We have
p (fuu (=)) N-11)
p p(uux)k=))
Jen [IKE(N) P (ux)
P p(ux) duN
IkEE[N p()))))))))))))))))))))))))))p
Σ
Σ
w/TTE (N)
pluN) duN
exp
2
Rd
NEIN (E(N)((M)
N-1
N-1
exp
Σ
ws w (N)
Tkd
Σ
un / KK
(N)
p (uex) dun
2 1= 1 LEIN I]
1=
N-1
= exp
Σ
Σ uk u (N)
exp
Σ uk
uN / / (2xx) dun
2
k=1
LE(N-1)M(A)
Rd
(19)
where (a) results from 17
Next, we aim to find simplified expression for A in 19). We have
N-1
cxp
UN
uk/
TRN
(N)
(uN) duN
R
exp
UN
uk/ u
TN,N
(N)
2
1
1lux11
duN
Rd
/
exp
1
Σ
uk TK.N
(N)
uN
Σ
ukk/
(N) T.NN
duN
R
2
2
k=
N-1
1)
N-1
exp
2
Σ
BAK/
T.NN
(N)
exp
1
un
Σ
/
(N) TK,N
dux
k=1
k=
1 NN-1N-1
exp
Σ
Σ
us wi/
(N) (N) T.N TLN
(20)
2
where to obtain (a), we assume that the marginal distribution of each uy. for k E [N]. is an isotropic Gaussian. Thus, we have p (uN)) o exp
(-2 ||u Also, (b) follows from the fact that the
integral Ses exp (- UN - EN- uu "
(N)
(N) duy evaluates to a constant.
19

---

Published as a conference paper at ICLR 2026
Image
Image
Image 3
exp(=]
aj [x) 0z)
Positives
Negatives
Applying aj
(b)
Figure 6: Illustration the impact of az, j E J on different terms our proposed Muscle loss function in (3) for a batch of three images: (a) when = 1, and (b) when az is determined based on the dependencies among outputs of representation models. For ease of visualization, Tu, is set to one for n, m 21,2.3).
By combining (19 and 20 we have
p (( (
exp
u m (I/TEN)
+1/ThN
(21)
IIRE[N P (ux)
Comparing
18 with 21
shows that the temperature parameters must satisfy the following equa-
tion:
1/) (N). 1/ 1/-(N== = + I/T..N T.N
[N-
(22)
Since temperature parameters are positive, equality must satisfy the following inequality: Tk.
22
implies that the temperature parameters
(N-1) (N)
Tkl
D EFFECT OF THE WEIGHTING COEFFICIENT Oj IN TOY EXAMPLE
Figure 6 illlstraates how aj j E J", impacts the terms in our proposed Muscle loss function. Con- sider batch of three images consisting of dog (image 1), cat (image 2), and a bench (image 3). The images are selected such that images I and 2 are more similar than images and 3, and images and are the least similar. Figure a and (b) correspond to the cases where aj 1, J T" (i.e., similar to the pairwise alignment) and where ag. J obtained based on proposed Muscle loss function in (3). respectively. =1 is the output of the representation model 1 for image 1. When 2 is considered the anchor vector, Figure 6 shows that the terms corresponding to the negatives originating from images 1 and (e.g., 23 + x have greater impact the loss function when a2 = 1, J". However, this is not the case Figure ((b). Specifically, ag tends to have a higher value when the negatives originate from images and (e.g.. 2 + 23), which are more dissimilar. Therefore, in the Muscle loss function, Oj can adjust the importance of the terms in the loss function based on the dissimilarity among the negatives.
20

---

Published as a conference paper at ICLR 2026
E
PROOF OF THEOREM 1
Based on equation 6
we have
[z] {=10) mm mpta)
E Muscee
(2,
Flog
ΣjeJ" f(2 12 Im motn
p( {x1t Di pl ) (1(" ) Iim
Flog
p (==", (2jm HM
Ejern pl ot ) (2%) ) mosn
mati
p( [=]" (= Jimi mat
p(2 )p( (2 matt
-Elog
p (2Th m/
2m 3/y
Jmm m//n
+ EGEF)
p )p(pp(( (=(") m=l, motin)
pl
(= ) mpinn
p p=()) {2,Th Jm Jmml, myn
LN
2 { mm m of
IN
Elog
1
= {2. Jmml,
P (2 JE
p
{2 m J mm m. mynh
(z.²) p
1m=1 mnn
{2 m- mfn
2
= =Elog
1
JT E
p z {-,Th W
J.N 3m Im mf
mm=1, mfn
(z )
{2) 1m= mfn
Elog
1
BNN-1
p 2 {2mhh I
mm=1 mmn
2 {27Th N
Imml of
(N 1) log((B) Eog
p p( { y mottn
(N - 1) log B3 1(2 ; [ Jmm mytn):
(23)
where equality (a) results from 10). Equality (b) is obtained by using [J"| = BN1 Inequal-
ity (c) results from the definition of [(2;; {2(" 1.
m==)) Finally, inequality (4) is concluded by
rearranging the terms in
F DISCUSSION ON THE MI RELATED TO MUSSCL (2?") AND LR CPAIRIISE (z.)
It is shown by Oord al. (2019) Tian al. (2020 that CIiONCE (2,") in 1 is related to the MI [[2;; = as follows:
[<=1:="" > log(B) B)
E.
InloNce (z(").
(24)
Using 24 for Chiiriise (2.² we have:
ECB (zx²)
E.C. (z!")
(
2 (N (B)
(25)
(E[
By rearranging the terms in 25 we have the following inequality for the pairwise alignment:
Σ
[(=]: 2 (N 1) log(B) IE.PPairiie
(26)
minimizing L (2,)) is equivalent to maximizing the MI Eme(N]y(n) 1(2;; 2. This pairwise MI formulation lacks conditional MI. other words. the
Based on inequality 266
other representation vectors 2!t where m' [N] \ {m, n), provide additional information
21

---

Published as a conference paper at ICLR 2026
[ x
[(
z?
x
22
(a)
Figure Venn diagrams illustrating the T(2 23, 2²²) among three variables z,, z?, and z/. (a) In
general, T(2,; z = =() + T(x] 2. |2.). Since z,, z;; and z are not independent of one does not equal [[]] =() + =(²); (b) A special case where [(2;; ; x²)
another. I(z
[(=); =²) [(2):=]) However, this case cannot occur because z!_ 2;, and z positives originating
from the same data sample æ., Since our proposed Muscle loss function related i =f). more
effectively captures dependencics between the representation vectors.
about the pair (z,² z; However, we know that all the positives z!' n E [M], originate from the same data sample æ; E DD and thus, they should share some joint dependencies. This MI perspective further confirms that pairwise alignment approaches cannot fully capture the dependencies among the representation vectors in multi-model scenarios.
On the other hand, as shown in Theorem [ minimizing our proposed Muscle loss func- tion is equivalent to maximizing the MI [[2;; (2(" (2 Jm=1, Figure 7 illustrates that {z,") 1 myn and EmglNtiin / [z,"; =²(") are not generally equal. In particular, by {
using the chain rule for MI (Thomas & Joy 2006), we have: 1(2;;; {2m"h Jmm=,, mynnn
1m-1 Σ [(=;; 2!Th (2 mm==1, m) m The conditional MI in this formulation captures the depen- dencies among all the representation vectors 2,, n E [N], thereby facilitating more effective knowl- edge transfer among them.
G DETAILS OF THE LOCAL TRAINING AND TEST DATASETS IN EXPERIMENTS
We consider that users employ heterogeneous pre-trained FMs tailored to their respective tasks. These pre-trained FMs can be fine-tuned for users' specific tasks, allowing the size of local datasets to remain small. Additionally, we assign different local dataset sizes to users based on the complexity of their tasks. Specifically:
Each user performing the IC10 task has 50 training samples and 4000 test samples.
Each user performing the IC100 task has 100 training samples and 2000 test samples.
Each user performing the MLC task has 200 training samples and 1000 test samples.
Each user performing the SS task has 1000 training samples and 3000 test samples.
Each user performing the TC task has 100 training samples and 5000 test samples.
The data samples are uniformly divided across users with the same task.
H DETAILS OF THE SELECTED MODELS AND ASSIGNED TASKS FOR USERS
In Setupl, we consider six users and assign their tasks and FMs as follows:
Users 1, 2, and 3 perform the MLC task using ViT-Base, VIT-Small, and ViT-Large, spectively.
Users 4 and 5 perform the IC100 task using ViT-Base and ViT-Small, respectively.
User 6 performs the IC10 task using ViT-Tiny.
22

---

Published as a conference paper at ICLR 2026
In Setup2, in addition the six users from Setup1, we consider four more with the following tasks and FMs::
Users 7 and 8 perform the SS task using SegFormer-BO and SegFormer-Bi respectively,
Users 9 and 10 perform the TC task using BERT-Base and SINBBERRTBBas, respectively.
For users performing the MLC, IC100, and IC10 tasks, we use two linear layers with a rectified linear unit (ReLU) in between as the task-specific prediction head. For users performing the SS task, we use the multi-layer perceptron (MLP) decoder proposed by Xie et al. (202I) as their task- specific prediction head. For users performing the TC task, we linear layer the task task-specific prediction head.
I IMPACT OF DECREASING MUSCLE LOss ON A
8.0
20.0
75
0.0
7.0
6.3
Musole
-20.0
6.0
-0..0
5.5
5.0
-60.0
30
90
120
150
Communication Round
Figure 8: Muscle loss and A evolution over communication rounds. Pascal VOC is used as the public dataset.
PERFORMANCE OF FEDMUSCLE USING A SYNTHETIC PUBLIC DATASET
To generate a synthetic public dataset, we follow existing works such as FedDF (Lin et al. 2020 and employ pre-trained generative models. In particular, we consider that each user generates 50 syn- thetic images using a pre-trained image-to-image diffusion pipeline based on Stable Diffusion (Rom- bach et al. 2022), , augmented with ControlNet image conditioning (Zhang al. 2023a) to enable controllable generation aligned with each user's 's local data. Figure 9 shows one the synthetic im- age samples generated by each user in Setup1. The 50 synthetic images generated by each user then combined to form a shared public dataset consisting of 300 unlabeled synthetic samples. The results in demonstatttt that FedMuscle can effectively benefit from a synthetic public dataset to align users' representation spaces. Thus, FedMuscle remains effective even when the samples in the shared public dataset are entirely synthetic. These results confirm that synthetic datasets can serve as viable and practical alternative to real public datasets for FedMuscle.
Table 6: Performance FedMuscle in Sctup1 when the public datasct is obtained by generating 300 synthetic data samples.
Fool Menic
FollMuacl (Ches)
Training
VrT-Bose
micre-Fi
44.97cm
42.17p414
mino-F
4187240
Large
MLC
mi
16300 cis
42.50 on
1C100
Accuracy
2850
24.7/ se:
Wir-Smal
10 100
Aceurcc
24.70 um
ViE-Tiny
1C10
Accuracy
61.77
45.77 us
(1))
+18.32
0.00
23

---

Published as a conference paper at ICLR 2026
ML.
User
IC
User 4
User 5
User
Figure 9; Illustraioon of a synthetically generated sample for each user in Setup
K IMPACT OF MUSCLE Loss FUNCTION IN FEDMUSCLE
Table 7: Performance FedMuscle Setupl using our proposed Muscle function compared
Gramian-bar based contrastive loss and pairwise alignment.
Public
FocMuule
Dateset
Metric
(Ours)
Gramisn-hood
Aligament
Training
ViT-Booe
MLC
mem-Fii
17
Pikca VOC
MTC
mem-Fi
Lang
MDC
mec
19
49.07
4137
02
Boee
JC100
Accuray
36.7
M483.on
24.7.00
Small
1C100
Accuaay
28.00 em
VIT Him
1C10
Accuracy
66371181
6133 cm
43.77
(831
+26.70
#34.01
+17.34
8000
VIT Buse
MLC
micro
48.17 u
42,27 em
42.17 oz
WTSSmoll
MLC
mem-F1
48.3
COLD
WiLLagge
MLC
mamn-1
44.07
42
WIT Base
IC100
Accuncy
37.27
20
15.13
24 24.77
WT-Smoll
IC100
Accuncy
5005
50g vm
10.20_ee
24 20
VIT Tim
1C10
Accuracy
03.25100
36,77 12
61.10
43.77 ow
(1521
+28.6
+22.11
+19.18
0,00
VIT Base
MLC
micro Ft
44
41.20 ez
42.17 ozr
CHARR-U 100
Small
MLC
micro 11
16
45.80.
46.00 60
43.87 cow
VereLange
MLC
memn-F
4144-
VIT Vir-Bee
IC100
Accuncc
53.43
31.70-004
31.20 ei
24.77-00
ViT-Small
IC100
Acoucyy
20 20
28.472m
25,00-1
24 M0mm
ViT-Tmy
IC10
Accuncy
5630ET
36.41-03
41.72n0
+16.8
+15.99
+10.18
0,00
L FEATURES OFFERED BY THE MUSCLE Loss
In this section, we provide detailed discussion of the features offered by our proposed Muscle loss that are not available the recently proposed tramian-basce contrastive loss by |Cichettii et al. 2025).
Theoretical foundation: Our proposed Muscle loss is derived from strong theoretical analysis, whereas the Gramian-based contrastive loss, though effective, based on intu- ition and lacks theoretical justification.
Consistency with InfoNCE: For two modalities models, the Muscle loss reduces the standard InfoNCE loss, demonstrating that it is well-structured extension. In contrast, the Cramian-based contrastive loss not equivalent to the InfoNCE loss in the two-modality (or two-model) case.
Well-definedness: The Gramian based contrastive loss is defined as follows (Cicchetti
exp(- Vol( "!..... 1.2 ..... x/7)/T et al. 2025): C (x = log موك exp( Voll= 23. =)) ......,
... "/.)) where Vol
denotes the volume the N -ddimensional parallelotope determined by the determinant of the Gramian matrix. However, in the case of N > d, the determinant of the Gramian matrix
24

---

Published as a conference paper at ICLR 2026
is zero. Although this condition may not hold in most realistic scenarios, it shows that the Gramian-base based contrastive loss is not a well-structured extension for multi-model settings. The Muscle loss does not have such limitations and remains well-defined regardless of the relationship between N and d.
Insights into temperature parameters: Our proposed Muscle loss provides new insights into the temperature parameters. Based on the analysis provided in Appendix [ we know, for example, that for three models, the temperature parameters for each par should be larger than those used for the same pair in two-model setting.
Control on negative transfer: Some tasks or modalities may be more semantically related to each other than others. Low semantic relevance can lead to negative transfer and degrade performance. The temperature parameters Ti.m, defined between each pair of models in the Muscle loss, provide flexible and principled mechanism to mitigate negative transfer among tasks with low semantic relevance. This is because semantically closer tasks or modalities can benefit from stronger alignment (lower temperature), while semantically distant ones require softer alignment (higher temperature). In contrast, the Gramian- contrastive loss uses single temperature parameter and lacks mechanism to mitigate negative transfer among tasks or modalities with low relevance.
Modularity and integration:
The structure of the Muscle loss enables modular integration into FL settings. As demonstrated in Section 5 when using the Muscle loss, each user transmits only its representation matrices and offloads to the server the part of the loss computation that depends on the representation matrices of other users. The Gramian- based contrastive loss, however, lacks such modularity, resulting in higher computational costs for users.
M INTEGRATION OF MUSCLE Loss INTO CREAMFL
In the CreamFL (Yu al. 2023) setup. the local models used by the users are as follows:
ResNet-18 (He et al.
2016 is used by the uni-modal image users.
GRU (Chung et al.
2014) is used by the uni-modal text users.
ResNet--1 and GRU serve as the image and text encoders, respectively, for the multi-modal users.
The global model at the server is designed to be larger than those used by the users. In particular, the server employs ResNet-101 and BERT as the image and text encoders, respectively. All models trained from scratch. In CreamFL, only the representations of the public data obtained from the models are communicated between the users and the server. Specifically, the server generates global image and text representations using the global model and sends them the users.
Each user first trains its local model using its local dataset and then applies a local contrastive regularization method. In this step, each user employs CL with the InfoNCE loss align the repre- sentations of the public data generated by its own model (e.g., local image representations) with the global representations of the other modality (c.g., global text representations). Additionally, each user aims to make the representations obtained from its own model for given modality (e.g., im- age) similar to the global representations of the same modality. To this end, each user employs loss function similar to the one proposed in MOON (Li et al., (2021a). Then, each user obtains the representations of the public data from its model and sends them to the server.
The server first trains the global model using the public dataset and then applies a global-local contrastive aggregation method. In this step, for each modality, the server assigns score to the local representations received from the users and aggregates them based on these scores. To compute the score for each local representation, the server calculates the InfoNCE loss for that representation as the anchor, using the global representations of the other modality as positives and negatives. Then, for each modality, the server minimizes the & distance between the global representations and the aggregated local representations.
The local contrastive regularization and global- contrastive aggregation methods in CreamFL heuristic approaches. For example, extending these methods to support more than two modalities
25

---

Published as a conference paper at ICLR 2026
would be challenging. We replace these heuristic approaches with our proposed Muscle loss, which can systematically capture dependencies among representations across any number of modalities. On the user side, replace local contrastive regularization with minimization of the Muscle loss to align the local representations of each modality (e.g.. image) with the global representations of all modalities (e.g., image and text). On the server side, we replace global-local contrastive aggregation with the Muscle loss to align the global representations of each modality with all local representations received from the users. For this experiment, we used the code provided by the CreamFL paper Yu et al. 2023) and retained all of their original hyperparameters.
N EFFECT OF LOCAL EPOCHS, COMMUNICATION ROUNDS, AND CL EPOCHS
Table 8: Performance FedMuscle under different numbers local epochs E, communication rounds R, and CL. epochs T. The product of local epochs and communication rounds is fixed at 150 (i.e., E x R = 150). The public dataset is derived from COCO.
User
Model
Tusk
Eval. Mieric
--1
FodMuscle
IC 10
R=-25
Local Training
T=3
7-5
7-1
T-3
T-5
7-3
VIT-Bise
micr-Fi
cB
88.00.a
47.33
48.00
ca
44.20.0
4900-co
49227
49.47.33
42.17.0z
WTSSmall
mim-Fii
45
4663pm
4677mm
50.8
3175gm
51.7209
51.1200
31.09am
5125
45.82om
Virraage
44
45.000
47.3g00
S120200
31.72400
47975mm
Biee
IC100
Accuracy
2507
26570w
31
3080
35.33w
0440.0
36372
37.00.3
24.7.oo
ViT-Small
1C:00
Accuncy
2077
3140.0
30
mm
10
24.0000
Wi-ling
IC10
Accuracy
32
38.20 1a
35.10-1
00.00 00
03000 on
0033-c0
03007 mm 6407.
43.77.ol
(16)
+5.71
+10.18
-1281
125.37
427.05
+27.3
+28.91
129.21
+29.06
0.00
O NON-IID DATA PARTITION AMONG USERS
DIMRI9
CEAR
105-0
300
92
82
72
60
53
65-
33
33
22
33
22
13
3-
12
Ulow c
Figure 10: Illustration of the non-IID data distribution among users.
P IMPACT OF THE OUTPUT DIMENSION AND TEMPERATURE
We investigate whether tuning the output dimension of the representation models (i.c., d) and the temperature parameters (i.e., Th.m and T as hyperparameters can further improve the perfor- mance of FedMuscle. Note that a higher results in increased communication costs. The results in Figure T[ a) show that the performance remains nearly constant regardless of the chosen output di- mension. Additionally, the results in Figure IT(b) show that lower temperature values lead to better overall performance. These findings are consistent with the results presented by Chen et al. (2020) and] Yu al. (2023) using the InfoNCE loss function.
We also investigate the impact of increasing the temperature parameter between tasks with lower semantic relevance in Setup2. The results in Table 9 show that assigning a very large temperature value to user pairs with SS and TC tasks improves the performance of both SS users and TC users. Beyond adjusting the temperature parameters, the semantic relevance between tasks can also be influenced by obtaining representations that enable meaningful knowledge transfer. Exploring other approaches, such as the method proposed in (Mukhoti et al. 2023). may serve as a potential solution for enhancing semantic relevance between SS and TCC tasks.
26

---

Published as a conference paper at ICLR 2026
15
4(%)
4(%)
0
128 256 512 1024 2048
(02,6.19) (04,03) 0.5 (06, 0.5)
10R 0663
11 0.5)
Output Dimension
(T mrr T
(b)
Figure 11: Impact of (a) the output dimension of the representation models and (b) the temperature parame- (4) (3) ters Tr.m and Trim on FedMuscle's performance. We set E 150, R 1, and T 5. The public dataset is derived from COCO.
Table 9: Performance of FedMuscle in Setup2 when large temperature parameter value (i.e., 100) is assigned between with SS and TC tasks.
User
Model
Tack
Bval. Metric
PooMinsee
Algorithm
(Ours)
Training
Wi--Bue
MLC
nicmo-F1
47.90
42.00
VIE Small
MLC
miceo F1
50.20
43,30
VIT Lange
MCL
mice F1
49.20
12:30
1010
Accuasy
31.20
20.90
Sm
10100
Accuayy
2790
20.20
VirTiny
1010
Accrasy
37:10
10700
1310
1250
Neghomer--1
53
12:30
1230
BERT Dase
TC
Accurac
48.20
10.10
10
DesIBERRT Base
TC
Accurscy
55380
51.00
A 1997
+13.
0.00
DISCUSSION ON POSSIBLE FUTURE WORK DIRECTIONS
Multi-Modal Representation Learning Although most existing self-supervised learning ap- proaches have focused on single-modality inputs (e.g.. images), performance can be enhanced by capturing the complementary information available across multiple modalities (Sükei et al. (024)) For example, in (Radford et al., 3 image captions are considered as an additional modality and CLIP is proposed to learn a multi-modal representation space by jointly training text and an image encoder. For two modalities, typical CL loss functions, such as InfoNCE, can be used. Multi-modal representation learning using image, text, and 3D point cloud modalities has recently been studied in (Xue et al. 2024a). However, they used the pairwise alignment approach, which cannot fully capture dependencies among the modalities. For more than two modalities, a systematic approach to learning multi-modal representation space that effectively captures dependencies among modal- ities is still needed. Recently, [Cicchetti et al. (2025) proposed a Gramian based contrastive loss. Although their approach is effective and intuitive, lacks theoretical justification and provides no mechanism to address negative transfer among modalities with low semantic relevance. While our primary motivation for deriving the Muscle loss function in was learn shared representation space across users' models in FMTL, we believe that the Muscle loss can also enhance performance multi-modal representation learning. particular, we can define positives and negatives based on the available modalities and then effectively capture the dependencies among them using our proposed Muscle loss function.
Performance Metric for Federated Multi Task Learning Users in an FMTL setting have differ- ent tasks with distinct performance metrics. Obtaining an overall performance metric across users is challenging due to task heterogeneity. In our experiments, we followed existing works (Lu et al.) 2024 Maninis et al. 2019) and used the average per-user performance improvement relative to the local training bascline (i.e., A) the overall performance metric. However, believe that this metric cannot fully capture the overall performance of an algorithm. This because a specific al- gorithm may enhance the performance of only one user and degrade the performance of all other users, yet may still be positive. For this reason, we presented detailed performance metrics for all users in experiments. Developing more informative overall performance evaluation metric for federated multi-task learning could be a valuable future research direction.
Selection of Representation Matrices in FedMuscle In Algorithm T randomly select M representation matrices from other users to compute the aggregated matrix and the weighting coeffi- cient vector for each user. Replacing this random selection with systematic approach that provides
27

---

Published as a conference paper at ICLR 2026
more relevant knowledge based each user's task and those of other users would be interesting direction for future research.
Convergence Analysis of FedMuscle In FedMuscle, there is no global model. Thus, new con- vergence definitions are needed to account for multiple objectives and heterogeneous tasks across users. These definitions may consider the convergence of each user's 's model individually or single objective that encompasses all users' models, similar to the Sheaf-FMTTL formulation proposed by Issaid et al. (2025). Furthermore, in FedMuscle, each user minimizes two losses: a task-specific loss using its local dataset and CL loss using a shared public dataset. The convergence of the Muscle loss and its generalization properties should be studied. Overall, our observation is that the CL loss FedMuscle can be viewed as regularizer that guides users' local models toward capturing com- mon, tank-agi information from other users' models, thereby improving their generalization capability. Exploring these theoretical aspects would be an interesting direction for future work. We have provided more insights FedMuscle's 's convergence in Appendix
R LIMITATIONS
FedMuscle relies a public dataset. Legal and copyright issues related the public data should be carefully considered in real-world applications. Similar to typical feder- ated learning algorithms, where users adhere to the global model architecture and training guidelines imposed by the server, in FedMuscle, users should also adhere to the public dataset provided by the server.
FedMuscle utilizes the Muscle loss the representations from users' models. Since users have heterogeneous models and diverse tasks, it is important to identify which part of each model should serve the representation model and which part should function as the task-specific prediction head. This distinction allows users to generate representations that are meaningful and suitable for CL. facilitating the alignment of representations across different users' models.
We considered multiple CV and NLP tasks in our experimental setup and provided com- parison with multi-modal federated learning algorithms based on their experimental set- tings. However, we believe that the experimental setup can be further expanded by in- corporating more challenging tasks. We will continue to explore this direction future work.
S DISCUSSION ON CONVERGENCE OF FEDMUSCLE
In this section, we show that the Muscle loss in FedMuscle be considered regularizer. We then show that, with some simplifications, and for the special case linear representation setting, FedMuscle is optimizing an objective nearly identical to the Sheaf-FMTL algorithm (Issaid al., 2025). Note that the Sheaf-FMTL formulation provides a unified view of several existing FL and FMTL algorithms with guaranteed convergence bounds.
In FedMuscle, we can formulate the overall optimization problem encompassing all user's models follows:
min (w" NE(N]
Σ (n((0) + (ww"), NEIN]
(27)
where we have (0)(0) (O") and the expectation is taken over the local dataset D for each user n E []]. We also have (ww")) E| "[ where the expectation taken over the
public dataset D. Note that x obtained from the output of user n's representation model. We have z"" = ex (w", æz)), where e" is user n's representation model and x; E D. In the following, step-by-step, simplify the expression (ww").
28

---

Published as a conference paper at ICLR 2026
Employing simplified version of the contrastive loss has been used in previous studies to facilitate analysis (Xue et al., 2024b. Wang & Liu 2021). D. Based on (3), we have
ai exp(2" Emen' 2," /Tmmm)
g" (0 log
Ejern aj exp(z]) EmEN' 3mm /T.mm
log ((......))
Blz
Σ zi
meN'
E log
Σ JEJ
aj exp (z
Σ
Jm h /
mEN"
E log (r(......))
- R[2
12
(NN)
(2,m
mEN"
E max jEJ
log ag + (z."
Σ MEN"
jim ((111)
(28)
where (a) follows from the log-sum-exp approximation. Let j E T be the index tuple that
maximizes the last term in (28). The optimization problem (7 is minimized by each user with
respect to its local model parameters l Thus, in the provided simplified contrastive loss in (28)
log (i.....)) and E log Crjuas are constants and can be ignored when optimizing (27). We have
g"(wn) ~ -B(2P
Σ mEN"
(2,
m jun
(N) m.
(29)
From 29 we observe that g" ("w") serves as a regularizer in the optimization problem (7 We
also observe that, by choosing large values for the temperature parameters (NN) the impact of this regularization term in (27p decreases. In the extreme case where the temperature parameters are set to co, (27) reduces to local training. In this case, each user optimizes its model solely using its local dataset without collaborating with other users.
Now, let us focus on linear representation setting in which e" xx)) = W". for W Rdxxd" where d" denotes the input dimension of user n's representation model based on its task. Note that focusing on linear representation setting has been widely adopted in FL (Collins al. 2021| and CI. (Nakada et al. 2021: Xue et al. 2024b) analysis. We have W"r; (æ: 3 Ta)
where @ denotes the Kronecker product and w" E IRd.d" is the vectorized form of the parameters of user nt's representation model. Thus, we can rewrite (ww") in (29 as follows:
("w"") ~ Ex( & Ia)w".
Σ
(N)
[(x,
I Dyme
(30)
meNa
Tn,m
By defining p" = E0 & La] and nmm = 1B(
E )m )T
X La], we have
9" (ww"))) =p" Σ (N)
n.m "n
mEA' Tn.m
(p" wt nm w mh)
(31)
(N)
meNa
7n,m
By combining (27 and 31
we have
min [
Er
D NE(N mEN*
(N) (p"w" pm.."""
(32)
DEIN]
Now, let us compare the optimization problem (32 with the Sheaf-FMTL optimization problem in (lssaid et al. 2025). The Sheaf-FMTL optimizatin problem is defined follows for ized setting in which each user is able to collaborate with its neighboring users in N (Issaid et al.) 2025):
a decentral-
min 6" ( ) =E[M]
Σ 10
A
Σ
Σ
pn 0" pm.. 0m
(33)
=E(M]
E [] meAa
29

---

Published as a conference paper at ICLR 2026
where Pnm and pm.. are learnable matrices and the term p".m pm..""" captures the dis-
crepancy or dissimilarity between local models 0 and 0m. We have p pmm"" p gr | + |mm. 6T (" 2| 2(Pn.m 9" pm..@m) By ignoring the regularization terms
ppm.or and pm f 1|?, the Sheaf-FMTI. optimization problem can be reformulated as fol-
lows:
min MM[M]
Σ ("
Σ
Σ
pn.m0".. pmm..0"
(34)
REIN
nE[M] mE.N"
By comparing the simplified version of the FedMuscle optimization problem in 32) with the Sheaf- FMTL optimization problem in 334, observe that
Due to the similarities between these optimization problems, the convergence of FedMuscle would be comparable that of the Sheaf-FMTL algorithm under the simplified assump- tions considered.
Due to the use of a public dataset in FedMuscle, p" and pm can obtained from the public dataset based the hard negative term in the contrastive loss. Thus, unlike in the Sheaf-FMTL algorithm, where pu,m and pm. are learnable matrices, p" and pm in FedMuscle are determined by the contrastive loss and the public dataset.
As in the Sheaf-FMTL algorithm, the term -p" "" in 32) can be viewed as capturing the discrepancy or dissimilarity between wn and wn. However, in FedMuscle, this term focuses on the representation models rather than the full model parameters ( and 0m
The Sheaf-FMTL algorithm uses a single hyperparameter to control the impact of other users' models on each user's local model training. However, in FedMuscle, the temperature parameters r arising from the contrastive loss control this impact between each pair of users and can be determined based on the semantic similarity among users' tasks.
T THE USE OF LARGE LANGUAGE MODELS (LLMs)
We used LLMs to edit the paper, including grammar, spelling, and word choice. Thus, LLMs did not play significant role in the research ideation writing of this paper.
30


Reviews:
[{"review_id": "qFzXmSHXdF", "reviewer": "Reviewer_VhyE", "summary": "This paper proposes to learn a shared realistic representation space across tasks in Federated multi-task learning (FMTL) rather than shared model parameters, as most existing approaches assume the use of fully or partially homogeneous models across users, which limits their applicability in realistic settings.", "questions": "See above", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper indicates a common limiting assumption among federated multitask\nlearning approaches (model congruity) and proposes a new method\n(FedMuscle) to overcome this limitation. The paper justifies its method\nvia theoretical results showing that their approach maximizes the mutual\ninformation among the models. The paper conducts experiments on both\nimage (ViT, SegFormer) and text (BERT, DistilBERT) domains to justify\nthe performance of FedMuscle compared to various baseline algorithms.", "weaknesses": "The requirement of a shared public dataset is very strong, particularly for\nfederated learning scenarios. On page 4, lines 163 to 165, this issue is addressed\nsuggesting using publicly available datasets or synthetic data samples but there are concerns of model collapse with synthetic data (though said concerns mostly focus on the recursively generated data by the model or\nmodel family itself reinforcing its own biases, and said issue is resolved via\nadding non-synthetic data which the local users’ private database would be\nnon-synthetic). The authors did not focus on any results on their method’s performance on synthetic data, so this is unclear.", "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": "A2zmyggPqc", "reviewer": "Reviewer_BYRu", "summary": "The authors propose an interesting idea of learning a shared representation space in federated multi-task learning settings. The proposed Muscle loss effectively captures dependencies across tasks, and the FedMuscle framework outperforms state-of-the-art baselines, demonstrating the effectiveness of the proposed method.", "questions": "Please see weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors provide a solid theoretical analysis of the proposed Muscle loss function.\n\n2. The proposed FedMuscle algorithm can be applied to various tasks, including computer vision and natural language processing, demonstrating its capability to handle model and task heterogeneity in federated learning.\n\n3. The contrastive Muscle loss can be seamlessly integrated into multimodal approaches.", "weaknesses": "From the perspective of federated learning theory analysis, the current work lacks proofs of convergence and generalization. Of course, this would be a substantial undertaking, and perhaps more in-depth theoretical analysis in this regard can be considered and refined in future research.", "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": "WyzyM8aBbL", "reviewer": "Reviewer_L4FP", "summary": "This study offers FedMuscle, a new federated multi-task learning (FMTL) approach that deals with model and task heterogeneity by aligning latent representations rather than exchanging model parameters.  The key concept is Muscle loss, a theoretically justified contrastive learning objective that aims to jointly align representations from several diverse clients rather than doing pairwise alignments like standard InfoNCE-based approaches.", "questions": "How sensitive is FedMuscle's performance to the number of clients (N) and sample parameter (M)?\n\nCan muscle loss deal with missing modalities or imbalanced representation spaces across tasks?\n\nHow plausible is the premise of a common unlabeled dataset in medical or financial FL contexts?\n\nHow does the approach perform under highly skewed non-IID distributions or in asynchronous settings?\n\nCould the authors compare the runtime (in hours or GPU-days) to FedHeNN or FedRCL?\n\nIs there a theoretical bound on the communication cost-performance trade-off?\n\nIs there a demonstrable association between the MI limit and Δ performance in experiments?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper covers a wide range of model structures and task types.\n\nThe MI lower-bound perspective is nicely integrated with the contrastive framework.\n\nCovers both unimodal and multimodal configurations; the results are consistent.\n\nClear improvement over strong baselines (FedHeNN, CoFED, CreamFL, and Muscle Loss).", "weaknesses": "The Muscle loss is simply a weighted multi-view InfoNCE variation; comparable theories exist, such as Gramian losses and multi-view MI maximization.\n\nDependence on a shared public dataset undermines the privacy argument and limits usefulness in confined contexts.\n\nThere are no theoretical assurances for FedMuscle's convergence or stability under customer heterogeneity.\n\nWeak ablations include no sensitivity analysis on critical hyperparameters (τ, M, B) or comparison to adaptive optimizers (FedAdam, FedNova).\n\nLack of interpretability of \"representation alignment\" and it is unclear what semantic information is transmitted between tasks or how this affects per-task specialization.", "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": "1EVJ6Oa73J", "reviewer": "Reviewer_pgBB", "summary": "This paper addresses federated multi-task learning (FMTL) with heterogeneous models and tasks. The authors propose Muscle loss, a novel contrastive learning objective that aligns representations from multiple models simultaneously by capturing dependencies among all models' representations. Building on this, they develop FedMuscle, a practical FMTL algorithm where users transmit representation vectors of public data to a server, which computes aggregated matrices and weighting coefficients. The paper demonstrates that minimizing Muscle loss is equivalent to maximizing a lower bound on mutual information among models' representations. Experiments on image and language tasks show consistent improvements over baselines.", "questions": "Refer to the Weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The shift from parameter-sharing to representation-alignment is conceptually appealing and well-motivated. The paper addresses a genuine gap in FMTL: most existing methods assume model congruity, while this work handles arbitrary heterogeneous architectures. The focus on learning a shared representation space rather than shared parameters is a valuable perspective.\n- The derivation of the Muscle loss from first principles through probability density ratios (Appendix A) demonstrates rigor. Theorem 1 establishing the MI lower bound is meaningful and provides principled justification for the approach. The analysis comparing Muscle loss to pairwise alignment from an MI perspective (Appendix F) clearly shows why capturing joint dependencies matters. \n- FedMuscle is well-engineered for federated settings. The communication cost reduction via random selection of M representation matrices (Algorithm 1) is pragmatic. Users don't reveal local model parameters (only representations on public data) which addresses privacy concerns better than some alternatives. The modularity of computation (server computes aggregations while users perform gradient updates) is efficient.\n- The paper evaluates across diverse tasks (image classification, semantic segmentation, text classification) with multiple datasets. The comparison with numerous baselines (FedRCL, SAGE, CoFED, FedDF, FedHeNN) is thorough. Setup2's multi-modal evaluation strengthens claims about generality. Integration into CreamFL demonstrates broader applicability. Ablation studies on E, R, T, M, and d are helpful.", "weaknesses": "- A fundamental limitation inadequately addressed is how to decompose each heterogeneous model into representation model $w_n$ and task-specific head $\\phi_n$. The paper assumes \"the output features of all representation models have the same dimension d\" and notes this \"requirement does not preclude model heterogeneity\" and \"can be relaxed by appending a lightweight, learnable projection head.\" This is vague. For practitioners: What if optimal representation dimensions differ by task?, How do you choose which layers constitute the encoder vs. decoder for arbitrary architectures?, Is the projection head frozen or jointly trained? How is its initialization determined?\nThis choice fundamentally affects what \"shared representation space\" means and is never empirically investigated.\n- FedMuscle critically relies on a shared public dataset D. The paper notes this dataset \"may be uni-modal or multi-modal, depending on the users' tasks\" and can be \"easily obtained.\" However, in Setup1, using CIFAR-100 (lower performance) vs. COCO (higher performance) shows ~12% performance swing in $\\Delta$. For multi-modal Setup2 with mixed CV/NLP tasks, using Flickr30K captions for both modalities seems contrived. Are image captions truly representative of what text classification models need? Real federated scenarios may not have access to high-quality public data relevant to all tasks. The paper doesn't analyze sensitivity to public data quality/task relevance rigorously.\n- The multi-modal evaluation in Setup2 mixes semantically distant tasks. Users 1-6 use vision tasks on CV data, while users 9-10 do text classification. The alignment mechanism forces all to learn representations in the same space using Flickr30K captions as the public dataset. This seems artificial. How does aligning a semantic segmentation model with a text classifier via image captions facilitate knowledge transfer? The small improvements for some users (SS tasks in Table 2) may reflect this semantic distance. The paper claims the approach \"handles both model and task heterogeneity,\" but doesn't sufficiently explore when heterogeneity becomes too extreme for representation alignment to be beneficial.\n- Some recent SOTA works on FMTL are missing, for instance:\n\n   - Yipan We, Yuchen Zou, Yapeng Li, Bo Du, “Towards Unified Modeling in Federated Multi-Task Learning via Subspace Decoupling”, arXiv preprint arXiv:2505.24185, 2025.\n   - Yuxiang Lu, Suizhi Huang, Yuwen Yang, Shalayiding Sirejiding, Yue Ding , Hongtao Lu, “FEDHCA2 : Towards Hetero-Client Federated Multi-Task Learning”, Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.\n   - Chaouki Ben Issaid, Praneeth Vepakomma, and Mehdi Bennis. \"Tackling Feature and Sample Heterogeneity in Decentralized Multi-Task Learning: A Sheaf-Theoretic Approach.\" Transactions on Machine Learning Research, 2025.", "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)