Federated Learning with Unlabeled Clients: Personalization Can Happen in Low Dimensions

## Abstract

Abstract Personalized federated learning has emerged as a popular approach to training on devices holding statistically heterogeneous data, known as clients.
However, most existing approaches require a client to have labeled data for training or finetuning
in order to obtain their own personalized model.
In this paper we address this by proposing FLowDUP, a novel method that is able to generate a
personalized model using only a forward pass with unlabeled data.
The generated model parameters reside in a low-dimensional subspace, enabling efficient communication and computation.
FLowDUP’s learning objective is theoretically motivated by our new transductive multi-task
PAC-Bayesian generalization bound, that provides performance guarantees for unlabeled clients.
The objective is structured in such a way that it allows both clients with labeled data and
clients with only unlabeled data to contribute to the training process.
To supplement our theoretical results we carry out a thorough experimental evaluation of FLowDUP, demonstrating
strong empirical performance on a range of datasets with differing sorts of statistically heterogeneous clients.
Through numerous ablation studies, we test the efficacy of the individual components of the method.

## 1 Introduction

Federated learning (FL) (McMahan et al., 2017) is a widely adopted approach to enable
privacy-preserving machine learning in distributed environments.
Data holding devices (clients) cooperatively train predictive models under the coordination of a central server, without sharing their local data.
Differing underlying behaviors often result in client data following different distributions, potentially rendering single global model training ineffective (Kairouz et al., 2021).
Personalized federated learning (Smith et al., 2017) has emerged as a popular approach to address this
challenge of statistical heterogeneity.
It enables clients to learn personalized models while still leveraging collective knowledge from the broader network.
However, since most approaches to personalization rely on some form of finetuning on the client local data they require
a client to have labeled data in order to obtain a personalized model.
This poses a major challenge for personalization in applications where active participation from the client (i.e. user) is required to create labels for their local data.
In such scenarios, even though some clients will likely be *labeled* (have data with labels), the majority of potentially participating clients will likely be *unlabeled* (i.e. have no labels for their data).
Similarly, most future clients will potentially also be unlabeled.

In this paper we address this issue with FLowDUP: Federated Low Dimensional Unlabeled Personalization.
FLowDUP learns to generate a personalized model for any client, even if it has only unlabeled data.
To do so, it trains a hypernetwork (Ha et al., 2017) that takes as input an unlabeled dataset and outputs the parameters of a personalized predictive model.
Crucially, rather than outputting all the parameters of a personalized model, the hypernetwork instead
outputs a parametrization in a subspace of much smaller dimension than the underlying personalized model.
The full model is obtained by multiplying the subspace parameters with a fixed (random) *expansion matrix*.
The use of a low-dimensional subspace is a crucial contribution of our work. It allows
for the generation of large model architectures and means that the hypernetwork is small
and efficient enough to be transmitted to and run on the client devices.
This allows FLowDUP to adhere strictly to the federated paradigm that clients data should not be transferred away from the device.
Without it, hypernetwork-based methods are only able to generate very small models and
are only practical if the hypernetwork runs on the central
server (Shamsian et al., 2021; Amosy et al., 2024; Scott et al., 2024), which is at odds with the FL
principle that clients’ data never leaves the clients’ device.

We propose a learning objective for FLowDUP that consists of two parts.
The first part measures the quality of the generated client models, for which it exploits that some of the
clients at training time do have labels.
The second part prevents overfitting by penalizing large deviations between the generated models
and a learned regularization term.
Evaluating it requires only unlabeled data, so it can be computed from all clients available during training. FLowDUP and its learning objective are motivated by theoretical results derived in a multi-task framework.
Specifically, we prove a generalization bound that provides performance guarantees on unlabeled clients. Our learning objective is derived by optimizing terms that appear in this bound.

In addition to our theoretical contributions we carry out an experimental evaluation over a range
of datasets, exhibiting various types of statistical heterogeneity, and at a range of proportions
of labeled clients present during training.
We conduct additional experiments and ablation studies to better understand the
different components of FLowDUP.

To summarize, our contributions are as follows:

- •
We propose FLowDUP, a method that generates low-dimensional personalized model parameters using only an on-device forward pass on unlabeled data.
- •
We derive generalization bounds in a multi-task framework, which, in our FL setting, provide guarantees on FLowDUP’s performance on unlabeled clients.
- •
Based on these bounds we propose a theoretically motivated training objective that is able to also leverage unlabeled clients during training.
- •
We conduct an experimental evaluation demonstrating strong empirical performance of FLowDUP and
illustrate the efficacy of the individual components with ablation studies.

## 2 Background

#### Notation

We assume a federated setting with $n$ clients participating in training. Each client $i\in\{1,\dots,n\}$ possesses a data distribution $D_{i}$ over shared inputs and output sets, $\mathcal{X}\times\mathcal{Y}$. The clients are statistically heterogeneous, meaning that $D_{i}\neq D_{j}$ for $i\neq j$ is possible.
We call $n_{L}$ the number of clients that hold labeled data and $n_{U}$ the number that hold only unlabeled data, i.e. $n_{L}+n_{U}=n$.
Writing $S_{i}$ for the dataset of client $i$, and $m_{i}=|S_{i}|$, we then have

$$ $\displaystyle S_{i}\coloneqq\begin{cases}(X_{i},Y_{i})=\big{(}x^{j}_{i},y^{j}_{i}\big{)}_{j=1}^{m_{i}}\sim D_{i}&\text{if }i\in\mathcal{I}_{L},\\ X_{i}=\big{(}x^{j}_{i}\big{)}_{j=1}^{m_{i}}\sim D_{i|\mathcal{X}}&\text{if }i\in\mathcal{I}_{U},\end{cases}$ (1) $$

where $D_{i|\mathcal{X}}$ denotes the marginal distribution of $D_{i}$ over $\mathcal{X}$.
We denote by $f(\cdot\,;\,\theta):\mathcal{X}\rightarrow\mathcal{Y}$ a predictive model parameterized by $\theta\in\mathbb{R}^{d}$,
and by $\ell:\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}$ a loss function, such that $\ell(y,f(x;\,\theta))$
measures the prediction quality of $f(\cdot\,;\,\theta)$.

#### Learning in a subspace

Modern neural networks use models with many parameters,
often more than the number of training examples.
However, it has been shown that their *intrinsic dimensionality* is much smaller than their number of parameters (Li et al., 2018),
i.e. it is possible to learn strong models in a random subspace of much smaller dimension than the full dimension of the model parameter space.
Formally, to learn model parameters $\theta\in\mathbb{R}^{d}$, given an initialization model $\theta_{0}\in\mathbb{R}^{d}$, and a random expansion matrix $P\in\mathbb{R}^{d\times k}$ (which describes the basis of a random subspace), we can describe a model in the generated random subspace by learning a vector $v\in\mathbb{R}^{k}$ as

$$ $\displaystyle\theta=\theta_{0}+Pv.$ (2) $$

Prior work in standard learning Li et al. (2018) and recently multi-task learning Zakerinia et al. (2025)
showed that $k$ can typically be chosen orders of magnitude smaller than $d$ while still allowing high
accuracy models to be trained.

In this work, we keep the matrix $P$ and the initialization $\theta_{0}$ fixed, such that $\theta$ is
completely determined by $v$.
With a slight abuse of notation, we then also write $f(\cdot\,;\,v)$ as shorthand for $f(\cdot\,;\,\theta_{0}+Pv)$.

## 3 Personalized federated learning with unlabeled clients

### 3.1 Training objective

Our goal is to generate personalized models for clients that possess only unlabeled data. We assume we have access to labeled and potentially also unlabeled clients during training, and that future clients we will encounter might be unlabeled.

The primary object we work with is a hypernetwork $h:\mathcal{P}(\mathcal{X})\rightarrow\mathbb{R}^{k}$,
where $\mathcal{P}(\mathcal{X})$ denotes the power set of $\mathcal{X}$.
The goal of $h$ is take in an unlabeled client dataset and
output the (low dimensional subpace) parameters of a personalized model
that works for the underlying client data distribution.
Formally, if $X\subset\mathcal{X}$ then $h$ is defined as

$$ $h(X)\coloneqq h_{2}\Bigg{(}\frac{1}{|X|}\sum_{x\in X}h_{1}(x)\Bigg{)},$ (3) $$

where $h_{1}:\mathcal{X}\rightarrow\mathbb{R}^{e}$ is feature extraction module, which is followed by an average across the (batch of) features, and a fully-connected module $h_{2}:\mathbb{R}^{e}\rightarrow\mathbb{R}^{k}$.

