Personalised Federated Learning On Heterogeneous Feature Spaces
Authors: Alain Rakotomamonjy
Maxime Vono
Hamlet Jesse Medina Ruiz
Liva Ralaivola
Venue: Accepted by TMLR ()

Abstract: Personalised federated learning (FL) approaches assume that raw data of   all clients are defined in a common space \emph{i.e.} all clients    store their data according to the same schema.   For real-world applications, this assumption is restrictive as clients, having their own systems to collect and then store data, may use {\em heterogeneous}    data representations. To bridge the gap between the assumption of a shared subspace   and the more realistic situation of client-specific spaces, we propose a general framework   coined FLIC that maps client's data onto a common feature space via local embedding functions,   in a federated manner. Preservation of class information in the latent space   is ensured by a distribution alignment with respect to a learned reference distribution.  We provide the algorithmic details of FLIC as well as   theoretical insights supporting the relevance of our methodology.   We compare its performances against FL benchmarks involving heterogeneous input features spaces.   Notably, we are the first to present a successful application of FL to Brain-Computer Interface signals  acquired on a different number of sensors.

Published in Transactions on Machine Learning Research (04/2024)
Personalised Federated Learning On Heterogeneous Feature Spaces
Alain Rakotomamonjy
Mlan..aaktoo@knso-roeee
Criteo AT Lab Paris, Fronce
Maxime Vono
m., m.vooo@enitto.ccommm com
Oriteo AI Lab
Poris, Pronce
Hamlet Jesse Medina Ruiz
hj. wy.meeiinaruis@eitee..oom
Criteo Al Lab Paris, Prance
Liva Ralaivola
Lralaivola@eriteoo.om
Criteo AI Lab
Puris, France
Reviewed on OpenReview: https: // openreview. net/ forun? id= JCZJaq
Abstract
Personalised federated learning (FL) approaches assume that raw data all clients are defined a common space iee all clients store their data according to the same schema. For real-world applications, this assumption is restrictive as clients, having their own systems to collect and then store data, may use heterogeneous data representations. To bridge the gap between the assumption of shared subspace and the more realistic situation of client-specific spaces, we proposc general framework coined FLIC that maps client's data onto a common feature space via local embedding functions, in a federated manner. Preservation of class information the latent space is ensured by a distribution aligmment with respect to a learned reference distribution. We provide the algorithmic details of FLIC as well thcoretical insights supporting the relevance of our methodology. We compare its performances against FL benchmarks involving heterogeneous input features spaces. Notably, we the first to present successful application of FL to Brain- Compuer Interface signals acquired on differont number of sensors.
1 Introduction
Federated learning (FL) is a machine learning paradigm where models are trained from multiple isolated data sets owned by individual agents/clients, where raw data need be transferred to a central server, nor even shared in any way (Kairouz et al., 2021). FL ensures data ownership, and structurally incorporates the principle of data exchange minimisation by only transmitting the required updates of the models being learned. Recently, FL works have focused personalised FL to tackle statistical data heterogeneity and used local models to fit client-specific data (Tan et al., 2022; Jiang et al., 2019; Khodak et al., 2019; Hanzely & Richtárik, 2020). However, most existing personalised FL works assume that the raw data on all clients share the same structure and are defined on a common feature space. Yet, in practice, data collected by clients may use differing structures: they may not capture the same information, some features may be missing or not stored, or some might have been transformed (e.g. via normalization, scaling, or linear combinations). An illustrative example of this, related to Brain-Computer Interfaces (Yger et al., 2016; Lv et al., 2021)

---

Published in Transactions on Machine Learning Research (04/2024)
and tackled in this paper, is the scenario where eletroenceephllorappyy signals are recorded from different subjects, with varying numbers of electrodes and a diverse range of semantic information (e.g. motor imagery tasks and resting state). To tackle the challenge of making FL possible in situations where clients have heterogencous feature spaces such as disparate dimensionalities or differing semantics of vector coordinates we present the first personalised FL framework specifically designed to address this learning scenario.
Proposed Approach. The key idea of our proposal is driven by two objectives: (i) clients' data have to be embedded in common latent space, and (ii) data related to the same semantic information (e.g. label) have to be embedded in the same region of this latent space. The first objective is prior necessary step before FL since it allows to define relevant aggregation scheme on the central server for model parameters (e.g. via weighted averaging). The second one is essential for a proper federated learning of the model parameters as FL approaches are known to struggle when data across clients follow different probability distributions (?). As shown later, this second objective is not guaranteed by performing client-indeeendet learning of embedding functions, such as via low- con-dimensional embeddings or autoencoders. To cope with this issue, we align elients' embedded feature distributions with latent. anchor distribulion that is shared across clients. The learning of the anchor distribution happens in federated way, which means it is updated locally on cach client and then combined on a central server through barycenter computation (Veldhuis, 2002; Banerjee et al., 2005). Then, we seamlessly integrate this distribution alignment mechanism, using local embedding functions and anchor distribution, into a personalised federated learning framework that is similar to the approach proposed by Collins et al. (2021), without any loss of generality.
Contributions. To help the reader better grasp the differences of our approach with respect to the existing literature, we here spell out our contributions:
1. We are the first to formalise the problem of personalised FL heterogeneous client's feature spaces. In contrast to existing approaches, the proposed general framework, referred to as FLIC, allows each client to leverage other clients' data even though they do not have the same raw representation.
2. We introduce a distribution alignment framework and an algorithm that learns the feature embedding functions along with the latent anchor distribution in a local and global foderated manner, respectively. We also show how those essential algorithmic components are integrated into a personalised FL algorithm, easing adoption by practitioners.
3. We provide algorithmic and theoretical support to the proposed methodology. In particular, we show that for a simpler but insightful learning scenario, FLIC is able to recover the true latent subspace underlying the FL problem.
Beyond competitive experimental analyses on loy and real-world problems, we stand out as a pioneer in Brain-Compute Interfaces (BCI) by being the first to learn from heterogeneous BCI datasets using federated learning. The proposed methodology and class nunbers, feat not achieved by any other methodology to our knowledge, and can have a strong impact on other medical domains with similar data heterogeneity.
handle data with different sensor counts
Conventions and Notations. The Euclidean norm on Rd is || |S] denotes the cardinality of the set S and N = N \ {0}. For n E N we refer to {1....,m] n) with [m]. N(m, Σ) is the Gaussian distribution with mean vector m and covariance matrix Σ and X v means that the random variable X is drawn from the probability distribution v. The Wasserstein distance of order between any probability measures μr, v on Rd with finite 2-moment is W2(u, (infçeT(mw)) S 110 0'))=a((0,00)) 02)2/2, where T(u, v) is the set of transference plans of μ and
2 Related works
As far as our knowledge gocs, the proposed methodology is the first one to tackle the problem of FL from heterogeneous feature spaces. However, some related ideas have been proposed in the literature. The idea of using distribution alignement has been considered in the FL literature but only for addressing distribution shifts
2

---

Published in Transactions on Machine Learning Research (04/2024)
Table 1: Related works. PFL refers horizontal personalised FL, VFL to vertical FL and FTL to federated transfer lcarning.
METHOD
TYPE
* PEATURE SPACES
MULTI-PARTY
NO SHARED ID
NO SHARED FEATURE
(ZHANG ET AT.,, 2021A)
PFT.
X
(DIAO AT.,, 2021)
PFT
(COLLINS ET AL., 2021)
PFL
X
X
X
(SHAMSIAN ET AL., 2021)
PFL
(HONG ET AL., 2022)
PFL
X
X
(MAKIIJA ET AL., 2022)
PFL
FLIC (TIIS PAPER)
PFL
J
(HARDY AL., 2017)
VFL
V
(YANG ET AL., 2019)
VFL
(GAO ET AL., 2019)
FTL
X
(SHARMA ET AL., 2019)
FTL
X
(LIU ET AL., 2020)
FTL
(Mont ET AL.., 2022)
FTL
X
on clients (Zhang al., 2021b; Ye al., 2022). Other mothodological works on autoencoders (Xu al., 2020), word embeddings (Alvarez-Melis Melis & Jaakkola, 2018: Alvarez- Melis et al., 2019) or FL under high statistical heterogeneity (Makhija et al., 2022; Luo et al., 2021; Zhou et al., 2022) similar ideas of distribution aligmment for calibrating feature extractors and classifiers. Comparing distributions from different spaces has also been considered in (non-FL) centralised manner using approaches like the Gromov- Wasserstein distance or related distances (Mémoli, 2011; Bunne et al., 2019; Alaya et al., 2022).
Several other works can also be broadly related to the proposed methodology. Loosely speaking, we can divide these related approaches into three categories namely (i) heteroonneuus-archhieeeuuu personalised FL, (ii) vertical FL and (iii) federated transfer learning.
Compared to traditional horizontal personalised FL (PFL) approaches, so-called aeterogeneous-arccttteetee ones are mostly motivated by local heterogeneity regarding resource capabilities of clients e.g. computation a nd storage (Zhang et al., 2021a; Diao et al., 2021; Collins et al., 2021; Shamsian et al., 2021; Hong et al., 2022; Makhija et al., 2022). Nevertheless, they never consider features defined on heterogencous subspaces, which is our main motivation. In vertical federated learning (VFL), clients hold disjoint subsets of features. However, a restrictive assumption is that large number of users are common across the clients (Yang et al., 2019; Hardy et al., 2017; Angelou et al., 2020; Romanini et al., 2021). In addition, up our knowledge, vertical personalised FL approach has been proposed far, which is restrictive if clients have different business objectives and/or tasks. Finally, some works have focused on adapting standard transfer learning approaches with heterogeneous feature domains under the FL paradigm. These federated transfer learning (FTL) approaches (Gao et al., 2019; Mori el al., 2022; Liu et al., 2020; Sharma et al., 2019) stand for FL variants of heterogeneeuu--eatheee feature transfer learning where there are b source clients and target client with a target domain. However, these methods do not consider the same sctting as ours and assume that clients share common subset of features. We compare the most relevant approaches among the previous ones in Table 1.
3 Proposed Methodology
Problem Formulation. We consider the problem where b E No clients want to solve a learning task within the centralised personalised FL paradigm (Yang et al., 2019; Kairouz et al., 2021), where a central entity orchestrates the collaborative solving of a common machine learning problem by the b clients, without requiring raw data exchanges. The clients are assumed to possess local data sets (Dihees such that, for any
E [b), D;
{(x)() Mi
(") ())en) where af!") stands for feature vector, (1) is label and n; Dz}- In contrast.
3

---

Published in Transactions on Machine Learning Research (04/2024)
Local
Global
Global
Local
Client with 28 x 28 Images
1
Φ1
(1)
A1
1
Shared representation
Local classification
Client 2 with 16 x 16 images
USPS
2
X2
F/2
L19
a
Bi
9
0
images
0
3
SVHN
X3
Figure 1: Illustration of part of the proposed methodology for b = clients with heterogencous digit images coming from three different data sets namely MNIST (Deng, 2012), USPS (Hull, 1994) and SVHN (Netzer et al., 2011). The circles with digits inside stand for a group of samples, of given class, owned by client and the size of the circles indicates their probability mass. In the subspace Φ {ui)em (and their level sets) refer to some learnable reference measures to which we seek to align the transformed version Vo, of Personalised FL then occurs in the space and aims at learning local models (0r)ie()) for cach client as well as (or, (4: Jej Non-personalised FL could also be considered and naturally embedded in the proposed distribution alignement framework.
to existing FL approaches, we assume that the raw input features (x((²) Joe[na] of clients live in helerogeneous spaces i.e. for any i E [b], x!!)) X, where X; is a client-specific measurable space. More precisely, for any E [b] and j E [n,], we assume that 2.") E X4 C R& such that {K] are not part of a common ground metric. This setting is challenging since standard FL approaches (MeMahan et al., 2017; Li et al., 2020) and even personalised FL ones (Collins al., 2021: Hanzely et al., 2021) cannot be directly applied. For simplicity, we assume that all clients want to solve multi-class classification task with C E N classes. We discuss later how regression tasks can be encompassed in the proposed framework.
Methodology. The goal of the proposed methodology, coined FLIC, is to learn a personalised model for each client while leveraging the information stored by other clients' data sets despite the heterogeneity issue. To address this feature space heterogeneity, we propose to map clicnt's features into a fixed-dimension common subspace C R by resorting learned local embedding functions {Φ; : X; + +) ]eep- order to preserve feature vector) from the original data distribution,
semantical information (such as the class associated seek at learning the functions for)ie(t) such that they are aligned with (i.e. minimise their distance to) some learnable latent auchor distributions that are shared across all clients. These anchor distributions act as universal "calibrators" for clients, preventing similar semantic information from different clients from scattering across the subspace &. This scattering would otherwise impede proper subsequent federated learning of the classification model. As depicted in Figure 1, local embedding functions are learned by aligning (c) conditioned on the class E [C], with C learnable anchor measures (Mc)eEjo)- This alignment is achieved by minimising their distance.
the mapped probability distributions, denoted as v
Remark We want to stress the signéfieance of aligning the class-conditional probability distributions v with respect to the anchor distributions pe- Local and independent learning of embedding functions di by each client does not guarantee alignment of resulting probability distributions in the common subspace Q. As in unsupervised multilingual embeddings (Grave et al., 2019), alignments crucial for preserving semantic similarity of class information. Misalignment occurs when projecting class-conditionals in lower-dimensional space using an algorithm, like t-sne, that seeks al presercing only local similarilies. Indeed, dala from different clients are projected in a subspace in which different ciass-conditionlls may overlop. This is also the case when using a neural network with random weights as projector or an auto-encoder. Examples of such phenomena are illssrated in Figure 3. Notably, this figure shows that the alignment with respect to the anchor distribution is crucial to ensure that the class-condittonal distributions are aligned in the common subspace T.
4

---

