Vertical Federated Learning with Missing Features During Training and Inference

## Abstract

Abstract Vertical federated learning trains models from feature-partitioned datasets across multiple clients, who collaborate without sharing their local data. Standard approaches assume that all feature partitions are available during both training and inference. Yet, in practice, this assumption rarely holds, as for many samples only a subset of the clients observe their partition. However, not utilizing incomplete samples during training harms generalization, and not supporting them during inference limits the utility of the model. Moreover, if any client leaves the federation after training, its partition becomes unavailable, rendering the learned model unusable. Missing feature blocks are therefore a key challenge limiting the applicability of vertical federated learning in real-world scenarios. To address this, we propose LASER-VFL , a vertical federated learning method for efficient training and inference of split neural network-based models that is capable of handling arbitrary sets of partitions. Our approach is simple yet effective, relying on the strategic sharing of model parameters and on task-sampling to train a family of predictors. We show that LASER-VFL achieves a 𝒪 ⁢ ( 1 / T ) 𝒪 1 𝑇 \mathcal{O}(\nicefrac{{1}}{{\sqrt{T}}}) caligraphic_O ( / start_ARG 1 end_ARG start_ARG square-root start_ARG italic_T end_ARG end_ARG ) convergence rate for nonconvex objectives in general, 𝒪 ⁢ ( 1 / T ) 𝒪 1 𝑇 \mathcal{O}(\nicefrac{{1}}{{T}}) caligraphic_O ( / start_ARG 1 end_ARG start_ARG italic_T end_ARG ) for sufficiently large batch sizes, and linear convergence under the Polyak-Łojasiewicz inequality. Numerical experiments show improved performance of LASER-VFL over the baselines. Remarkably, this is the case even in the absence of missing features. For example, for CIFAR-100, we see an improvement in accuracy of 21.4 % percent 21.4 21.4\% 21.4 % when each of four feature blocks is observed with a probability of 0.5 and of 12.2 % percent 12.2 12.2\% 12.2 % when all features are observed.

## 1 Introduction

In federated learning (FL), a set of clients collaborates to jointly train a model using their local data without sharing it (Kairouz et al., ). In horizontal FL, data is distributed by samples, meaning each client holds a different set of samples but shares the same feature space. In contrast, vertical FL (VFL) involves data distributed by features, where each client holds different parts of the feature space for overlapping sets of samples. Whether an application is horizontal or vertical FL is dictated by how the data arises, as it is not possible to redistribute the data. This work focuses on VFL.

In VFL (Liu et al., ), the global dataset $\mathcal{D}\coloneqq\{\bm{x}^{1},\dots,\bm{x}^{N}\}$, where $N$ is the number of samples, is partitioned across clients $\mathcal{K}\coloneqq\{1,\dots,K\}$. Each client $k\in\mathcal{K}$ typically holds a local dataset $\mathcal{D}_{k}\coloneqq\{\bm{x}_{k}^{1},\dots,\bm{x}_{k}^{N}\}$, where $\bm{x}_{k}^{n}$ is the block of features of sample $n$ observed by client $k$. We have that $\bm{x}^{n}=(\bm{x}_{1}^{n},\dots,\bm{x}_{K}^{n})$. Unlike horizontal FL, in VFL, different clients collect local datasets with distinct types of information (features). Such setups—e.g., an online retail company and a social media platform holding different features on shared users—typically involve entities from different sectors, reducing competition and increasing the incentive to collaborate. To train VFL models without sharing the local datasets, split neural networks (Ceballos et al., ) are often considered.

In split neural networks, each client $k$ has a representation model $\bm{f}_{k}$, parameterized by $\bm{\theta}_{\bm{f}_{k}}$. The representations extracted by the clients are then used as input to a fusion model $g$, parameterized by $\bm{\theta}_{g}$. This fusion model can be at one of the clients or at a server. We thus want to learn the parameters $\bm{\theta}_{\mathcal{K}}\coloneqq(\bm{\theta}_{\bm{f}_{1}},\dots,\bm{\theta}_
{\bm{f}_{K}},\bm{\theta}_{g})$ of the resulting predictor $h$. To achieve this, we aim to solve:

$$ $\min_{\bm{\theta}_{\mathcal{K}}}\quad\frac{1}{N}\sum_{n=1}^{N}\ell(h(\bm{x}^{n };\bm{\theta}_{\mathcal{K}}),y^{n})\quad\text{where}\quad h(\bm{x}^{n};\bm{ \theta}_{\mathcal{K}})\coloneqq g\left(\left\{\bm{f}_{k}(\bm{x}_{k}^{n};\bm{ \theta}_{\bm{f}_{k}})\right\}_{k=1}^{K};\bm{\theta}_{g}\right),$ (1) $$

with $\ell$ denoting a loss function and $y^{n}$ the label of sample $n$, which we assume to be held by the same entity (client or server) as the fusion model. We illustrate this family of models in Figure 1(a).

#### Generalization and availability in standard VFL.

We see in (1) that, even for a single sample $n$, predictor $h$ depends on all the blocks of features, $\{\bm{x}^{n}_{k}:\forall k\}$. Thus, for both training and inference, (1) requires the observations of all the clients to be available.(^4^44This contrasts with horizontal FL, where the local data of one client suffices to estimate mini-batch gradients and perform inference.) Building on our example of an online retailer and a social media company sharing users, each company is also likely to have unique users. This applies to both training and test data. Further, if any client drops from the federation during inference, its block will permanently stop being observed. In both the case of nonshared users and of clients leaving the federation, standard VFL predictors become unusable. Figure 1(b) illustrates this phenomenon: for both training and inference, if each of $K=3$ clients observes its block of features with an (independent) probability of $0.5$, only $0.5^{3}=12.5\%$ of the original data is usable. Thus, restricting training to fully-observed samples hinders generalization, while restricting inference to such samples limits the availability and utility of the model.

#### Dealing with missing features.

To handle missing features in training data, some approaches expand the dataset, filling the missing features before collaborative training (Kang et al., ), while others use their partition of partially-observed samples for local representation learning but exclude them from collaboration (He et al., ). These methods can outperform standard VFL when partially observed samples are present, but they add new training stages and auxiliary modules, making them more complex. Moreover, they train a joint predictor that requires the collaboration of all clients during inference. On the other hand, to improve robustness against missing blocks at inference, each client can train a local predictor using only its own features, avoiding collaboration altogether. This also addresses missing features in the training data. Alternatively, collaborative methods can employ techniques such as knowledge distillation (Huang et al., ) and data augmentation (Gao et al., ) to train local predictors, also adding complexity with multiple stages. However, while robust to missing features, these local predictors cannot utilize additional feature blocks if available at test time, resulting in wasted data and reduced predictive power. We illustrate this in Figure 1(b).

Figure: (a) Split neural network.
Refer to caption: extracted/5963893/figures/splitnn.png

Therefore, there is a gap in the literature when it comes to leveraging all the available data without either dropping incomplete samples or ignoring existing features. This raises the following question:

*Can we design an efficient VFL method that is robust to missing features at both training and inference with provable convergence guarantees, without wasting data?*

In this work, we answer this question in the affirmative. We propose LASER-VFL (Leveraging All Samples Efficiently for Real-world VFL), a novel method that enables both training and inference using any and all available blocks of features. To the best of our knowledge, this is the first method to achieve this. Our approach avoids the multi-stage pipelines that are common in this area, providing a simple yet effective solution that leverages strategic sharing of model parameters alongside a task-sampling mechanism. By fully utilizing all data (see Figure 1(b)), our method leads to significant performance improvements, as demonstrated in Figure 1(c). Remarkably, our approach even surpasses standard VFL when all samples are fully observed. We attribute this to a dropout-like regularization effect introduced by task sampling.

#### Our contributions.

The main contributions of this work are as follows.

- •
We propose LASER-VFL, a simple, hyperparameter-free, and efficient VFL method that is flexible to a varying number of feature blocks during both training and inference. To the best of our knowledge, this is the first method to achieve such flexibility and avoid wasting either training data or test data.
- •
We show that LASER-VFL converges at a $\mathcal{O}(\nicefrac{{1}}{{\sqrt{T}}})$ rate for nonconvex objectives in general, and at a $\mathcal{O}(\nicefrac{{1}}{{T}})$ rate for a sufficiently large batch size. Further, under the Polyak-Łojasiewicz (PL) inequality, we achieve linear convergence.
- •
Numerical experiments show that LASER-VFL consistently outperforms baselines across multiple datasets and varying data availability patterns. It demonstrates superior robustness to missing features and, notably, when all features are available, it still outperforms even standard VFL.

#### Related work.

VFL shares challenges with horizontal FL, such as communication efficiency (Liu et al., , Valdeira et al., ) and privacy preservation (Yu et al., ), but also faces unique obstacles, such as missing feature blocks. During training, these unavailable partitions render the observed blocks of other clients unusable, and at test time, they can prevent inference altogether. Therefore, most VFL literature assumes that all features are available for both training and inference—an often unrealistic assumption that has hindered broader adoption of VFL (Liu et al., ).

To address the problem of missing features during VFL training, some works use nonoverlapping, or nonaligned, samples (that is, samples with missing feature blocks) to improve generalization. In particular, Feng () and He et al. () apply self-supervised learning to leverage nonaligned samples locally for better representation learning, while using overlapping samples for collaborative training of a joint predictor. Alternatively, Kang et al. (), Yang et al. (), and Sun et al. () employ semi-supervised learning to take advantage of nonaligned samples. Although these methods enable VFL to utilize data that conventional approaches would discard, they still train a joint predictor and thus require all features to be available for inference, which remains a limitation.

A recent line of research leverages information from the entire federation to train local predictors. This can be achieved via transfer learning, as in Feng and Yu (), Kang et al. () for overlapping training samples and in Liu et al. (), Feng et al. () for handling missing features. Knowledge distillation is another approach, as used by Ren et al. (), Li et al. () for overlapping samples and by Li et al. (), Huang et al. () which leverage nonaligned samples when training local predictors (only the latter considers scenarios with more than two clients). Xiao et al. () recently proposed using a distributed generative adversarial network for collaborative training on nonoverlapping data and synthetic data generation. These methods yield local predictors that outperform naive local approaches but fail to utilize valuable information when other feature blocks are available during inference.

Two recent works enable a varying number of clients to collaborate during inference. Sun et al. () employs party-wise dropout during training to mitigate performance drops from missing feature blocks at inference, while Gao et al. () introduces a multi-stage approach with complementary knowledge distillation, enabling inference with different client subsets. Both methods, however, require fully-observed training data and lack convergence guarantees.

In contrast to prior work, LASER-VFL can handle any subset of feature blocks being present during both training and inference without wasting data, while requiring only minor modifications to the standard VFL approach. It differs from conventional VFL solely in its use of strategic sharing of model parameters and task-sampling mechanisms, without the need for additional stages.