To generate a client model $f$ given $X$, we first compute $v=h(X)$,
then the full-dimensional parameters for $f$
are obtained by random expansion: $\theta=\theta_{0}+Pv$ as defined in ([2](#S2.E2)).
The final personalized model is $f(\cdot\,;\,\theta):\mathcal{X}\rightarrow\mathcal{Y}$.
Note that $P$ is a random matrix and the initialization $\theta_{0}$ is
typically random, and both are fixed before training.
Therefore, these matrices can be generated on the client using an
agreed-on random seed, and do not need to be transmitted via the
network.

We denote the trainable parameters of $h$ by $\psi_{h}$.
Additionally, FLowDUP training also uses a learnable regularization, $\psi_{r}\in\mathbb{R}^{k}$,
in the low dimensional model subspace to prevent overfitting.
The learnable parameters of FLowDUP are therefore $\psi\coloneqq(\psi_{h},\psi_{r})$.

Training FLowDUP means to learn parameters $\psi$ that, given unlabeled data $X$, are able to output personalized client model parameters that work on the distribution that $X$ was drawn from.
Our proposed training objective for doing this has two parts, a loss $\mathcal{L}$, and a (learnable)
regularizer, $\Omega$.

$$ $\min_{\psi}\ \mathcal{L}(\psi)+\lambda\Omega(\psi),$ (4) $$

The loss measures the quality of a model that is generated for a client, and can therefore
be evaluated only on clients that have at least some labeled data,

$$ $\displaystyle\mathcal{L}(\psi)$ $\displaystyle\coloneqq\sum_{i\in\mathcal{C}}\ \sum_{(x^{\prime},y^{\prime})\in(X^{\prime}_{i},Y^{\prime}_{i})}\ell(y^{\prime},f(x^{\prime};\,h(X_{i}\,;\,\psi_{h}))),$ (5) $$

where, $\mathcal{C}$ is a cohort (subset) of clients, $X_{i}$ is a batch of data without labels
from client $i$ and $(X^{\prime}_{i},Y^{\prime}_{i})$ is a batch of data with labels from client $i$.
In case $(X_{i},Y_{i})$ are empty, for example because client $i$ has no labeled data, the value of the sum is simply taken to be 0.
Intuitively, $\mathcal{L}$ penalizes the hypernetwork parameters $\psi_{h}$ for outputting
models that perform poorly on the clients’ data distributions.
Notice that different data batches are used to generate the client model and evaluate it. This is because we want $f$ to be good on the client distribution, and not just on the actual batch used to generate $f$.

The regularizer ensures that the learned models do not diverge too much from each other.
It prevents overfitting to individual clients by penalizing large deviations between the client
model parameters and a global learned regularizer. It can be computed on all clients, because
only unlabeled data is required for its evaluation,

$$ $\displaystyle\Omega(\psi)$ $\displaystyle\coloneqq\sum_{i\in\mathcal{C}}\|h(X_{i}\,;\,\psi_{h})-\psi_{r}\|^{2},$ (6) $$

where, again, $\mathcal{C}$ is a cohort of clients and $X_{i}$ is an unlabeled batch of data from client $i$.
Note that, as required for federated learning, the participating clients and data batches can be
different in every update step. The regularization strength, $\lambda\in\mathbb{R}$, is a hyperparameter.
We provide a more in depth and formal motivation for our loss function $\mathcal{L}$, based on our generalization
bound, in Section [4](#S4).

### 3.2 Training procedure

Figure: Algorithm 1 FLowDUP

Figure: Algorithm 2 ClientUpdate

FLowDUP is designed to be trainable by standard federated learning frameworks.
Algorithms [1](#alg1) and [2](#alg2) show the necessary steps.
Algorithm [1](#alg1) shows the server steps of FLowDUP that follow
a standard Federated Averaging (McMahan et al., 2017) pattern.
Training proceeds over a number of rounds. Each round, the server samples a subset of clients with the restriction that a certain fraction, $\alpha$, are clients with labeled data.
Each client in the subset receives the current parameters $\psi$ from the server, computes an update to $\psi$ and shares this update with the server. The server then aggregates the client updates and uses the aggregate to update the learnable parameters using an update rule of their choice.

The most important step, namely the ClientUpdate, is shown in Algorithm [2](#alg2).
The client receives the latest iteration of the parameters $\psi$ from the server and computes
an update direction to those parameters by training on their local data for $E$ local epochs.
At each epoch, the client samples a batch of data (line [5](#alg2.l5))
and randomly splits the batch into two equal parts (lines [6](#alg2.l6), [7](#alg2.l7)).
The first part is used to generate the subspace parameter $v$ (line [8](#alg2.l8)),
note that no labels are required for this step.
The client then computes its contribution to the regularizer value $\Omega$ (line [9](#alg2.l9)).
Next, the labeled clients compute the labeled portion of the loss $\mathcal{L}$ (line [11](#alg2.l11))
of the generated client model on the second half of the batch.
The total loss is computed (line [12](#alg2.l12)) and the client updates the learnable parameters
$\psi$ with a single gradient step (line [13](#alg2.l13)).
After all local epochs have been run, the client computes the difference between the final and initial parameters (line [16](#alg2.l16)) and returns this update to the server.

### 3.3 Model generation

After training, inference with FLowDUP is simple and lightweight, in particular it requires no model training
or finetuning, just a single forward pass through $h$.
Given a client, with unlabeled data $X=(x)_{j=1}^{m}$, the client gets $\psi$ from the server and generates a personalized model by passing $X$ through the hypernetwork: $v=h(X\,;\,\psi_{h})$ and expanding to full dimensionality:
$\theta=\theta_{0}+Pv$. The client then has a personalized classifier $f(\cdot\,;\,\theta)$ that it can use
(either on $X$ and/or on future data).

## 4 Theory

In this section, we provide a generalization bound for *transductive multi-task learning* in the PAC-Bayesian framework, which provides a theoretical justification for our algorithm.
In contrast to *inductive learning*, which uses a labeled training dataset to learn a model for inference on future samples, *transductive learning* (Vapnik, 1998), involves access to a set of unlabeled data for which prediction are desired. Similarly, in transductive multi-task learning, there are both labeled and unlabeled tasks, and the goal is to learn models for all tasks jointly. As our main theoretical contribution, we prove a PAC-Bayesian generalization bound for transductive multi-task learning for stochastic algorithms that generate a model given an unlabeled dataset. The proof and the general results are provided in Appendix [A](#A1). Here we provide a result tailored for our algorithm. Specifically, define a Gaussian distribution over the parameters of the hypernetwrk, $\rho_{h}=\mathcal{N}(\psi_{h}\,;\,\alpha_{h}\text{Id})$ for a fixed $\alpha_{h}$, and $n$ posterior models for $n$ clients as $Q_{i}=\mathcal{N}(h(S_{i}\,;\,\psi_{h})\,;\,\alpha_{\theta}\text{Id})$, and a regularization distribution as $\mathcal{Q}=\mathcal{N}(\psi_{r}\,;\,\alpha_{r}\text{Id})$. Our result bounds the gap between the true risk, $\mathcal{R}$, of all clients

$$ $\displaystyle\mathcal{R}(\rho_{h})$ $\displaystyle=\mathbb{E}_{\psi_{h}^{\prime}\sim\rho_{h}}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{(x^{\prime},y^{\prime})\sim\mathcal{D}_{i}}\ell\big{(}y^{\prime},f(x^{\prime};\,h(S_{i}\,;\,\psi_{h}^{\prime}))\big{)},$ (7) and the training risk, $\widehat{\mathcal{R}}$ , of the labeled clients $\displaystyle\widehat{\mathcal{R}}(\rho_{h})$ $\displaystyle=\mathbb{E}_{\psi_{h}^{\prime}\sim\rho_{h}}\frac{1}{n_{L}}\sum_{i=1}^{n_{L}}\frac{1}{m}\sum_{j=1}^{m}\ell\big{(}y^{j}_{i},f(x^{j}_{i};\,h(S_{i}\,;\,\psi_{h}^{\prime}))\big{)}.$ $$

theoremgaussianbound
For all $\delta>0$, and any loss function $\ell:\mathcal{Y}\times\mathcal{Y}\rightarrow[0,1]$, the following statement holds
with probability at least $1-\delta$ over the sampling of $n_{L}$ clients of $n$ clients and randomness of the dataset.
For all parameter vectors, $\psi=(\psi_{h},\psi_{r})$:

$$ $\displaystyle\mathcal{R}(\rho_{h})\leq\widehat{\mathcal{R}}(\rho_{h})$ $\displaystyle\ +\ \sqrt{\Big{(}1-\frac{n_{L}}{n}\Big{)}\frac{\frac{1}{2\alpha_{h}}\|\psi_{h}\|^{2}+c_{1}}{2n_{L}}}$ $\displaystyle\ +\ \operatorname*{\mathbb{E}}_{\psi_{h}^{\prime}\sim\rho}\sqrt{\frac{\frac{1}{2\alpha_{\theta}}\sum_{i=1}^{n}\operatorname*{\mathbb{E}}_{\psi_{r}^{\prime}}\|h(S_{i}\,;\,\psi_{h}^{\prime})-\psi_{r}^{\prime}\|^{2}+\frac{1}{2\alpha_{r}}\|\psi_{r}\|^{2}+c_{2}}{2mn}}$ (8) $$

where $c_{1}$ and $c_{2}$ are logarithmic terms in $n$ and $n_{L}$.

#### Discussion.

[Section˜4](#S4) provides a direct mathematical justification for
FLowDUP.
Our overall goal is to achieve high accuracy across all tasks, i.e. minimize the left hand side of ([8](#S4.E8)).
That is not a computable quantity, though, so as a proxy we instead minimize its upper
bound on the right hand side of ([8](#S4.E8)).
That consists of $\widehat{\mathcal{R}}(\rho_{h})$, which corresponds to the training loss across labeled tasks, i.e. FLowDUP’s loss $\mathcal{L}$, and some complexity terms. Besides $\|\psi_{h}\|$ and $\|\psi_{r}\|$,
which are automatically minimized when using optimization steps with weight decay, the latter in
particular contains $\sum_{i=1}^{n}\operatorname*{\mathbb{E}}_{\psi_{r}^{\prime}}\|h(S_{i}\,;\,\psi_{h}^{\prime})-\psi_{r}^{\prime}\|^{2}$,
which corresponds to FLowDUP’s
regularizer $\Omega$.
The difference between the theoretical viewpoint and the practical algorithm
is that in practice we use deterministic neural networks instead of stochastic models,
so no expected value operations are required. Also, we treat the regularization strength as
a hyper-parameter that can be model-selected, instead of relying on the (often overly
pessimistic) values provided by generalization theory.

On a technical level, [Section˜4](#S4) has a number of desirable properties.
Firstly, it directly reflects the transductive setting, as it controls the
risk across all clients in terms of the training loss of just the labeled ones.
The complexity term, however, is computed from all clients and its denominator scales with
the total number of available samples, which justifies the use of unlabeled clients
during training, which contribute to the learning of a shared regularizer.
Finally, the sample complexity of the first complexity term is improved by $\sqrt{1-n_{L}/n}$
compared to standard inductive bounds, which holds the promise of better between-client
generalization in transductive compared to inductive learning.

## 5 Experiments

Here we present our empirical results. In [5.1](#S5.SS1)
we explain our experimental setup, in [5.2](#S5.SS2)
we include our main results and in [5.3](#S5.SS3) we provide
additional experiments and ablation studies to better understand FLowDUP.
Implementation details and further ablation studies can be found in Appendix [B](#A2).
The values in all results tables are given as the mean and standard deviation of runs over 3 random seeds.
We implement our experiments using the pfl-research library (Granqvist et al., 2024).

### 5.1 Experimental Details

#### Baselines

We are interested in producing a personalized model for an unlabeled client and consider only baselines that
are capable of doing this.
The simplest to consider are those that train a single (full dimensional) global model $f$ on just
the labeled clients.
This of course results in a predictive model that can be used by any unlabeled client. In this category we include
Federated Averaging (FedAvg) (McMahan et al., 2017) and FedProx (Li et al., 2020).
We also include a baseline that we call LD-FedAvg that trains the low dimensional subspace parameterization
of $f$, using federated averaging.
This baseline therefore has the same client model $f$ parameterization as the personalized models
generated by FLowDUP.
Finally, we include FedTTA (Ye et al., 2024) a federated test time adaption method. It personalizes a
globally trained model $f$ to an unlabeled client with a gradient update obtained by a forward pass
with unlabeled data through an adaptation model.

#### Datasets

We include a range of datasets with different types of statistically heterogeneous clients.
Following prior work in personalized FL we simulate label heterogeneity with a non-IID partitioning of
CIFAR-10 (Krizhevsky, 2009), where each client receives 100 datapoints from only 2 classes.
We simulate heterogeneity in the features $\mathcal{X}$ by partitioning and rotating Fashion-MNIST
(Xiao et al., 2017) and MNIST (LeCun et al., 1998), as first done by Ghosh et al. (2020).
Specifically, we do this by creating an IID partitioning of the dataset into clients, where each client
receives 100 randomly IID sampled datapoints.
Then each client randomly samples a rotation from $\{0^{\circ},90^{\circ},180^{\circ},270^{\circ}\}$ and rotates all of their
images by that amount.
Finally, we also include results for FEMNIST (Caldas et al., 2018), a federated dataset with 62 classes and
an existing partition into around 3500 clients, holding in total $\approx 800000$ datapoints.
The clients in FEMNIST exhibit both feature (different writing styles) and label (different
class frequencies per client) heterogeneity.
For each of the above datasets we simulate a range of different levels of label prevalence
in the clients.
Specifically, for $p\in\{0.1,0.2,1.0\}$ we randomly choose $pn$ of the $n$ total clients
to have labels and the rest are fully unlabeled.

#### Network architectures

We run experiments using two different architectures for the client model $f$, a CNN following prior work (McMahan et al., 2017) and a ResNet18 (He et al., 2016).
On partitioned CIFAR10 we present results for both the CNN and the ResNet18, for the remaining datasets
we use a CNN.
For FLowDUP recall the structure of $h$ is:
$h(X)=h_{2}(\frac{1}{|X|}\sum_{x\in X}h_{1}(x))$.
For our main results in Section [5.2](#S5.SS2) we always implement $h_{1}$ using the
same architecture (CNN or ResNet18) as $f$ for that particular setting.
In Section [5.3](#S5.SS3) we explore the case when they differ.
We implement $h_{2}$ as a fully connected network with a single hidden layer.
Both FLowDUP and LD-FedAvg work with subspace parameterizations of the client models $f$.
For the results in Section [5.2](#S5.SS2) we set the subspace dimension
to $k=10^{4}$.
This represents an approximately $15\times$ reduction in the number of client model parameters for the CNN and $1100\times$ for ResNet18. We examine the impact the choice of $k$ has in [5.3](#S5.SS3).
For FedTTA the global model architecture is the same as that used for FedAvg and the
adaptation model is the same 3-layer network as used in (Ye et al., 2024).

### 5.2 Experimental results

**Table 1: Class-partitioned CIFAR10 for varying fractions, $p$, of the training clients having labeled data. Accuracy on unlabeled test clients.**
| Model | CNN | ResNet18 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| $\bm{p}$ | 0.1 | 0.2 | 1.0 | 0.1 | 0.2 | 1.0 |
| FedAvg | 47.9 $\pm$ 0.1 | 53.5 $\pm$ 0.4 | 63.6 $\pm$ 0.4 | 63.6 $\pm$ 0.9 | 68.9 $\pm$ 0.7 | 75.9 $\pm$ 0.3 |
| FedProx | 47.6 $\pm$ 0.3 | 53.3 $\pm$ 0.3 | 63.7 $\pm$ 0.6 | 63.5 $\pm$ 0.8 | 68.8 $\pm$ 0.5 | 75.9 $\pm$ 0.3 |
| LD-FedAvg | 41.6 $\pm$ 0.6 | 47.8 $\pm$ 1.1 | 56.2 $\pm$ 0.7 | 54.6 $\pm$ 0.8 | 60.0 $\pm$ 0.5 | 65.4 $\pm$ 0.1 |
| FedTTA | 57.2 $\pm$ 1.0 | 60.2 $\pm$ 1.1 | 69.7 $\pm$ 3.8 | 69.3 $\pm$ 2.3 | 77.2 $\pm$ 1.5 | 81.9 $\pm$ 2.1 |
| FLowDUP | 66.1 $\pm$ 1.5 | 75.3 $\pm$ 1.8 | 86.6 $\pm$ 0.1 | 72.7 $\pm$ 1.3 | 82.9 $\pm$ 1.8 | 92.3 $\pm$ 0.4 |

**Table 2: Rotated Fashion-MNIST (left) and FEMNIST (right) for varying fractions, $p$, of the training clients having labeled data. Accuracy on unlabeled test clients.**
| Dataset | Rotated Fashion-MNIST | FEMNIST |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| $\bm{p}$ | 0.1 | 0.2 | 1.0 | 0.1 | 0.2 | 1.0 |
| FedAvg | 77.7 $\pm$ 0.1 | 79.2 $\pm$ 0.2 | 81.4 $\pm$ 0.6 | 76.5 $\pm$ 0.3 | 78.0 $\pm$ 0.3 | 84.3 $\pm$ 0.1 |
| FedProx | 77.3 $\pm$ 0.7 | 78.7 $\pm$ 0.4 | 80.0 $\pm$ 0.6 | 72.3 $\pm$ 0.4 | 74.8 $\pm$ 0.4 | 83.3 $\pm$ 0.0 |
| LD-FedAvg | 76.1 $\pm$ 1.1 | 78.8 $\pm$ 0.3 | 81.3 $\pm$ 0.1 | 61.4 $\pm$ 1.4 | 74.4 $\pm$ 0.3 | 80.6 $\pm$ 0.4 |
| FedTTA | 78.3 $\pm$ 0.5 | 80.2 $\pm$ 0.1 | 81.2 $\pm$ 0.2 | 71.1 $\pm$ 2.4 | 72.6 $\pm$ 3.5 | 74.3 $\pm$ 3.1 |
| FLowDUP | 81.5 $\pm$ 0.3 | 83.3 $\pm$ 0.5 | 87.3 $\pm$ 0.2 | 84.5 $\pm$ 0.4 | 86.4 $\pm$ 0.3 | 89.3 $\pm$ 0.1 |

Our main results are shown in Tables [1](#S5.T1) and [2](#S5.T2).
Table [1](#S5.T1) shows class-partitioned CIFAR10
as we vary the fraction of labeled clients present in the training set $(p)$.
The left part of the table shows the results for when the model architecture is a CNN
and the right part for ResNet18.
The results show that FLowDUP performs best across the board.
As expected all methods improve as the fraction of labeled clients increases and when we
switch from using a CNN to a ResNet18.
Notably, even when using the CNN architectures, FLowDUP is able to outperform the
non-personalized baselines using a ResNet18.
We also observe that LD-FedAvg has significantly lower performance than FedAvg. This shows that
in fact for a single global predictive model, lower subspace dimension reduces the power of the model.
However, FLowDUP’s strong performance shows that, given a well learned representation of client
heterogeneity in $h$, a low dimensional subspace model can obtain good personalized performance.

Table [2](#S5.T2) shows the results for Rotated Fashion-MNIST (left)
and FEMNIST (right). The results for Rotated MNIST can be found in the Appendix in Table
[5](#A2.T5).
For these datasets all methods use the CNN architecture.
When comparing to CIFAR10, for both these datasets FedTTA performs quite poorly.
This is to be expected given that FedTTA adapts the global model based on the output
logits of the unlabeled data.
When statistical heterogeneity is present only in the labels (as is the case for class-partitioned
CIFAR10), the logits are a good summary of the client personalization requirements.
However, in the presence of feature heterogeneity the logits alone provide a poor adaptation signal.
FLowDUP performs better as it personalizes based on all the clients feature data.

### 5.3 Additional Studies

Here we provide results from additional experiments investigating the components
of FLowDUP.
In addition to these we include in the appendix a study of the effect that training with unlabeled clients has (Tables [6](#A2.T6) and [7](#A2.T7)) as well
as the benefits of the learnable regularizer $\psi_{r}$ (Table [8](#A2.T8)).

**Table 3: Accuracy of FLowDUP for architecture combinations on class-partitioned CIFAR10.**
| Model Architecture | Fraction Labeled $\bm{(p)}$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| $\bm{h_{1}}$ | $\bm{f}$ | 0.1 | 0.2 | 0.5 | 1.0 |
| CNN | CNN | 66.1 $\pm$ 1.5 | 75.3 $\pm$ 1.8 | 83.3 $\pm$ 1.5 | 86.6 $\pm$ 0.1 |
| CNN | ResNet18 | 71.4 $\pm$ 0.5 | 80.0 $\pm$ 3.3 | 88.0 $\pm$ 0.9 | 90.7 $\pm$ 0.2 |
| ResNet18 | CNN | 67.0 $\pm$ 2.4 | 78.0 $\pm$ 3.1 | 87.3 $\pm$ 1.3 | 88.5 $\pm$ 1.1 |
| ResNet18 | ResNet18 | 72.7 $\pm$ 1.3 | 82.9 $\pm$ 1.8 | 91.0 $\pm$ 0.4 | 92.3 $\pm$ 0.4 |

#### Differing client model architectures

We train $h$ to output personalized model parameters in a subspace of dimension $k$.
However, for any given $k$ we are still free to choose the architecture of the client model $f$.
Moreover, the choice of architecture of $f$ does not affect the communication cost of the
training procedure (which is critical and typically the primary bottleneck when running federated training
in practice).
It does, however, affect the compute cost on the client side.
On the flip side, the choice of hypernetwork architecture does affect the communication cost as $h$ is transmitted
from server to client. This gives FLowDUP some flexibility when trading off communication vs
computational cost.
In Table [3](#S5.T3) we show results for CIFAR10 for different combinations of hypernetwork and
client model architectures.
The results show that for a fixed subspace dimension $k$, changing the architecture of $h_{1}$ or $f$ can have a substantial impact on performance.
Most notably, when comparing row 1 with 2, and row 3 with 4, we see
that changing $f$ from a CNN to a ResNet18 gives a performance boost of 4-5%
without incurring any additional communication cost.

**Table 4: Accuracy of FLowDUP with different subspace dimensions, $k$, on class-partitioned CIFAR10.**
| Fraction Labeled $\bm{(p)}$ | 0.1 | 0.2 | 0.5 | 1.0 |
| --- | --- | --- | --- | --- |
| FLowDUP ($k=200$) | 56.6 $\pm$ 1.9 | 65.7 $\pm$ 2.9 | 70.0 $\pm$ 1.8 | 71.2 $\pm$ 0.7 |
| FLowDUP ($k=500$) | 60.0 $\pm$ 1.9 | 69.8 $\pm$ 2.8 | 75.9 $\pm$ 1.2 | 76.9 $\pm$ 0.8 |
| FLowDUP ($k=2000$) | 64.5 $\pm$ 0.3 | 72.5 $\pm$ 1.7 | 80.3 $\pm$ 1.5 | 83.5 $\pm$ 0.9 |
| FLowDUP ($k=10000$) | 66.1 $\pm$ 1.5 | 75.3 $\pm$ 1.8 | 83.3 $\pm$ 1.5 | 86.6 $\pm$ 0.1 |

#### Varying the subspace dimension

While the choice of client model architecture affects only the computational cost of FLowDUP, the
subspace dimension $k$ affects both communication and computational cost.
This is because the final layer of $h$ has dimension $k$ and the matrices in [2](#S2.E2) scale with $k$.
Table [4](#S5.T4) shows results for FLowDUP on CIFAR10 as we vary $k$. In these
experiments we fix the architectures of $h_{1}$ and $f$ to be CNN.
We observe that performance increases monotonically with $k$.
Most notably we observe that, with the exception of $p=0.1$, FLowDUP outperforms
the baselines, even for very low values of $k$ which in turn incur low compute costs on the client.

Figure: Figure 1: Visualization of the client embeddings on Rotated Fashion-MNIST. Projection to two dimensions using PCA (left) and t-SNE (right).
Refer to caption: /html/2505.15579/assets/x1.png

#### Understanding dataset embeddings

Recall the hypernetwork architecture from ([3](#S3.E3)).
We can interpret the input to $h_{2}$ as an embedding representation of the client dataset $X$, which we call

$$ $\displaystyle r(X)\coloneqq\frac{1}{|X|}\sum_{x\in X}h_{1}(x).$ (9) $$

To better understand the inner workings of this mechanism we visualize the space of dataset embeddings.
Concretely, for each validation client in the Rotated Fashion-MNIST dataset, we compute $r(X_{i})$ for the $h$
that was trained on the training clients.
For our experiments these representations are 256-dimensional, therefore to create a visualization we apply dimensionality
reduction. We separately apply both PCA and t-SNE (van der Maaten & Hinton, 2008) to project the representations into 2 dimensions.
The scatter plots (Figure [1](#S5.F1))
show that FLowDUP learns to organize the space of client dataset
representations according to the present statistical heterogeneity, namely that the
representations are clustered according to which rotation the data has.
Note that nowhere in the learning objective is this hard coded.
Rather FLowDUP has learned the type of heterogeneity present in the client datasets, and organized the
representation space so that in this case, roughly speaking, clients with the same rotation receive the same model.

## 6 Related work

#### Personalized federated learning for unlabeled clients

First proposed by Smith et al. (2017), personalized FL arose from the observation that
statistical heterogeneity of client data distributions can make training a single global model suboptimal (Kairouz et al., 2021).
However, most works in personalized FL assume that all clients hold labeled data that they
can use to obtain a personalized model, see Appendix [C](#A3).
A small number of works have looked at the problem of personalizing models to unlabeled clients.
Ye et al. (2024) apply a popular test time adaptation method to a FL setting. A global predictive model is personalized on a client
using gradient updates obtained with forward passes through an adaptation model.
Closest to our own work are Amosy et al. (2024); Scott et al. (2024) which use a combination of an embedding
network and hypernetwork to generate personalized models using only unlabeled data.
Their approaches differ from ours in two key aspects.
Firstly, while they are able to produce personalized models for unlabeled clients, they are
not able to use the unlabeled clients during training. In contrast, FLowDUP also leverages
unlabeled clients during training, which leads to improved performance.
Secondly, and more importantly, they generate all the parameters of the personalized model,
rather than a low dimensional parameterization as FLowDUP does.
Due to the high computational cost of this approach they are only able to generate
very small personalized models for the clients.
Moreover, even for such small models
the hypernetwork is too large to send to the clients and instead remains on the server. This leads to complicated multi-step communication schemes that are
inefficient and break the federated learning paradigm of working only with aggregate client statistics in order to maintain privacy.
Finally, several works have explored Federated Representation Learning in the presence of statistically
heterogeneous clients (Zhang et al., 2020; Jang et al., 2022).
These methods typically aim to learn a feature extractor(s) using unlabeled client data.
However, obtaining a personalized predictive model on the client would still require
the client to possess labels, and as such these methods are not applicable to the task
we aim to solve.

#### Theory

Beyond standard single-task learning, it has been shown that in *multi-task learning (MTL)* (Caruana, 1997; Baxter, 1995), when learning multiple related tasks, sharing information between them can provably improve the performance. Specifically, different works had studied the generalization behavior of multi-task learning(Maurer, 2006; Crammer & Mansour, 2012; Pontil & Maurer, 2013; Pentina & Lampert, 2017; Yousefi et al., 2018; Du et al., 2021) using notions like VC-dimensions(Vapnik & Chervonenkis, 1971) or Rademacher complexity(Bartlett & Mendelson, 2002).
Recently, with the success of PAC-Bayesian bounds for studying generalization behavior of neural networks (Dziugaite & Roy, 2017), and due to their strength in parameterizing multi-task learning and meta-learning(Schmidhuber, 1987), PAC-Bayes multi-task learning has become an active line of research  (Pentina & Lampert, 2014; Amit & Meir, 2018; Rothfuss et al., 2023; Guan & Lu, 2022; Zakerinia et al., 2024), However, the focus of these works is *inductive multi-task learning*, where all tasks have access to labeled data. In contrast in this work, we introduce *transductive multi-task learning*, where only a subset of the tasks have labeled data.

## 7 Conclusion

We presented FLowDUP, a method that can generate personalized models for clients
using only a single forward pass with unlabeled data.
FLowDUP’s training objective is theoretically motivated by our new generalization bound
in a transductive multi-task framework and is able to make use of both labeled and unlabeled
clients during training for improved performance.
Our empirical evaluation demonstrated that FLowDUP can achieve strong performance
in the presence of both feature and label heterogeneity.
In additional studies, we investigated the effect of model architecture choices,
subspace dimension, the benefit of using unlabeled clients and how
FLowDUP is able to capture dataset heterogeneity.

FLowDUP’s main strength is also its main limitation: because it uses only the unlabeled
data to produce a personalized model, it will struggle in a setting where the conditional
distributions, $(D_{i})_{Y|X}$, differ across clients in a way such that knowledge of the
marginals, $(D_{i})_{X}$, are insufficient to infer good predictive models.
Such situations could be possible, for example, in subjective prediction tasks, such as
product recommendations or sentiment analysis.
This limitation is not specific to FLowDUP though: all methods that rely only on unlabeled
client data would fail in this setting, so additional information sources would be required,
such as meta-data or user participation.

## Acknowledgements

This work was supported in part by the Austrian Science Fund (FWF) [10.55776/COE12].
This research was also supported by the Scientific Service Units (SSU) of ISTA through resources provided by Scientific Computing (SciComp).

## References

- Aghajanyan et al. (2021)
Aghajanyan, A., Gupta, S., and Zettlemoyer, L.
Intrinsic dimensionality explains the effectiveness of language model fine-tuning.
In *Association for Computational Linguistics (ACL)*, 2021.
- Amit & Meir (2018)
Amit, R. and Meir, R.
Meta-learning by adjusting priors based on extended PAC-Bayes theory.
In *International Conference on Machine Learning (ICML)*, 2018.
- Amosy et al. (2024)
Amosy, O., Eyal, G., and Chechik, G.
Late to the party? On-demand unlabeled personalized federated learning.
In *IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)*, pp.  2173–2182. IEEE, 2024.
- Arivazhagan et al. (2019)
Arivazhagan, M. G., Aggarwal, V., Singh, A. K., and Choudhary, S.
Federated learning with personalization layers.
2019.
URL [http://arxiv.org/abs/1912.00818](http://arxiv.org/abs/1912.00818).
- Bartlett & Mendelson (2002)
Bartlett, P. L. and Mendelson, S.
Rademacher and Gaussian complexities: Risk bounds and structural results.
*Journal of Machine Learning Research (JMLR)*, 2002.
- Baxter (1995)
Baxter, J.
Learning internal representations.
In *Conference on Computational Learning Theory (COLT)*, 1995.
- Baxter (2000)
Baxter, J.
A model of inductive bias learning.
*Journal of Artificial Intelligence Research (JAIR)*, 12:149–198, 2000.
- Bégin et al. (2014)
Bégin, L., Germain, P., Laviolette, F., and Roy, J.-F.
PAC-Bayesian Theory for Transductive Learning.
In *International Conference on Artificial Intelligence and Statistics (AISTATS)*, 2014.
- Caldas et al. (2018)
Caldas, S., Wu, P., Li, T., Konečný, J., McMahan, H. B., Smith, V., and Talwalkar, A.
LEAF: A benchmark for federated settings.
2018.
URL [http://arxiv.org/abs/1812.01097](http://arxiv.org/abs/1812.01097).
- Caruana (1997)
Caruana, R.
Multitask learning.
*Machine Learning*, 28:41–75, 1997.
- Chen et al. (2023)
Chen, D., Yao, L., Gao, D., Ding, B., and Li, Y.
Efficient personalized federated learning via sparse model-adaptation.
In *International Conference on Machine Learning (ICML)*, 2023.
- Collins et al. (2021)
Collins, L., Hassani, H., Mokhtari, A., and Shakkottai, S.
Exploiting shared representations for personalized federated learning.
In *International Conference on Machine Learning (ICML)*, 2021.
- Crammer & Mansour (2012)
Crammer, K. and Mansour, Y.
Learning multiple tasks using shared hypotheses.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2012.
- Dinh et al. (2020)
Dinh, C. T., Tran, N. H., and Nguyen, T. D.
Personalized federated learning with moreau envelopes.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2020.
- Du et al. (2021)
Du, S. S., Hu, W., Kakade, S. M., Lee, J. D., and Lei, Q.
Few-shot learning via learning the representation, provably.
In *International Conference on Learning Representations (ICLR)*, 2021.
- Dziugaite & Roy (2017)
Dziugaite, G. K. and Roy, D. M.
Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.
In *Uncertainty in Artificial Intelligence (UAI)*, 2017.
- Fallah et al. (2020)
Fallah, A., Mokhtari, A., and Ozdaglar, A. E.
Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2020.
- Ghosh et al. (2020)
Ghosh, A., Chung, J., Yin, D., and Ramchandran, K.
An efficient framework for clustered federated learning.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2020.
- Granqvist et al. (2024)
Granqvist, F., Song, C., Áine Cahill, van Dalen, R., Pelikan, M., Chan, Y. S., Feng, X., Krishnaswami, N., Jina, V., and Chitnis, M.
pfl-research: simulation framework for accelerating research in private federated learning, 2024.
URL [https://arxiv.org/abs/2404.06430](https://arxiv.org/abs/2404.06430).
- Guan & Lu (2022)
Guan, J. and Lu, Z.
Fast-rate PAC-Bayesian generalization bounds for meta-learning.
In *International Conference on Machine Learning (ICML)*, 2022.
- Ha et al. (2017)
Ha, D., Dai, A. M., and Le, Q. V.
Hypernetworks.
In *International Conference on Learning Representations (ICLR)*, 2017.
- Hanzely et al. (2020)
Hanzely, F., Hanzely, S., Horváth, S., and Richtárik, P.
Lower bounds and optimal algorithms for personalized federated learning.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2020.
- He et al. (2016)
He, K., Zhang, X., Ren, S., and Sun, J.
Deep residual learning for image recognition.
In *IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)*, pp.  770–778. IEEE Computer Society, 2016.
- Jang et al. (2022)
Jang, J., Ha, H., Jung, D., and Yoon, S.
FedClassAvg: Local representation learning for personalized federated learning on heterogeneous neural networks.
In *International Conference on Parallel Processing*. ACM, 2022.
- Jiang et al. (2019)
Jiang, Y., Konečný, J., Rush, K., and Kannan, S.
Improving federated learning personalization via model agnostic meta learning.
2019.
URL [http://arxiv.org/abs/1909.12488](http://arxiv.org/abs/1909.12488).
- Kairouz et al. (2021)
Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., Bonawitz, K., Charles, Z., Cormode, G., Cummings, R., D’Oliveira, R. G. L., Eichner, H., Rouayheb, S. E., Evans, D., Gardner, J., Garrett, Z., Gascón, A., Ghazi, B., Gibbons, P. B., Gruteser, M., Harchaoui, Z., He, C., He, L., Huo, Z., Hutchinson, B., Hsu, J., Jaggi, M., Javidi, T., Joshi, G., Khodak, M., Konecný, J., Korolova, A., Koushanfar, F., Koyejo, S., Lepoint, T., Liu, Y., Mittal, P., Mohri, M., Nock, R., Özgür, A., Pagh, R., Qi, H., Ramage, D., Raskar, R., Raykova, M., Song, D., Song, W., Stich, S. U., Sun, Z., Suresh, A. T., Tramèr, F., Vepakomma, P., Wang, J., Xiong, L., Xu, Z., Yang, Q., Yu, F. X., Yu, H., and Zhao, S.
Advances and open problems in federated learning.
*Foundations and Trends in Machine Learning*, 14(1–2), 2021.
- Krizhevsky (2009)
Krizhevsky, A.
Learning multiple layers of features from tiny images.
*Master’s thesis, Department of Computer Science, University of Toronto*, 2009.
- LeCun et al. (1998)
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P.
Gradient-based learning applied to document recognition.
*Proc. IEEE*, 1998.
- Li et al. (2018)
Li, C., Farkhoor, H., Liu, R., and Yosinski, J.
Measuring the intrinsic dimension of objective landscapes.
In *International Conference on Learning Representations (ICLR)*, 2018.
- Li et al. (2020)
Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V.
Federated optimization in heterogeneous networks.
In *Proceedings of the Third Conference on Machine Learning and Systems, MLSys*, 2020.
- Li et al. (2021)
Li, T., Hu, S., Beirami, A., and Smith, V.
Ditto: Fair and robust federated learning through personalization.
In *International Conference on Machine Learning (ICML)*, 2021.
- Lin et al. (2022)
Lin, S., Han, Y., Li, X., and Zhang, Z.
Personalized federated learning towards communication efficiency, robustness and fairness.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2022.
- Lotfi et al. (2022)
Lotfi, S., Finzi, M., Kapoor, S., Potapczynski, A., Goldblum, M., and Wilson, A. G.
PAC-Bayes compression bounds so tight that they can explain generalization.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2022.
- Marfoq et al. (2021)
Marfoq, O., Neglia, G., Bellet, A., Kameni, L., and Vidal, R.
Federated multi-task learning under a mixture of distributions.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2021.
- Marfoq et al. (2022)
Marfoq, O., Neglia, G., Vidal, R., and Kameni, L.
Personalized federated learning through local memorization.
In *International Conference on Machine Learning (ICML)*, 2022.
- Maurer (2004)
Maurer, A.
A note on the PAC Bayesian theorem.
*CoRR*, cs.LG/0411099, 2004.
URL [http://arxiv.org/abs/cs.LG/0411099](http://arxiv.org/abs/cs.LG/0411099).
- Maurer (2006)
Maurer, A.
Bounds for linear multi-task learning.
*Journal of Machine Learning Research (JMLR)*, 7:117–139, 2006.
- McAllester (1998)
McAllester, D. A.
Some PAC-Bayesian theorems.
In *Conference on Computational Learning Theory (COLT)*, 1998.
- McMahan et al. (2017)
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A.
Communication-efficient learning of deep networks from decentralized data.
In *International Conference on Artificial Intelligence and Statistics (AISTATS)*, 2017.
- Park et al. (2025)
Park, S., Jeong, J., Kim, Y., Lee, J., and Lee, N.
ZIP: An efficient zeroth-order prompt tuning for black-box vision-language models.
In *International Conference on Learning Representations (ICLR)*, 2025.
- Pentina & Lampert (2014)
Pentina, A. and Lampert, C. H.
A PAC-Bayesian bound for lifelong learning.
In *International Conference on Machine Learning (ICML)*, 2014.
- Pentina & Lampert (2017)
Pentina, A. and Lampert, C. H.
Multi-task learning with labeled and unlabeled tasks.
In *International Conference on Machine Learning (ICML)*, 2017.
- Pontil & Maurer (2013)
Pontil, M. and Maurer, A.
Excess risk bounds for multitask learning with trace norm regularization.
In *Conference on Computational Learning Theory (COLT)*, pp.  55–76. PMLR, 2013.
- Rothfuss et al. (2023)
Rothfuss, J., Josifoski, M., Fortuin, V., and Krause, A.
Scalable PAC-Bayesian meta-learning via the PAC-Optimal hyper-posterior: From theory to practice.
*Journal of Machine Learning Research (JMLR)*, 2023.
- Schmidhuber (1987)
Schmidhuber, J.
Evolutionary principles in self-referential learning.
Diploma thesis, Technical University of Munich, Germany, 1987.
- Scott et al. (2024)
Scott, J., Zakerinia, H., and Lampert, C. H.
Pefll: Personalized federated learning by learning to learn.
In *International Conference on Learning Representations (ICLR)*, 2024.
- Seldin et al. (2012)
Seldin, Y., Laviolette, F., Cesa-Bianchi, N., Shawe-Taylor, J., and Auer, P.
PAC-Bayesian inequalities for martingales.
*IEEE Transactions on Information Theory*, 58(12):7086–7093, 2012.
- Shamsian et al. (2021)
Shamsian, A., Navon, A., Fetaya, E., and Chechik, G.
Personalized federated learning using hypernetworks.
In *International Conference on Machine Learning (ICML)*, 2021.
- Smith et al. (2017)
Smith, V., Chiang, C., Sanjabi, M., and Talwalkar, A.
Federated multi-task learning.
In *Conference on Neural Information Processing Systems (NeurIPS)*, 2017.
- Thrun & Pratt (1998)
Thrun, S. and Pratt, L. (eds.).
*Learning to Learn*.
Kluwer Academic Press, 1998.
- van der Maaten & Hinton (2008)
van der Maaten, L. and Hinton, G.
Visualizing data using t-sne.
*Journal of Machine Learning Research*, 9(Nov):2579–2605, 2008.
- Vapnik (1998)
Vapnik, V. N.
*Statistical Learning Theory*.
Wiley, 1998.
ISBN 978-0-471-03003-4.
- Vapnik & Chervonenkis (1971)
Vapnik, V. N. and Chervonenkis, A. Y.
On the uniform convergence of relative frequencies of events to their probabilities.
*Theory of Probability & Its Applications*, 16(2):264–280, 1971.
- Wu et al. (2023)
Wu, Y., Zhang, S., Yu, W., Liu, Y., Gu, Q., Zhou, D., Chen, H., and Cheng, W.
Personalized federated learning under mixture of distributions.
In *International Conference on Machine Learning (ICML)*, 2023.
- Xiao et al. (2017)
Xiao, H., Rasul, K., and Vollgraf, R.
Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.
2017.
URL [http://arxiv.org/abs/1708.07747](http://arxiv.org/abs/1708.07747).
- Ye et al. (2023)
Ye, R., Ni, Z., Wu, F., Chen, S., and Wang, Y.
Personalized federated learning with inferred collaboration graphs.
In *International Conference on Machine Learning (ICML)*, 2023.
- Ye et al. (2024)
Ye, T., Chen, C., Wang, Y., Li, X., and Gao, M.
UPFL: unsupervised personalized federated learning towards new clients.
In *Proceedings of the 2024 SIAM International Conference on Data Mining, SDM*, 2024.
- Yousefi et al. (2018)
Yousefi, N., Lei, Y., Kloft, M., Mollaghasemi, M., and Anagnostopoulos, G. C.
Local Rademacher complexity-based learning guarantees for multi-task learning.
*Journal of Machine Learning Research (JMLR)*, 19(38):1–47, 2018.
- Zakerinia et al. (2024)
Zakerinia, H., Behjati, A., and Lampert, C.
More flexible PAC-Bayesian meta-learning by learning learning algorithms.
In *International Conference on Machine Learning (ICML)*, 2024.
- Zakerinia et al. (2025)
Zakerinia, H., Ghobadi, D., and Lampert, C. H.
From low intrinsic dimensionality to non-vacuous generalization bounds in deep multi-task learning.
*arXiv preprint arXiv:2501.19067*, 2025.
- Zhang et al. (2020)
Zhang, F., Kuang, K., You, Z., Shen, T., Xiao, J., Zhang, Y., Wu, C., Zhuang, Y., and Li, X.
Federated unsupervised representation learning.
2020.
URL [https://arxiv.org/pdf/2010.08982](https://arxiv.org/pdf/2010.08982).
- Zhang et al. (2023)
Zhang, H., Li, C., Dai, W., Zou, J., and Xiong, H.
FedCR: Personalized federated learning based on across-client common representation with conditional mutual information regularization.
In *International Conference on Machine Learning (ICML)*, 2023.

## Appendix A Transductive multi-task learning

In this section, we provide our general theoretical contributions and their proof.

#### PAC-Bayesian theory

(McAllester, 1998) studies the generalization behavior of stochastic models (a distribution over a set of models $f\in\mathcal{F}$). Formally, when training a posterior distribution $Q\in\mathcal{M}(\mathcal{F})$ using a dataset $S$ consisting of $m_{L}$ i.i.d. labeled samples from a distribution $D$ over $\mathcal{X}\times\mathcal{Y}$ , PAC-Bayes bounds guarantee upper-bounds for the true risk of a posterior $Q$ i.e. $\mathcal{R}(Q)=\operatorname*{\mathbb{E}}_{f\sim Q}\operatorname*{\mathbb{E}}_{(x,y)\sim D}\ell(f,x,y)$ based on its training risk ($\widehat{\mathcal{R}}(Q)=\operatorname*{\mathbb{E}}_{f\sim Q}\frac{1}{m_{L}}\sum_{i=1}^{m_{L}}\ell(f,x_{i},y_{i})$), and a complexity term based on $\operatorname{\mathbf{KL}}(Q\|P)$ where $P\in\mathcal{M}(\mathcal{F})$ is a data-independent prior over the space of models, and $\operatorname{\mathbf{KL}}$ is the Kullback-Leibler divergence. As an example, from (Maurer, 2004) we have that for any prior $P$, and any $\delta>0$, over the sampling of the dataset $S$ for all posteriors $Q$ it holds:

$$ $\displaystyle\mathcal{R}(Q))$ $\displaystyle\leq\widehat{\mathcal{R}}(Q)+\sqrt{\frac{\operatorname{\mathbf{KL}}(Q\|P)+\log(\frac{2\sqrt{m_{L}}}{\delta})}{2m_{L}}},$ (10) $$

This bound holds in an inductive learning setting, i.e. we have a distribution $D$, and the goal is to generalize from training data to the unseen data. On the other hand, in transductive learning (Vapnik, 1998), there is a set of labeled data, and a set of unlabeled data, and the goal is to generalize from labeled to unlabeled data. Transductive learning has also been studied in the PAC-Bayes literature. Therefore, if we have $m$ samples which $m_{L}$ of them is labeled we want to generalize from $\widehat{\mathcal{R}}(Q)$ to $\operatorname*{\mathbb{E}}_{f\sim Q}\frac{1}{m}\sum_{i=1}^{m}\ell(f,x_{i},y_{i})$. For this problem,
Bégin et al. (2014) have proved for any prior $P$, and any $\delta>0$, over the sampling of $m_{L}$ labeled samples out of $m$ samples (without replacement) for all posteriors $Q$ it holds:

$$ $\displaystyle\operatorname*{\mathbb{E}}_{f\sim Q}\frac{1}{m}\sum_{i=1}^{m}\ell(f,x_{i},y_{i})\leq\widehat{\mathcal{R}}_{L}(Q)+\sqrt{(1-\frac{m_{L}}{m})\frac{\operatorname{\mathbf{KL}}(Q\|P)+\log(\frac{t(m_{L},m)}{\delta})}{2m_{L}}},$ (11) $$

where $t(m_{L},m)=3\log(m_{L})\sqrt{m_{L}(1-\frac{m_{L}}{m})}$. Since the dominant term in both bounds is the $\operatorname{\mathbf{KL}}$ terms, the transductive setting has an improved generalization guarantee by a factor of $\sqrt{1-\frac{m_{L}}{m}}$.

#### PAC-Bayes multi-task learning

Beyond standard single-task learning, it has been shown theoretically that when learning multiple related tasks, sharing information between them can improve performance. Formally, in *multi-task learning (MTL)* (Caruana, 1997), there are $n$ tasks with different distributions $D_{1},\dots,D_{n}$ and respective datasets $S_{1},\dots,S_{n}$, and the goal is to learn individual models (or Posteriors) jointly. *Meta-learning* (Schmidhuber, 1987), or *learning to learn* (Thrun & Pratt, 1998), extends multi-task learning to the setting where there will also be future tasks that are not observed yet, with the assumption that the observed tasks and future tasks are all i.i.d. samples from a distribution $\tau$ over an environment of tasks (Baxter, 2000). In multi-task learning, the goal is to generalize from the average training error to the average true risk of observed clients, and in meta-learning, the goal is to generalize to the expected true risk over $\tau$. Given that in the training step, we can not learn a model for future tasks, meta-learning is formalized as learning an algorithm to apply to a future task. However, past works focused on multi-task learning and meta-learning in an inductive way, i.e. the model suggested by Baxter (2000). In this work, we focus on a transductive version where there are $n$ tasks, out of which $n_{L}$ tasks are selected randomly to have a labeled dataset, and the remaining $n-n_{L}$ tasks have only unlabeled data.

Following Pentina & Lampert (2014), PAC-Bayes is a popular framework to study multi-task learning and meta-learning, given its ability to share information between tasks through the concept of a prior.
In this work, we follow the framework introduced in Zakerinia et al. (2024). As our theoretical contribution, we prove new bounds for the transductive multi-task learning scenario we described.

Formally, we have access to $n_{L}$ labeled datasets and $n-n_{L}$ unlabeled datasets, and our goal is to generate models with good performance on all tasks. Since we have unlabeled tasks, we work with a family of algorithms $\mathcal{A}$ that can create posteriors from unlabeled data i.e. $A:\mathcal{P}(\mathcal{X})\to\mathcal{M}(\mathcal{F})$ is a mapping from the set of unlabeled datasets to models. We aim to learn a stochastic algorithm i.e. a *meta-posterior* over a set of algorithms ($\rho\in\mathcal{M}(\mathcal{A})$) that have good performance on all tasks. For each algorithm, $A(X_{1}),\dots,A(X_{n})$ are the task posteriors, and we denote to $\mathcal{Q}(A)$ as a hyper-posterior, a distribution over priors to capture the similarities between the outputs of the algorithm, and to have a data-dependent prior for our PAC-Bayes bounds. For a detailed explanation of the role of hyper-posterior, we refer the reader to Zakerinia et al. (2024).

For each meta-posterior, our goal is to minimize the average true risk of all tasks:

$$ $\displaystyle\mathcal{R}(\rho)$ $\displaystyle=\mathbb{E}_{A\sim\rho}\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}_{(x_{i},y_{i})\sim\mathcal{D}_{i}}\operatorname*{\mathbb{E}}_{f\sim A(X_{i})}\ell(y_{i},f(x_{i}))$ (12) $$

However, the true distribution of tasks is unknown, and we can not compute the training risk of unlabeled clients; therefore, the computable object is the average training risk of labeled clients:

$$ $\displaystyle\widehat{\mathcal{R}}(\rho)$ $\displaystyle=\mathbb{E}_{A\sim\rho}\frac{1}{n_{L}}\sum_{i=1}^{n_{L}}\frac{1}{m}\sum_{j=1}^{m}\operatorname*{\mathbb{E}}_{f\sim A(X_{i})}\ell(y_{i,j},f(x_{i,j}))$ (13) $$

We now state our main theoretical results:

###### Theorem A.1 .

For any fixed meta-prior $\pi$, fixed hyper-prior $\mathcal{P}$ and any $\delta>0$
with probability at least $1-\delta$ over the sampling of the datasets,
for all distributions $\rho\in\mathcal{M}(\mathcal{A})$ over algorithms,
and for all hyper-posterior functions $\mathcal{Q}:\mathcal{A}\rightarrow\mathcal{M}(\mathcal{M}(\mathcal{F}))$
it holds

$$ $\displaystyle\mathcal{R}(\rho)$ $\displaystyle\leq\widehat{\mathcal{R}}(\rho)+\sqrt{(1-\frac{n_{L}}{n})\frac{\operatorname{\mathbf{KL}}(\rho||\pi)+\log(\frac{t(n_{L},n)}{\delta})}{2n_{L}}}+\operatorname*{\mathbb{E}}_{A\sim\rho}\sqrt{\frac{C(A,\mathcal{Q},\mathcal{P})+\log(\frac{8mn}{\delta})+1}{2mn}}$ (14) $$

where,

$$ $\displaystyle C(A,\mathcal{Q},\mathcal{P})$ $\displaystyle=\operatorname{\mathbf{KL}}(\mathcal{Q}(A)\|\mathcal{P})\quad+\operatorname*{\mathbb{E}}_{P\sim\mathcal{Q}(A)}\sum_{i=1}^{n}\operatorname{\mathbf{KL}}(A(S_{i})||P)$ (15) $$

###### Proof.

For the proof, we define the following intermediate objective, which quantifies the training risk of all clients. Note that this object exists (There is a labeling for the unlabeled data, however, we do not have access to them, i.e. the loss function for unlabeled clients is well-defined, but we can not compute it.)

$$ $\displaystyle\widetilde{\mathcal{R}}(\rho)$ $\displaystyle=\mathbb{E}_{A\sim\rho}\frac{1}{n}\sum_{i=1}^{n}\frac{1}{m}\sum_{j=1}^{m}\operatorname*{\mathbb{E}}_{f\sim A(X_{i})}\ell(y_{i,j},f(x_{i,j}))$ (16) $$

We divide the proof into bounding $\mathcal{R}(\rho)-\widetilde{\mathcal{R}}(\rho)$, and bounding $\widetilde{\mathcal{R}}(\rho)-\widehat{\mathcal{R}}(\rho)$ separately. For stating our results, we use the following notations as in Zakerinia et al. (2024):

- •
*Posterior $\mathfrak{Q}(A,\mathcal{Q}(A))$:* given as input an algorithm $A\in\mathcal{A}$ and
a hyper-posterior mapping $\mathcal{Q}:\mathcal{A}\to\mathcal{M}(\mathcal{F})$ as input, it outputs the
distribution over $\mathcal{M}(\mathcal{F})\times\mathcal{F}^{\otimes n}$ with the following generating
process:
*i)* sample a prior $P\sim\mathcal{Q}(A)$,
*ii)* for each task, $i=1,\dots,n$, sample a model $f_{i}\sim A(X_{i})$.
- •
*Prior $\mathfrak{P}$:* For the hyper-prior $\mathcal{P}\in\mathcal{M}(\mathcal{F})$ as input, it outputs the distribution over $\mathcal{M}(\mathcal{F})\times\mathcal{F}^{\otimes n}$ with the following generating process: *i)* sample a prior $P\sim\mathcal{P}$, *ii)* for each task, $i=1,\dots,n$, sample a model $f_{i}\sim P$.

First part is bounding $\mathcal{R}(\rho)-\widetilde{\mathcal{R}}(\rho)$ i.e. multi-task generalization bound for all $n$ tasks. The change of objective and intermediate step we have here compared to Zakerinia et al. (2024) allows us to consider the generalization behavior of all tasks (labeled and unlabeled) in contrast to only the labeled tasks. For the proof, for any task $i$ and any model $f_{i}$ we define:

$$ $\displaystyle\Delta_{i}(f_{i})=\operatorname*{\mathbb{E}}_{x\sim D_{i}}\ell(y,f_{i}(x))-\frac{1}{m}\sum_{j=1}^{m}\ell(y_{i,j},f_{i}(x_{i,j}))$ (17) $$

By this definition and the definitions of $\widetilde{\mathcal{R}}$ and $\mathcal{R}$ we have:

$$ $\displaystyle\operatorname*{\mathbb{E}}_{(P,f_{1},f_{2},...,f_{n})\sim\mathfrak{Q}(A,\mathcal{Q}(A))}\Big{[}\frac{1}{n}\sum_{i=1}^{n}\Delta_{i}(f_{i})\Big{]}=\mathcal{R}(A)-\widetilde{\mathcal{R}}(A)$ (18) $$

By *change of measure inequality* (Seldin et al., 2012), for any $\lambda>0$, any $A\in\mathcal{A}$ and any $\mathcal{Q}:\mathcal{A}\to\mathcal{M}(\mathcal{M}(\mathcal{F}))$, we have:

$$ $\displaystyle\mathcal{R}(A)-\widetilde{\mathcal{R}}(A)-\frac{1}{\lambda}\operatorname{\mathbf{KL}}\big{(}\mathfrak{Q}(A,\mathcal{Q}(A))||\mathfrak{P}\big{)}$ $\displaystyle\leq\frac{1}{\lambda}\log\operatorname*{\mathbb{E}}_{(P,f_{1},f_{2},...,f_{n})\sim\mathfrak{P}}\prod_{i=1}^{n}e^{\frac{\lambda}{n}\Delta_{i}(f_{i})}$ (19) $$

Because $\mathfrak{P}$ is independent of $S_{1},...,S_{n}$, we have

$$ $\displaystyle\operatorname*{\mathbb{E}}_{S_{1},\dots,S_{n}}\operatorname*{\mathbb{E}}_{(P,f_{1},f_{2},...,f_{n})\sim\mathfrak{P}}\prod_{i=1}^{n}e^{\frac{\lambda}{n}\Delta_{i}(f_{i})}=\operatorname*{\mathbb{E}}_{P\sim\mathcal{P}}\prod_{i=1}^{n}\operatorname*{\mathbb{E}}_{S_{i}}\operatorname*{\mathbb{E}}_{f_{i}\sim P}e^{\frac{\lambda}{n}\Delta_{i}(f_{i})}$ (20) $$

By Hoeffding’s lemma for $\Delta_{i}(f_{i})\in[0,1]$, we have

$$ $\displaystyle\operatorname*{\mathbb{E}}_{S_{i}}\operatorname*{\mathbb{E}}_{f_{i}\sim P}e^{\frac{\lambda}{n}\Delta_{i}(f_{i})}\leq e^{\frac{\lambda^{2}}{8n^{2}m}}.$ (21) $$

Therefore by combining ([20](#A1.E20)) and ([21](#A1.E21)) we have:

$$ $\displaystyle\operatorname*{\mathbb{E}}_{S_{1},\dots,S_{n}}\operatorname*{\mathbb{E}}_{(P,f_{1},f_{2},...,f_{n})\sim\mathfrak{P}}\prod_{i=1}^{n}e^{\frac{\lambda}{n}\Delta_{i}(f_{i})}\leq e^{\frac{\lambda^{2}}{8nm}}.$ (22) $$

By Markov’s inequality, for any $\epsilon>0$ we have

$$ $\displaystyle\mathbb{P}_{S_{1},\dots,S_{n}}\Big{(}\operatorname*{\mathbb{E}}_{(P,f_{1},f_{2},...,f_{n})\sim\mathfrak{P}}\prod_{i=1}^{n}e^{\frac{\lambda}{n}\Delta_{i}(f_{i})}\geq e^{\epsilon}\Big{)}\leq e^{\frac{\lambda^{2}}{8nm}-\epsilon}$ (23) $$

Hence by combining ([19](#A1.E19)) and ([23](#A1.E23)) we get

$$ $\displaystyle\mathbb{P}_{S_{1},\dots,S_{n}}\Big{(}\exists{A,\mathcal{Q}}:\mathcal{R}(A)-\widetilde{\mathcal{R}}(A)-\frac{1}{\lambda}\operatorname{\mathbf{KL}}(\mathfrak{Q}(A,\mathcal{Q}(A))||\mathfrak{P}(\mathcal{P}))\geq\frac{1}{\lambda}\epsilon\Big{)}\leq e^{\frac{\lambda^{2}}{8nm}-\epsilon},$ (24) $$

or, equivalently, it holds for any $\delta>0$ with probability at least $1-\frac{\delta}{2}$:

$$ $\displaystyle\forall{A,\mathcal{Q}}:\quad\mathcal{R}(A)-\widetilde{\mathcal{R}}(A)\leq\frac{1}{\lambda}\operatorname{\mathbf{KL}}(\mathfrak{Q}(A,\mathcal{Q}(A))||\mathfrak{P})+\frac{1}{\lambda}\log(\frac{2}{\delta})+\frac{\lambda}{8nm}$ (25) $$

With a union bound for $\lambda\in\{1,\dots,4mn\}$, and choosing the best $\lambda$ we get:

$$ $\displaystyle\mathbb{P}_{S_{1},\dots,S_{n}}\Big{(}\forall{A,\mathcal{Q}}:\mathcal{R}(A)-\widetilde{\mathcal{R}}(A)$ $\displaystyle\leq\sqrt{\frac{\operatorname{\mathbf{KL}}\big{(}\mathfrak{Q}(A,\mathcal{Q}(A))||\mathfrak{P}\big{)}+\log(\frac{8mn}{\delta})+1}{2mn}}\Big{)}\geq 1-\frac{\delta}{2}$ (26) $$

With probability at least $1-\frac{\delta}{2}$, the bound holds for all $A\in\mathcal{A}$, therefore, we can take the expectation for $A\in\rho$. Given that
$\operatorname*{\mathbb{E}}_{A\sim\rho}[\mathcal{R}(A)-\widetilde{\mathcal{R}}(A)]=\mathcal{R}(\rho)-\widetilde{\mathcal{R}}(\rho)]$, the following holds with probability at least $1-\frac{\delta}{2}$, for all $\rho$, and $\mathfrak{Q}$:

$$ $\displaystyle\forall{\rho,\mathcal{Q}}:\widetilde{\mathcal{R}}(\rho)-\widehat{\mathcal{R}}(\rho)\leq\operatorname*{\mathbb{E}}_{A\sim\rho}\sqrt{\frac{\operatorname{\mathbf{KL}}(\mathfrak{Q}(A,\mathcal{Q}(A))||\mathfrak{P})+\log(\frac{8mn}{\delta})+1}{2mn}}.$ (27) $$

For $\operatorname{\mathbf{KL}}(\mathfrak{Q}(A,\mathcal{Q}(A))||\mathfrak{P})$, we have:

$$ $\displaystyle\operatorname{\mathbf{KL}}(\mathfrak{Q}(A,\mathcal{Q}(A))||\mathfrak{P})$ $\displaystyle=\operatorname*{\mathbb{E}}_{P\sim\mathcal{Q}(A)}\Big{[}\operatorname*{\mathbb{E}}_{f_{i}\sim A(X_{i})}\ln\frac{\mathcal{Q}(A)(P)\prod_{i=1}^{n}A(X_{i})(f_{i})}{\mathcal{P}(P)\prod_{i=1}^{n}P(f_{i})}\Big{]}$ $\displaystyle=\operatorname*{\mathbb{E}}_{P\sim\mathcal{Q}(A)}\Big{[}\ln\frac{\mathcal{Q}(A)(P)}{\mathcal{P}(P)}\Big{]}+\operatorname*{\mathbb{E}}_{P\sim\mathcal{Q}(A)}\Big{[}\sum_{i=1}^{n}\operatorname*{\mathbb{E}}_{f_{i}\sim A(S_{i})}\ln\frac{A(X_{i})(f_{i})}{P(f_{i})}\Big{]}\Bigg{]}$ $\displaystyle=\operatorname{\mathbf{KL}}(\mathcal{Q}(A)\|\mathcal{P})+\operatorname*{\mathbb{E}}_{P\sim\mathcal{Q}(A)}\sum_{i=1}^{n}\operatorname{\mathbf{KL}}(A(S_{i})||P)$ (28) $$

Combining equations ([27](#A1.E27)) and ([28](#A1.E28)) gives that with probability $1-\frac{\delta}{2}$, for all $\rho\in\mathcal{M}(\mathcal{A}),\mathcal{Q}:\mathcal{A}\to\mathcal{M}(\mathcal{M}(\mathcal{F}))$:

$$ $\displaystyle\widetilde{\mathcal{R}}(\rho)-\widehat{\mathcal{R}}(\rho)\leq\operatorname*{\mathbb{E}}_{A\sim\rho}\sqrt{\frac{C(A,\mathcal{Q},\mathcal{P})+\log(\frac{8mn}{\delta})+1}{2mn}}$ (29) $$

For the second part, we apply the transductive bound ([11](#A1.E11)) to generalization between $\widetilde{\mathcal{R}}(\rho)$ and $\widehat{\mathcal{R}}(\rho)$, we get:

$$ $\displaystyle\mathcal{R}(\rho)$ $\displaystyle\leq\widehat{\mathcal{R}}(\rho)+\sqrt{(1-\frac{n_{L}}{n})\frac{\operatorname{\mathbf{KL}}(\rho||\pi)+\log(\frac{t(n_{L},n)}{\delta})}{2n_{L}}}$ (30) $$

Combining ([29](#A1.E29)) and ([30](#A1.E30)) proves the theorem.
∎

From this general theorem, we can now prove the generalization bound for our algorithm FLowDUP.

*

###### Proof.

For each hypernetwork with trainable parameters $\psi_{h}$, consider an algorithm $A$ that for task $i$ generates $A(S_{i})=\mathcal{N}(h(S_{i}\,;\,\psi_{h})\,;\,\alpha_{\theta}\text{Id})$.

For each $v\in\mathbb{R}^{k}$, the corresponding prior is $P=\mathcal{N}(v;\alpha_{\theta}\text{Id})$, and $\mathcal{P}=\mathcal{N}(0;\alpha_{r}\text{Id})$ is the hyper-prior over priors i.e. over their mean. Define the meta-prior as $\pi=\mathcal{N}(0;\alpha_{h}\text{Id})$ over $\mathcal{A}$ i.e. over hypernetwork parameters $\psi_{h}$.

Define the meta-posterior as a Gaussian distribution over the parameters of the hypernetwork, $\rho_{h}=\mathcal{N}(\psi_{h}\,;\,\alpha_{h}\text{Id})$ for a fixed $\alpha_{h}$. Define the hyper-posterior as $\mathcal{Q}=\mathcal{N}(\psi_{r}\,;\,\alpha_{r}\text{Id})$.

By applying the theorem [A.1](#A3.EGx13) to the defined distributions, we complete the proof.
∎

###### Theorem A.1 .

###### Proof.

###### Proof.

## Appendix B Experiments

Here we include additional experiments (Section [B.1](#A2.SS1))
and implementation details (Section [B.2](#A2.SS2)).

### B.1 Additional experiments

#### Rotated MNIST

We include here results for Rotated MNIST
The observations here are broadly similar to those discussed in Section [5.2](#S5.SS2)
with FLowDUP exhibiting the best performance overall.

**Table 5: Accuracy on Rotated MNIST for varying fractions ($p$) of the training clients having labeled data.**
| Fraction Labeled $\bm{(p)}$ | 0.1 | 0.2 | 0.5 | 1.0 |
| --- | --- | --- | --- | --- |
| FedAvg | 92.9 $\pm$ 0.2 | 94.8 $\pm$ 0.2 | 96.0 $\pm$ 0.2 | 96.4 $\pm$ 0.1 |
| FedProx | 92.8 $\pm$ 0.4 | 94.7 $\pm$ 0.1 | 95.8 $\pm$ 0.2 | 96.3 $\pm$ 0.1 |
| LD-FedAvg | 90.5 $\pm$ 0.3 | 92.4 $\pm$ 0.3 | 94.1 $\pm$ 0.0 | 95.0 $\pm$ 0.1 |
| FedTTA | 93.2 $\pm$ 0.3 | 94.8 $\pm$ 0.1 | 96.6 $\pm$ 0.1 | 97.3 $\pm$ 0.1 |
| FLowDUP | 94.0 $\pm$ 1.4 | 96.6 $\pm$ 0.4 | 97.8 $\pm$ 0.1 | 98.5 $\pm$ 0.1 |

#### The effect of unlabeled clients

One of the advantages of FLowDUP is that it’s learning objective is structured in such a
way that it allows unlabeled clients that are present during training to contribute to
regularizing $h$.
Specifically, unlabeled clients are able to compute the regularizer $\Omega$ in
[6](#S3.E6) and obtain from this a gradient update for $\psi$.
Here we investigate the effect that unlabeled clients have.
To do this we run FLowDUP on partitioned CIFAR10 both with and without using the unlabeled clients.
For FLowDUP without unlabeled clients we set the cohort size to 100 and sample only clients with labeled data.
For FLowDUP with labeled clients we set the cohort size to 200 and the labeled client sampling rate to $\alpha=0.5$
so that the number of labeled clients present during each round is the same in both cases.
We vary the total fraction of labeled clients, with $p\in\{0.05,0.1,0.2,0.5\}$.
As always we set $k=10^{4}$.
Tables [6](#A2.T6) and [7](#A2.T7) show the results.
As we can see, overall using unlabeled clients leads to an increase in performance.
This is most pronounced for lower values of $p$ (in particular at $p=0.5$ the effect is small
or negligible) and occurs for both architectures.

**Table 6: Partitioned CIFAR10 Accuracy for CNN.**
| Fraction Labeled $\bm{(p)}$ | 0.05 | 0.1 | 0.2 | 0.5 |
| --- | --- | --- | --- | --- |
| FLowDUP (without unlabeled) | 51.5 $\pm$ 1.8 | 65.3 $\pm$ 1.3 | 74.5 $\pm$ 1.3 | 82.5 $\pm$ 1.0 |
| FLowDUP (with unlabeled) | 52.3 $\pm$ 1.4 | 67.5 $\pm$ 0.4 | 76.6 $\pm$ 2.3 | 83.8 $\pm$ 0.7 |

**Table 7: Partitioned CIFAR10 Accuracy for ResNet18.**
| Fraction Labeled $\bm{(p)}$ | 0.05 | 0.1 | 0.2 | 0.5 |
| --- | --- | --- | --- | --- |
| FLowDUP (without unlabeled) | 54.8 $\pm$ 1.3 | 71.7 $\pm$ 1.2 | 83.3 $\pm$ 3.5 | 90.5 $\pm$ 0.4 |
| FLowDUP (with unlabeled) | 57.5 $\pm$ 1.6 | 72.6 $\pm$ 0.6 | 83.0 $\pm$ 2.7 | 90.6 $\pm$ 0.7 |

#### The learnable regularizer

FLowDUP prevents overfitting by penalizing large deviations between the generated client subspace
parameters and a learned regularizer $\psi_{r}$.
Intuitively, we can think of $\psi_{r}$ as a global subspace model that
clients should not deviate too much from.
Through the lens of our theoretical results, [4](#S4), $\psi_{r}$ is the mean of a learnable
prior distribution on the client models.
Here we investigate the efficacy of learning the regularizer.
To do this we replace $\psi_{r}$ by 0, so that the regularization term in the loss becomes instead

$$ $\displaystyle\Omega=\sum_{i\in\mathcal{C}}\|h(X\,;\,\psi_{h})\|^{2}.$ (31) $$

Note, this is of course still a reasonable regularizer more in line with classic $\ell_{2}$
regularization.
We train FLowDUP using this new non-learnable regularizer and compare it our proposed version
using the learnable regularizer [6](#S3.E6).
As in the previous section we again take a cohort size of 100 labeled clients and
100 unlabeled clients each round, i.e. $\alpha=0.5$. We set $k=10^{4}$ and
vary $p\in\{0.1,0.2,0.5,1.0\}$.
The results can be seen in Table [8](#A2.T8).
They indeed show that the extra flexibility afforded by learning the regularizer
leads to a modest boost in performance for all values of $p$.

**Table 8: The effect of learning the regularizer $\psi_{r}$. Accuracy on partitioned CIFAR10 for CNN.**
| Fraction Labeled $\bm{(p)}$ | 0.1 | 0.2 | 0.5 | 1.0 |
| --- | --- | --- | --- | --- |
| FLowDUP (Without Learnable Reg) | 65.5 $\pm$ 0.4 | 75.2 $\pm$ 0.8 | 83.3 $\pm$ 1.7 | 85.6 $\pm$ 0.8 |
| FLowDUP (With Learnable Reg) | 67.5 $\pm$ 0.4 | 76.6 $\pm$ 0.3 | 83.8 $\pm$ 0.7 | 86.6 $\pm$ 0.1 |

### B.2 Implementation details

#### Hyperparameters

There are number of general federated learning hyperparameters that are shared across methods.
We train all methods for the same number of global rounds $T$. For class partitioned
CIFAR10 and FEMNIST we set $T=1000$ while for Rotated Fashion-MNIST and Rotated MNIST
we set $T=500$.
All methods use a client cohort size each round of size 100,
a local number of epochs set to $E=1$ and a local
batch size of $B=20$ for FEMNIST and $B=50$ on all other datasets.
For all methods we tune the local learning rate $\eta_{l}$ on validation clients.

Regarding method specific hyperparameters.
For FedProx, we set $\mu$ following Li et al. (2020), that is $\mu=1$ initially
and is updated over the course of training as described in Li et al. (2020).
For FedTTA we tune the adaptive learning rate.
For FLowDUP, unless otherwise stated, we set the labeled client sampling rate to $\alpha=0.9$
and the subspace dimension to $k=10^{4}$. We tune the regularization strength $\lambda$.

#### Model architectures

When used for prediction the CNN follows the architecture used in prior work
(McMahan et al., 2017) while the ResNet18 follows the standard architecture with the
final classification layer having output dimension 10 for the 10 classes
present in CIFAR10.
When used for $h_{1}$, we instead replace the final linear layers of the CNN and
ResNet18 with another with output dimension 256.
For $h_{2}$ we always use a fully connected network with a single hidden layer
and ReLU non-linearity.

## Appendix C Additional Related Work

#### Personalized federated learning

Approaches to personalized FL typically fall into one of the following categories: Meta-learning based approaches (Jiang et al., 2019; Fallah et al., 2020) which learn a single global model that can be easily personalized using a small number of gradient steps on the client.
Parameter decomposition-approaches (Arivazhagan et al., 2019; Collins et al., 2021; Marfoq et al., 2022; Chen et al., 2023; Wu et al., 2023) which divide the learnable
parameters into some that are shared across clients (such as a feature extractor) and some that are specific individual to each client (such as a classification head).
Federated multi-task approaches (Smith et al., 2017; Dinh et al., 2020; Hanzely et al., 2020; Marfoq et al., 2021; Li et al., 2021; Lin et al., 2022; Ye et al., 2023; Zhang et al., 2023)
learn separate models for each client while still sharing some information across clients, for instance by regularizing towards some global model.
All of the above approaches, however, require a client to have labeled data
in order to obtain a personalized model.

#### Learning in a subspace

This formulation of intrinsic dimensionality and learning in a subspace has been studied in different contexts in the literature. Li et al. (2018) introduced the formulation and showed that for real-world problems the intrinsic dimension is much smaller than the total number of parameters $k\ll d$. Aghajanyan et al. (2021) showed that pretraining reduces the intrinsic dimensionality of fine-tuning. Lotfi et al. (2022) used this parametrization to achieve non-vacuous generalization bounds for neural networks. Zakerinia et al. (2025) extended the definition of the intrinsic dimension to multi-task learning. Park et al. (2025) used the formulation for black-box prompt tuning. Lin et al. (2022) used it to reduce the communication of the global model in personalized federated learning.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. What is the detailed form of the regularization term $\psi_r$ in implementation, will the authors make their code public?
2. The paper claims that FLowDUP supports training with unlabeled data, but all experimental datasets are public labeled datasets. How were the results in Tables 1 and 2 obtained, were the training data labeled, unlabeled, or partially labeled?
3. What was the data splitting scheme between labeled and unlabeled clients?
soundness: 3
strengths: 1. Addressing personalization for unlabeled clients is a relatively unexplored area in PFL.
2. The paper introduces a new transductive multi-task PAC-Bayesian bound that motivates the objective and gives its proof.
3. The low-dimensional subspace formulation has the potential to reduce communication and computation costs in federated settings.
summary: The paper proposes a PFL method called FLowDUP that generates a client-specific model using only unlabeled data. FLowDUP employs a hypernetwork to produce low-dimensional subspace parameters, which are expanded via a fixed random matrix to yield the full model. Its training objective combines a labeled-client loss and an unlabeled-client regularizer, theoretically supported by a new transductive multi-task PAC-Bayesian generalization bound. Experiments across serveral widely-used datasets show improved performance compared with baselines such as FedAvg, FedProx, and ATP.
weaknesses: 1. The proposed method is built upon the foundations of regulation-based and hypernetwork-based personalized federated learning (PFL). However, the related work section does not provide an adequate review or a clear comparison between prior methods and the proposed approach.
2. The paper uses the hypernet for personalization which is not a new idea. Besides, the authors emphasize that the proposed method can handle federated learning with labeled data, yet there is no direct comparison with existing hypernetwork-based PFL approaches, such as [1] and [2].
3. If FedAvg cannot be trained without labels, how was it used for comparison? If FedAvg was trained with labeled data but the proposed method used unlabeled data, how is this comparison fair and meaningful?
4. The paper claims that the proposed approach is computationally and communication efficient, yet no experiments or quantitative evidence are presented to substantiate this claim.
5. The writing and structural focus of the paper could be improved. While substantial space (including Appendix A) is devoted to PAC-Bayes theory, its direct connection to the proposed method is not clearly articulated and should be elaborated. Conversely, the discussion of basic federated learning concepts is overly detailed and could be condensed.

[1] Shamsian, A. et al. “Personalized Federated Learning Using Hypernetworks.” ICML, 2021.

[2] Yang, Z. et al. “Hypernetwork-Based Physics-Driven Personalized Federated Learning for CT Imaging.” IEEE Transactions on Neural Networks and Learning Systems, 2023.

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: 1-FLowDUP fails when conditional distributions \(D_{i|Y|X}\) (label given input) differ across clients but marginal distributions \(D_{i|X}\) (input alone) are insufficient to infer predictive models—e.g., subjective tasks like product recommendations or sentiment analysis. The authors acknowledge this but provide no mitigation strategies, restricting the method’s real-world applicability.

2-The default \(k=10^4\) is chosen empirically, with no theoretical analysis of how k relates to dataset heterogeneity, model architecture, or performance. While experiments show k correlates with accuracy, the lack of a principled way to select k (e.g., via the generalization bound) reduces FLowDUP’s usability.

3-All experiments use the \(pfl-research\) simulation library—there is no validation on actual edge devices (e.g., smartphones, IoT sensors). The authors claim low computational cost, but simulated efficiency may not translate to real hardware (e.g., memory constraints, latency).
soundness: 3
strengths: 1-FLowDUP directly addresses a well-identified pain point of PFL (unlabeled client incompatibility) with a non-trivial solution. Unlike prior hypernetwork-based methods (e.g., Amosy et al. 2024; Scott et al. 2024), it uses low-dimensional subspaces to enable on-device hypernetwork deployment and leverages unlabeled clients for training—two key innovations that distinguish it from related work.

2-It explicitly connects the regularizer \(\Omega\) and loss L to generalization performance, providing a mathematical justification for FLowDUP’s objective. This contrasts with many PFL works that lack theoretical guarantees for unlabeled clients.

3-Low-dimensional parameters reduce communication costs, the hypernetwork is lightweight for edge devices, and no client data is transmitted. This aligns with the core goals of federated learning (privacy, efficiency).
summary: This paper addresses a critical limitation of existing personalized federated learning (PFL) methods—their reliance on labeled client data to generate personalized models—by proposing FLowDUP (Federated Low Dimensional Unlabeled Personalization), a framework enabling personalized model generation for unlabeled clients via a single forward pass on unlabeled data. FLowDUP trains a hypernetwork h that takes unlabeled client data as input and outputs low-dimensional subspace parameters \(v_i\). Full model parameters \(\theta_i\) are derived via \(\theta_i = \theta_0 + Pv_i\) (where P is a fixed random expansion matrix and \(\theta_0\) is an initial parameter), avoiding high-dimensional parameter transmission. The objective combines a loss L (evaluated on labeled clients to measure model quality) and a learnable regularizer \(\Omega\) (computed on all clients, including unlabeled ones, to prevent overfitting).
weaknesses: 1-FLowDUP fails when conditional distributions \(D_{i|Y|X}\) (label given input) differ across clients but marginal distributions \(D_{i|X}\) (input alone) are insufficient to infer predictive models—e.g., subjective tasks like product recommendations or sentiment analysis. The authors acknowledge this but provide no mitigation strategies, restricting the method’s real-world applicability.

2-The default \(k=10^4\) is chosen empirically, with no theoretical analysis of how k relates to dataset heterogeneity, model architecture, or performance. While experiments show k correlates with accuracy, the lack of a principled way to select k (e.g., via the generalization bound) reduces FLowDUP’s usability.

3-All experiments use the \(pfl-research\) simulation library—there is no validation on actual edge devices (e.g., smartphones, IoT sensors). The authors claim low computational cost, but simulated efficiency may not translate to real hardware (e.g., memory constraints, latency).

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: Please address the weaknesses.
soundness: 2
strengths: - The paper investigates an important problem in FL where clients without labels can often exist and not be able to gain anything out of federated learning. 
- The paper provides theoretical insights into their proposed algorithm's loss function regarding using labeled clients to benefit the unlabeled client's personalized models. 
- The paper is clearly written.
summary: This paper proposes FLowDUP to train personalized models for clients participating in federated learning with and for even clients that have unlabeled data using only a forward pass. FLowDUP trains a hypernetwork that takes the unlabeled dataset and outputs parameters of a personalized predictive model that allows them to train in a subspace that is much smaller than the actual underlying personalzied model. They provide intuition behind their method by giving a generalization bound with multi-task PAC-Bayesian. The paper includes experimental evaluation of FLowDUP
weaknesses: - FLowDUP relies largely on the labeled clients' data, and due to this I have a couple of concerns. First, if the labeled clients data does not match well of the unlabeled clients' data, which is often the case in federated learning due to data heterogeneity and cannot really be controlled since they are clients' private data, the performance I believe will largely suffer. Moreover, it is rather unfair for the labeled clients to provide their training resource/data for the benefit of training personalized models for these unlabeled clients. In the perspective of incentives and fairness of clients the method also is not really convincing. 
- In the evaluation, the model used and the dataset used seems rather outdated, and easy ones and not sure how convincing the results actually are. Moreover, as mentioned above it seems like the data heterogeneity across clients will matter quite a lot here and the fact that this is not really explored in either in the theory or experiment section is concerning.


Decision: Reject