Published in Transactions on Machine Learning Research (04/2024)
deploy Once data from the heterogencous spaces are embedded in the same latent subspace Φ, federated learning methodology for training from this novel representation space. While any standard FI. approach e.g. FedAvg (McMahan et al., 2017) can be used, we consider a personalised FL where each client has a local model tailored to its specific data distribution as statistical heterogeneitics that are still present in & (Tan et al., 2022). Hence, given the aforement ioned local embedding functions {6t), the model parameters {0, E [d and some non-negative weights associated to each client (wsJeE such that Σ W/ = 1, we consider the following empirical risk minimisation problem:
min f(01:b, (1.)) Σ (1., 010.0118
(1)
and for any i E [b],
n2
f((0, 0i) = 74
Σ j=1
Φ " 90.
(2) (o))
(2)
where ((.-)) stands for a classification loss function between the true label vf and the predicted one (0
(f) 90, (d, ( where 9ø, is the local model that admits personalised architecture parameterised by 0, and taking as input an embedded feature vector (xx))) Φ
Objective Function. At this stage, we are able to integrate the FL paradigm and the local embedding function learning into a global objective function we want to optimise, see (1). Remember that we want to learn the parameters {Oshep of personalised FL models, in conjuction with some local embedding functions (dihe()) and shared anchor distributions {Ne}. In particular, the latter have to be aligned with class-conditional distributions ) We enforce this aligummet via a Wasserstein regularisation term leading to regularised version of the empirical risk minimisation problem defined in (1), namely
Pis, Pi.c
arg min 010 1=1
Σ
FilOs, Pz, (100,,,.)r..),,
and for any i E [b],
J 1 Σ Σ
Fr dG (1.)) = wif((0, + A1ws
Σ W3 (pow//o) + A2wss Σ cE>;
(c.90.
(
(3)
eE);
j=1
where ;2 ;j
E [J]tetto] stand for samples drawn from {Meteetc), and λ1 A2 > 0 are regularisation parameters. The second term in (3) aims at aligning the conditional probability distributions of the transformed features to the anchors. The third one is an optional term aspiring to calibrate the anchor distributions with the classifier in cases where two more classes are still ambiguous after mapping onto the common feature space; it has also some benefits to tackle covariate shift in standard FL (Luo et al., 2021).
Design Choices and Justifications. In the sequel, we consider Gaussian anchor measures He = N(ve, Ee) where V. R and c E [C]. One of the key advantages of this Gaussian assumption that, under mild assumptions, it guarantees the existence of a transport map T(() such that T (v:) = μt, owing to Brenier's theorem (Santambrogio, 2015) since a mixture of Gaussians admits a density with respect to the Lebesgue measure. Hence, in our case, learning the local embedding functions boils down to approximating this transport map T(() by Φi. We also approximate the conditional probability mcasures {u/"; Vihej by using Gaussian measures ( N(m (0) ((0) ) E Dihee such that for any i E [6] and e E [C], m. (c) and (c) stand for empirical mcan vector and covariance matrix. The relevance of this approximation is detailed in Appendix $1.2.
These two Gaussian choices (for the anchor distributions and the class-conditional distributions) allow us to have a closed-form expression for the Wasserstein distance of order 2 which appears in (3), see e.g. Gelbrich (1990); Dowson Landau (1982). More precisely, we have for any i E [b] and E [C],
W (Mor v.)) = | m 12 + 222 (Ior
5

---

Published in Transactions on Machine Learning Research (04/2024)
Table 2: Current personalised FL techniques that can be embedded in the proposed framework. The parameters and 3 respectively stand for model weights while w E [0, 1]. Typically, stand for the shared weights associated to the shared layers of neural network architecture and B; for local ones aiming at performing personalised classification.
Algorithm
Local model
Local weights
Fe
98. 9o:
(6)
0,
L2GD
9o.
(1) = go.
0, wa + (1 w)B
FedRep
90. = 9b. o 9a
0; = |a, 3,1
where (:).)) denotes the Bures distance between two positive definite matrices (Bhatia et al., 2019).
Remark 2. Instead of Gaussian distribution approzimations, we can consider more complez probobility distributions. For instance, we can use a Gaussian mixture model (GMM) and still be able to compute cheaply the Wasserstein distance (Chen al., 2018). We can even make no hypothesis on the data distribution and compule the dislance for aligument using the linear programmingg based OT-formalation (Flamary el al., 2021) use any other IPM such as the marimum mean discrepancy (MMD), see Gretton et al. (2012). However, in practice, we found that the closed-form Wasserstein distance achieves slightly better performance thon MMD.
Regarding the parametrization the local embedding functions di, we consider a neural network architecture which takes into account the nature the input data and outputs to a vector of dimension k. For instance, for the digit datasets, we consider a convolutional neural network (CNN), followod by pooling and lincar layers whereas for the BCI datasets, we have used a fully connected neural network as the data has been preprocessed. The choice of the architecture is left to the user and can be adapted to the problem at hand.
Reference Distribution for Regression. For regression problem, except the change in the loss function, we also need to define properly the reference anchor distributions. Since our goal is to map all samples for all clients into a common latent subspace, in which some structural information about the regression problem is preserved. As such, in order to reproduce the idea of using a Gaussian mixture model as anchor distribution, we propose to use an infinite number of Gaussian mixtures which the distribution of a transformed feature vector (xx) R associated to response y is going to be mapped on a unit-variance Gaussian distribution whose mcan depends uniquely on y. Formally, define the anchor distribution as
= -(m))) IT)
where (v) is a vector of dimension k that is uniquely defined. In practice, we consider m( yo + (1 (1- (b-y)b where and b are two distinct vectors in R...
When training FLIC, this means that for a client i, we can compute W3 Myg v.) based on the set of training samples ((x, v))- In practice, if for a given batch of samples we have a single sample of value x, then the Wasserstein distance boils to [[dx(z) - m(v) [[, which means that we are going to map x to its corresponding vector on the segment [a, b).
4 Algorithm
As detailed in (3), we perform personalisation under the FL paradigm by considering local model architectures (i)
{go. heej and local weights 01:3- As an example, we could resort to federated averaging with fine-tuning (e.g. Fe (Collins et al., 2022)), model interpolation (e.g. L2GD (Hanzely & Richtárik, 2020; Hanzely et al., 2020)) or partially local models (e.g. FedRep (Collins et al., 2021) or the works of Oh et al. (2022); Singhal et al. (2021)). Table details how these methods can be embedded into the proposed methodology and explicitly shows the local model and local weights associated each algorithm as well as how the global weights are taken into account locally. Tn Algorithm 1, we detail the pseudo-code associated to specific
6

---

Published in Transactions on Machine Learning Research (04/2024)
instance of the proposed methodology when FedRep is resorted to learn model parameters (0shee(6) under the FL paradigm. For FedRep, each 0; includes local and global weights. Besides these learnable parameters, the algorithm also learns the local embedding functions 61 and the anchor distributions M1.C-
In practice at a given epoch t of the algorithm, a subset A S [6] of clients are selected to participate to (t) the training process. Those clients receive the current latent anchor distribution Pic and the current shared representation part of g() Then, each client locally updates Φi, O, and (1 The number of local steps M for each client is a hyperparameter of the algorithm. Afterwards, clients send back to the server their updated (t) version of g(+) and P1.C Updated global parameters g(++1) and M1.C' (++1)) are then obtained by weighted averaging of client updates on appropriate manifolds. The use of the Wasserstein loss in (3) naturally leads to perform averaging of the local anchor distributions via a Wasserstein barycenter; algorithmic details are provided in the next paragraph. In Algorithm 1, we use for the sake of simplicity the notation DescStep(F (((tee) to .1ic)) with respect to a subset of paramcters in (0;, oi, /1.C). This subset is specified the socond argument of DescStep. fully detailed version of Algorithu 1 is provided in the supplementary material (see Algorithm S2).
denote a (stochastic) gradient descent step on the function p(t.m) = F, (A((.m)) .6 ((.mm)
Updating and Averaging Anchor Distributions. In this paragraph, provide algorithmic details regarding steps 14 and 19 in Algorithm 1. For any e [C], the anchor distribution Me involves two Icarnable parameters namely the mean vector Un and the covariance matrix Eg. For the mean, step 14 stands for a (stochastic) gradient descent step aiming to obtain a local version of ve and step 19 boils down to compute the averaging of the mean vectors: ve ((+1)) = (6/ Ac+l) EiEA,,, W
((+1), For updating the covariance matrix, we resort to the Cholesky dccomposition to enforce the positive semi-definite constraint of the covariance matrix. Hence, we rewrite it as E. = LoLe where Le E R and optimise in step 14 with respect to the factor Le instead of Σe. Then, can handle the gradient computation of the Bures distance in step 14 using the work of Muzellee & Cuturi (2018); and obtain a local factor (((++) at itcration t. In step 19, we compute I((+1) = (b/lAz+1l) ΣieA, (([((+)) and set x((+1) = L. [z( When X2 = in (3), these mcan vector and covariance matrix updates exactly boil down to perform one stochastic (because of partial participation) gradient descent step to solve the Wasserstein barycenter problem arg min Ei= wi W (
Pre-training i1.0. Owing to the introduction of reference anchor distributions which carry semantical infor- mation, each local feature embedding function can be pre-trained by optimising the loss Eeer. W3 (ee o )
While in theory, pre-training may not be necessary, we believe that it helps reach a better solution of the federated learning problem parameters 0, are optimised starting from better latent representation of vø:- This is a phenomenon that has been observed in the context of fine-tuning (Kunar et al., 2022) or domain generalisation (Rame et al., 2022). Note that in practice, we initialize the mean vectors de of the anchor distributions by randomly sampling a Gaussian distribution with a very large variance. By doing so, we ensure a good separability of the classes and the anchor distributions in the latent space. The covariance matrices are initialized as identity matrices.
Remark 3. Regarding privacy, FL ensures that raw data never leaves the client device. Compared to the (t+1) base FL algorithm such as FedRep, FLIC exchanges between clients and server the anchor distribution Pi,1. This can be seen a further privacy risk. Howweerr the anchor distribution is a set of C Gaussian distributions, providing only aggregated information about the data and it needs to be clarified based on further work the level of privacy it brings.. A possible solution to miligate potential privacy leak is to consider differential privacy (Dwork & Roth, 2014) and use differentially privote version of the Wasserstein barycenter algorithm based on a differentially private Wasserstein distance (Lê Tien et al., 2019; Goldfeld & Greenewald, 2020; Rakotomamonjy & Ralaivola, 2021).
5 Aon-Asymptotic Convergence Guarantees in a Simplified Setting
Deriving non-asymptotic convergence bounds for Algorithm 1 in the general case is challenging since the considered C-class classification problem leads to jointly solving personalisod FL and foderated Wasserstein barycenter problems. Regarding the latter, obtaining non-asymptotic convergence results is still an active research area in the centralised learning framework (Altschuler et al., 2021). As such, we propose to analyse
7

---

Published in Transactions on Machine Learning Research (04/2024)
Algorithm FLIC for FedRep
Require: initialisation Pi:c: P1.b ((0) ,3160
for t= 0 to - do Sample set of Ar-1 of active clients.
for i E Ar+1 do
4:
// Communication to clients
The central server sends the global parameters (9 and PIC to A4t1 clients. // Update local porameters
(t)
for 0 to M - do
(l.+11) f- DescStep p(l.m) :p (l.m)) p(t.m)) g(t.m)
8 g( + DescStep
9:
(1+10)) 0; (1.,M)
10:
8. g(c+1,0) =3((.,0)
11:
12:
// Update global parameters
(t-1) + DescStep (((M (1)
13:
M.1:C + DescStep plaaa (2)
(I+1)
14:
15:
// Communication with the server
Send (++1) and Mi,1.C to central server.
(++1)
16:
17:
// //ueraggng global parameters
(e+1) |Aril ExEre, waa (aa(t+1)
18:
(++1) + Vassersteinarreenttrr (++1)
19:
M1C
Ensure: parameters (T)), Mi.0 Pib P1b
(T,0) g(T.D)
simpler regression framework where the anchor distribution is known beforehand and not learned under the FL paradigm. While we acknowledge that this theoretical analysis is based simplified setting of our approach, it still offers insightful perspective and we leave the general case for future work.
More precisely, we assume that x. N(m,, Mx) with m, E R& and M, g for i [b],] [nd]. In addition, we consider that the continuous scalar labels are generated via the oracle model ( (5) ob (a!²)) where A E Rxd & IRd and (+)- are the ground-truth parameters and the feature embedding functions, respectively. We make the following assumptions the ground truth, which are inherited from those of FedRep.
H1. (i) For any [b], k; 2 k.
(ii) For any i E [b), j E [nı], embedded features ot (x.,)) (z.( are distributed according to N(0,, Tg)-
(Mi Ground-truti truth model parameters satisfy || 1l dd for i [b] and A has orthonormal columns.
(iv) For any t E "O....,T-11. |Ar+1 with 1 < b' < b, and if we select b' clients, their ground-truth head parameters {8 hiEAu span Idd.
(u) In (2), (((.)) the l2 norm, Wi = 1/b, 0, [A, B.) and 96. (x) (AB,) x for x RR
Under HI-(i), Delon et al. (2022, Theorem 4.1) show that of can expressed as non-unique affine map with closed-form expression. To align with the true latent anchor distribution μ == N(O,, Ix); we propose estimate &e by leveraging this closed-form mapping between N(m,, E,) and μ.. Because of the non-unicity of of, we show in Theorem that we can only recover it up to a matrix multiplication. Interestingly, Theorem also proves that the global representation A(7) learnt via FedRep (see Algorithm S3 in Appendix) is able correct this feature mapping indetermination. Associated convergence behavior is illustrated in Figure 2 toy example whose details are postponed to Appendix $2.
Theorem Assume HI. Then, for any x; E R we have d1(2,) = QØ7 (2+) where Q E R of the form diag, (+1). Under additional technical assumplions delavild in Appendix S2, have for any E (0,..., T T-13
8

---

Published in Transactions on Machine Learning Research (04/2024)
10
Principal angle distance
10 -2
dist(A A') with FedRep
10-4
dist(A (4) AA*) with FLIC-FedRep
dist(A (0) QA") with FLIC-FedRep
10
10-8 8
100
200
300
400
500
Epoch t
Figure 2: Illustration of the convergence behavior of FLIC toy linear regression problem. The plot shows the principal angle distance bet ween the estimated linear global map A and the true linear global map A (up to a matrix multiplication). (blue) A is estimated based on the latent space feature mapped obtained using of and FedRep. (orange) A(t) is estimated based the latent space feature mapped obtained using á,. (green) A is estimated based the latent space feature mapped obtained using 6 but compared to QA*. Here, we show FLIC is able to recover the true global map A up to matrix multiplication.
and with high probability,
dist(A(t+)) ,(A)) S (1 (iis (AAA)
where (0,1) is detailed explicitly in Theorem 83 and dist denotes the principal angle distance.
6 Numerical Experiments
For numerically validating the benefits associated to the proposed methodology FLIC, we consider toy problems with different characteristics of heterogeneity; as well as experiments on real data, namely (i) a digit classification problem from images of different sizes, (ii) an object classification problem from either images or text captioning on clients, and (iii) a Brain-Computer Interfaces problem. Code for reproducing part of the experiments is available at https: tttp//gttttuu..tt/// com wwt.///rratttom//l/..
Baselines. Since the problem we are addressing novel, no FL competitor in the literature can serve as a baseline beyond local learning. However, propose to modify the methodology proposed by Makhija et al. (2022) and Collins et al. (2022) to make them applicable to clients with heterogeneous feature spaces. The method proposed by Makhija et al. (2022) handles local representation models with different architectures. Since their key idea to align the latent representations of fixed-dimeessonallty inputs shared by the server to all clients, we propose an alternative approach called FedHeNN, that works for clients with different feature spaces, where we build a Representation Alignment Dataset (RAD) based on the largest feature space and then prune to obtain lower-dimensional RAD for each client. We can also adapt the FedRep approach (Collins al., 2022) to our setting by considering local feature transformation followed by shared global representation model and local classification head. This approach, denoted as HetFedRep, maps the input data to fixed dimensionality before the shared global representation model. We can understand this approach as a special case of FLiC where the local feature transformation are not enforced to align with the anchor distributions. We adapted our network architecture to match the baselines by considering two variants. Following Makhija al. (2022), we treated all layers except the last as the representation learning module for a fair comparison. Therefore, our approach, the alignment applies the penultimate layer, and the last layer is the classifier layer. We call this model FlassClass. Additionally, introduced another model, called

---

Published in Transactions on Machine Learning Research (04/2024)
Class
Class 0
Class
Class
4-
0-
-10-
6-
-15
-8-
2.5 5.0
Figure 3: Projecting the Gaussian coasiconniisonnls stored on 20 clients obtained from the "added features' toy problem. For readability, the problem has been made easier by considering only 3 classes with all classes on clients. Means of the class-condittonall are drawn from a zero-mean Gaussian with a covariance matrix I'I with o = 3. Added features range from to 10. (left) t-sne projection. (right) Multi Dimensional Scaling (MDS) projection. Different tones of color represent the same class-conditionals of a given problem. From this figure, we remark the overlapping of classes regardless of the algorithm used for projection (t-sne uses random initialization and while MDS uses a deterministic ones). This emphasizes the need for proper alignement of the conditonn during projection.
FLIC-HLL similar to FedRep, but with an extra trainable global hidden layer with and 3, parameters for respectively the shared representation and classification layers.
Data Sets. We consider four different classification problems to assess the performances of our approach. For all simulations, we assume prior probability shift e.g each client will have access to data of only specific classes. The first problem is a toy problem with 20 classes and Gaussian class- cass-conniiional distributions, where we conduct two sub-experiments: adding random spurious features and applying a random linear mapping, both of random dimensionality on each client. Section $3.1 in the supplementary material provides more details about these data sets. The second problem involves digit classification using MNIST and USPS datasets, with dimensions of 28 x 28 and 16 x 16, respectively. Each client hosts a subset of either MNIST or USPS dataset. The third experiment addresses a multimodal problem using a subset of the TextCaps dataset (Sidorov et al., 2020), an image captioning dataset. We converted it into a 4-class classification problem with 12, 000 and 3, 000 examples for training and lesting, respectively, based on caption text or image. We used pre-trained models (Bert and ResNet) to embed the caption and image into 768-dimensional and 512-dimensional vectors. To create heterogeneity, Each client hosts either image text embeddings. Finally, the fourth problem is real medical problem denoted as Brain-Computer Interface (BCT) which consists in classifying mental imagery EFEG datasets into five classes. The datasets we have considered is based on six datasets from the mental imagery MOABB data repository (Jayaram & Barachant, 2018) (details are given in Section $3.1 in the supplement). Each those EEG datasets have been acquired on different subjects, have different number of channels and classes. We used a vectorised channel-deenndet covariance matrices representations of cach EEG signals as a feature (Yger et al., 2016; Barachant et al., 2022). Hence, the dimensionality of the feature space is different for each dataset. We have considered each subject in all the experiments client owning his own dataset. In practice, the number of training examples on client ranges from 30 to 600 while the dimensionality of the features goes from to 1,830.
randomly pruned 10% of features each client.
Illustrating the need for anchor distributions. The main bottleneck for applying FL algorithms to heterogencous feature spaces the lack of common space. However, one can argue that this common space can be created by projecting the data onto joint common space. As we have claimed, we illustrated here that this is not sufficient. To do so, we have considered the "added noisy features' toy problem with
10

---

Published in Transactions on Machine Learning Research (04/2024)
85.0
75.0
Local
82.5
FedHeNN
72.5
80.0
HetFedRep FLic Class
70.0
Accuracy (%)
Accuracy (36
77.5
FLic
67.5
75.0
65.0
Local
62.5
72.5
FedHeNN
60.0
HetFedRep
70.0
FLic Class
57.5
FLic
55.0
25
50
75
125 # Clients
100
150
175
200
25
50
75
100 125
150
175
200
# Clients
Figure Performance of FLIC and competitors on the toy data sets with respect to the number of clients. (left) Gaussian classes in dimension k = with added noisy feature. (right) Gaussian classes in dimension k 30, transformed by random linear mapping. Only 3 classes are present each client among the 20 possible ones.
the 3 class-conditionals over 20 clients. We have projected those class-conditionals onto a common space using different projection algorithms, namely I-sne (Van der Maaten & Hinton, 2008) and multi-dimensional scaling (MDS) (Kruskal, 1964). The results are shown in Figure 3. We can note that just projecting into a common subspace without taking into account the somantic of the class-conditionals leads to overlapping of classes. This emphasizes the need for a proper alignment of the class-conditionals during projection, based reference distribution for proper alignment of the class-conditionals by exhibiting the poor performance of FedRep trained on data from this common space.
propose in FLIC. In Appendix $3.7, we provide quantitative evidence of the need
Experimental Setting. For the experimental analysis, we use the codebase of Collins et al. (2021) with some modifications to meet our setting. For experiments, we consider T communication rounds for all algorithms; and at each round, client participation rate of 0.1. The number of local epochs for training has been set to M 10. As optimisers, have used Adam optimiser with a learning rate of 0.001 for all problems and approaches. Further details given in Section S3.3 in the supplement. For each component of the latent anchor distribution, we consider Gaussian with learnable mean vectors and fixed Identity covariance matrix. As such, the Wasserstein barycenter computation boils down to simply average the mean of client updates and for computing the third term (3), we just sample from the Gaussian distribution. Accuracies are computed as the average accuracy over all clients after the last epoch which local models are trained.
Results Toy Data Sets. Figure depicts the performance, averaged over 5 runs, of the different algorithms with respect to the number of clients and when only 3 classes are present in each client. Since the fixed amount of training data (2000 classes) are equally shared across clients holding a given class, we increase the number of clients, we expect the problem to be harder and the performance to decrease. For both data sets, we can note that for the noisy feature setting, FLIC improves on FedHeNN of about 3% of accuracy across the setting, performs better than local learning and is comparable to HetFedRep. For the linear mapping setting, FLIC achieves better than other approaches with gain of performance of about 4% while the gap tends to decrease as the number of clients increases. For this problem, HelFedRep fails to achieve good performance, highlighting the need for proper alignment of the dass-conditional. Since the local embedding functions are handled as any local layers and not enforced to align with the anchor distributions. Interestingly for both problems, FLIC-HL performs slightly better than FLIC-Class showing the benefit. of adding a shared representation layer a in addition to the local embedding functions and the classification layer.
Figure also illustrates how samples embedded with Φi evolve during training towards the anchor distributions P.C- At start, they are clustered client-wise (same marker with different. colors grouped together.) As
11

---

Published in Transactions on Machine Learning Research (04/2024)
Class
Class 1
Class Cless Class Class 11
Class
Class 5
Class Class 11
Class 6 Class
Mean vector μ4
Figure 5: 2D t-sme projection of 5 classes partially shared by 3 clients for the toy linear mapping dataset after learning the local transformation functions for (left) 10 epochs, (middle) 50 epochs, (right) 100 epochs. The three different markers represent the different clients while classes are represented by different color tones. The marker represents the class-conditional mean of the reference anchor distribution. We note that training set converges lowards those means.
Table 3: Performance over 3 runs of our FLIC model and the competitors real-data problems (Digits and TextCops data set).
Data sets (setting)
Local
FedHeNNf
HetFedRep
FLIC-Class
FLIC-HL
Digits (b 100, 3 Classes/cliee /
97.49 + 0.4
97.45 + 0.5
57.85 + 1.4
97.83 + 0.3
97.70 = 0.2
Digits (b = 100, 5 Classes /client)
96.15 + 0.3
96.15 + 0.2
54.87 + 5.0
96.46
96.54 0.7
Digits (b 200, 3 Classes cclient)
93.33 + 0.2
93.40 + 0.4
67.99 + 2.2
94.50 + 0.3
94.51 0.3
Digits (b 200, 5 Classes/ /client
87.48 = 1.5
87.22 + 1.8
48.88 + 3.0
91.11 + 0.6
91.10 = 0.7
TextCaps (b 100, 2 Classes /client))
84.19 + 0.8
83.99 + 0.7
87.05 + 0.7
89.14 + 1.1
89.68 = 0.7
TextCaps 100, Classes (casses/client)
76.04 0.8
75.39 + 0.9
77.99 0.6
81.27 0.2
81.50 = 0.2
TextCaps (b 200, 2 Classes
Classes/cliens cclient)
83.78 + 1.8
83.89 + 1.7
85.48 + 1.5
87.73 + 0.8
87.74 = 1.3
Tex1Caps 200, Classes /clsent/client)
74.95
74.77 1 1.0
75.73 0.8
79.08 0.7
78.49 1 0.7
BCI (b=54)
73.51 + 0.8
70.84 + 1.0
75.03 + 0.6
75.17
0.9
76.27 = 0.2
BCI (b=40)
73.98 + 0.2
71.48 + 0.6
74.23 + 0.7
75.09 1 1.0
75.82 0.3
iterations go, samples converge towards the relevant class wise anchor distributions, For instance, we note that after 50 epochs, samples from the class 5 are converging towards the mean of the related anchor distribution. After 100 epochs, they are almost all concentrated around the mean.
Results on Digits and TextCaps Data Sets. Performance, averaged over 3 runs, of all algorithms on the real-word problems are reported in Table For the Digits data set problem, we first remark that in all sitnations, FT ahgo performs bit better than local learning. In addition, both variants of FLIC achieve better accuracy than competitors. Difference in performance favor of FLIC reaches 3% for the most difficult. problem. For the TerlCaps data set, gains in performance of FLIC-HL reach about 4% across settings. While FedHeNN and FLIC algorithms follow the same underlying principle (alignment of representation in a latent space), our framework benefits from the use of the latent anchor distributions, avoiding the need of sampling from the original space. Instead, FedHeNN may fail as the sampling strategy of their RAD approach sulfers from the curse of dimensionality and does not properly lead to a successful feature alignment.
Results on Brain-Computer interfaces. For the BCI problem, performances are Table 3. In this case, the Local model performance corresponds also to the usual BCI performance measure as models are usually subject-specific We can note that. FLIC-HI. achieves better performance than all competitors with gain of
12

---

Published in Transactions on Machine Learning Research (04/2024)
Table 4: Performance over 3 runs of our FLIC model and the competitors on the same real-world problems when processed so as to have same input sizes.
Data sets (setting)
Local
FedHellNN
HetFedRep
FedRep
FLIC-Class
FLIC-HL
Digits- Resize 100, Classes)
97.70 0.1
97.62 L
92.63 1.8
95.76 1
98.14
98.05 0.1
Digits-Rest Resize (b 100, Classes)
96.36 0.1
96.37 = 0.2
94.90 0.7
96.07 = 0.5
96.94 0.1
96.91 0.1
Digits- Digits-Resize Resize 200, Classes)
93.62
0.3
93.56 + 0.2
69.21 3.0
92.93 + 0.8
94.73
0.4
94.54 0.2
Digits Resize (h- = 200, Classes)
87.74 = 1.3
87.49 + 1.0
69.27 1.8
94.57 + 0.9
91.40 + 0.4
91.16 0.3
Text Clip (b 100, Classes)
96.55 0.5
96.44 + 0.5
96.45
0.4
82.31 = 2.6
96.59 + 0.5
96.65 0.4
Text-aps (b 100, Classes
94.38
0.2
94.21 1
94.13 0.2
76.47 1
94.34 0.3
94.21 0.3
TextCaps--lipp (b 200, Classes)
96.27
96.16
L 0.4
96.36 0.2
83.33
1 1.7
96.55 0.2
96.44 0.2
TextCaps-Clip (6 200, Classes)
93.78 0.2
93.55 1 0.2
93.74 0.4
73.23 1 0.5
94.04 0.2
93.88 0.3
BCT-common (b=40)
71.24 0.6
70.63 + 0.4
72.09 0.1
71.50 + 0.3
72.17 + 0.3
72.12 + 0.6
about 3% of accuracy compared to BCT baseline. In addition, we pave the way to learning BCT models with heterogeneous datasets.
Preprocessing datasets to same dimensionality. Some preprocessing can be applied to the above datasets so that standard 'same dimensionality" FL methods can be considered. We can apply simple resizing of image. For the TextCaps classification problem can extract features from multimodal embedder such CLIP Radford et al. (2021b) in which embeddings between text and images are aligned and of dimension 512. For the BC1 problem, we used EEGs acquired from the set of common electrodes. This experiment also serves at proving that FLIC not only useful for hetcrogencous feature spaces but can be applied when feature spaces are identical. In such a case, we expect to benefit from the aligument of the class-conditionals to enhance performances. We have run the same experiments as above but using these preprocessings and also compared to plain FedRep and report the results in Table 4. FLIC still achieves slightly better performance than competitors 7 out of 9 settings, but more importantly, one should highlight the gain in performance on BCT problems when considering all sensors at disposal.
7 Conclusion
We have introduced novel and general framework, referred to as FLIC, for personalised FL when clients have heterogeneous feature spaces. Under this framework, we proposed a FL algorithm involving two key components: (i) local feature embedding function; and (ii) a latent anchor distribution which allows to match similar semantical information from each client. Experiments on relevant data sets have shown that FLIC achieves better performances than competing approaches. Finally, we provided theoretical support to the proposed methodology, notably via a oon-asympootic convergence result.
Limitations and Broader impacts of FLIC. One main limitation of FLIC is that it requires common feature space to be defined with ad-hoc dimensionality. While this dimensionality can be chosen by the user, it is not clear how to select it in practice and has to be set to default value (in our case 64). In addition, the proposed approach has a computational overhead due to the need of learning the local embedding functions. In addition, it is worth noting that the proposed approach is a stateful algorithm that is applicable in a cross-silo silo setting. FLIC has the potential to widen the scope of privacy-aware FL applications by allowing clients to have heterogeneous feature spaces. This is particularly relevant for medical applications where data are collected from different sources and may have different formats.
13

---

Published in Transactions on Machine Learning Research (04/2024)
References
Mokhtar Z Alaya, Maxime Bérar, Gilles Gasso, and Alain Rakotomamonjy. Theoretical guarantees for bridging metric measure embedding and optimal transport. Neurocomputing, 468:416-430, 2022.
Jason Altschuler, Sinho Chewi, Patrik Robert Gerber, and Austin J Stromme. Averaging on the Bures- Wasserslein manifold: dimension-free convergence of gradient descent. In Advances in Neural Information Processing Systems, 2021.
David Alvarez-Melis and Tommi Jaakkola. Gromov-Wassersten Alignment of Word Embedding Spaces. In Conference on Empirical Methods in Natural Language Processing, pp. 1881 1890, 2018.
David Alvarez-Melis, Stefanie Jegelka, and Tommi S. Jaakkola. Towards Optimal Transport with Global Invariances. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), International Conference on Artificial Intelligence and Statistics, volume 89, pp. 1870-1871) 2019.
Nick Angelou, Ayoub Benaissa, Bogdan Cebere, William Clark, Adam James Hall, Michael A Hoeh, Daniel Liu, Pavlos Papadopoulos, Robin Rochm, Robert Sandmann, et al. Asymmetric private set intersection with applications to contact tracing and private vertical federated machine learning. ar preprinl arXiv:2011.09550, 2020.
Arindam Banerjee, Srujana Merugu, Inderjit S. Dhillon, and Joydeep Ghosh. Clustering with Bregman Divergences. Journal of Machine Learning Research, 6738):1705-1749, 2005.
Alexandre Barachant, Quentin Barthélemy, Jean- King, Alexandre Gramfort, Sylvain Chevallier, Pedro L. C. Rodrigues, Emanuele Olivetti, Vladislav Goncharenko, Gabriel Wagner vom Berg, Ghiles Reguig, Arthur Lebeurrier, Erik Bjäreholt, Maria Sayu Yamamoto, Pierre Clisson, and Marie- Corss-CCnsttecee Corsi. pyrriemamm/prriemannt vO.3, July 2022. URL https: pttpp::/doi..oo wtt.or/////..//eeeeeee 7547583.
httpp::/da...rrr//0.5///////
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the Bures-Wee Wasserstsn distance between positive definite matrices. Expositiones Mathematicae, 37 2019.
Charlotte Bunne, David Alvarez-Melis, Andreas Krause, and Stefanie Jegelka. Learning Generative Models across Incomparable Spaces. In International Conference on Machine Learning, volume 97, pp. 851 861, 2019.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konevený, H Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812.01097, 2018.
Wei-Ning Chen, Christopher A Choquette Choo, Peter Kairouz, and Ananda Theertha Suresh. The Fundamental Price of Secure Aggregation in Differentially Private Federated Learning. In International Conference on Machine Learning, volume 162, pp. 3056 3089, 2022.
Yongxin Chen, Tryphon T Georgion, and Allen Tannenbaum. Optimal transport for gaussian mixture models. IEEE Access, 7:6269 6278, 2018
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting Shared Representations for Personalized Federated Learning. In International Conference on Machine Learning, pp. 2099-2099, 2021.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. FedAvg with Fine Tuning: Local Updates Lead to Representation Learning. In Advances in Neural Information Processing Systems, 2022.
Julie Delon, Agnes Desolneux, and Antoine Salmona. Gromov Wasserstein distances between Gaussian distributions. Journal of Applied Probability, 59(4):1178 1198, 2022.
Li Deng. The MNIST Database of Handwritten Digit Images for Machine Learning Research. IEEE Signal Processing Magazine, 4144:1111,44, 2012.
14

---

Published in Transactions on Machine Learning Research (04/2024)
Enmao Diao, Jie Ding, and Vahid Tarokh. HeteroFL: Computation and Communicttion Efficient Federated Learning for Heterogeneous Clients. In International Conference. on Learning Regresenlations, 2021.
D.C Dowson and B.V Landau. The Fréchet distance between multivariate normal distributions. Journal of Multivariate Analysis, 12(3):450 455, 1982.
Paul-Ambroisc Duquenne, Hongyu Gong, and Holger Schwenk. Multimodal and Multilingual Emboddings for Large-Scale Speech Mining. In Advances in Neural Information Processing Systems, volume 34, pp. 15748- 141411 2021.
Cynthia Dwork and Aaron Roth. The Algorithmic Foundations of Differential Privacy. Foundations and Trends in Theoretical Computer Science, 9(3 4):211 407, 2014.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fourmier, Léo Gautheron, Nathalie T.H. Gayrand, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1 2021. URL http::/jmmiii org/ 2ttp/ppprrs/22222055555 html.
Dashan Gao, Yang Liu, Anbu Huang, Ce Ju, Han Yu, and Qiang Yang. Privacy-preserving Helerogeneous Federated Transfer Learning. In JEEE International Conference on Big Data (Big Data), pp. 2552-2556. 2019.
Matthias Gelbrich. On Formula for the L2 Wasserstein Metric between Measures on Euclidean and Hilbert Spaces. Mathematische Nachrichten, 147(():185-208 1990.
Ziv Goldfeld and Kristjan Greenewald. Gaussian-smoottee optimal transport: Metric structure and statistical efficiency. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 3327 3337. PMLR, 26 html.
Aug 2020. URL https: https://prott //ppc//ee/oc/ngsss ..s.......... mww. wat..preessvv001. got.///olleeld..ae..mo
Edouard Grave, Armand Joulin, and Quentin Berthet. Unsupervised alignment of embeddings with wasserstein procrustes. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 1880 1890. PMLR, 2019.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):7237777 7733-773, 2012. URL http://jmlr.... /ttp:////m/rrr////////// papers/v13/grettontt html.
Filip Hanzely and Peter Richtárik. Federated learning of mixture of global and local models. ai preprint 0556:00000. 2020.
Filip Hanzely, Slavomír Hanzely, Samnel Horváth, and Peter Richtárik. Lower bounds and optimal algorithms for personalized federated learning. arXiv preprint arXiv:2010.02372 2020.
Filip Hanzely, Boxin Zhao, and Mladen Kolar. Personalized Federated Learning: A Unified Framework and Universal Optimization Techniques. arXiv preprint arXiv: 2102.09749. 2021.
Stephen Hardy, Wilko Henecka, Hamish Ivey-Law, Richard Nock, Giorgio Patrini, Guillaume Smith, and Brian Thore. Private federated learning vertically partitioned data via entity resolution and additively homomorphic encryption. arXiv preprint arxiv:1711. 10667, 2017.
Junyuan Hong, Haotao Wang, Zhangyang Wang, and Jiayu Zhou. Efficient Split-Mix Federated Learning for On-Demand and In-Situ Customization. International Conference on Learning Representations, 2022.
J. J. Hull A database for handwritten text recognition research. IEEE Transactions on Pattern Analysis and Machine Intelligence, 16(5):550 15445545-540-.- 1994.
15

---

Published in Transactions on Machine Learning Research (04/2024)
Prateek Jain, Pranceth Netrapalli, and Sujay Sanghavi. Low-Rank Matrix Completion Using Alternating Minimization. In ACM Symposium on Theory of Compuling, pp. 665 674, 2013.
Vinay Jayaram and Alexandre Barachant. Moabb: trustworthy algorithm benchmarking for beis. Journal of neural engineering, 1566)106 2018.
Yihan Jiang, Jakub Konevený, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint 19X.::1999.1.// 2019.
Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mchdi Bennis, Arjun Nitin Bhagoji, K. A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G.T. D'Oliveira, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang Hc, Lie Hc, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konevený, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramer, Praneeth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and Open Problems in Federated Learning. Foundations and Trends in Machine Learning, 14 (1 21:112100, 2021.
Mikhail Khodak, Maria-Florina F Balcan, and Ameet S Talwalkar. Adaptive gradient-based meta-learning methods. Advances in Neural Information Processing Systems, 32:5917-5928, 2019.
Joseph B Kruskal. Multidimensional scaling by optimizing goodness of fit to nonmetric hypothesis. Psychometrika, 29(1):1-27, 1964.
Ananya Kumar, Aditi Raghunathan, Robbic Matthew Jones, Tengyu Ma, and Percy Liang. Finc-tuning can distort pretrained features and underperform out-of-distribution. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL
https://openreww httpp://pppparreeeeaaaa/c/oaaaad net//forum/dd-UUnee222233
Nam Lê Tien, Amaury Habrard, and Marc Scbban. Differentially private optimal transport: Application to domain adaptation. In IJCAI, pp. 2852 2858, 2019.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated Optimization in Heterogeneous Networks. In Machine Learning and Systems, volume 2, pp. 429-450, 2020.
Yang Liu, Yan Kang, Chaoping Xing, Tianjian Chen, and Qiang Yang. A Secure Federated Transfer Learning Framowork. IEEE Intelligent Systems, 35 2020.
Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, and Jiashi Feng. No Fear of Hetcrogeneity: Classifier Calibration for Federated Learning with Non-IID Data. In Advances in Neural Informalion Processing Systems, volume 34, 2021.
Zhihan Lv, Liang Qiao, Qingjun Wang, and Francesco Piccialli. Advanced Machine-Learning Methods for Brain-Computer Interfacing. IEEE/ACM Transactions on Computational Biology and Bioinformatics, 18 (5):1688- 1698, 2021.
Disha Makhija, Xing Han, Nhat Ho, and Joydeep Ghosh. Architecture Agnostic Federated Learning for Neural Networks. In International Conference on Machine Learning, volume 162, pp. 14860-14870 2022.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication- Efffcient Learning of Deep Networks from Decentralized Data. In International Conference on Artificial Intelligence and Statistics, volume 54, pp. 2273-1282, 2017.
Facundo Mémoli. Gromor-Wassersteh distances and the metric approach object matching. Foundations of compulational mathemalics, 144444117-487, 2011.
16

---

Published in Transactions on Machine Learning Research (04/2024)
Fan Mo, Hamed Haddadi, Kleomenis Katevas, Eduard Marin, Diego Perino, and Nicolas Kourtellis. PPFL: Federated Learning with Trusted Execution Environments. In International Conference on Mobile Systems, Applications, and Services, pp. 94 108, 2021.
Junki Mori, Isamu Teranishi, and Ryo Furukawa. Continual Horizontal Federated Learning for Heterogeneous Data. Xáv preprint arx:::2003. 02100,,02110..00 2022.
Boris Muzellec and Marco Cuturi. Generalizing Point Embeddings using the Wasserstein Space of Elliptical Distributions. In Advances in Newral Information Processing Systems, volume 31, 2018.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading Digits in Natural Images with Unsupervised Feature Learning- In NIPS Workshop on Deep Learning ond Unsupercised Feature Learning, 2011.
Jachoon Oh, SangMook Kim, and Se- Young Yun. FodBABU: Toward Enhanced Representation for Federated Image Classification. In International Conference on Learning Representations, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Tlya Sutskever. Learning Transferable Visual Models From Natural Language Supervision. In International Conference on Machine Learning, volume 139, pp. 8748 8763, 2021a.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning, pp. 8748 8763. PMLR, 2021b.
Alain Rakotomamonjy and Liva Ralaivola. Differentially private sliced wasserstein distance. In International Conference on Machine Learning, pp. 8810 8820. PMLR, 2021.
Alexandre Rame, Matthicu Kirchmeyer, Thibaud Rahier, Alain Rakotomamonjy, Patrick Gallinari, and Matthieu Cord. Diverse weight averaging for out-of-distribution generalization. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 10821-10836. Curran Associates, Inc., 2022. URL https::/proceeeeeeee wew.reeeddddnneeennn..eeeeee paper paper_filess
Thomas Rippl, Axel Munk, and Anja Sturm. Limit laws of the empirical Wasscrstein distance: Gaussian distributions. Journal of Multivoriate Analysis, 151:90-109, 2016.
Daniele Romanini, Adam James Hall, Pavlos Papadopoulos, Tom Titcombe, Abbas Ismail, Tudor Cebere, Robert Sandmann, Robin Roehm, and Michael A Hoeh. PyVertical: A Vertical Federated Learning Framcwork for Multi-headed SplitNN. arXiv preprint arXiv:2104.00488 2021.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized Federated Learning using Hypernetworks. In International Conference on Machine Learning, volume 139, pp. 9489 9502 2021.
Shreya Sharma, Chaoping Xing, Yang Liu, and Yan Kang. Secure and Efficient Federated Transfer Learning. In IEEE International Conference on Big Dala (Big Dala), pp. 1569-2576. 2019.
Oleksäi Sidorov, Roughang Hu, Marcus Rohrbach, and Amanpreet Singh. Textcaps: a dataset for image captioning with reading comprehension. In European Conference on Computer Vision, pp. 742-758, 2020.
Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, John Rush, and Sushant Prakash. Federated Roconstruction: Partially Local Federated Learning. In Advances in Neural Information Processing Systems, volume 34, pp. 11220 11232, 2021.
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards Personalized Federated Learning. IEEE Transactions on Neural Networks and Learning Systems, pp. 1-17, 2022.
17

---

Published in Transactions on Machine Learning Research (04/2024)
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9111), 2008.
Veldhuis. The centroid of the symmetrical Kullback-Leible distance. IEEE Signal Processing Letters, (3):96-99, 99, 2002.
Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Dala Science. Cambridge University Press, 2018.
Cedric Villani. Optåmal Transport: Old and New. Springer Berlin Heidelberg, 2008.
Hongteng Xu, Dixin Luo, Ricardo Henao, Svati Shah, and Lawrence Carin. Learning Autoencoders with Relational Regularization. In International Conference on Machine Learning, volume 119, pp. 10576 10586, 2020.
Qiang Yang, Yang Liu, Tiajjaan Chen, and Yongxin Tong. Federated Machine Learning: Concept and Applications. Transactions on Intelligent. Systems and Technology. 10(2), 2019.
Rui Ye, Zhenyang Ni, Chenxin Xu, Jianyu Wang, Siheng Chen, and Yonina C Eldar. Fedfm: Anchor-based feature matching for data heterogeneity in federated learning. ar preprind orXiv:2210.0761. 2022.
Florian Yger, Maxime Berar, and Fabien Lotte. Riemannian approaches in brain-computer interfaces: review. IEEE Transactions Neural Systems and Rehabilitation Engineering, 25(10):1753 1-6217-7---12,, 2016.
Jie Zhang, Song Guo, Xiaosong Ma, Haozhao Wang, Wenchao Xu, and Feijie Wu. Parameterized Knowledge Transfer for Personalized Federated Learning. In A. Beygelzimer, Y. Dauphin, Liang, and Wortman Vaughan (cds.), Advances in Neural Information Processing Systems, 2021a.
Lin Zhang, Yong Luo, Yan Bai, Bo Du, and Ling- Yu Duan. Fodcrated learning for non-iid data via unified feature learning and optimization objective alignment. In Proceedings of the IEEE/CVE international conference on computer vision, pp. 4420 4428, 2021b.
Tailin Zhou, Jun Zhang, and Danny Tsang. FedFA: Federated Learning with Feature Anchors to Align Feature and Classifier for Heterogeneous Data. arXiv preprint Xix: 22211 1.09299, 2022.
18

---

Published in Transactions on Machine Learning Research (04/2024)
Appendices
Personalised Federated Learning On Heterogeneous Feature Spaces
Notations and conventions. We denote by B(R)) the Borel a-field of R2, M(R)) the set of all Borel measurable functions f on I and |I-II the Euclidean norm on R For μ a probability mcasure on B(R")) and ʃ M(R*) a y-integrable function, denote by plf) the integral of f with respect to (w.r.t.) μ. Let p and v be two sigma- finite measures on (RR B(RR)). Denote by μ < v if μ is absolutely continuous w.r.t. and du/dv the associated density. We say that Ç is transference plan of μ and if it is a probability measure on (Rd x R2, B(Rd x ())) such that for all measurable set A of I2, <(A R =(A) and C(R4 x A)=v/h v(A). We denote by T(u, v) the set of transference plans of J and v. In addition, we say that a couple of -a-ddoo variables XX,Y) is a coupling of and if there exists Ç E T(p, p) such that (X, Y) are distributed according to Ç. We denote by (R()) the sot of probability mcasures with finite I-moment: for all 1 E Pi [Rd): [==]dd[(()) < oo. We denote by (2((R)) the set of probability measures with finite 2-moment: for all E (2 [l=]]ddp(x) < co. We define the squared Wasserstein distance of order associated with ||- for any probability measures J4, A.V.E P2(R by
W3(m, = inf [ (ET(p.)) JRIxRd
| - y/12d((x. y)
By Villani (2008, Theorem 4.1), for all p. V probability measures on Rd, there exists transference plan Ç* E T(p, v) such that for any coupling (X,Y) distributed according to 5" W2(u, Bi||= y//3|//2, This kind of transference plan (respectively coupling) will be called an optimal transference plan (respectively optimal coupling) associated with W2. By Villani (2008, Theorem 6.16), P2 (2 equipped with the Wasserstein distance W2 is a complete separable metric space. For the sake of simplicity, with little abuse, we shall use the same notations for a probability distribution and its associated probability density function. For n 2 1, we refer to the set of integers between and n with the notation [n]. The d-multiddmensional Gaussian probability distribution with mean j Rd and covariance matrix Σ E Raxd is denoted by N(J,, E).. Given two matrices M, N Rxxd the principal angle distance between the subspaces spanned by the columns of M and N is given by N), = [ Nl [| M12 where M, N are orthonormal bases of (pan(M) and Span( respectively. Similarly, N are orthonormal bases of orthogonal complements Sp)n(M)+ and Span(N)* respectively. This principal angle distance is upper bounded by 1, see Jain et al. (2013, Definition 1). Outline. This supplementary material aims at providing the interested reader with further understanding of the statements pointed out in the main paper. More precisely, in Appendix S1, we support the proposed methodology FLIC with algorithmic and thcoretical details. In Appendix S2, we prove the main results stated in the main paper. Finally, in Appendix S3, we provide further experimental design choices and show complementary numerical results.
S1 Algorithmic and Theoretical Insights
In this section, we highlight alternative but limited ways to cope with feature space heterogeneity; and justify the usage, in the objective function (3) of the main paper, of Wasserstein distances with empirical probability distributions instead of true ones. In addition, we detail the general steps depicted Algorithm 1.
S1.1 Some Limited but Common Alternatives to Cope with Feature Space Heterogeneity
Depending on the nature of the spaces (K)e the feature transformation functions {or can be either known beforchand or more difficult to find. As an example, for any i E [b], X S X, then we can set mask functions as feature transformation functions in order to only consider features that are shared across all the clients. Besides, we could consider multimodal embedding models to perform feature transformation
19

---

Published in Transactions on Machine Learning Research (04/2024)
each client (Duquenne et al., 2021). For instance, if clients own either pre-processed images text of titles, descriptions and tags, then we can use the Contrastive Language-Image Pre- Tre-Trainnn (CIIP) model as feature transformation function (Radford et al., 2021a). These two examples lead to the solving of a classical (personalised) FL problem which can be performed using existing state-of-the-art approaches. However, when the feature translormation functions cannot be easily found beforehand, solving the FL. problem at stake becomes more challenging and has never been addressed the federated learning literature so far, up to the authors' knowledge.
S1.2 Use of Wasserstein Losses Involving Empirical Probability Distributions
Since the true probability distributions {u!o) ;e E Ji)se( are unknown priori, we propose in the main paper to estimate the latter using [ Mile(()) and to replace W3 (Pe, v by W (Me ( in the objective function (3) in the main paper. As shown in the following result, this assumption theoretically grounded when the marginal distributions of the input features are Gaussian.
Thcorem S2. For any E [b] and c E [C), let ( = -D/ where D( denotes the subset of the local data set D; only involuing observations associated to the label Besides, assume that vid is Gaussian with mean vector m E IR and full-rank covariance matrix 21(0) Rk Then, have in the limiting case n! 00,
(Me (Me vo.)) w3 Mes 0/00) V n (c) (W2
in distribution 2
E!! rjj with (Kg.rahsems standing for (eigenvvalue, eigenvector) pairs of the symmetric covoriance matrix 1
where z/0) ~ N(O, (0) and s = ( -m)
Proof. The proof follows from Rippl et al. (2016, Theorem 2.1) with the specific choices = 1.0 He and i (c) which are defined in Section 3 in the main paper.
S1.3 Detailed Pseudo-Code for Algorithm
In Algorithm $2 we provide algorithmic support Algorithm the main paper by detailing how to perform each step. Note that we use the decomposition LL to enfore the positive semi-definite constraint for the covariance matrix E.
20

---

Published in Transactions on Machine Learning Research (04/2024)
Algorithm S2 Detailed version of FLIC when using FedRep
.1%00 and step-size vy fj for fj Require: initialisation a Pi.c [ " %.ol with Ee ( 1
(0.0),
1: for 0 to T 1 do
2: Sample
of Ac41 of active clients.
3: for E Ar41 do
4: The central server sends a and (t) Act1.
5:
// Update local parameters
6:
for = 0 to M -1 do
Sample a fresh batch (6.m)) of nf samples with nf [ma]- ((.m) and CE 3% via z(retim) (1) + 18 where <(C.m) ~ Ne).
T:
Sample z(z...m) Me for j E t+1
8:
of (f++l))
(tmm)
Σ
Voll (1 g ofumm (o)(m) (x. (x.2))
9:
JE
1
Σ V. W
M
(4.me)
cE
10:
2 gft.m+1) x- g.(.m)
Bj
sel (fim)
11:
with Bj V
Vi
()) (1)
g
((tim))
pi
(2(0)] (201) 1722 Eneo, V (i
.9 ((0 p(tes [Zcilooo)
( ((t+110) (1.
8! p((41,u) g(:.br)
1/ Update global perameters
((+1) + (1)
nh
Σ Aj
15:
141
JE((LMM)
with A, ( Val Me
())
x ))
17 E.Es: Vol Mi oawso 22hcaao))3
( (0
16:
for 1 to C do
//m0 (c.1) (.(e.))
are the empirical mean and covariance of dlc.M) (240)) for j l+1
(1.M)
19:
1/ is the mean of anchor distribution. ple
Update ((00)) $1 using d.t,M).
20:
((41)) 1 - V.. oo!' m. m//00/1'
ndo Ecer
n,
Σ Voll (ice(1))
11 (alt) [2(uaamm)
jeilhm)
((41) ( Adi VI,
23 23
(IL) ,,(00)))))))) s((r1)
1112 ExEr,
The
Σ Vinl 10(0)
(ac0 s(ttae) (20u
(5.m))
jelffim
23:
// Communication with the server (t+1) and l((+1) to central server. Send af ((+1) 1 1.1
24:
1,C // Averaging global parameters
26:
(++1) |A,1l Eiere wia; wia
27:
to C do for
28:
( ((+1) (b/lAr+1) Σ Eiere (1t(t+1) (t+1)
29:
[(--)) = (b/lAc+1l) ExeAch w:(((t+)) and set s(++1) (0+1) 1 11+1)TT
Ensure: parameters aT) (T)
11:0: ( ((T.0)
S1.4 Additional Algorithmic Insights
Scalability. When the number of classes C' is large, both local computation and communication costs are increased. In this setting, we propose to partition all the classes into Cineta < meta-classes and consider reference mcasures {Je)oE(Cmenul associated to these meta-classes. As an example, if we are considering a dataset made of features associated to animals, the meta-class refers to an animal (e.g. dog) and the class refers to specific breed (e.g. golden retriever).
21

---

Published in Transactions on Machine Learning Research (04/2024)
Privacy Consideration. As other standard (personalised) FL algorithms, FLIC satisfies first-order privacy guarantees by not allowing raw data exchanges but rather exchanges of local Gaussian statistics. Note that FLIC stands for a post-hoc approach and can be combined with other privacy/confin /confideitti techniques such as differential privacy (Dwork & Roth, 2014), sccure aggregation via secure multi-party computation (Chen et al., 2022) or trusted execution environments (Mo et al., 2021).
Inference on New Clients. When client who has not participated to the training procedure appears, there is need to re-launch potentially costly federated learning procedure. Instead, the server sends the shared parameters {a("), (100) } the new client and the latter only needs to learn the local parameters
(os.
S2 Proof of Theorem 1
This section aims at proving Theorem in the main paper. To this end, we first provide in Appendix S2.1 a closed-form expression for the estimated embedded features based the features embedded by the oracle. Then, in Appendix $2.3, we show technical lemmata that will be used in Appendix $2.2 to show Theorem 1.
To prove our results, we consider the following set of assumptions.
H1. (i) For any i [b], E [ma], ground-truth truth embedded features of (z!?) distributed according to N(Ok, I&).
(ii) Ground-truth model parameters satisfy || vd for i E [b] and A has orthonormal columns.
(iii) For any t E fo....,T 1}. AAtzl = [rb] with < [bb] S b, and if we select [rb] clients, their
ground-truth head parameters {B) NEALT spon RRd
(in) In (2) in the main paper, :) is the & norm, w, 1/b, 0, (.) and ga. (x) (AB,) x for R.
$2.1 Estimation of the Feature Transformation Functions
As in Section 5 in the main paper, we assume that z/()) N(m,, Σ.) with m, E Rº and E, R& for E [ [n:]- In addition, we consider that the continuous scalar labels generated via the oracle model ()) = (A*B;)) oj (x(")) where A E Rx , { E TRd and (-) are ground-truth parameters and feature transformation function, respectively. Under H1-(i), the oracle feature transformation functions {ør heep are assumed to map k-ddimensional Gaussian distributions N(m,, Σ,) to common k-dimension Gaussian N(O,, Tx)- As shown in Delon et al. (2022, Theorem 4.1), there exist closed-form expressions for { hepp) Gromov-Wasserster plans. More precisely, these oracle feature transformation stand for affine maps and are of the form, for any e [b].
which can be shown to stand for solutions of Gromov- Wasserstein problem restricted to Gaussian transport
of (2. (1) = j(
(D())) (
Oc (x) m
where f(i..) = diag, (+1) is a k-dimensional diagonal matrix with diagonal elements in (-1,1), E, = P.D.P. is the diagonalisation of Σ, and (a)) stands for the restriction of D, to the first & components. In the sequel, we assume that all oracle feature transformation functions share the same randomness, that is p(t.*) = I = (aaag((#))
For the sake of simplicity, we assume that we know the true latent distribution of oj (x(")) and such consider pre-fixed reference latent distribution that equals the latter, that is μ N(O,, Tx). Since we know from Delon et al. (2022, Thcorem 4.1) that there exist mappings between Gaussian distributions with supports associated to different metric spaces, propose estimate for the ground-truth feature transformation functions defined by for any i E [b].
a. ( (1) [2 0k {
(a((0) m,)
where 7 = (+1g)(++), By noting that 7 = =îī;, where Q E blck is a diagonal matrix of the form diag, (+1), it follows that
O. (2 x = -0 x
(S1)
22

---

Published in Transactions on Machine Learning Research (04/2024)
Algorithm S3 FLIC-FedRep for linear regression and Gaussian features
Require: step size п, number of outer iterations T', participation rate (0, 1), diagonalizations >; P.D,P." sorting cigenvalues in docreasing order.
// Estimation of embedded features
For each client E [b], set 6s (z)0)= ) 0k. (z. m.) // Initialisation (
Each client E sends Z; (1/n;) ΣØ= (u: )2 (o!") 1ox (5) IT
to the central server.
The central server computes T rank-d SVD (1/b) E Z,)
6: The central server initialises A = U. for t= 0 1 do Sample
of Ar-1 of active clients such that |A4+1| [rb]. for E Ar+1 do
10:
The central server sends too Act1.
11:
// Update local parameters
B((+1) arg ming, Σj= y - 3, [A(0)) (z(1))
12:
13:
// Update global parameters
A (1+1) = A - V Ej= (1,0 (3((t AT te 11 ( 2
)²
14:
15:
// Communication with the server
Send ( to the central server.
16:
17:
// Averaging and orthogonalisation of global parameter ((+1) [rb] ΣiEA,h 1
18:
1(+1), R(++1) QR (A(+t)).
19:
Ensure: parameters A B(t)
In Appendix S2.2, the equation (S1) will allow us to relate the ground-truth labels n, with estimated predictions g.5) = (( Tā.(x))) via Algorithm S3 starting from the same embedded features.
(A*);) o (x!
S2.2 Proof of Theorem
Let B E gxd the matrix having local model parameters (Bi as columns and denote by BA.L. Rirb) xd its restriction to the row defined by A41 where |At1l [rb] for some r E (0,1), For the sake of simplicity, assume in the sequel that all clients have the same number of data points that is for any E [b], n; = n. For random batches samples {xx!") ". " E [π]] we define similarly to Collins et al. (2021); Jain et al. (2013), the random linear operator A RRrb]xd - BRrbjn for any M E Rlrbjxd A(M) = (a(())))))) MMMsa(a sisssssaa where e; stands for the i-th standard vector of R.rb] Using these notations, it follows from Algorithm S3 that for any {O...., 1}, the model parameters
g(++1) - [A(44)),3 ((tt),,,((+)) (((+1)) are computed as follows:
B((-1) arg min
1
[tt+))
B [A BA [1()][))))) [)
(S2)
[rb] n
BA
[rb] ( (Bh (A B(ttt) (4(0)))((1)0)))) [T) QB((t1) (t+1)
((+11) A
(t++1) R(++1) QR (A(+t)).
(S3)
23

---

Published in Transactions on Machine Learning Research (04/2024)
where A(++1) stands for specific instance of A depending on the random subset of active clients available at each round and AT is the adjoint operator of A defined by AT(M) Eie(traj] (o (o(
The update (S2) admits a closed-form expression as shown in the following lemma.
Lemma S1. For any E... .........,,T- T-.., 1, we have
Bhtt1) BA QA -F(O)
where F()) is defined in (S12), A is defined in (S3) and B is defined in (S2).
Proof. The proof follows from the same steps as in Collins et al. (2021, Proof of Lemma 1) using (S2).
Under HI, we have the following non-asymptotic convergence result.
Theorem S3. Assume H1. Then, for any I; E R we have (x;)) - (x,) where Q E Bock is of the form diag, (+1). Define En = dist(A (0), QA*). Assume that n 2 c(d log( (rb]))/E - x/ [rb]) for some absolute constant > O. Then, for any E (O...., 11, n S ( and with high probability at least -110k -110d log(Lrb)), we have
(itt( ,QA)) S (x)(=))()//2221)))))) dist(AA (dss)(t QA*),
where A is computed via Algorithem 83, dist denotes the principal angle distance and k. E (0, is defined
EE /
Proof. The proof follows first by plugging Lemma $3, Lemma S8 and Lemma S9 into Lemma S2. Then, we use the same technical arguments and steps as in Collins et al. (2021, Proof of Lemma 6).
S2.3 Technical Lemmata
In this section, we provide of useful technical lemmata to prove our main result in Appendix S2.2. Notations. We begin by defining notations that will be used in the sequel. For any E {0, ..... T 1},
define
(
B(++1) Ac+1
[A
Bi [A']T
(S4)
In addition, let
G(f)
clfl
c
D
D
G
.c()
,D(1)
(
(
c
((t)
D
D
d
dd
where for p.q [d].
((t)
Σ Σ
E (oT(a())) Qa41a/9TTQ0001
(S5)
pq
n
iEAe+1 j=1
(1)) = Σ Σ
(0i(2(1))
(z)
(S6)
iEA:
(1)) = (a T
(S7)
24

---

Published in Transactions on Machine Learning Research (04/2024)
with up" E R standing for the p-th column of At) B and og R standing for the p-th column of A E gexd, Finally, we define for any E Act1.
n = o7 (0)
(S8)
(G ) = "TTTQQQ
(S9)
(C(O)) - "QU'Q'" QUIQAA
(S10)
( QAA"
(S11)
Using these notations, we also define 3 1(8)) (1 E lrbbbd and D -
) Fl0) ", (C" -(())))))
(S12)
Technical results. To prove main result in Theorem S3, we begin by providing first upper bound on the quantity of interest namely dist (A(t+1), QA*). This is the purpose of the next lemma. Lemma S2. For any E t0...., T } and 17 > 0, we have
dist (1(+t),004 s Ci + C2,,
where
1 (R(tt)) -
C1 = [A QA(0)
Ia " (1(++2) T (((+1)
(S13)
[nb]
C2 7 [rb]] (H (QA (Qquthy ((+11) ((+1)) (2(4zt) Q (-(ttt)
T
b(++1)
|| (
A
(S14)
where A is defined in (S3), B is defined in (S2), z(t) is defined in (S4) and R(t) comes from the QR factorisation of A(t), step 19 in Algorithm S3.
Proof. The proof follows from the same steps as in Collins et al. (2021, Proof of Lemma 6) and by noting that dist (A) QA*) = dist (Q) A*) for I E (0....,7-))
1}. 11.
We now have to control the terms Ci and Cz. For the sake of clarity, split technical results siming to upper bound of C₁ and C2 in two different paragraphs.
Control of Ct.
Lemma S3. Assume HI. Let 8 = cds/2 (log([obg]rrb//oo 1/2 for some absolute constant c > 0. Then, for any I E (0, (O....,77 T - -}, with probability at least -111k log((rd]), we have for da 1/2 and n 1/(46 m.
C a 5 11 (1 11 (1- dist (A (A4)) -2 miin. + 21
8
5 max dist (1 QA+)) || ( -1
where Z T2 are defined in (S15)-(S16), C1 is defined in (S13), A(t) is defined in (S3) and R(t) comes from the QR factorisation of A0), see step 19 in Algorithm $3.
Proof. Using Cauchy-Schward inequality, we have
Ia " (
(R(+)))
C $
[rb]
Apt1
2
= dist (A0,0 Q.4))
La- [rb]
n
[B
T [(++))
(R(++))
Act1
25

---

Published in Transactions on Machine Learning Research (04/2024)
Define the following minimum and maximum singular values:
T min,*
min AC(1 [n)
Mmi
BA
(S15)
[rb]
d max,*
min ACbb//
mmax
BA
(S16)
V[rb]
Using Collins et al. (2021, Proof of Lemma 6, equations (67)-(68)), we have for dd S 1/2 where da is defined in Lemma S4 and S 1444 )
La
" P [bb]
Blttt) Ae41
S 77 (1 dist (A (44)) 0 min, + 2m
8a J
mmax,
6
with probability at least 1-e 1116" log(Ibb) The proof is concluded by combining the two previous bounds.
Control of C2. We begin by showing four intermediary results gathered in the next four lemmata.
Lemma S4. Assume H1. Let & = cd3/2 vlog( (rb])/nn/ for some absolute constant > 0. Then, for any [ 111k
E {O...., T 1}, with probability at least e111k log((bb)), we have
((a0)-111 - 8d"
1
where G(t) is defined in (55).
Proof. The proof stands as a straightforward extension of Collins (2021, Proof of Lemma 2) by noting that the random variable QØi (or = (xx is nab-Gaussian under H1-(i); and as such is omitted. J
Lemma S5. Assume HI. Let 8 = cd [log(Lbb))
[rb]]////// 1/2 for some absolute constant c > 0. Then, for any E (O....,Te 1}, with probabilily at least -111k log(Lrd]), we have
Bx 1,
S da ||B dist Q4+)
where GLD) is defined in (S5), D(t) is defined in (S7), ((t) is defined in (S6) and A in (S3).
Proof. Without loss generality and ease notation, we remove the superscript (1) in the proof and re-index the indexes of clients in Ar- Let H GD C. From (S8), (S9), (S10) and (S11), it follows, for any e [[rb]], that
H'= A A - 1)QA
Hence, by using the definition of H, we have
rb
[n B
[l CCD
Σ
H
|
1
H
2
[b
where the last inequality follows almost surely from H1-(ii). As in Collins et al. (2021, Proof of Lemma 3), then define for any [n], the vectors
( 1 [AJ] (AAT - ())
Vn
60 A (a!))
Vn"
26

---

Published in Transactions on Machine Learning Research (04/2024)
Let St-1 denotes the d-dimensional unit spheres. Then, by Vershynin (2018, Corollary 4.2.13), define Na, the 1/4-net over Sd-1 such that [Mal s 9 Therefore, by using Vershynin (2018, Equation (4.13)), we have
11 22
Σ (2, u (1) (0))) .v). ==1
max N.VEN
Since 07 (x) is standard Gaussian vector, is sub- Gaussian and therefore (z, ( and (o y) are sub-Gaussian with norms addition, we have
+ [A (AAT 11l2 (I/VssdistAA, QA*) and (1/vn) respectively. In
E (2, u; ()) (0 .0]
= E 1 (ATT (000(("")
π
Vn
1 T [AT (AA A4y
-
where we have used the fact that Elo MoT
Q2= and (AAT-1A 44A== 0. The rest of the
proof is concluded by using the Bernstein incquality by following directly the steps detailed in Collins et al. (2021, Proof of Lemma 3, see equations (35) to (39)).
Lemma S6. Assume HI. Let 8a = cd3/2 / (bb])/n for some absolute constant > 0. Then, for any e [T], with probability at least 1 e -111A log((=bl), have
p(0
5d
|B l1 dist QA+)
<
1 da
where F(I) is defined in (S12) and A(O in (S3).
Proof. By the Canchy-Schwarz incquality, we have ||](|| 8 B2 l9. S ||[ (G()) D() C( ll, s da | 11 The proof is concluded by combining the upper bounds given in Lemma S4 and Lemma $5.
= || - ) l1 S
Lemma S7. Assume H1 and let 8 = cdv/k V [rb] for some absolute positive constant c. For any E [T] and whenever 8' I d, we have with probability al least -110k -10df log([rb])
(Coca1)a (2 Q 2(0)
T
((0)
S 8/ dist (A0),AA))
[bb]
2
where Bl) is defined in (S2) and z is defined in (S4).
Proof. Let L [T]. Note that we have
(2(4109441) (2(0) (Z4)e² Q z n
m
200(4)))2000011 [1n0]] 2 (1) [B!0]
Let 1 and 1d-1 denote the k-dimensional and d-dimensional unit spheres, respectively. Then, by Vershynin (2018, Corollary 4.2.13), we can define Ns. and Na, 1/4-nets over Sk-1 and 8d-1, respectively, such that W&l < g* and (al < gd. Therefore, by using Vershynin (2018, Equation (4.13)), we have
(1((41))4))) (2 Q
z10) Q=(())
((t)
27

---

Published in Transactions on Machine Learning Research (04/2024)
=2 max T
m Σ
Σ
QQ (2) 24 (. (o)(( [3!0]] - - [3(0] c
WEN,
n SEA. j=1
(x.", Σ Σ (Q0) (x. ( (u, Qo7 (o) v) (u, <0!" .0).
2 DEN j=1
max
(S17)
In order to control (S17) using Bernstein inequality as in Lemma S5, we need to characterise, in particular, the sub-Gaussianiy, of (u, = and (8/ ,u) which require a bound on ||>]"|| and || respectively. From Lemma S1, we have [ T = (A4) A
(A*)
T
(z(")T which leads to
(A4) QA //0 A 3/
(2 1A)A 18811 + {(00
s 2d (A QA") + 2 5!001²
Using (812) and the Cauchy-Schward inequality, we have
12
Di00 (C40
|| 2 (4.100 co.(0))
2
||
S 1
(i, -+.(0)
c+((1)
2
(S18)
where the last inequality follows from H1-(ii..
Using Lemma S4 and Lemma $5 and similarly to Collins et al. (2021, Equation (45)), it follows for any E A that
2
| s 4d dist(A "Q*),,
with probability at least 1 el 10d (og((+b]),
Similarly, using Lemma S1 and (S18), we have with probability at least 1 - el10d" log([rb)) and for any i E Ar,
that
|| s $2||14001 QA
3 3:1 +2|15(0) 1²=
s 4d.
Besides, note we have
E (1 ) (o m5] (u, (u,,,,)) 200)(()())))
The proof is then concluded by applying the Bernstein inequality following the same steps as in the final steps of Collins et al. (2021, Proof of Lemma 5).
We are now ready control C2.
Lemma S8. Assume H1 and let 8'2 cdy// [rb] n. for some absolute positive conslant For any E (O....,T 1}, >> and whenever 8y, S d, we have with probability at least 1 - e
110k e 110d2 log((rh])
C2 < no' dist (.0, || (R(C+L) -1
where C2 is defined in (S14), A is defined in (S3) and R(t) comes from the QR factorisation of A(t, see step 19 in Algorithm $3.
28

---

Published in Transactions on Machine Learning Research (04/2024)
Proof. Let t E {0.... ................................ T - 1} and 7 > 0. Then, whenever 84 s d, we have with probability least e 110k c -100d log([rb]), we have
1
[11I"(QA44t)) ((++1)
((4+)) (z(+t1) Q 2 BAree
(
C2
[rb]
S " [rb]
(1 A (2(ttl) Q- Q- (1-20t1) Bl(t2)
(n 7
A1t1
S n99s dist (A", Q4+) | (R(tt))
where we used the Cauchy-Schward inequality in the second inequality and Lemma S7 for the last one.
Control of| || (R(+))) -1 11.. To finalise our proof, it remains to bound (R((+))) -1 1l2. The associated result is depicted in the next lemma.
Lemma S9. Define óa = ba + 8's where da and 8' are defined in Lemma S4 and Lemma $5, respectively. Assume H1. Then, we have with probabillly at least -110k -110d log([rb]),
(
8 0 max,
1/2
$ 1 2
4n (1 0a)²
Proof. The proof follows from Collins et al. (2021, Proof of Lemma 6).
S3
Experimental Details
S3.1
Data Sets
We provide some details about the datasets we used for our numerical experiments
S3.1.1 Toy data sets
The first toy dataset, denoted noisy features, is 20-class classification problem in which the features for a given class is obtained by sampling a Gaussian distribution of dimension 5, with random mean and Identity covariance matrix. Mean for each class is sampled from a Gaussian distribution with zero mean and diagonal covariance matrix I with o = 0.8.
For building the training set, we sample 2000 examples for each class and equally share those examples among clients who hold that class. Then, in order to generate some class imbalances on clients, we randomly subsample examples on all clients. For instance, with 100 clients and 2 classes per clients, this results in problem with total of about 16k samples with a minimal number of samples of 38 and maximal one of 400. In order to get different dimensionality, we randomly append on each client dataset some Gaussian random noisy features with dimensionality varying from 1 to 10.
The second toy dataset, denoted as linear mapping, is a 20-class classification problem where each class- conditional distribution is Gaussian distribution of dimension 5, with random mcan and identity covariance matrix. Mean for each class is sampled from a Gaussian distribution with zero mean 0 and diagonal covariance matrix I with 0.5. matrix. As above, we generate 2000 samples per class, distribute and subsample them across clients in the similar way, leading to total number of samples of about 15k. The dimensionality perturbation is modelled by random (Gaussian) linear transformation that maps the original samples to a space which dimension goes from 3 to 100.
Examples from given class are equally shared among clients after random permutation of their index. For subsampling samples on a client, we randomly select a percentage of samples to keep ranging from 5% 100%, and them uniformly randomly select among all classes the clients the ones we keep.
29

---

Published in Transactions on Machine Learning Research (04/2024)
Table S1: Summary of the Brain- Comp-Co Interfaces dataset we used. We report the number of subjects (#Subj), the number of channels (#Chan), the number of classes (# Classes), the number of trials per class (Tiilsss class) and the number of features (#features) on the covariance representation has been vectorized.
NAME
#SUBJ
# CCAAN
# #CLASSES
# TRIALS CLASS
# FEATURES
ALEXMI
8
136
BNC12014001
9
22
144
253
ENCT2014002
14
120
BNC12014004
9
360
6
WEIBO2014
10
60
80
1830
7Hou2016
4
14
160
105
For these problems, the test sets are obtained in the same way as the training sets but 1000 examples for each class shared across clients.
S3.1.2 MNIST-USPS
We consider a digit classification problem with the original MNIST and USPS data sets which are respectively of dimension 28 x 28 and 16 x 16 and we assume that client hosts either a subset of MNIST or USPS data set. We use the natural train/lest split of those datasets and randomly share them accross clients.
S3.1.3 TextCaps data set
The TextCaps data set (Sidorov et al., 2020) is an Image captioning dataset for which goal to develop a model able to produce a text that captions the image. The dataset composed of about 21k images and 110k captions and cach image also comes with an object class. For our purpose, we have extracted pair of 14977 images and captions from the following four classes Bottle, Car, Food and Book. At each run, those pairs are separated in 80% train and 20% test sets. Examples from the TextCaps datasets are presented in Figure S7. Images and captions are represented by vectors by feeding them respectively to pre-trained ResNet18 and a pretrained Bert, leading to vectors of size 512 and 768.
Each client holds either the image or the text representation of subset of examples and the associated vectors are randomly pruned of up to 10% coordinates. As such, all clients hold dataset with different dimensionality.
For this dataset, we randomly split the samples for the train and test and share them across clients.
S3.2 Brain-Computer Interfaces data set
The Brain-Computer Interfaces dataset we used are summarized in Table S1. Each dataset description can be obtained from the MOABB library (Jayaram & Barachant, 2018) and at the following URL: http: //moabb. neurotechx. com/ doca/ /ataaaattssss.... html. For each subject, we select the predefined train/lesi splits or used 75% of the trials for training and the remaining 25% for testing. We used a bandpass prefiltering between 8 and 30 Hz of the EEG signals and extracted a covariance matrix for each trial using all available channels. These covariance matrices are vectorized and used as feature. The classes that we used for the classification problem are the following ones: [fleft hand', 'right hand', 'eet', 'tongue', 'rest") and subset of them available for each dataset. For these datasets, if a subject has more than one session, then use the last one for generating the test set and the remaining ones for the training set. If we have only single session, we randomly split the trials into training and test sets.
S3.3 Models and Learning Parameters
For the toy problems, the TextCaps data set and the BCI one, as a local transformation functions we used a fully connected neural network with one input, one hidden layer and one output layers. The number of units in hidden layer has been fixed to 64 and the dimension of latent space been fixed to 64. ReL.U activation
30

---

Published in Transactions on Machine Learning Research (04/2024)
3.0
clients 10
clients 10
7,70
clients 50
2.9
000
clients 50
285
2,000
clients 87
2.0
clients 87
2.8
2.7
50
2.3795
2.6
2.3790
2.3
780
500
2.4
2.4
2.2
2.3
2.2
200
800
#local epochs
#local epochs
clients 10
clients 50
clients 87
180
200
220
loss
35
3.0
200
#local epochs
Figure S1: Evolution of the local loss curve of three different clients for three different learning situations. See text for details.
has been applied after the input and hidden layers. For the digits dataset, we used a CNN model with 2 convolutional layers followed by a max-pooling layer and a sigmoid activation function. Once flattened, we have a one fully-connected layer and ReLU activation. The latent dimension is fixed to 64.
For all datasets, as for the local model ge,- in order to be consistent with competitors, we first considered single layer lincar model implementing the local classifier well as model with one input layer (linear units followed by LeakyReLU activation funcion) denoting the shared representation layer and an output linear layer.
For training, all methods use Adam with a default learning rate of 0.001 and a batch size of 100. Other hyperparameters have been set as follows. Unless specified, the regularization strength À and X2 have been fixed to 0.001. Local sample batch size is to 100 and the participation rate to 0.1. For all experiments, have set the number of communication round T to 50 and the number of local epochs to respectively 10 and 100 for the real-world and toy datasets. For FLIC, in FedRep those local epochs is followed by epoch for representation learning. We have trained the local embedding functions for 100 local epochs and batch size of 10 for toy datasets and TextCaps and while of 100 for MNIST-USPS and BCI, Reported accuracies are computed after local training for all clients.
S3.4 Ablating Loss Curves
In order to gain some understanding on the learning mechanism that involves local and global training respectively due to the local embedding functions, the local classifier and the global representation learning, we propose to look at local loss curves across different clients.
Here, we have considered the linear mapping toy dataset as used in the toy problem analysis. However, the learning parameters have chosen different from those have used to produce the results to
31

---

Published in Transactions on Machine Learning Research (04/2024)
highlight some specific features. The number of epochs (communication rounds) is set to 100 with client activation ration of 0.1. Local epochs are shared for either training the local parameters or the global ones (note that in our reference Algorithm 1, the global parameter is updated only once for each client) Those latter are trained starting after the 20-th commmication round and in this case, the local epochs are equally shared between local and global parameter updates. Note that because of the randomness in the client. selection at each epoch, the total number of local epochs is different from client to client. We have evaluated three learning situations and plotted the loss curves for each client.
the local embedding functions and the global models are kept fixed, and only the local classifier is Irained. Examples of loss curves for 3 clients are presented in the left plot of Figure $1. For this learning situation, there no shared global parameters that are trained locally. Hence, the loss curve is typical of those obtained by stochastic gradient descent with smooth transition, at multiple of 100 local epochs, when given client is retrained after communication round.
the local embedding functions are kept fixed, while the classifier and global parameters are updated using half of the local epochs each. This situation is interesting and reported in middle plot in Figure S1. We can see that for some rounds of 100 local epochs, strong drop in the loss occurs at starting at the 50th local epoch because the global parameters are being updated. Once the local update of a client is finished the global parameter is sent back to the server and all updates of global parameters are averaged by the server. When a client is selected again for local updates, it is served with a new global parameter (hence a now loss value which causes the discontinuity in the loss curve at the beggining of each local update.
all the part (local embedding functions, global paramcter and the classifier) of the models are trained. Note at first that the loss value for those curves (bottom plot in Figure S1) is larger than for the two first most left plots as the Wasserstein distance to the anchor distribution is now taken into account and tends to dominate the loss. The loss curves are globally decreasing with larger drops in loss at the beginning of local epochs.
S3.5 On Importance of Alignment Pre- Training and Updates.
We have analyzed the impact of pre-training the local emebedding functions and their updates during learning for fixed anchor reference distributions. For this, we have kept the same global experimental settings as for the performance comparisons except that we fixed the number of users to 100. We have compared two learning situations the first one in which the local embedding functions are pre-trained for K epochs and updated during local training and the second in which they are also pre-trained for K epochs and kept fixed during local training. We have reported the classification performance when the number of pre-training epochs K varies from 1 to 200. Results, averaged over 5 runs are shown in Figure S2, At first, we can note that there exists a performance gap between and 200 epochs of pre-training. Nonetheless, this gap is not that large except for the toy added noisy feature dataset. For the three datasets, increasing the number of pre-training epochs up to a certain number tends to increase performance, but overfitting may occur. The latter is mostly reflected in the toy linear mapping dataset for which 10 to 50 epochs is sufficient for good pre-training. Examples of how classes evolves during pre-training are illustrated in Figure 5, through t-sne projection. We also illustrate cases of how pre-training impact on the test set and may lead to overfitting shown in Figure S4.
S3.6 On the Impact of the Participation Rate
We have analyzed the effect of the participation rate of each client into our federated learning approach. Figure S3 reports the accuracies, averaged over 3 runs, of our approach for the toy datascts and the TextCops problem with respect to the partication rale at each round. We can note that the proposed approach is rather robust to the participation rate but may rather suffer from overfitting due to overtraining of local models. On the left plot, performances, measured after the last communication round, for Tept is stable over participation rate while those performances tend to decrease for the toy problems. We associate these decrease to overfitting since when we report (see right plot.) the best performance over communication rounds
32

---

Published in Transactions on Machine Learning Research (04/2024)
13
TextCaps
Toy
TextCaps -ixed-- fieee-
Toy LM fixed
Accuracy
65-
100
125
#pretralning apochs
Figure S2: Impact of the number of pre-training epochs for de the model accuracy. We compared the case where o, is (plain) updated during local training or (dashed) kept fixed. Results for three different datasets are reported.
00
Accuracy (80
Accuracy 1%)
TextCaps Toy
TextCaps
Toy
Toy LM
Toy LM
6m
0.20
8:20
Participation rate
Participation rate
Figure S3: Evolution of the performance of our FLIC-Class algorithm with respects to the participation rate of clients, using the same experimental setting as in Figure 4. (left) evaluating performance after last communication rounds, (right) best performance across communication rounds.
(and not the last onc), they are stable for all problems. This suggests that number of local epochs may be dependent to the task on each client and the client participation rate.
S3.7 Comparing FLIC with t-sne and FedRep
We have claimed several times the need for a proper aligument of the class- consicoonditonal distributions when mapping into a common latent subspace. We have illustrated this need Figure 3 where we have projected the class-conditional distributions of the toy linear mapping dataset without alignment and using classical dimensionality reduction algorithm like I-sne oracle multidimensional scaling (MDS). Without aligument, mapped class-conditional distributions can get mixed up and learning accurate classification model becomes hard.
In order to back up our claim with quantitative results, we have compared the performance of FLIC with the baseline that consists in projecting data on clients into a common latent subspace (here R84) using t-sne and then applying FedRep on the projected data. We have followed the same protocol as for the results on toy dataset and reported the results in Figure $5. As expected, the perfomance of this bascline in terms of classification accuracy is lower than the one of FLIC and even lower than the of local classifier.
33

---

Published in Transactions on Machine Learning Research (04/2024)
Cass Cass Caeell
Caos Cavos
Cavos
Clack
Caus
Closs
Cass Cam Caa 11
Class
Class
Class
Class
Figure S4: 2D t-sne projection of classes partially shared by 3 clients for the toy linear mapping dataset after learning the local embedding functions for (left) 10 epochs, (middle) 50 epochs, (right) 100 epochs. Original dimensions clients vary from 5 to 50. Top row shows the projection the training set while bottom row plots show both training and test set. Star * markers represent the projection of the mean of each class-conditonnal The three different marker styles represent the different clients. Classes are denoted by colors and similar toncs of color distinguish train and test sets. We see that each class from the training set from each client converges lowards the mean of its anchor distribution, represented by the star marker. Interestingly, we also remark that unless convergence is reached, empirical class-conditional distributions on each clients not equal making necessary the learning of joint representation. From the bottom plots, understand that distribution alignment impacts mostly the training set but this alignment does not always generalize properly to the test sets.
80
70
Local
Accuracy (%6
Accuracy (%)
60
FedHeNN
HetFedilep
Local
Class
FedHeNN
50
FLic
HetFedRep
t-sne FedRep
FLic Class
45
FLic HL
Lsnne FedRep
30
35
50
75
100 # Clients
125
150
175
200
25
50
125 # Clients
150
175
200
Figure $5: Comparing the baseline t-sne projection + FedRep to FLIC and other baselines We report the classification accuracy for the (left) toy noisy features datasct and (right)) toy linear mapping datasct.
34

---

Published in Transactions on Machine Learning Research (04/2024)
300
500
Local
Local
250
FedHeNN
FedHeNN
HetFedRep
400
HetFedRep
200
FLic Class
Flic Class
FLic HL
FLic HL
Time
150
200
100
50
100
50
150 # Clients
175
150
# Clients
Figure S6: Comparing the computational cost for training all the models over 50 epochs on the toy dataset. (left) toy noisy fealures dataset and (right) loy linear mapping dataset.
S3.8 On the computational cost of FLIC
The approach we propose has computational overhead compared to the classical federated learning approach and compared the basclines that are able handle heterogencous feature spaces. For instance, in addition to the local version of the model on each client, need to update the local embedding function and the anchor distributions parameters.
In order to properly quantify the computational cost of FLIC, report in Figure S6 the running time needed for training each model over 50 epochs and for increasing number of clients, for the toy dataset problems. Remind that as the number of clents increase, the number of samples on each client decreases. We can see that roughly, the computational cost of FLIC is about two times higher than for other competitors, across the range of clients and for the two lcarning problems.
S3.9 Experiments on FEMNIST
We have also evaluated our approach on the FEMNIST dataset. It is part of the LEAF (LEAF: A Benchmark for Federated Seutings) project (Caldas al., 2018), which aims provide benchmarking framework for learning in federated settings. It 28 by 28 pixel image dataset that built by partitioning the data in the Extended MNIST dataset based on the writer of the digit/chand cigit/character. It contains 62 different classes, including 10 digits, 26 lowercase letters, and 26 uppercase letters. As in Collins et al. (2021), subset of 50 and 150 clients (one writer, one client) and limited ourselves to the 10 digits samples available for each writer. Samples in FEMNIST are all of the same dimensionality so our goal here is to compare our approach to some baselines such as FedRep for a well-known federated learning problem. For training, we used the same parameters as for the other datasets except that the number of epochs has been set to 1000, as the learning problem more complex and needs more epochs to converge.
have considered
3 runs, are reported in Table $2. We first note that the performance of FedRep is in line with the ones reported in Collins et al. (2021). Then, remark that our approach outperforms the competitors on this datasct by a large margin with a performance gain of about 5% over the best competitor. This is a strong result that shows the potential of our approach on real-world datasets even when the feature spaces are homogeneous. We can also highlight the impact of pre-training the local embedding functions when comparing the performance FLIC and HetFedRep (the latter is a variant of FLIC where the local embedding functions are not pre-trained). Here, pre-training the local embedding functions play the role of domain adaptation mechanism that allows to better align the class-conditional distributions of each client. with common shared representation.
Results, averaged
35

---

Published in Transactions on Machine Learning Research (04/2024)
Table S2: Performance comparison on the FEMNIST dataset. We report the classification accuracy for the FEMNIST dataset for 50 and 150 clients.
METHOD
50 CLIENTS
150 CLIENTS
LoCAL
77.29
0.6
79.59 + 0.7
FEDHENN
71.01 + 0.8
74.77 + 0.2
HETFEDREP
52.90
2.0
57.78 + 2.3
FEDREP
70.29 + 3.0
74.55 + 1.3
CLIC-CLASS
82.25
0.4
84.81
+ 0.4
FLIC-HL
82.37 + 0.6
84.30 + 0.4
pan sits on a hob with bearing the logo Hamilton Beach Stay n Go
Deciding on whether to drink spring water or a 7UP.
ana
STAVE
LIBRA
Jass
bookshalf consists about 20 diffirrnt types
The yellow convertible is display somewhere in California.
Nesrtti CRUNEELD
'HER COMPURII
A.NALS
Figure S7: Examples of some Text pairs of image/caption from the classes we considered of (top-left) Food, (top-right) Bottle, (bottom-left) Book (bottom-right) Car. We can how difficult some examples can be, especially from the caption point of view since few hint about the class is provided by the text.
36


Reviews:
[{"review_id": "MUMV8U8eYM", "reviewer": "Reviewer_5DVB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed FLIC framework is interesting -- in particular, the alignment of representations of same-class examples with a latent anchor distribution. In numerical experiments, the proposed methods FLIC-HL and FLIC-Class perform well on these problems.\n\nThe FLIC methods have only marginal gains over competing methods. The theoretical result (Theorem 1) appears unsatisfactory.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a federated learning algorithm to handle cases where the datasets and features are  heterogeneous across clients. The key ingredients of the proposed FLIC framework are the following:\n* Map client data to a common representation space with learnable mapping $\\phi_i$ for client $i$.\n* In the common representation space, encourage examples from class $c$ from all clients to be close to an anchor Gaussian distribution $\\mu_c$.\n* On top of these representations, learn a classifier parameterized by global parameters $\\alpha$ and local parameters $\\beta_i$ for clients $i \\in [b]$.\n\nThe authors apply their methods to \n* Two toy datasets with Gaussian data\n* MNIST and USPS digit classification task\n* An image, caption to class task \n* A Brain Computer Interface data task.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Details for experiments should be fully specified. For example, in the toy dataset problems, the following details should be added.\n* How are the means of the class-conditional distributions generated? How separated are the means compared to the noise?\n* How are examples subsampled in each client?\n* How is the random matrix formed?\n\nFor the plot in Figure 3 (left), more details are required for reproducibility. For example, the Gaussian distributions used to generate data for the 10 classification problems.\n\nIn Figure 4 (left), why is the accuracy going down as we increase the number of clients? In Figure 4 (right), why does it look very different from 4 (left)? Why is HetFedRep so poor the right plot, but good on the left plot?\n\nIn Theorem 1:\n* The result is on regression, whereas the focus of FLIC is classification.\n* Are you assuming that $k_i \\ge k$ for $i \\in [b]$? If so, consider making it explicit.\n* Suppose $k_i = k$ for simplicity. Assume $\\Sigma = P D P^T$. May be I am missing something, but should the mapping from $N(0, \\Sigma)$ to $N(0, I)$ be\n$f(x) = D^{-1/2} P^T x$ instead of $f(x) = D^{-1/2} x$?\nNote that, for $y=Bx$, \n$$E[yy'] = E[Bxx'B'] = B \\Sigma B' = D^{-1/2} P' PDP' (D^{-1/2}P')' = I$$\nYou can apply another orthonormal transformation and still the covariance would be identity. So, there is nothing special about the diagonal matrices with $\\pm 1$, I think.\n\n\n**Following comments are minor in nature**\n\nWhy don't we explicitly enforce separation between means $\\mu_c$ of different classes? Is it automatically done by the classification loss minimization?\n\nWhat does FLIC stand for?\n\nIt should be mentioned in the main text how $\\phi_i$ can be parametrized.\n\n\nTypo: In page 7, right below Algorithm 1, the description of the algorithm, the step numbers seems to be wrong.\n\nTypo: Correct the typo \"FedRed\".\n\nWhat is principal angle distance between two orthonormal matrices A, B?", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JXATQwwZuA", "reviewer": "Reviewer_beo7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* Defines a strategy for learning client-specific embeddings across heterogeneous feature spaces in a federated way.\n* The need for such an approach is well-motivated, especially in the case of cross-institution FL where dataset formats may not be standardized a priori.\n* This learned feature representation method is broadly applicable to existing FL algorithms. \n\nWeaknesses:\n* There is a need for specifying in advance the latent dimensionality and having some prior knowledge of the latent distribution.\n* The theoretical analysis is for a simplified setting in which the anchor distribution is known beforehand, which may not be the case in practice.\n* Lacks comparison with alternative ways of embedding distinct data in a common subspace. \n* The empirical results are quite simple. The chosen datasets and settings do not have prior work on them to have benchmarks and points of reference that could be used to put the numbers in context in terms of expected performance. The difference in performance across methods are quite minimal with the top accuracy falling within the confidence intervals of other approaches (e.g. Table 3).\n* The paper anchors the contribution and discussion on extending FedRep to the heterogeneous feature space setting, making it difficult to understand for anyone unfamiliar with the prior work, even though the contributions are more broadly applicable.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a method for applying federated learning (FL) to heterogeneous datasets, which may have different schema or feature spaces. In order to enable FL, which typically assumes all client data has been drawn from the same feature space, across clients with distinct feature spaces, the authors propose having each client learn a local embedding function that maps client data to a common subspace. This framework \"FLIC\" is compatible with existing FL algorithms such as FedRep and FedAvg with fine-tuning. The authors include some empirical results on a range of datasets - from toy to brain-computer interface data - and some theoretical analysis of the proposed approach.\n\nPrimary contributions:\n* First FL framework for operating on client data with heterogeneous feature spaces.\n* First empirical results for using FL on brain computer interface data.\n* Distribution alignment algorithm for learning local embeddings functions to map data to a common subspace.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical changes:\n* Additional experimental results would strengthen the work.\n    * Include baseline performance of FLIC on federated datasets commonly used in the literature, where the partitions are based on user and all data has the same format. This would indicate how well the method works on data that is non-iid but not completely disjoint, and offer a comparison point to common FL approaches.\n    * Demonstrate FedAvg + FT with learned feature embeddings for comparison (simpler baseline for personalized FL).\n    * Compare FLIC to t-sne and multi-dimensional scaling in the experimental results. Depicting clusters visually using these methods on a toy set does not demonstrate the effect of these methods in practice.\n\n* Additional description of figures and discussion of what the results show is necessary.\n    * Figure 2: It is not clear how the two plots included in the figure relate. Either the relation should be discussed or the plots should be separated. What should be understood from the leftmost plot? There is no discussion and no metrics to indicate whether the projected features are far from the ground truth. The caption of the rightmost plot includes numbers that are not matching what is shown and requires further explanation.\n    * Figure 4: What is the effect of increasing the number of clients on a toy problem? Is this a more or less challenging problem? Include analysis of the results and interpretation of how or why each method differs.\n\n* Section 4 claims that the presented algorithm avoids the client drift phenomenon. It is unclear how this is possible, as the federated algorithm presented has clients train locally for some number of steps between rounds of parameter averaging. The communication cost of using this approach is strictly higher than other federated algorithms without it, such as FedAvg or FedRep, due to the additional parameters that are needed to define the learned anchor measures. Claiming FLIC is communication-efficient is misleading and incorrect.\n\nMinor changes:\n* Section 4 “Averaging Anchor Distributions” references incorrect steps in Algorithm 1.\n* The experimental section could use further clarification. Is the reported accuracy on a held out portion of each client’s data?\n* It is not clear what should be expected when plotting performance with respect to the number of clients.\n* Table 4 accuracies should include confidence intervals over the 3 runs.\n* Simplifying the notation through the manuscript and algorithms would help with the readability. As this approach is compatible with different model definitions it would be useful to use a simple base model \\theta in the algorithms and throughout the text, rather than splitting the model into \\alpha and \\beta components, since that is not the contribution of this work. The key change this work presents is in adding a local embedding layer, which could be expressed much more simply.\n* Typo “jsut” on page 10.\n* A limitation that should be noted is that this is a stateful algorithm applicable to the cross-silo setting.", "broader_impact_concerns": "Learning on heterogeneous data formats has the potential to enable learning across data collected by different institutions without standardizing the format a priori. Making the claim of expanded privacy-aware FL scope, however, is misleading without adding a privacy analysis to strengthen \"Remark 3\" and demonstrating that there is no further privacy leakage than what is given by FedAvg. Stating how the proposed approach might be compatible with privacy techniques such as differential privacy or secure aggregation would strengthen the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "H6dset5lrv", "reviewer": "Reviewer_eR4X", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper proposes a novel framework for federated learning on heterogeneous feature spaces via the mapping of input features into a common embedding space\n2. Theoretical analysis for a simple regression framework is presented.\n3. Experimental results are presented for a toy example and real-world datasets such as Digits, TextCaps, and BCI datasets.\n4. The paper discusses the limitations and privacy impacts of the proposed framework.\n5. The appendix includes interesting ablation studies.\n\nWeaknesses/questions:\n1. The paper was hard to follow and can use some improvement in writing.\n2. It is not clear how the anchor distributions are initialized. In particular, what is $v_{1:C}^0$ in Algorithm S2? Is it equal to $\\hat{m}_{1:C}^0$?\n3. FLIC has memory and communication overheads stemming from $v_{i, 1:C}, L_{i, 1:C}$.\n4. FLIC is computationally heavy at the client end as it needs to compute gradients with respect to multiple parameters and loss functions.\n5. It would be helpful to present some quantitative results comparing the overheads of FLIC to the other baselines. \n6. The Averaging anchor distribution paragraph on page 7 mentions \"We provide algorithmic details regarding steps\n14 and 20 in Algorithm 1\". Algorithm 1 only has 18 steps and that paragraph doesn't seem to be consistent with the algorithm.\n7. Section S3.5 and the caption for Figure S2 are not clear.  It mentions \"We have considered two learning situations: one in which they are\nupdated during local training (as usual) and another one in they are kept fixed all along the training\". Is the transformation $\\phi$ pre-trained in these two cases? The discussion on the impact of pre-training the embedding space is interesting and needs to be more clearly written.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper formulizes personalized Federated Learning on heterogeneous client’s feature spaces and proposes a framework FLIC. FLIC learns a function locally at each client to map the input feature space to a common embedding space. To preserve semantical information from the original data distribution, the locally learned functions at each client are aligned with some learnable latent anchor distributions that are shared across all clients. A  personalized Federated Learning algorithm is then employed on the common embedding space to learn the desired task. The paper shows that for a simpler regression framework where the anchor distribution is known beforehand and not learned under the FL paradigm, FLIC can recover the true latent subspace underlying the FL problem. The effectiveness of FLIC is shown through experiments on toy and real-world problems including heterogeneous BCI datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See Weaknesses/questions", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose FLIC, a framework that enables personalized FL in scenarios where different clients may have different data representations (heterogeneous feature spaces). This involves clients learning functions to map their data to a common feature space, and the embedding functions being optimized such that the learned feature space represents similar data/same classes across clients similarly. The authors provide simplified theoretical analysis showing that their approach recovers the true latent subspace. The experiment shows some gains, even if marginal, over other embedding methods. The authors incorporated presentation improvements (clarified concepts, algorithm steps, and provided missing details on the experimental setup, removed some statements not supported by the experiments), and ran additional experiments comparing to other baselines as suggested by the reviewers, further improving the paper.
Decision: Accept (poster)