## 2 Definitions and Preliminaries

Our global dataset $\mathcal{D}$ is drawn from an input space $\mathcal{X}$ which is partitioned into $K$ feature spaces $\mathcal{X}=\mathcal{X}_{1}\times\dots\times\mathcal{X}_{K}$ such that $\mathcal{D}_{k}\subseteq\mathcal{X}_{k}$, where, recall, $\mathcal{D}_{k}$ is the local dataset of client $k$. Ideally, we would train a general, unconstrained predictor $\tilde{h}\colon\mathcal{X}\mapsto\mathcal{Y}$, where $\mathcal{Y}$ is the label space, by solving
$\min_{\tilde{\bm{\theta}}}\frac{1}{N}\sum_{n=1}^{N}\ell(\tilde{h}(\bm{x}^{n};
\tilde{\bm{\theta}}),y^{n})$. However, VFL brings the additional constraint that this must be achieved without sharing the local datasets. That is, for all $k$, the local dataset $\mathcal{D}_{k}$ must remain at client $k$. As mentioned in Section 1, standard VFL methods approximate $\tilde{h}$ by $h\colon\mathcal{X}\mapsto\mathcal{Y}$, as defined in (1), allowing for collaborative training without sharing the local datsets, but they cannot handle missing features during training nor inference.

Another way to train predictors without sharing local data is to learn local predictors. In particular, each client $k\in\mathcal{K}$ can learn the parameters $\bm{\theta}_{k}$ of a predictor $h_{k}:\mathcal{X}_{k}\mapsto\mathcal{Y}$ to approximate $\tilde{h}$:

$$ $h_{k}(\bm{x}^{n}_{k};\bm{\theta}_{k})\coloneqq g_{k}(\bm{f}_{k}(\bm{x}^{n}_{k} ;\bm{\theta}_{\bm{f}_{k}});\bm{\theta}_{g_{k}})\quad\text{where}\quad\bm{ \theta}_{k}\coloneqq(\bm{\theta}_{\bm{f}_{k}},\bm{\theta}_{g_{k}}).$ (2) $$

The representation models $\bm{f}_{k}:\mathcal{X}_{k}\mapsto\mathcal{E}_{k}$, where $\mathcal{E}_{k}$ is the representation space of client $k$, are as in (1), yet the fusion models $g_{k}:\mathcal{E}_{k}\mapsto\mathbb{R}$ differ from $g:\mathcal{E}_{1}\times\cdots\times\mathcal{E}_{K}\mapsto\mathbb{R}$ in (1) and $h_{k}$ differs from $h$. This approach is useful in that $h_{k}$ allows client $k$ to perform inference (and be trained) independently from all the other clients, but it does not make use of the features observed by clients $j\not=k$.

More generally, to achieve robustness to missing blocks in the test data while avoiding wasting other features, we wish to be able to perform inference based on any possible subset of blocks, $\mathscr{P}(\mathcal{K})\setminus\{\varnothing\}$, where $\mathscr{P}(\mathcal{K})$ denotes the power set of $\mathcal{K}$. Standard VFL allows us to obtain a predictor for the blocks $\mathcal{K}$ and the local approach provides us with predictors for the singletons $\{\{i\}\colon i\in\mathcal{K}\}$ in the power set. However, none of the other subsets of $\mathcal{K}$ is covered by either approach.

A naive way to achieve this would be to train a predictor for each set in $\mathscr{P}(\mathcal{K})\setminus\{\varnothing\}$ in a decoupled manner. That is, we could train each of the following predictors $\{h_{\mathcal{J}}\colon\prod_{k\in\mathcal{J}}\mathcal{X}_{k}\mapsto\mathcal{Y}\}$ independently, using the collaborative training approach of standard VFL for $\mathcal{J}$:

$$ $\left\{h_{\mathcal{J}}(\bm{x}^{n}_{\mathcal{J}};\bm{\theta}_{\mathcal{J}}) \coloneqq g_{\mathcal{J}}((\bm{f}_{k}(\bm{x}_{k}^{n};\bm{\theta}_{\bm{f}_{k}( \mathcal{J})})\colon k\in\mathcal{J});\bm{\theta}_{g_{\mathcal{J}}})\colon \mathcal{J}\in\mathscr{P}(\mathcal{K})\setminus\{\varnothing\}\right\},$ (3) $$

where $\bm{\theta}_{\mathcal{J}}=((\bm{\theta}_{\bm{f}_{k}}(\mathcal{J})\colon k\in
\mathcal{J}),\bm{\theta}_{g_{\mathcal{J}}})$ and $\bm{x}^{n}_{\mathcal{J}}=(\bm{x}^{n}_{k}\colon k\in\mathcal{J})$. These predictors can either be at one of the clients or at the server. The set of predictors in (3) includes the local predictors $\{h_{k}\}$ and the standard VFL predictor $h$, but also all the other nonempty sets in the power set $\mathscr{P}(\mathcal{K})$.(^5^55 The notation in (3) differs slightly from (1) and (2), which use a simpler, more specific formulation.)
This approach addresses the issue of limited flexibility and robustness in prior methods, which struggle with varying numbers of available or participating clients during inference. However, by requiring the independent training of $2^{K}-1$ distinct predictors, it introduces a new challenge: the number of models would grow exponentially with the number of clients, $K$. Consequently, the associated memory, computation, and communication costs would also increase exponentially.

Further, if the predictors in (3) are all held by a single entity, this setup introduces a dependency of all clients on that entity. To enhance robustness against clients dropping from the federation, we would like each client to be able to ensure inference whenever it has access to its corresponding block of the sample. That is, we want each client $k\in\mathcal{K}$ to train a predictor for every nonempty subset of $\mathcal{K}$ that includes $k$, defined as $\mathscr{P}_{k}(\mathcal{K})\coloneqq\{\mathcal{J}\in\mathscr{P}(\mathcal{K})
\colon k\in\mathcal{J}\}$. This allows $k$ to make predictions regardless of whether the information on the blocks observed by other clients is available and can be leveraged to improve performance.

In the next section, we present our method, LASER-VFL, which enables the efficient training of a family of predictors that can handle scenarios where features from an arbitrary set of clients are missing. We have included a table of notation in Appendix A.

## 3 Our method

#### Key idea.

When striving to have a predictor at each client to perform inference using any subset of blocks that includes its own, the main challenge is to circumvent the exponential complexity that can ensue from the exponential number of possible combinations of available blocks. To address this, in LASER-VFL, we share model parameters across the predictors leveraging different subsets of blocks and train them so that the representation models allow for good performance across the different combinations of blocks. Further, we train the fusion model at each client to handle any combination of representations that includes its own. To train the predictors on an exponential number of combinations of missing blocks while avoiding an exponential computational complexity, we employ a sampling mechanism during training, which allows us to essentially estimate an exponential combination of objectives with a subexponential complexity.

### 3.1 A tractable family of predictors sharing model parameters

In our approach, each client $k\in\mathcal{K}$ has a set of predictors $\{h_{(k,\mathcal{J})}\colon\mathcal{J}\in\mathscr{P}_{k}(\mathcal{K})\}$.
Each predictor $h_{(k,\mathcal{J})}\colon\prod_{k\in\mathcal{J}}\mathcal{X}_{k}\mapsto\mathcal
{Y}$ has the same target, but a different domain. Therefore, we say that each predictor performs a distinct task. We define the predictors of client $k$ as follows:

$$ $\left\{h_{(k,\mathcal{J})}\left(\bm{x}^{n}_{\mathcal{J}};\bm{\theta}_{(k, \mathcal{J})}\right)\coloneqq g_{k}\left(\frac{1}{|\mathcal{J}|}\sum_{{j\in \mathcal{J}}}\bm{f}_{j}(\bm{x}_{j}^{n};\bm{\theta}_{\bm{f}_{j}});\bm{\theta}_{ g_{k}}\right)\colon\mathcal{J}\in\mathscr{P}_{k}(\mathcal{K})\right\},$ (4) $$

where $\bm{\theta}_{(k,\mathcal{J})}=((\bm{\theta}_{\bm{f}_{j}}\colon j\in\mathcal{J}
),\bm{\theta}_{g_{k}})$ and $|\mathcal{J}|$ denotes the cardinality of set $\mathcal{J}$.

#### Model parameter sharing.

Note that, although we can see (4) as $|\mathscr{P}_{k}(\mathcal{K})|=2^{K-1}$ different predictors, the predictors are made up of different combinations of only $K$ representation models and $K$ fusion models.
That is, we only require parameters $\bm{\theta}\coloneqq(\bm{\theta}_{\bm{f}_{1}},\dots,\bm{\theta}_{\bm{f}_{K}},
\bm{\theta}_{g_{1}},\dots,\bm{\theta}_{g_{K}})$. In particular, in contrast to (3), where the predictors used different representation models for each task, in (4), we share the representation models across predictors. Similarly to the representation models, we have a single fusion model per client.

It is also important to note that each fusion model takes the specific form
$g_{k}\left((\bm{f}_{i}(\bm{x}_{i}^{n};\bm{\theta}_{\bm{f}_{i}})\colon i\in
\mathcal{J});\bm{\theta}_{g_{k}}\right)=g_{k}(\frac{1}{|\mathcal{J}|}\sum_{{i
\in\mathcal{J}}}\bm{f}_{i}(\bm{x}_{i}^{n};\bm{\theta}_{\bm{f}_{i}});\bm{\theta
}_{g_{k}})$. By employing a nonparameterized aggregation mechanism on the extracted representations whose output does not depend on the number of aggregated representations, the same fusion model can handle different sets of representations of different sizes. In particular, we opt for an average (rather than, say, a sum) because neural networks perform better when their inputs have similar distributions (Ioffe and Szegedy, ).

With this approach, we have avoided an exponential memory complexity by sharing the weights across an exponential number of predictors such that we have $K$ representation models and $K$ fusion models. In the next subesection, we will go over the efficient training of this family of predictors.

### 3.2 Efficient training via task-sampling

The family of predictors introduced in Section 3.1 can efficiently perform inference for an exponential number of tasks. We will now discuss how to train this family of predictors while avoiding exponential computation and communication complexity during the training process. The key to this approach lies in carefully sampling tasks at each gradient step of model training.

#### Our optimization problem.

As mentioned in Section 1, we are not only interested in dealing with missing blocks of features during inference, but also during training. In particular, we assume that the availability of each block $k$ in sample $n$ follows a random distribution and denote by $\mathcal{K}_{o}^{n}$ the (random) subset of blocks of features of sample $n$ that is observed in the training data. Let the loss corresponding to fusion model $k$ using blocks of features $\mathcal{I}$ for prediction be denoted as $\tilde{\mathcal{L}}_{k,\mathcal{I}}(\bm{\theta};n)\coloneqq\ell\left(h_{(k,
\mathcal{I})}\left(\bm{x}^{n}_{\mathcal{I}};\bm{\theta}_{(k,\mathcal{I})}
\right),y^{n}\right)$, to train our predictors (4), we consider the following optimization problem:

$$ $\min_{\bm{\theta}}\,\left\{\mathcal{L}(\bm{\theta})\coloneqq\mathbb{E}\left[ \tilde{\mathcal{L}}(\bm{\theta})\coloneqq\frac{1}{N}\sum_{n=1}^{N}\tilde{ \mathcal{L}}_{n}(\bm{\theta})\right]\right\}\;\;\text{where}\;\;\tilde{ \mathcal{L}}_{n}(\bm{\theta})\coloneqq\sum_{k\in\mathcal{K}_{o}^{n}}\sum_{ \mathcal{I}\in\mathscr{P}_{k}(\mathcal{K}_{o}^{n})}\frac{1}{|\mathcal{I}|} \cdot\tilde{\mathcal{L}}_{k,\mathcal{I}}(\bm{\theta};n).$ (5) $$

where the expectation is over the set of available feature blocks $\{\mathcal{K}_{o}^{n}\colon\forall n\}$. In (5), we employ the normalization ${1}/{|\mathcal{I}|}$ to avoid having a larger weight being assigned to prediction tasks concerning a larger number of blocks, as these will show up at $|\mathcal{I}|$ different clients. Further, note that, if different blocks of features have different probabilities of being observed, we can also normalize the loss with respect to the different probabilities which can be computed during the entity alignment stage of VFL (see below), before taking the expectation in (5). We omit this here for simplicity.

Looking at (5), we see that, while we are interested in tackling $K\times 2^{K-1}$ different tasks, which falls in the realm of multi-objective optimization and multi-task learning (Caruana, ), we opt for combining them into a single objective, rather than employing a specialized multi-task optimizer. This corresponds to a weighted form of unitary scalarization (Kurin et al., ).

Figure: Algorithm 1 LASER-VFL Training

#### Our optimization method.

We learn parameters $\bm{\theta}$ using a gradient-based method, summarized in Algorithm 1, that we now describe. At the start of step $t$, we have $\bm{\theta}^{t}\coloneqq(\bm{\theta}^{t}_{\bm{f}_{1}},\dots,\bm{\theta}^{t}_{
\bm{f}_{K}},\bm{\theta}^{t}_{g_{1}},\dots,\bm{\theta}^{t}_{g_{K}})$ and choose a mini-batch $\mathcal{B}^{t}\subseteq[N]\coloneqq\{1,\dots,N\}$ such that $\mathcal{K}_{o}^{n}$ is the same for all $n$ in the same mini-batch $\mathcal{B}^{t}$ (we elaborate on this below) and define $\mathcal{K}^{\mathcal{B}^{t}}_{o}\coloneqq\mathcal{K}^{n}_{o}$. Then, each client $k\in\mathcal{K}_{o}^{\mathcal{B}^{t}}$ broadcasts its representation $\bm{f}_{k}(\mathcal{B}^{t}_{k};\bm{\theta}^{t}_{\bm{f}_{k}})$, where $\mathcal{B}^{t}_{k}\coloneqq\{\bm{x}_{k}^{n}\colon\forall n\in\mathcal{B}^{t}\}$, to the other clients in $\mathcal{K}_{o}^{\mathcal{B}^{t}}$.

At this point, each client $k\in\mathcal{K}_{o}^{\mathcal{B}^{t}}$ holds the representations corresponding to the set of observed features of the current mini-batch, $\mathcal{K}^{\mathcal{B}^{t}}_{o}$. Now, to train for all the possible combinations of feature blocks while avoiding an exponential complexity, each client $k$ samples $K_{o}^{t}\coloneqq|\mathcal{K}^{\mathcal{B}^{t}}_{o}|$ tasks (the number of blocks observed for the samples in $\mathcal{B}^{t}$) from $\mathscr{P}_{k}\left(\mathcal{K}^{\mathcal{B}^{t}}_{o}\right)$ to be used to update the model parameters at iteration $t$ (note that $K_{o}^{t}\leq K$) and weights their losses, so that we get an unbiased estimator of the gradient. This allows us to train the model in such a way that it can perform inference for any subset of available blocks of features.

More precisely, the task-defining feature blocks sampled at client $k$ are given by $\mathcal{S}_{k}^{t}\coloneqq\left\{\mathcal{S}_{k1}^{t},\dots,\mathcal{S}_{kK_
{o}^{t}}^{n}\right\}$, where $\mathcal{S}_{kj}^{t}$ is a set containing $j$ feature blocks. Each set $\mathcal{S}_{kj}^{t}$ is sampled from a uniform distribution over a subset of $\mathscr{P}_{k}\left(\mathcal{K}_{o}^{\mathcal{B}^{t}}\right)$ that contains its sets of size $j$. That is:

$$ $\mathcal{S}_{kj}^{t}\sim\mathcal{U}\left(\mathscr{P}_{k}^{j}\left(\mathcal{K}_ {o}^{\mathcal{B}^{t}}\right)\right)\quad\text{where}\quad\mathscr{P}_{k}^{j} \left(\mathcal{K}_{o}^{\mathcal{B}^{t}}\right)\coloneqq\left\{\mathcal{J}\in \mathscr{P}_{k}(\mathcal{K}_{o}^{\mathcal{B}^{t}})\colon|\mathcal{J}|=j\right\}.$ $$

This task-sampling mechanism provides us with $K_{o}^{t}\leq K$ representations that correspond to the blocks in $\mathcal{S}_{k}^{t}$. This allows us to efficiently estimate the loss (linear complexity), whose exact computation would lead to an exponential complexity. Based on these sampled representations, we compute the following estimate for the loss:

$$ $\tilde{\mathcal{L}}_{k,\mathcal{S}_{k}^{t}}(\bm{\theta}^{t};\mathcal{B}^{t}) \coloneqq\frac{1}{|\mathcal{B}^{t}|}\sum_{n\in\mathcal{B}^{t}}\sum_{j\in[K_{o} ^{t}]}a_{j}^{n}\cdot\tilde{\mathcal{L}}_{k,\mathcal{S}_{kj}^{t}}(\bm{\theta}^{ t};n)\quad\text{where}\quad a_{j}^{n}={\binom{K_{o}^{t}-1}{j-1}}/j.$ (6) $$

The weighting $a_{j}^{n}$ counteracts the probability of each task being sampled, allowing us to obtain an unbiased estimate of the gradient. Now, we use this to minimize (5) by doing the following update:

$$ $\bm{\theta}^{t+1}=\bm{\theta}^{t}-\eta\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{ t}\mathcal{S}^{t}}(\bm{\theta}^{t}),$ (7) $$

where $\mathcal{S}^{t}\coloneqq\left\{\mathcal{S}_{1}^{t},\dots,\mathcal{S}_{K}^{t}\right\}$ and

$$ $\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm{\theta}^{t}) \coloneqq\frac{1}{|\mathcal{B}^{t}|}\sum_{n\in\mathcal{B}^{t}}\sum_{k\in \mathcal{K}_{o}^{n}}\sum_{j\in[|\mathcal{K}_{o}^{n}|]}a_{j}^{n}\cdot\nabla \tilde{\mathcal{L}}_{k,\mathcal{S}_{kj}^{t}}(\bm{\theta}^{t};n).$ $$

After completing the training in Algorithm 1, we can perform inference, as described in Algorithm 2. This algorithm consists simply of sharing of representations of the observed blocks test sample $\bm{x}^{\prime}$ by the clients which observe their partitions for this sample, $\mathcal{K}^{\prime}_{o}$, following by each of these clients $k\in\mathcal{K}^{\prime}_{o}$ using $h_{(k,\mathcal{K}^{\prime}_{o})}$ to predict the label.

Figure: Algorithm 2 LASER-VFL Inference

#### Aligning the samples based on block availability.

Entity alignment is a process the precedes training in VFL where the samples of the different local datasets are matched so that their features are correctly aligned, thus allowing for collaborative model training (Liu et al., ). The identification of the available blocks for each sample can be performed during this prelude to VFL training, allowing for a batch selection procedure that ensures that each batch contains samples with the same observed blocks, which we use in our method.

#### Extensions.

It is important to highlight that our method can also be easily employed in situations where different clients hold different labels, as it naturally fits multi-task learning problems. Further, it can easily be applied to setups where any subset of $\mathcal{K}$ holds the labels, sufficing to drop fusion models from the clients which do not hold the labels or have them tackle unsupervised or self-supervised learning tasks instead.

## 4 Convergence guarantees

In this section, we present convergence guarantees for our method. Let us start by stating the assumptions used in our results.

###### Assumption 1 ( $L$ -smoothness and finite infimum) .

Function $\tilde{\mathcal{L}}\colon\bm{\Theta}\mapsto\mathbb{R}$ is differentiable and there exists a constant $L\in(0,\infty)$ such that:

$$ $\forall\;\bm{\theta}_{1},\bm{\theta}_{2}\in\bm{\Theta}\colon\quad\lVert\nabla \tilde{\mathcal{L}}(\bm{\theta}_{1})-\nabla\tilde{\mathcal{L}}(\bm{\theta}_{2} )\rVert\leq L\lVert\bm{\theta}_{1}-\bm{\theta}_{2}\rVert.$ (A 1 ) $$

The inequality above holds if $\mathcal{L}$ is $L^{\prime}$-smooth.
We assume and define $\mathcal{L}^{\star}\coloneqq\inf_{\bm{\theta}}\mathcal{L}(\bm{\theta})>-\infty$.

###### Assumption 2 (Unbiased) .

The mini-batch gradient estimate is unbiased. That is, let $\tilde{\mathcal{L}}_{\mathcal{B}^{t}}(\bm{\theta})\coloneqq\frac{1}{|\mathcal{
B}^{t}|}\sum_{n\in\mathcal{B}^{t}}\tilde{\mathcal{L}}_{n}(\bm{\theta})$, we have for any mini-batch $\mathcal{B}^{t}\subseteq[N]$ that:

$$ $\forall(\bm{\theta},t)\in\bm{\Theta}\times\{0,1,\dots,T-1\}\colon\quad\mathbb{ E}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}}(\bm{\theta})\right]=\nabla \tilde{\mathcal{L}}(\bm{\theta}),$ (A 2 ) $$

where the expectation is with respect to the mini-batch $\mathcal{B}$.

###### Assumption 3 (Bounded variance) .

There exists a constant $\sigma\in(0,\infty)$ such that, for any mini-batch $\mathcal{B}^{t}\subseteq[N]$ of size $B$ and any set of task-defining blocks $\mathcal{S}^{t}$:

$$ $\forall(\bm{\theta},t)\in\bm{\Theta}\times\{0,1,\dots,T-1\}\colon\quad\mathbb{ E}\left\lVert\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm{ \theta}^{t})-\nabla\tilde{\mathcal{L}}(\bm{\theta})\right\rVert^{2}\leq\frac{ \sigma^{2}}{B},$ (A 3 ) $$

where the expectation is with respect to $\mathcal{B}^{t}$ and $\mathcal{S}^{t}$.

Using the assumptions above, we can derive the results below, where we use the following filtration:

$$ $\mathcal{F}^{t}\coloneqq\sigma(\mathcal{B}^{0},\mathcal{S}^{0},\mathcal{B}^{1} ,\mathcal{S}^{1},\dots,\mathcal{B}^{t-1},\mathcal{S}^{t-1}).$ $$

We also define a random variable $I_{k}^{n}$ that is equal to one if block $k$ of sample $n$ is observed and zero otherwise. Further, we let $\bm{I}^{n}=(I^{n}_{1},\dots,I^{n}_{K})$ and $\bm{I}=(I^{1},\dots,I^{N})$.

We use the following lemma, showing the unbiasedness of our update vector, in our proof of Theorem 1.

###### Lemma 1 (Unbiased update vector) .

If the mini-batch gradient estimate is unbiased (A2), then, for all $t\geq 0$:

$$ $\mathbb{E}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm {\theta})\mathrel{|}\sigma(\bm{I},\mathcal{F}^{t})\right]=\nabla\tilde{ \mathcal{L}}(\bm{\theta}^{t}),$ (8) $$

where the expectation is with respect to mini-batch $\mathcal{B}^{t}$ and the sampled set of tasks $\mathcal{S}^{t}$.

###### Theorem 1 (Main result) .

Let $\{\bm{\theta}^{t}\}$ be a sequence generated by Algorithm 1, if $\mathcal{L}$ is $L$-smooth and has a finite infimum (A1), the mini-batch estimate of the gradient is unbiased (A2), and our update vector has a bounded variance (A3), we then have that, for $\eta\in(0,1/L]$:

$$ $\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left\lVert\nabla\mathcal{L}(\bm{\theta}^ {t})\right\rVert^{2}\leq\frac{2\Delta}{\eta T}+\frac{\eta L\sigma^{2}}{B},$ (9) $$

where $\Delta\coloneqq\mathcal{L}(\bm{\theta}^{0})-\mathcal{L}^{\star}$ and the expectation is with respect to $\{\mathcal{B}^{t}\}$, $\{\mathcal{S}^{t}\}$, and $\bm{I}$.

If we take the stepsize to be $\eta=\sqrt{\frac{2\Delta B}{L\sigma^{2}T}}$, we get from (9) that

$$ $\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left\lVert\nabla\mathcal{L}(\bm{\theta}^ {t})\right\rVert^{2}\leq\sqrt{\frac{8\Delta L\sigma^{2}}{2BT}},$ $$

thus achieving a $\mathcal{O}(\nicefrac{{1}}{{\sqrt{T}}})$ convergence rate. Yet, if we assume that the batch size is sufficiently large, $B=\Omega(\nicefrac{{\sigma^{2}}}{{\epsilon}})$, we get from (9) that we can achieve an error level $\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left\lVert\nabla\mathcal{L}(\bm{\theta}^
{t})\right\rVert^{2}\leq\epsilon$ in $\mathcal{O}(\nicefrac{{1}}{{T}})$ iterations.

We can further establish linear convergence under the Polyak-Łojasiewicz inequality Polyak ().

###### Assumption 4 (PL inequality) .

We assume that there exists a positive constant $\mu$ such that

$$ $\forall\bm{\theta}\in\bm{\Theta}\colon\quad\lVert\nabla\mathcal{L}(\bm{\theta} )\rVert^{2}\geq 2\mu(\mathcal{L}(\bm{\theta})-\mathcal{L}^{\star}).$ (A 4 ) $$

We use the expected suboptimality $\delta^{t}\coloneqq\mathbb{E}\mathcal{L}(\bm{\theta}^{t})-\mathcal{L}^{\star}$, where the expectation is with respect to $\{\mathcal{B}^{t}\}$, $\{\mathcal{S}^{t}\}$, and $\bm{I}$, as our Lyapunov function in the following result.

###### Theorem 2 (Linear convergence) .

Let $\{\bm{\theta}^{t}\}$ be a sequence generated by Algorithm 1, if $\mathcal{L}$ is $L$-smooth and has a finite infimum (A1), the mini-batch estimate of the gradient is unbiased (A2), our update vector has a bounded variance (A3), and the PL inequality (A4) holds for $\mathcal{L}$, we then have that, for $\eta\in(0,1/L]$:

$$ $\delta^{T}\leq(1-\mu\eta)^{T}\delta^{0}+\frac{\eta L\sigma^{2}}{2\mu B}.$ $$

It follows from $\mu\leq L$ that $1-\mu\eta\in(0,1)$. Therefore, we have linear convergence to the global optimum for the full-batch case ($\sigma=0$) and, more generally, we have linear convergence to a $\mathcal{O}(\sigma^{2})$ neighborhood around it.

We defer the proofs of Lemma 1, Theorem 1, and Theorem 2 to Appendix B.

###### Assumption 1 ( L 𝐿 L italic_L -smoothness and finite infimum) .

###### Assumption 2 (Unbiased) .

###### Assumption 3 (Bounded variance) .

###### Lemma 1 (Unbiased update vector) .

###### Theorem 1 (Main result) .

###### Assumption 4 (PL inequality) .

###### Theorem 2 (Linear convergence) .

## 5 Experiments

In this section, we describe the experiments and analyze the results to empirically evaluate the performance of LASER-VFL and compare it against baseline approaches. We first provide a brief overview of the baseline methods and tasks considered. Next, we discuss the metrics used for evaluation and present the experimental results, followed by a brief discussion.

We compare our method to the following baselines:

- •
Local: as can be seen in (2), the local approach leverages its block for training and inference whenever it is observed, but ignores the remaining blocks.
- •
Standard VFL (Liu et al., ): as can be seen in (1), the standard VFL model can only be trained on and used for inference for fully-observed samples. When the model is unavailable for prediction, a random prediction is made.
- •
VFedTrans (Huang et al., ): VFedTrans introduces a multi-stage VFL training pipeline to collaboratively train local predictors. (We give more details about VFedTrans below.)
- •
Ensemble: in the ensemble approach, we train local predictors as in Local. However, instead of performing decoupled inference, the clients share their predictions and select a joint prediction by majority vote. Ties are broken at random.

Our experiments focus on the following tasks:

- •
HAPT (Reyes-Ortiz et al., ): a human activity recognition dataset.
- •
Credit (Yeh and Lien, ): a dataset for predicting credit card payment defaults.
- •
MIMIC-IV (Johnson et al., ): a time series dataset concerning healthcare data. We focus on the task of predicting in-hospital mortality from ICU data corresponding to patients admitted with chronic kidney disease. We resort to the data processing pipeline in (Gupta et al., ).
- •
CIFAR-10 & CIFAR-100 (Krizhevsky et al., ): are two popular image datasets, the former containing 10 classes and the latter containing 100.

For all datasets, we split the samples into $K=4$ feature blocks. For example, for the CIFAR datasets, each image is partitioned into four quadrants, each assigned to a different client. Given our interest in ensuring that all clients perform well at inference time, we average the metrics across clients. In Appendix C, we provide detailed descriptions of how we compute accuracy and $F_{1}$-score.

**Table 1: Test metrics for $K=4$ clients across different datasets and varying probabilities of missing blocks during training and inference, averaged over five seeds ($\pm$ standard deviation).**
| Training $p_{\text{miss}}$ | 0.0 | 0.1 | 0.5 |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Inference $p_{\text{miss}}$ | 0.0 | 0.1 | 0.5 | 0.0 | 0.1 | 0.5 | 0.0 | 0.1 | 0.5 |
| HAPT (accuracy, %) |  |  |  |  |  |  |  |  |  |
| Local | $82.1\pm 0.6$ | $82.2\pm 0.7$ | $81.9\pm 1.6$ | $82.3\pm 0.4$ | $82.4\pm 0.4$ | $81.8\pm 1.0$ | $80.7\pm 1.0$ | $80.8\pm 1.3$ | $80.3\pm 1.7$ |
| Standard VFL | $92.7\pm 0.4$ | $61.9\pm 3.8$ | $17.8\pm 5.1$ | $91.3\pm 1.0$ | $61.1\pm 3.7$ | $17.6\pm 5.0$ | $81.7\pm 2.2$ | $54.9\pm 3.4$ | $16.4\pm 4.2$ |
| VFedTrans | $82.5\pm 0.4$ | $82.5\pm 0.4$ | $82.3\pm 0.7$ | $82.9\pm 0.3$ | $82.9\pm 0.2$ | $82.8\pm 0.7$ | $82.2\pm 0.5$ | $82.2\pm 0.5$ | $82.2\pm 0.8$ |
| Ensemble | $90.5\pm 0.9$ | $90.0\pm 0.7$ | $84.3\pm 1.4$ | $90.9\pm 0.7$ | $90.2\pm 1.0$ | $84.4\pm 1.1$ | $88.7\pm 1.3$ | $88.5\pm 1.6$ | $82.6\pm 1.7$ |
| LASER-VFL (ours) | $\textbf{92.9}\pm 0.6$ | $\textbf{92.2}\pm 0.6$ | $\textbf{87.0}\pm 1.2$ | $\textbf{93.6}\pm 0.5$ | $\textbf{93.1}\pm 0.5$ | $\textbf{88.2}\pm 1.2$ | $\textbf{92.3}\pm 1.1$ | $\textbf{91.7}\pm 0.9$ | $\textbf{86.0}\pm 1.8$ |
| Credit ($F_{1}$-score$\times 100$) |  |  |  |  |  |  |  |  |  |
| Local | $37.7\pm 3.1$ | $37.6\pm 3.1$ | $37.5\pm 3.2$ | $37.6\pm 3.0$ | $37.6\pm 2.9$ | $37.5\pm 3.2$ | $36.2\pm 3.5$ | $36.2\pm 3.6$ | $36.0\pm 3.6$ |
| Standard VFL | $45.7\pm 2.6$ | $38.8\pm 2.5$ | $31.0\pm 0.7$ | $42.4\pm 1.2$ | $38.3\pm 0.7$ | $31.0\pm 0.7$ | $32.4\pm 2.4$ | $32.5\pm 1.3$ | $30.3\pm 0.5$ |
| VFedTrans | $37.7\pm 1.5$ | $37.6\pm 1.4$ | $37.2\pm 1.6$ | $39.5\pm 0.8$ | $39.4\pm 0.8$ | $39.2\pm 0.9$ | $35.7\pm 0.3$ | $35.6\pm 0.3$ | $35.6\pm 0.4$ |
| Ensemble | $42.1\pm 1.0$ | $42.1\pm 1.0$ | $40.1\pm 0.8$ | $41.4\pm 1.3$ | $40.8\pm 0.9$ | $39.2\pm 1.1$ | $\textbf{42.4}\pm 2.1$ | $\textbf{41.8}\pm 1.7$ | $40.1\pm 1.1$ |
| LASER-VFL (ours) | $\textbf{46.5}\pm 2.8$ | $\textbf{45.0}\pm 2.5$ | $\textbf{43.7}\pm 1.2$ | $\textbf{43.1}\pm 4.2$ | $\textbf{41.9}\pm 4.0$ | $\textbf{41.3}\pm 1.9$ | $41.5\pm 4.2$ | $40.9\pm 4.0$ | $\textbf{41.4}\pm 1.7$ |
| MIMIC-IV ($F_{1}$-score$\times 100$) |  |  |  |  |  |  |  |  |  |
| Local | $74.9\pm 0.5$ | $74.9\pm 0.5$ | $74.8\pm 0.5$ | $75.0\pm 0.3$ | $75.0\pm 0.2$ | $75.0\pm 0.3$ | $73.0\pm 0.5$ | $73.0\pm 0.5$ | $73.0\pm 0.5$ |
| Standard VFL | $91.6\pm 0.2$ | $77.0\pm 1.3$ | $52.5\pm 2.2$ | $90.9\pm 0.3$ | $76.8\pm 1.3$ | $52.5\pm 2.2$ | $81.1\pm 0.4$ | $70.2\pm 0.8$ | $51.8\pm 1.8$ |
| Ensemble | $84.2\pm 0.3$ | $83.9\pm 0.3$ | $77.9\pm 1.6$ | $84.1\pm 0.4$ | $83.7\pm 0.5$ | $78.1\pm 1.1$ | $82.1\pm 0.3$ | $81.8\pm 0.5$ | $76.4\pm 1.1$ |
| LASER-VFL (ours) | $\textbf{92.0}\pm 0.2$ | $\textbf{91.2}\pm 0.3$ | $\textbf{85.5}\pm 1.0$ | $\textbf{91.1}\pm 0.5$ | $\textbf{90.2}\pm 0.3$ | $\textbf{84.7}\pm 0.9$ | $\textbf{87.7}\pm 0.2$ | $\textbf{86.9}\pm 0.2$ | $\textbf{82.0}\pm 1.0$ |
| CIFAR-10 (accuracy, %) |  |  |  |  |  |  |  |  |  |
| Local | $72.4\pm 0.1$ | $72.5\pm 0.2$ | $72.5\pm 0.9$ | $71.8\pm 0.2$ | $72.0\pm 0.3$ | $72.0\pm 0.7$ | $67.6\pm 0.3$ | $67.7\pm 0.4$ | $68.0\pm 1.1$ |
| Standard VFL | $87.8\pm 0.2$ | $59.5\pm 12.0$ | $11.6\pm 3.5$ | $85.6\pm 0.8$ | $58.0\pm 11.2$ | $11.5\pm 3.4$ | $52.3\pm 10.2$ | $37.0\pm 9.0$ | $10.9\pm 2.1$ |
| Ensemble | $83.8\pm 0.2$ | $82.5\pm 0.6$ | $74.9\pm 0.5$ | $83.2\pm 0.3$ | $82.0\pm 0.8$ | $74.3\pm 0.7$ | $79.5\pm 0.5$ | $78.1\pm 0.5$ | $70.3\pm 0.3$ |
| LASER-VFL (ours) | $\textbf{90.3}\pm 0.2$ | $\textbf{89.3}\pm 0.3$ | $\textbf{81.6}\pm 1.2$ | $\textbf{89.7}\pm 0.2$ | $\textbf{88.6}\pm 0.3$ | $\textbf{81.1}\pm 1.1$ | $\textbf{85.5}\pm 0.3$ | $\textbf{84.6}\pm 0.4$ | $\textbf{77.0}\pm 1.1$ |
| CIFAR-100 (accuracy, %) |  |  |  |  |  |  |  |  |  |
| Local | $46.4\pm 0.2$ | $46.5\pm 0.1$ | $46.5\pm 0.1$ | $45.5\pm 0.3$ | $45.4\pm 0.4$ | $45.3\pm 0.8$ | $37.4\pm 0.6$ | $37.7\pm 0.5$ | $37.7\pm 0.9$ |
| Standard VFL | $62.1\pm 0.3$ | $40.0\pm 9.3$ | $2.3\pm 2.8$ | $55.6\pm 1.8$ | $35.6\pm 7.6$ | $2.3\pm 2.7$ | $15.6\pm 6.5$ | $10.2\pm 4.3$ | $1.4\pm 0.9$ |
| Ensemble | $58.5\pm 0.5$ | $56.3\pm 0.9$ | $48.6\pm 0.3$ | $57.5\pm 0.6$ | $55.3\pm 0.7$ | $47.4\pm 0.7$ | $48.2\pm 0.9$ | $46.3\pm 0.7$ | $40.1\pm 0.7$ |
| LASER-VFL (ours) | $\textbf{69.7}\pm 0.4$ | $\textbf{68.2}\pm 0.8$ | $\textbf{58.4}\pm 1.5$ | $\textbf{68.3}\pm 0.2$ | $\textbf{66.8}\pm 0.5$ | $\textbf{57.0}\pm 1.5$ | $\textbf{58.5}\pm 0.9$ | $\textbf{57.3}\pm 0.7$ | $\textbf{48.7}\pm 0.8$ |

We do not run VFedTrans for MIMIC-IV, CIFAR-10, and CIFAR-100 due to its greater complexity. In particular, the training pipeline of VFedTrans includes three stages: federated representation learning, a local representation distillation, and training a local model. The first two stages are performed per pair of clients and are thus repeated $K-1$ times. This means that experiments with VFedTrans take significantly longer to run (and to tune the hyperparameters of each stage of the training pipeline). Thus, given how VFedTrans performs similarly to Local in the first two experiments, we limit our experiments with VFedTrans to these simpler tasks.

#### Analyzing the empirical results.

In Table 1, we see that the methods with local predictors (Local and VFedTrans) are robust to missing blocks during both training and inference, yet they fall behind when all the blocks are observed. In contrast, Standard VFL performs well when all the blocks are observed, yet its performance degrades very quickly in the presence of missing blocks. The Ensemble method improves upon Standard VFL significantly in the presence of missing blocks, as it leverages the robustness its local predictors. Yet, when all the blocks are observed, although Ensemble improves over the local predictors significantly, it still cannot match the performance of Standard VFL. LASER-VFL consistently outperforms the baselines across the different probabilities of missing blocks during training and inference.(^6^66With the exception of two setups for the Credit dataset where Ensemble outperforms our method slightly, but this minimal difference falls well within the standard deviation.) In particular, even when the samples are fully-observed, LASER-VFL outperforms Standard VFL. We believe this is due to the regularization effect of the task-sampling mechanism in our method, which effectively behaves as a form of dropout.

## 6 Conclusions

In this work, we introduced LASER-VFL, a novel method for efficient training and inference in vertical federated learning that is robust to missing blocks of features without wasting data. This is achieved by carefully sharing model parameters and by employing a task-sampling mechanism that allows us to efficiently estimate a loss whose computation would otherwise lead to an exponential complexity. We provide convergence guarantees for LASER-VFL and present numerical experiments demonstrating the improved performance of LASER-VFL, not only in the presence of missing features but, remarkably, even when all samples are fully observed. For future work, it would be interesting to explore applications of our method to multi-task learning, seeing how our method behaves when different clients tackle different tasks, as this fits very naturally into our setup.

## Acknowledgements

This work is supported in part by the Fundação para a Ciência e a Tecnologia through the Carnegie Mellon Portugal Program and by the grants U.S. National Science Foundation ECCS-2318441.

## Appendix A Table of Notation

**Table 2: Table of Notation**
| Symbol | Description |
| --- | --- |
| $K$ | Number of clients |
| $N$ | Number of samples |
| $\bm{x}^{n}$ | Sample $n$ |
| $\bm{x}^{n}_{k}$ | Block $k$ of sample $n$ |
| $\mathcal{D}$ | Global dataset |
| $\mathcal{D}_{k}$ | Local dataset |
| $\bm{\theta}$ | Parameters of our model |
| $\bm{\theta}_{\mathcal{K}}$ | Parameters of the standard VFL model |
| $\bm{\theta}_{\bm{f}_{k}}$ | Parameters of the representation model at client $k$ |
| $\bm{\theta}_{\bm{g}_{k}}$ | Parameters of the fusion model at client $k$ |
| $\ell$ | Loss function |
| $\mathcal{L}$ | Objective function |
| $h$ | Predictor of the standard VFL model |
| $h_{(k,\mathcal{J})}$ | Predictor at client $k$ using feature blocks $\mathcal{J}$ |
| $\mathscr{P}(\mathcal{S})$ | Power set of $\mathcal{S}$ |
| $\mathscr{P}_{j}(\mathcal{S})$ | Subsets in the power set of $\mathcal{S}$ containing $j$ |
| $\mathscr{P}^{i}_{j}(\mathcal{S})$ | Subsets in the power set of $\mathcal{S}$ of size $i$ containing $j$ |
| $\bm{f}_{k}$ | Representation model at client $k$ |
| $g_{k}$ | Fusion model at client $k$ |
| $h_{k}$ | Predictor at client $k$ |
| $g$ | Standard VFL fusion model |
| $h$ | Standard VFL predictor |
| $h^{\star}$ | Ideal predictor (without FL constraints) |
| $\mathcal{K}_{o}^{n}$ | Set of observed features blocks of sample $n$ |
| $\bm{x}^{\prime}$ | Test sample |

## Appendix B Proofs

### B.1 Preliminaries

The following quadratic upper bound follows from the $L$-smoothness of $\tilde{\mathcal{L}}$ (A1):

$$ $\forall\;\bm{\theta}_{1},\bm{\theta}_{2}\in\bm{\Theta}\colon\quad\tilde{ \mathcal{L}}(\bm{\theta}_{1})\leq\tilde{\mathcal{L}}(\bm{\theta}_{2})+\nabla \tilde{\mathcal{L}}(\bm{\theta}_{2})^{\top}(\bm{\theta}_{1}-\bm{\theta}_{2})+ \frac{L}{2}\lVert\bm{\theta}_{1}-\bm{\theta}_{2}\rVert^{2}.$ (10) $$

Let $X$ denote a random variable and let $F$ be a convex function, Jensen’s inequality states that

$$ $F(\mathbb{E}(X))\leq\mathbb{E}(F(X)).$ (11) $$

### B.2 Proof of Lemma 1

We define the following shorthand notation for the conditional expectations

$$ $\mathbb{E}_{t}\left[\cdot\right]\coloneqq\mathbb{E}_{\mathcal{B}^{t}\mathcal{S }^{t}}\left[\cdot\mathrel{|}\sigma(\bm{I},\mathcal{F}^{t})\right]$ $$

and

$$ $\mathbb{E}_{t^{+}}\left[\cdot\right]\coloneqq\mathbb{E}_{\mathcal{B}^{t} \mathcal{S}^{t}}\left[\cdot\mathrel{|}\sigma(\bm{I},\mathcal{F}^{t},\mathcal{B }^{t})\right].$ $$

From the definition of our update vector $\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm{\theta}^{t})$, in (7), we have that

$$ $\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}} (\bm{\theta}^{t})\right]=\mathbb{E}_{t}\left[\frac{1}{|\mathcal{B}^{t}|}\sum_{ n\in\mathcal{B}^{t}}\sum_{k\in\mathcal{K}_{o}^{n}}\sum_{j=1}^{|\mathcal{K}_{o} ^{n}|}a_{j}^{n}\cdot\nabla\tilde{\mathcal{L}}_{k,\mathcal{S}_{kj}^{t}}(\bm{ \theta}^{t};n)\right]$ $$

and, using the law of total expectation, we arrive at

$$ $\begin{split}\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t} \mathcal{S}^{t}}(\bm{\theta}^{t})\right]&=\mathbb{E}_{t}\left[\frac{1}{| \mathcal{B}^{t}|}\sum_{n\in\mathcal{B}^{t}}\mathbb{E}_{t^{+}}\left[\sum_{k\in \mathcal{K}_{o}^{n}}\sum_{j=1}^{|\mathcal{K}_{o}^{n}|}a_{j}^{n}\cdot\nabla \tilde{\mathcal{L}}_{k,\mathcal{S}_{kj}^{t}}(\bm{\theta}^{t};n)\right]\right] \\ &=\mathbb{E}_{t}\left[\frac{1}{|\mathcal{B}^{t}|}\sum_{n\in\mathcal{B}^{t}} \sum_{k\in\mathcal{K}_{o}^{n}}\sum_{j=1}^{|\mathcal{K}_{o}^{n}|}a_{j}^{n}\cdot \mathbb{E}_{t^{+}}\left[\nabla\tilde{\mathcal{L}}_{k,\mathcal{S}_{kj}^{t}}(\bm {\theta}^{t};n)\right]\right].\end{split}$ $$

Now, it follows from the fact that $\mathcal{S}_{kj}^{t}\sim\mathcal{U}(\mathscr{P}_{k}^{j}(\mathcal{K}_{o}^{n}))$ that

$$ $\begin{split}\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t} \mathcal{S}^{t}}(\bm{\theta}^{t})\right]&=\mathbb{E}_{t}\left[\frac{1}{| \mathcal{B}^{t}|}\sum_{n\in\mathcal{B}^{t}}\sum_{k\in\mathcal{K}_{o}^{n}}\sum_ {j=1}^{|\mathcal{K}_{o}^{n}|}a_{j}^{n}\cdot\sum_{\mathcal{J}\in\mathscr{P}_{k} ^{j}(\mathcal{K}_{o}^{n})}\mathbb{P}(\mathcal{S}_{kj}^{n}=\mathcal{J})\cdot \nabla\tilde{\mathcal{L}}_{k,\mathcal{J}}(\bm{\theta}^{t};n)\right]\\ &=\mathbb{E}_{t}\left[\frac{1}{|\mathcal{B}^{t}|}\sum_{n\in\mathcal{B}^{t}} \sum_{k\in\mathcal{K}_{o}^{n}}\sum_{j=1}^{|\mathcal{K}_{o}^{n}|}\frac{a_{j}^{n }}{|\mathscr{P}_{k}^{j}(\mathcal{K}_{o}^{n})|}\cdot\sum_{\mathcal{J}\in \mathscr{P}_{k}^{j}(\mathcal{K}_{o}^{n})}\nabla\tilde{\mathcal{L}}_{k,\mathcal {J}}(\bm{\theta}^{t};n)\right]\\ &=\mathbb{E}_{t}\left[\frac{1}{|\mathcal{B}^{t}|}\sum_{n\in\mathcal{B}^{t}} \sum_{k\in\mathcal{K}_{o}^{n}}\sum_{j=1}^{|\mathcal{K}_{o}^{n}|}\frac{a_{j}^{n }}{\binom{|\mathcal{K}_{o}^{n}|-1}{j-1}}\cdot\sum_{\mathcal{J}\in\mathscr{P}_{ k}^{j}(\mathcal{K}_{o}^{n})}\nabla\tilde{\mathcal{L}}_{k,\mathcal{J}}(\bm{ \theta}^{t};n)\right].\end{split}$ $$

Now, from the definition of $a_{j}^{n}$, we have that

$$ $\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}} (\bm{\theta}^{t})\right]=\mathbb{E}_{t}\left[\frac{1}{|\mathcal{B}^{t}|}\sum_{ n\in\mathcal{B}^{t}}\sum_{k\in\mathcal{K}_{o}^{n}}\sum_{j=1}^{|\mathcal{K}_{o} ^{n}|}\frac{1}{j}\cdot\sum_{\mathcal{J}\in\mathscr{P}_{k}^{j}(\mathcal{K}_{o}^ {n})}\nabla\tilde{\mathcal{L}}_{k,\mathcal{J}}(\bm{\theta}^{t};n)\right],$ $$

which we can rewrite as

$$ $\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}} (\bm{\theta}^{t})\right]=\mathbb{E}_{t}\left[\frac{1}{|\mathcal{B}^{t}|}\sum_{ n\in\mathcal{B}^{t}}\sum_{k\in\mathcal{K}_{o}^{n}}\sum_{\mathcal{J}\in\mathscr {P}_{k}(\mathcal{K}_{o}^{n})}\frac{1}{|\mathcal{J}|}\cdot\nabla\tilde{\mathcal {L}}_{k,\mathcal{J}}(\bm{\theta}^{t};n)\right].$ $$

Now, from the definition of $\nabla\tilde{\mathcal{L}}_{n}(\bm{\theta})$, in (5), we have that

$$ $\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}} (\bm{\theta}^{t})\right]=\mathbb{E}_{t}\left[\frac{1}{|\mathcal{B}^{t}|}\sum_{ n\in\mathcal{B}^{t}}\nabla\tilde{\mathcal{L}}_{n}(\bm{\theta})\right].$ $$

Finally, from (A2), we have that:

$$ $\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}} (\bm{\theta}^{t})\right]=\nabla\tilde{\mathcal{L}}(\bm{\theta}^{t}),$ $$

arriving at the result that we set out to prove.

### B.3 Proof of Theorem 1

Let us show the convergence of our method. It follows from the quadratic upper bound in (10) which ensues from the $L$-smoothness assumption (A1), that

$$ $\tilde{\mathcal{L}}(\bm{\theta}^{t+1})-\tilde{\mathcal{L}}(\bm{\theta}^{t}) \leq\nabla\tilde{\mathcal{L}}(\bm{\theta}^{t})^{\top}(\bm{\theta}^{t+1}-\bm{ \theta}^{t})+\frac{L}{2}\left\lVert\bm{\theta}^{t+1}-\bm{\theta}^{t}\right \rVert^{2}.$ $$

Using the fact that $\bm{\theta}^{t+1}=\bm{\theta}^{t}-\eta\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{
t}\mathcal{S}^{t}}(\bm{\theta}^{t})$, we get that:

$$ $\tilde{\mathcal{L}}(\bm{\theta}^{t+1})-\tilde{\mathcal{L}}(\bm{\theta}^{t}) \leq-\eta\nabla\tilde{\mathcal{L}}(\bm{\theta}^{t})^{\top}\nabla\tilde{ \mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm{\theta}^{t})+\frac{\eta^{2}L }{2}\left\lVert\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm{ \theta}^{t})\right\rVert^{2}.$ $$

Now, we take the conditional expectation
$\mathbb{E}_{t}\left[\cdot\right]=\mathbb{E}_{\mathcal{B}^{t}\mathcal{S}^{t}}
\left[\cdot\mathrel{|}\sigma(\bm{I},\mathcal{F}^{t})\right]$, defined in Appendix B.2,
arriving that:

$$ $\mathbb{E}_{t}\left[\tilde{\mathcal{L}}(\bm{\theta}^{t+1})\right]-\tilde{ \mathcal{L}}(\bm{\theta}^{t})\leq-\eta\nabla\tilde{\mathcal{L}}(\bm{\theta}^{t })^{\top}\mathbb{E}_{t}\left[\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t} \mathcal{S}^{t}}(\bm{\theta}^{t})\right]+\frac{\eta^{2}L}{2}\mathbb{E}_{t} \left[\left\lVert\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}( \bm{\theta}^{t})\right\rVert^{2}\right].$ $$

We can now use the result from Lemma 1 to arrive at the following inequality:

$$ $\mathbb{E}_{t}\left[\tilde{\mathcal{L}}(\bm{\theta}^{t+1})\right]-\tilde{ \mathcal{L}}(\bm{\theta}^{t})\leq-\eta\left\lVert\nabla\tilde{\mathcal{L}}(\bm {\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L}{2}\mathbb{E}_{t}\left[\left \lVert\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm{\theta}^{ t})\right\rVert^{2}\right].$ (12) $$

Now, since it follows from Lemma 1 that

$$ $\mathbb{E}_{t}\left[\left\lVert\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t} \mathcal{S}^{t}}(\bm{\theta}^{t})-\nabla\tilde{\mathcal{L}}(\bm{\theta}^{t}) \right\rVert^{2}\right]=\mathbb{E}_{t}\left[\left\lVert\nabla\tilde{\mathcal{L }}_{\mathcal{B}^{t}\mathcal{S}^{t}}(\bm{\theta}^{t})\right\rVert^{2}\right]- \left\lVert\nabla\tilde{\mathcal{L}}(\bm{\theta}^{t})\right\rVert^{2},$ $$

Thus, we have from the bounded variance assumption (A3) that

$$ $\mathbb{E}_{t}\left[\left\lVert\nabla\tilde{\mathcal{L}}_{\mathcal{B}^{t} \mathcal{S}^{t}}(\bm{\theta}^{t})\right\rVert^{2}\right]\leq\left\lVert\nabla \tilde{\mathcal{L}}(\bm{\theta}^{t})\right\rVert^{2}+\frac{\sigma^{2}}{B}.$ $$

Using the inequality above in (12), we arrive at

$$ $\begin{split}\mathbb{E}_{t}\left[\tilde{\mathcal{L}}(\bm{\theta}^{t+1})\right] -\tilde{\mathcal{L}}(\bm{\theta}^{t})&\leq-\eta\left\lVert\nabla\tilde{ \mathcal{L}}(\bm{\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L}{2}\left\lVert \nabla\tilde{\mathcal{L}}(\bm{\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L \sigma^{2}}{2B}\\ &=-\eta\left(1-\frac{\eta L}{2}\right)\left\lVert\nabla\tilde{\mathcal{L}}(\bm {\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L\sigma^{2}}{2B}.\end{split}$ $$

Therefore, for a stepsize $\eta\in(0,1/L]$, we have that

$$ $\mathbb{E}_{t}\left[\tilde{\mathcal{L}}(\bm{\theta}^{t+1})\right]-\tilde{ \mathcal{L}}(\bm{\theta}^{t})\leq-\frac{\eta}{2}\left\lVert\nabla\tilde{ \mathcal{L}}(\bm{\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Taking the conditional expectation $\mathbb{E}[\cdot\mathrel{|}\sigma(\bm{I})]$, we get from the inequality above that have:

$$ $\mathbb{E}\left[\tilde{\mathcal{L}}(\bm{\theta}^{t+1})-\tilde{\mathcal{L}}(\bm {\theta}^{t})\mathrel{|}\sigma(\bm{I})\right]\leq-\frac{\eta}{2}\cdot\mathbb{E }\left[\left\lVert\nabla\tilde{\mathcal{L}}(\bm{\theta}^{t})\right\rVert^{2} \mathrel{|}\sigma(\bm{I})\right]+\frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Now, taking the unconditional expectation (over $\bm{I}$) and using the law of total expectation and the fact that
$\mathcal{L}(\bm{\theta})=\mathbb{E}_{\bm{I}}[\tilde{\mathcal{L}}(\bm{\theta})]$, we get that

$$ $\mathbb{E}\mathcal{L}(\bm{\theta}^{t+1})-\mathbb{E}\mathcal{L}(\bm{\theta}^{t} )\leq-\frac{\eta}{2}\cdot\mathbb{E}\left[\mathbb{E}\left[\left\lVert\nabla \tilde{\mathcal{L}}(\bm{\theta}^{t})\right\rVert^{2}\mathrel{|}\sigma(\bm{I}) \right]\right]+\frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Therefore, we have from Jensen’s inequality (11) that

$$ $\mathbb{E}\mathcal{L}(\bm{\theta}^{t+1})-\mathbb{E}\mathcal{L}(\bm{\theta}^{t} )\leq-\frac{\eta}{2}\cdot\mathbb{E}\left\lVert\mathbb{E}\left[\nabla\tilde{ \mathcal{L}}(\bm{\theta}^{t})\mathrel{|}\sigma(\bm{I})\right]\right\rVert^{2}+ \frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Using the dominated convergence theorem, we can interchange the expectation and the gradient operators, arriving at the following descent lemma in expectation:

$$ $\mathbb{E}\mathcal{L}(\bm{\theta}^{t+1})-\mathbb{E}\mathcal{L}(\bm{\theta}^{t} )\leq-\frac{\eta}{2}\cdot\mathbb{E}\left\lVert\nabla\mathcal{L}(\bm{\theta}^{t })\right\rVert^{2}+\frac{\eta^{2}L\sigma^{2}}{2B}.$ (13) $$

Taking the average of the inequality above for $t\in\{0,1,\dots,T-1\}$, we get that:

$$ $\frac{\mathbb{E}\mathcal{L}(\bm{\theta}^{T})-\mathcal{L}(\bm{\theta}^{0})}{T} \leq-\frac{\eta}{2T}\sum_{t=0}^{T-1}\mathbb{E}\left\lVert\nabla\mathcal{L}(\bm {\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L\sigma^{2}}{2}.$ $$

Lastly, rearranging the terms and using the existence and definition of $\mathcal{L}^{\star}$ (A1), we have that

$$ $\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}\left\lVert\nabla\mathcal{L}(\bm{\theta}^ {t})\right\rVert^{2}\leq\frac{2(\mathcal{L}(\bm{\theta}^{0})-\mathcal{L}^{ \star})}{\eta T}+\frac{\eta L\sigma^{2}}{B},$ $$

thus arriving at the result in (9).

### B.4 Proof of Theorem 2

Following the same steps as in the proof of Theorem 1, we can arrive at the same descent lemma in expectation (13),

$$ $\mathbb{E}\mathcal{L}(\bm{\theta}^{t+1})-\mathbb{E}\mathcal{L}(\bm{\theta}^{t} )\leq-\frac{\eta}{2}\cdot\mathbb{E}\left\lVert\nabla\mathcal{L}(\bm{\theta}^{t })\right\rVert^{2}+\frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Rearranging the terms and subtracting the infimum on both sides, we get that

$$ $\mathbb{E}\mathcal{L}(\bm{\theta}^{t+1})-\mathcal{L}^{\star}\leq\mathbb{E} \mathcal{L}(\bm{\theta}^{t})-\mathcal{L}^{\star}-\frac{\eta}{2}\cdot\mathbb{E} \left\lVert\nabla\mathcal{L}(\bm{\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L \sigma^{2}}{2B}.$ $$

Now, using the definition of our Lyapunov function, $\delta^{t}=\mathbb{E}\mathcal{L}(\bm{\theta}^{t})-\mathcal{L}^{\star}$, we arrive at the following inequality:

$$ $\delta^{t+1}\leq\delta^{t}-\frac{\eta}{2}\cdot\mathbb{E}\left\lVert\nabla \mathcal{L}(\bm{\theta}^{t})\right\rVert^{2}+\frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Now, from the PL inequality (A4), we have that:

$$ $\delta^{t+1}\leq\delta^{t}-\mu\eta\cdot\mathbb{E}\left[\mathcal{L}(\bm{\theta} ^{t})-\mathcal{L}^{\star}\right]+\frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Therefore, again using the definition of our Lyapunov function, we arrive at

$$ $\delta^{t+1}\leq(1-\mu\eta)\cdot\delta^{t}+\frac{\eta^{2}L\sigma^{2}}{2B}.$ $$

Recursing the inequality above, we get that

$$ $\delta^{T}\leq(1-\mu\eta)^{T}\delta^{0}+\frac{\eta^{2}L\sigma^{2}}{2B}\sum_{t= 0}^{T-1}(1-\mu\eta)^{t}.$ $$

Finally, using the sum of a geometric series, we arrive at

$$ $\delta^{T}\leq(1-\mu\eta)^{T}\delta^{0}+\frac{\eta L\sigma^{2}}{2\mu B},$ $$

which corresponds to the result that we set out to prove.

## Appendix C Experiment details

#### Notes on VFedTrans.

For the VFedTrans (Huang et al., ) pipeline, we use FedSVD LABEL:chai2022practicallosslessfederatedsingular for federated representation learning and an autoencoder for local representation learning, as these are the best performing methos in Huang et al. (). For the local predictor, we use a multilayer perceptron, as in Huang et al. ().

#### Notes on MIMIC-IV.

We use ICU data of MIMIC-IV v1.0 (Johnson et al., ) and follow the data processing pipeline in Gupta et al. (). We focus on the task of predicting in-hospital mortality for patients admitted with chronic kidney disease. We do feature selection with diagnosis data and the selected features of chart events (labs and vitals) used in the “Proof of concept experiments” Gupta et al. () as input features.

#### Computing metrics.

As mentioned in the main text of the paper, we want our metrics to capture that fact that we want all clients to perform well at inference time. Therefore we consider the average metrics across the clients. More precisely, we compute the metrics for Table 1 as follows:

- •
For the test accuracy, let $\hat{y}^{n}(k)$ denote the prediction of client $k$ for sample $n$, we compute
$\text{accuracy}=100\times\frac{1}{N}\sum_{n=1}^{N}\left(\frac{1}{|\mathcal{K}_
{o}^{n}|}\sum_{k\in\mathcal{K}_{o}^{n}}\bm{1}(\hat{y}^{n}(k)=y^{n})\right);$
- •
For the $F_{1}$-score ($\times 100$), we compute
$F_{1}\text{-score}=100\times\frac{1}{K}\sum_{k=1}^{K}\frac{P_{k}\cdot R_{k}}{P
_{k}+R_{k}},$
where $P_{k}$ and $R_{k}$ are the precision and recall of client $k$, with respect to its predictor and (only) the samples it observed.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. What do hyperparameters refer to in vertical FL? Is it the optimization parameters such as learning rate or some specific parameter in vertical FL?

2. Does the proposed algorithm allow multiple local updates in each machine like traditional FL?
soundness: 3
strengths: 1. This paper proposes a new algorithm to address the missing features in vertical FL.

2. Valid convergence analysis and experiments are provided.
summary: This paper studies vertical FL with missing features. The authors propose a new vertical FL method named **LASER-VFL** for efficient training and inference of split neural network-based models. This method relies on the strategic sharing of model parameters and on task-sampling to train a family of predictors. The authors also provide convergence rates for **LASER-VFL** under different objective functions.
weaknesses: 1. The notations are too complex, which is not friendly for readers. The authors should try to simplify notations.

2. I'm not an expert in vertical FL, so I'm confused about the **key idea** presented on page 4.  So it will be better to add workflow diagrams to explain the method.

3. The experiments are conducted on the data with synthetic missing features. It would be more convincing to show results on real-world vertical FL data with missing features.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: How would the performance of the aggregation mechanism differ when using weighted averaging rather than simple averaging? What if some clients have more powerful feature blocks? Giving uniform weights would alter the performance compared to other concatenation methods.

How would the performance be with even larger number of clients?

For small number of clients, the combinatorial approach becomes feasible, how does the combinatorial approach (which can be considered as an upper bound) differ from the proposed method?

How does the method perform against the other discussed related works which does solve the same problem but may have different training stages/modules?

How would the results be in a convex setting? Where does it converge? Does it guarantee optimal performance of VFL with missing clients?
soundness: 3
strengths: The authors propose an approach for VFL that enables both training and inference using all available blocks of features.

The authors provide theoretical guarantees for convergence under L-smoothness, unbiasedness, and bounded variance assumptions. 

The authors conduct empirical evaluations of the proposed methods over 4 datasets, showing good performance.
summary: This paper focuses on vertical federated learning (VFL), where data is partitioned by features rather than by samples. In this setting, each client possesses only a subset of the features for a specific sample. The paper addresses the challenge of missing features in the dataset and the heterogeneous availability of clients during the federation. Standard VFL approaches cannot handle these challenges, as it assumes that all features are available during both training and inference. To tackle these issues, this work proposes LASER-VFL, which relies on strategic sharing of model parameters and task sampling to train a family of predictors, allowing training and inference using any and all available blocks of features, leading to full utilization of the data. The paper provides theoretical guarantees for non-convex objectives and demonstrates empirical results on four different datasets, showing good performance.
weaknesses: For VFL models, the authors mention that split neural networks are often considered and explains that they seek to find multiple parameters for each block and for a fusion model. Then mentions that they aim to solve (1) which requires all the blocks, which contradicts the paper objective seeking VFL with missing blocks. The authors should be clear from the early stage whether they aim to solve (1), or other works aimed to solve (1). 

The authors mentions that LASER-VFL enables both training and inference
using all available blocks of features. However, in Figure 1 (b) (our method), shows that samples 4, 6, and 7, which has missing blocks, are only used for training (and not for testing), which is not clear why.

The authors go through some previous works very briefly in the introduction without going into details of their approaches and mention that these previous works add training stages and modules, making them more complex. While related work was provided later this paragraph seemed both redundant and unclear. 

The authors mentioned some works that dealt with missing features. Then surprisingly claim that to the best of their knowledge, LASER-VFL is the first method to enables both training and inference using any and all available blocks of features. (line 108: We propose LASER-VFL (Leveraging All Samples Efficiently for Real-world VFL), a novel method that enables both training and inference using any and all available blocks of features. To the best of our knowledge, this is the first method to achieve this.)

Citations should be \citep when discussing some paper (not the authors). For example in line 144 it should be in (Feng & Yu 2020; Kang et al. 2022a) not in Feng & Yu (2020); Kang et al. (2022a).

In VFL problems what cases would really requires more than a few clients to collaborate? In the empirical setting the authors considered 4 clients. If we have a few companies, why not simply do a combinatorial approach? For such few client cases, say 2 or 3, the authors could show the optimal performance of training the 2^k-1 models as an optimal target against the proposed solution. 

The authors introduce a lot of notations, which should be minimized to the least possible to make it easier to follow. (While the authors clearly mentions that the notation in (3) differs slightly from (1) and (2), this still should be avoided as much as possible). 

While the authors mention several methods that dealt with missing blocks, most of these approaches were not considered in the empirical evaluation. Being more complex and requiring multiple stages or modules should not be a reason to skip the comparison. It could be that these added stages or complexities lead to better accuracy and/or better efficiency in terms of rounds to accuracy. 

Giving uniform weights in the aggregation of the feature representations of different blocks of features for the different clients (Eq. (4)) would alter the performance compared to other concatenation methods. Averaging and sum can cancel some information which can be a drawback compared to concatenation methods.

While the authors claim that the method is suitable for a growing number of clients, avoiding the exponential growth in computation and memory. Empirical results should be provided to deal with more than 4 clients to see that the method indeed can achieve good performance. For example, would be nice to show for 8 and/or 16 clients. 

The theory shows convergence in the non-convex setting, under L-smoothness and other assumptions, which is appreciated. However, such results do not reveal the optimality of the method.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - How does this paper compare to the methods proposed in Ganguli et al., [2023]? They seem to be doing something very similar including the dropout mechanism during training. Perhaps, the authors missed this very related work.

- The model architecture resembles a single-layer GNN with summation aggregation. Each client has it's own fusion predictor. Is this correct or am I missing something?

- While the weighting method for learning an "exponential number of models" is interesting, it is perhaps a bit misleading to argue for the "compactness" of learning an exponential number of models. Simple dropout could also learn "an exponential number of models" but it is difficult to argue that these models would be anywhere near the naive solution of actually directly training an exponential number of models.

- There are other types of parametric aggregation methods that can handle arbitrary sizes. For example, an attention-based aggregation or similar is a parametric aggregation method but can handle arbitrary length inputs.

- The notation is a bit overwhelming and needs to be simplified and clarified. In Eq 5 alone, there are 4 different uses of the $\mathcal{L}$ with different subscripts.

- Choosing a mini-batch with the same missing features in a real-world VFL scenario would introduce bias in the training batch since it is unlikely for the missing features to be missing completely at random (MCAR). Rather, different groups of samples will likely have similar missing features. Can you explain how your method would handle this? This seems to be a severe limitation as the exact missingness pattern is unlikely to happen often unless the number of clients is very small.

- Related to above, it seems that this biasedness in mini-batches would cause Assumption 2 in your theory to be violated. Thus, it is unclear if your theory is reasonable for the practical algorithm and real-world situation.

- Relatedly, how do you handle the fact that there might be missingness patterns that are far more common than others? 

- Does this assume that all clients participate during the training process (i.e., no client leaves even temporarily)? It seems that this requires all clients to be available during training (even if they do not participate in the training round).

- How are the baseline methods trained when there is missingness at test time? Are the batches with high missingness simply thrown out? It seems there are some simple baselines to adapt standard VFL by simply imputing zeros at train and test time or something like that. Stronger baselines with some naive adaptations like this for missing values would be helpful.  Currently, it seems that standard VFL is unfairly penalized.

- Do your experiments assume missingness completely at random (MCAR)? I'm guessing they do but wanted to double check. This would align with your theory.

- What happens if you do not weight things but just use simple dropout when sampling your missing features?
soundness: 3
strengths: - The paper proposes a Monte-Carlo approach with weighting to train a split-NN architecture to be robust to missing features during training and inference. The method is relatively simple and avoids complex hyperparameters or multi-stage configurations.

- The paper provides theoretic convergence results on their proposed method.

- The empirical results show significant improvement over baselines methods used in the paper when there is missingness at inference time.
summary: This paper proposes to handle missing features in VFL during both training and testing by sharing model parameters and sampling different configurations to approximate the exponential number of possible predictors (with different possible non-missing features).
The intuition is that entirely local models are robust to missing features but poor predictors while global models are stronger predictors but not robust missing features.
The paper proposes to train a set of models for all possible subsets of missing features by sharing model parameters and sampling different missingness patterns, called "tasks" in this paper, during training time.
The sampled tasks are weighted to approximate the loss function of all missing value patterns using Monte Carlo sampling of the tasks.
The paper provides theoretic results on convergence and empirically demonstrates their method compared to several baseline methods.
weaknesses: ------ After reading author response ------
While I still have some concerns, the author response helped reduce them. Thus, I've increased my score by 1.
------ Original -------
- Either the setting is a small number of clients (say 4 as used in the experiments) and then it might be feasible to actually just train all possible models---this calls into question the motivation for learning "an exponential number of models" since there are only 24 models. Or, the setting is a large number of clients (say 10-100) and then finding large batches of samples such that the missing values are the same would result in either small or biased batches---which would violate assumption 2. Thus, the motivation of the paper or the requirement for large unbiased batches seem to be at odds with each other. 

- The assumption of unbiased batches would require the missingness to be Missing Completely at Random (MCAR) to ensure that the batches are unbiased. This seems unlikely to hold in real-world applications.

- The paper seems to overclaim a bit on the novelty or significance of the model sharing aspect. The model sharing aspect is a common approach in GNNs or other modern networks such as attention-based architectures, yet these do not claim significant novelty in terms of sharing model parameters across nodes.

- The paper lacks an empirical comparison to the method proposed in Ganguli et al. [2023] that can handle inference time faults. This seems to be a closely related method. They use dropout (like your task sampling), gossiping (like the ensemble-based baseline), and replication of the fusion/predictor (as you do as well).

[Ganguli et al., 2023] Ganguli, S., Zhou, Z., Brinton, C. G., & Inouye, D. I. (2023). Fault Tolerant Serverless VFL Over Dynamic Device Environment. arXiv preprint arXiv:2312.16638.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: By transforming the objective into a finite sum problem, is the convergence analysis almost the same as the classic convergence analysis of SGD?
soundness: 3
strengths: - The writing is clear in general;
- The paper is well-motivated, the missing feature during training and inference is a crucial problem in vertical federated learning from a practical perspective.
- The proposed method is simple and easy to understand.
summary: The author(s) proposed LASER-VFL to handle the missing feature challenge arising from the vertical federated learning problem. By transforming the objective into a finite sum problem, LASER-VFL can handle missing features (clients) via the task sampling technique, which is similar to client dropout. LASER-VFL is backed with theoretical analysis and some simple numerical experiments.
weaknesses: - The experiment section is rather weak:
    - Experiments details are missing, for example, the train/test split of dataset, the neural network architecture being used, the hyper-parameter tunning...
    - A baseline is missing, the naïve method described in section 2 (train exponentially many predictors).
    - The number of clients is set to 4 for all experiments, how does the result change when we increase the number of clients? If there are 4 clients, then it is not really computational prohibitive to train exponentially many predictors.
- Notations are quite heavy in line 293-342.
- The theoretical analysis assumes features are missed uniform randomly, which simplifies the problem significantly. By this assumption and transforming the objective into a finite sum problem, the author(s) could directly use classic convergence analysis of SGD without much technical modification. Therefore, the theoretical novelty is kind of limited.
- Minor: missing reference at line 855.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - Why are there K fusion models?  Is it because each client needs a specialized fusion model for its own task?
- Could the authors elaborate on line 2 in Algorithm 1 (Client K sample the same mini-batch .... via a shared seed)? What is the impact of this shared seed, and why is the same mini-batch important?
- Could the authors elaborate on how domain shift is addressed? Or how would the framework generalize if there were a significant level of domain shift across clients?
- Does the proposed framework have specific requirements for the fusion model and projectors? Also, what is the architecture of the fusion model and predictors in Section 5?
soundness: 2
strengths: I really enjoyed reading this paper. The presentation is good and easy to follow, and the experimental results are very convincing. More specifically:
- The paper addresses missing features problem, which is an important area of study in VFL.
- The proposed idea is simple yet effective. The authors create a fusion of arbitrary combinations of predictors on each client, allowing them to leverage all available features. To avoid sharing an exponential number of predictors, they propose a strategy in which only K predictors and K fusion models need to be shared. 
- They demonstrate improved performance compared to the existing works.
summary: This paper focuses on vertical federated learning. More specifically, the authors present a framework called LASER-VFL, which is more robust in handling missing features and any arbitrary combination of partitions. They provide a convergence rate analysis along with  experimental results.
weaknesses: - While the authors present related works, it is difficult to identify which components of the proposed framework are novel when reading Sections 2 and 3. I encourage the authors to provide a more detailed comparison with existing VFL approaches that address missing features.
- The results are convincing, especially given that they cover different levels of missing features. However, only two VFL methods are used for comparison. The authors are encouraged to include additional comparisons to better highlight the work's significance in the field of VFL.

Meta-review: The authors propose a new vertical federated learning (FL) method named LASER-VFL, designed for efficient training and inference of split neural network-based models in vertical FL settings with missing features. LASER-VFL addresses the challenge of missing labels by training a set of representation models (capable of performing well across different combinations of feature blocks) and a fusion model on each client (to handle any combination of representations).

The approach is interesting and compelling, and the numerical experiments verify its effectiveness. The paper is also complemented by a theoretical analysis under standard assumptions, such as smoothness and the PL condition. However, Theorems 1 and 2 provide standard convergence results for unbiased stochastic gradient descent, and their connection to the proposed algorithm remains underexplored. For instance, the presented results do not clarify how the number of clients $K$ affects the convergence guarantees, nor do they provide insights into the quality of the fusion or representation models. I believe the contribution would be significantly stronger if the theoretical analysis offered concrete guidelines and insights into the scalability of the method. Please add a discussion on this to the final version of the work.

Although the proposed method is not yet fully supported by the theory, it is intriguing and novel, and it has the potential to inspire further research within the vertical FL community. Thus I recommend acceptance.

Decision: Accept (Poster)