Vertical Federated Learning with Missing Features During Training and Inference

Thank you for your comments.

Re: weaknesses in the experiments section

We have significantly strengthened the experiments section. In particular, we have made the following improvements:

• We have added some details on the experiments in the supplementary material (given the space contraints) and we will be making our code publicly available upon acceptance, thus ensuring reproducibility.

• We have added the naive combinatorial method described in Section 2 to Table 1 and we have added new experiments comparing this combinatorial approach to our method and to the other baselines for a range of numbers of clients (see Section 5 of the revised manuscript). These experiments show that, even for a small number of clients, the combinatorial approach is significantly slower than our method, as seen in the following table, where we present the runtime of the combinatorial approach divided by the runtime of our LASER-VFL method for CIFAR-10:

  Number of Clients (K)	2	3	4	5	6	7
  Runtime Ratio	1.76	3.19	5.96	10.91	19.41	35.80

  This demonstrates that the combinatorial approach may become impractical even for a modest number of clients. Moreover, the combinatorial approach is not optimal in terms of performance: (1) in the presence of missing features during training, for each batch, it can only train the predictors using a set of features that is a subset of the set of observed features, significantly decreasing the number of training samples and thus achieving worse generalization; and (2) even when there is no missing training data, our LASER-VFL method still outperforms the combinatorial approach due to the dropout-like regularization effect introduced by task sampling.

• We have added a new baseline to the experiments in Section 5: PlugVFL [1]. PlugVFL is the most highly cited of the three works enabling inference in VFL from varying sets of feature blocks. After recently obtaining the code from the authors, we were able to include PlugVFL as a baseline in our experiments. While PlugVFL outperforms most baselines, it does so inconsistently and our LASER-VFL method ultimately outperforms PlugVFL, as shown in the revised Table 1 and in the experiments with a varying number of clients.

  Note that we have made the necessary extensions to PlugVFL to support missing training data and more than two clients, which were not considered in the approach described in the original paper. See Appendix D of the revised manuscript for details.

  We would also like to explain that the reason for not including more baselines in our original submission was the difficulty of reproducing results from the related methods, due to the unavailability of their source code. In this spirit, we are happy to share our code and make it open-source upon acceptance, to enable others to compare their methods with ours. In the meantime, we have made our code publicly available in an anonymos GitHub repository: https://anonymous.4open.science/r/robust-vfl-7A1E/README.md

Re: heavy notation

We have worked to simplify the notation in the revised manuscript and improve the readability of the section you mentioned, namely by reducing the use of notation. (See the "Our optimization method" paragraph in Section 3.2). However, given the topic of the paper, further simplification is challenging, as the use of indices for sets of models, feature blocks, and stochastic approximations is often essential for clarity and precision.

Re: theoretical analysis

Our key theoretical contribution lies in framing the problem of missing features in vertical federated learning as a finite-sum problem. This formulation allows us to leverage analytical tools developed for finite-sum problems to establish convergence guarantees for our method, thus confirming its validity and supporting its design under uniformly missing blocks of features—a desirable property. This stands in contrast to previous works on vertical federated learning enabling a varying number of clients to collaborate during inference, which do not provide convergence guarantees [1, 2].

Re: missing reference

Thank you for bringing the missing reference to our attention.

References

[1] Jingwei Sun, Zhixu Du, Anna Dai, Saleh Baghersalimi, Alireza Amirshahi, Qilin Zheng, David Atienza, & Yiran Chen. (2024). PlugVFL: Robust and IP-Protecting Vertical Federated Learning against Unexpected Quitting of Parties.

[2] Gao, Dashan, et al. "Complementary Knowledge Distillation for Robust and Privacy-Preserving Model Serving in Vertical Federated Learning." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. No. 18. 2024.

Dear Reviewer GeBH,

Thank you again for your valuable comments and feedback. We hope that our response has addressed your concerns and clarified your questions. As the discussion period comes to a close, please feel free to reach out if you have any additional questions.

Best regards,
The Authors

Dear Reviewer GeBH,

Thank you again for your time and effort in reviewing our paper. Since tomorrow is the final day for reviewers to post comments, we would appreciate your feedback on whether our responses to your comments and revisions to the manuscript are satisfactory. If you have any further questions, please let us know and we will be happy to provide additional clarifications. If not, we would be grateful if you could consider increasing your score to reflect the fact that we have addressed your concerns.

Best regards,
The Authors

Thank you for your comments.

Re: a more detailed comparison with existing approaches

The reason why we did not perform a more detailed comparison with individual existing VFL methods addressing either missing features during training or during inference is that these differ very significantly from our method (and often from each other); they have little more in common with our work than their objective. In fact, even in terms of the task being tackled, we believe our method is the first one addressing both of these challenges simultaneously (missing features during training and during inference). To the best of our knowledge, the whole idea of modeling predictions from different subsets of features as separate tasks, each with its own loss term, and then sampling from these tasks to make training feasible, is novel.

Re: additional baselines

We have added a new, strong baseline to Table 1: the combinatorial approach described in Section 2. These new results show that our method outperforms even the combinatorial approach. We attributed this to (1) the presence of missing features during training, meaning that in the combinatorial approach, for each batch, we can only train the predictors corresponding to a set of features that is a subset of the set of observed features, significantly decreasing the number of training samples and thus achieving worse generalization; and (2) even in the absence of missing training data, our LASER-VFL method still outperforms the combinatorial approach due to the dropout-like regularization effect introduced by task sampling.

Re: why $K$ fusion models

The main reason for having $K$ fusion models, one at each client, is to ensure that any client can continue performing the prediction task it was trained for during collaborative VFL (whether or not it is the same task as other clients) regardless of whether any other client leaves the federation. If a single fusion model were trained at a given client and that client were to leave the federation during inference, this would render the models of all the other clients unavailable during inference. Therefore, having a local fusion model gives more insurance of continued model availability during inference to each of the clients. Nevertheless, as you mentioned, this also allows each client to have a specialized fusion model for its own task, which can be very convenient.

Re: on the use of a shared seed

The use of a shared seed allows the different clients to randomly sample the same indices for each mini-batch without having to communicate the batch indices at each communication round; this is common in VFL [1, 2]. This matching mini-batch is important because otherwise we would train the model on invalid samples. To illustrate this, consider a dataset with only two samples $\lbrace \mathbf{x}_1, \mathbf{x}_2 \rbrace$ with each sample $\mathbf{x}_i$ further partitioned into two blocks of features $\mathbf{x}\_i=(\mathbf{x}\_{i1}, \mathbf{x}\_{i2})$. If a client holding $\lbrace\mathbf{x}\_{i1}:\forall i\rbrace$ samples a sample (mini-batch of size one) without ensuring it matches the sample of the another client holding $\lbrace\mathbf{x}\_{i2}:\forall i\rbrace$, we may end up computing the loss corresponding to a prediction based on $(\mathbf{x}\_{i1}, \mathbf{x}\_{j2})$, where $i\not=j$, which is not a valid sample.

Re: domain shift

As mentioned above, while having a fusion model at each client allows each client to tackle a different task, this is not the focus of this work. For this reason, we did not explore such a domain shift across clients, leaving it as future work instead.

Re: requirements on the model architecture

Our framework is very flexible regarding the types of models it can accommodate. In the experiments, we used a variety of representation models, including multilayer perceptrons, an LSTM, and ResNets. For the fusion models, we employed a simple linear transformation (a single linear layer neural network), as it was sufficient to achieve good performance; however, other fusion models can also be considered, as our LASER-VFL method does not restrict the fusion model to consist of a single layer. We have added more details about our experiments in the supplementary material of the revised manuscript.

References

[1] Castiglia, Timothy, Shiqiang Wang, and Stacy Patterson. "Flexible vertical federated learning with heterogeneous parties." IEEE Transactions on Neural Networks and Learning Systems (2023).

[2] Valdeira, Pedro, et al. "Communication-efficient Vertical Federated Learning via Compressed Error Feedback." arXiv preprint arXiv:2406.14420 (2024).

Dear Reviewer Cmzg,

Thank you again for your time and effort in reviewing our paper. Since tomorrow is the final day for reviewers to post comments, we would appreciate your feedback on whether our responses to your comments and revisions to the manuscript are satisfactory. If you have any further questions, please let us know and we will be happy to provide additional clarifications. If not, we would be grateful if you could consider increasing your score to reflect the fact that we have addressed your concerns.

Best regards,
The Authors

Re: comparison with the combinatorial approach

Regarding your comment that "it is difficult to argue that these models would be anywhere near the naive solution of actually directly training an exponential number of models", as mentioned above, we have added new experimental results in which we compare the combinatorial approach to our method and to the other baselines (see Section 5 of the revised manuscript).

These new results show that our method in fact outperforms the combinatorial approach. We attributed this to (1) the presence of missing features during training, meaning that in the combinatorial approach, for each batch, we can only train the predictors corresponding to a set of features that is a subset of the set of observed features, significantly decreasing the number of training samples and thus achieving worse generalization; and (2) even in the absence of missing training data, LASER-VFL still outperforms the combinatorial approach due to the dropout-like regularization effect introduced by task sampling.

Re: notation complexity

We have worked to simplify the notation in the revised manuscript (see the "Our optimization method" paragraph in Section 3.2). However, given the topic of the paper, further simplification is challenging, as the use of indices for sets of models, feature blocks, and stochastic approximations is often essential for clarity and precision.

Re: client participation

In this work, we did not consider communication faults nor clients leaving during training. In contrast to [2], which focuses on dynamic networks in cross-device VFL, our work focuses on the cross-silo setting, which is the most common setting in VFL [3], and tackle the fundamentally different problem of missing features in VFL, where the missing blocks of features are entirely absent from training, at all iterations.

In cross-silo settings, there are fewer clients than in cross-device settings, typically with more reliable communication links. For example, as mentioned in the paper, an online retail company and a social media platform may hold different features on shared users. In such applications, communication failures are less problematic than in cross-device settings.

Regarding your statement that, in our LASER-VFL method, all clients must be available during training, even if they do not participate in a round, we believe this is not the case. A client that does not observe its feature block in a batch does not need to join that communication round.

Re: dealing with missing features in the standard VFL baseline

The baseline standard VFL method does waste samples with missing features, making a random prediction instead. However, we did include other VFL-based baselines in the original submission, such as the ensemble method, which is still outperformed by our method, as is our newly added baseline: the combinatorial approach. While this is a strong baseline, it still performs worse than our LASER-VFL method, as mentioned above.

Re: further baselines

We have added a new baseline to the experiments in Section 5: PlugVFL [4]. PlugVFL is the most highly cited of the three works enabling inference in VFL from varying sets of feature blocks. After recently obtaining the code from the authors, we were able to include PlugVFL as a baseline in our experiments. While PlugVFL outperforms most baselines, it does so inconsistently and our LASER-VFL method ultimately outperforms PlugVFL, as shown in the revised Table 1 and in the experiments with a varying number of clients.

Note that we have made the necessary extensions to PlugVFL to support missing training data and more than two clients, which were not considered in the approach described in the original paper. See Appendix D of the revised manuscript for details.

We would also like to explain that the reason for not including more baselines in our original submission was the difficulty of reproducing results from the related methods, due to the unavailability of their source code. In this spirit, we are happy to share our code and make it open-source upon acceptance, to enable others to compare their methods with ours. In the meantime, we have made our code publicly available in an anonymos GitHub repository: https://anonymous.4open.science/r/robust-vfl-7A1E/README.md

References

[3] Liu, Yang, et al. "Vertical federated learning: Concepts, advances, and challenges." IEEE Transactions on Knowledge and Data Engineering (2024).

[4] Jingwei Sun, Zhixu Du, Anna Dai, Saleh Baghersalimi, Alireza Amirshahi, Qilin Zheng, David Atienza, & Yiran Chen. (2024). PlugVFL: Robust and IP-Protecting Vertical Federated Learning against Unexpected Quitting of Parties.

Thank you for your comments.

Re: combinatorial approach for a small number of clients

• We have added new experiments comparing the combinatorial approach to our method and to the other baselines for a range of numbers of clients (see Section 5 of the revised manuscript). These experiments show that, even for a small number of clients, the combinatorial approach is significantly slower than our method, as seen in the following table, where we present the runtime of the combinatorial approach divided by the runtime of our LASER-VFL method for CIFAR-10:

  Number of Clients (K)	2	3	4	5	6	7
  Runtime Ratio	1.76	3.19	5.96	10.91	19.41	35.80

  This demonstrates that the combinatorial approach may become impractical even for a modest number of clients. Further, as mentioned below in our response, the combinatorial approach is in fact outperformed by our LASER-VFL method.

Re: MCAR

Thank you for raising this point. We have made it clearer in the revised manuscript that we do assume that the blocks of features are Missing Completely at Random (MCAR). As pointed out in a recent review paper on handling missing data [1] (in general machine learning; not specific to federated learning), the MCAR assumption is the most common type of assumption when dealing with missing data (compared to Missing At Random (MAR) and Missing Not At Random (MNAR)). This being said, it would be interesting to extend our convergence guarantees to the more challenging MAR setting, which is often more realistic.

Our experiments in the original submission also assume MCAR and similar probabilities of missing for the different blocks of features. However, to address your concern, we are currently working on experiments with different missingness patterns to include in the revised manuscript.

Re: VFL and GNNs

Regarding your comments about the connection to GNNs and the use of parameter sharing, while we are not experts in GNNs (so please correct us if we interpreted your comment incorrectly), it seems to us that you are saying that we can interpret each client as a node in a complete graph and interpret the whole model as a GNN. If this is the case, we believe it may be possible to see our model as a single-layer GNN with average aggregation at each client’s own fusion predictor, yes. If our understanding of your comments is correct, a small correction on your comment about “sharing model parameters across nodes” would be that we do not share model parameters across nodes, but rather across tasks. Nevertheless, we do not claim that the idea of sharing parameters is novel. We simply claim that, by sharing model parameters across tasks and sampling a subset of tasks from the full set at each iteration, we obtain a good, novel solution to the problem of missing features in VFL. This solution is indeed simple, but we see that as a strength, not a weakness.

Re: missing reference

Thank you for bringing the work by Ganguli et al [2] to our attention, as indeed it is worth including in our paper as a related work. This being said, while [2] focuses on dealing with communication failure during inference for a network of edge devices in the cross-device FL (settings with a large number of clients), our work focuses on missing features during both training and inference.

Note however that, while task sampling and dropout are related, they are not exactly the same. Our task-sampling mechanism allows the loss computed at each client to take into account multiple sets of tasks at each iteration (each task corresponding to prediction from a different subset of blocks of features), whereas dropout considers a single combination at each iteration. Further, our task-sampling mechanism takes into account the presence of missing training data (which is not considered in [2]); while in [2] the authors "simulate a dynamic graph using dropout", this is not equivalent to having missing blocks of features during training, as a missing block is missing throughout training, rather than simply being absent from a given (set of) iteration(s).

References

[1] Zhou, Youran, Sunil Aryal, and Mohamed Reda Bouadjenek. "Review for Handling Missing Data with special missing mechanism." arXiv preprint arXiv:2404.04905 (2024).

[2] Ganguli, Surojit, et al. "Fault-tolerant vertical federated learning on dynamic networks." arXiv preprint arXiv:2312.16638 (2023).

We have spent additional time understanding the paper by Ganguli et al. [2] and running new experiments. As a result, we have some further responses to share below.

Re: comparison with Ganguli et al. [2]

We would like to highlight that, to the best of our understanding, for the special case where all clients can always communicate with each other (a static, complete graph, which is the setup considered in our work), the method proposed by Ganguli et al. [2] matches the PlugVFL method proposed by Sun et al. [4], in the absence of the privacy mechanism therein. While [2] addresses the challenge of extending [4] to dynamic networks in the cross-device setting and thus has a different focus from that of our paper, our LASER-VFL method outperforms PlugVFL [4] in the special case of static networks in the cross-silo setting and, consequently, also outperforms [2] in this special case.

Re: experiments with nonuniform missing features

To address your question regarding how our method handles cases where some feature missing patterns are more common than others, we have added an experiment to the revised manuscript. See the "Nonuniform missing features" paragraph in Section 5 and Table 3 in Appendix E. In this experiment, we assume that each block $k$ has a probability of not being observed, $p_k$. Instead of setting all $\{p_k\}$ to the same value as in the previous experiments, we sample each $p_k$ independently from a beta distribution: $p_k \sim \text{Beta}(2.0, 2.0)$. As an example, for one of the seeds of our experiments with $K=4$ clients, we have $(p_1, p_2, p_3, p_4)=(0.43, 0.24, 0.46, 0.74)$. Our LASER-VFL method demonstrates consistent performance without any degradation and, as in the previous experiments, outperforms the baselines.

References

[2] Ganguli, Surojit, et al. "Fault-tolerant vertical federated learning on dynamic networks." arXiv preprint arXiv:2312.16638 (2023).

[4] Jingwei Sun, Zhixu Du, Anna Dai, Saleh Baghersalimi, Alireza Amirshahi, Qilin Zheng, David Atienza, & Yiran Chen. (2024). PlugVFL: Robust and IP-Protecting Vertical Federated Learning against Unexpected Quitting of Parties.

Dear Reviewer wVzL,

Thank you again for your valuable comments and feedback. We hope that our response has addressed your concerns and clarified your questions. As the discussion period comes to a close, please feel free to reach out if you have any additional questions.

Best regards,
The Authors

Dear Reviewer wVzL,

Thank you again for your time and effort in reviewing our paper. Since tomorrow is the final day for reviewers to post comments, we would appreciate your feedback on whether our responses to your comments and revisions to the manuscript are satisfactory. If you have any further questions, please let us know and we will be happy to provide additional clarifications. If not, we would be grateful if you could consider increasing your score to reflect the fact that we have addressed your concerns.

Best regards,
The Authors

Thank you for your comments.

Re: clarity when introducing Problem (1)

We have made the sentence leading up to Problem (1) clearer.

Re: missing feature blocks in Figure 1 (b)

Samples 4, 6, and 7 in Figure 1(b) (our method) are not used by client 1 for inference because they correspond to samples where client 1 does not observe its own feature block. In the example in the paper of an online retailer and a social media company sharing users, this corresponds to the case where an online retailer does not perform the prediction task for an individual who is not one of its users.

Re: mentioning prior art before the "Related work" paragraph

The purpose of the "Dealing with missing features" paragraph in the introduction, which, as you noted, discusses previous works only at a high level, is to highlight that while there have been attempts to address this problem in VFL, none have fully solved it. This highlights the interest of the community on the topic and, more importantly, allows us to identify an important gap in the existing literature. In contrast, the related works section provides a more comprehensive review of the literature, including more detail on the approaches that these works have employed to address the issue of missing features in VFL.

Re: claim of novelty

We reiterate that, to the best of our knowledge, our LASER-VFL method is the first method to enable both training and inference using any and all available blocks of features. In contrast, the other works mentioned in the paper address either missing features during training or missing features during inference, but not both. Our approach addresses both challenges simultaneously.

Re: citations formatting

The following quote is from the "Formatting Instructions for ICLR 2025 Conference Submissions": "When the authors or the publication are included in the sentence, the citation should not be in parenthesis" (emphasis added by us).

Re: combinatorial approach for a small number of clients

We have added new numerical experiments (see Section 5 of the revised manuscript) comparing the combinatorial approach with our method and the other baselines. These experiments show that, even with a small number of clients, the combinatorial approach is significantly slower than our method, as seen in the following table, where we present the runtime of the combinatorial approach divided by the runtime of LASER-VFL for CIFAR-10:

Moreover, the combinatorial approach is not optimal in terms of performance: (1) in the presence of missing features during training, for each batch, it can only train the predictors using a set of features that is a subset of the set of observed features, significantly decreasing the number of training samples and thus achieving worse generalization; and (2) even when there is no missing training data, LASER-VFL still outperforms the combinatorial approach due to the dropout-like regularization effect introduced by task sampling.

Re: notation complexity

We have worked to simplify the notation in the revised manuscript. However, given the topic of the paper, further simplification is challenging, as the use of indices for sets of models, feature blocks, and stochastic approximations is often essential for clarity and precision. The choice to use slightly different notation in (3), compared to (1) and (2), was made to keep (1) and (2) as simple as possible. This approach helps ease the reader into the new notation before introducing the set notation, which is only necessary in (3). While consistency is important, we viewed this as a trade-off and prioritized simplicity where feasible.

Thank you for your comments.

Re: notation complexity

We have worked to simplify the notation in the revised manuscript (see the "Our optimization method" paragraph in Section 3.2). However, given the topic of the paper, further simplification is challenging, as the use of indices for sets of models, feature blocks, and stochastic approximations is often essential for clarity and precision.

Re: workflow diagram

Thank you for your suggestion of the workflow diagram. We have added a diagram illustrating our LASER-VFL method in Appendix B of the revised manuscript (as space constraints prevented its inclusion in the main paper) and referenced it just before introducing the key idea on page 4.

Re: synthetic missing features

We agree that running experiments on real vertical FL data with missing features would be ideal. However, such data is typically privately held by companies and is not publicly accessible. Consequently, as is standard practice in this area [1, 2, 3], we have used publicly available datasets, synthetically partitioned the data across clients, and simulated missing feature patterns.

Re: hyperparameters in vertical FL

If your question regarding hyperparameters refers to our statement in the contributions bullet points that our LASER-VFL method is "hyperparameter-free", we would like to clarify that we were not implying the removal of any hyperparameters specific to vertical FL. Rather, we meant that, unlike other works addressing missing features in vertical FL [1, 2, 3], LASER-VFL does not introduce any additional hyperparameters. Please let us know if we have misunderstood your question.

Re: multiple local updates

The proposed algorithm does not consider multiple local updates. Extending LASER-VFL to accommodate multiple local updates would be an interesting direction for future work. However, it is important to note that the use of multiple local updates in vertical FL, while possible [4], is not as prevalent as in horizontal FL. The main reason for this is that performing multiple local updates requires all clients to have access to the labels used for computing the loss, as well as the fusion model.

To illustrate this, consider a simplified vertical FL model represented as a composition $g \circ f$, where $g$ is the fusion model and $f$ represents a local representation model. After the first local update of $f$, further local updates would require additional forward and backward passes over both $f$ and $g$ at the client, including the computation of the loss $\ell(y^n,g(f(x^n)))$. In contrast, by using a single update, LASER-VFL enables clients to use their (possibly different) labels locally and communicate only the derivative of the loss with respect to the features extracted by each representation model.

References

[1] Kang, Yan, Yang Liu, and Xinle Liang. "FedCVT: Semi-supervised vertical federated learning with cross-view training." ACM Transactions on Intelligent Systems and Technology (TIST) 13.4 (2022): 1-16.

[2] He, Yuanqin, et al. "A hybrid self-supervised learning framework for vertical federated learning." IEEE Transactions on Big Data (2024).

[3] Huang, Chung-ju, Leye Wang, and Xiao Han. "Vertical federated knowledge transfer via representation distillation for healthcare collaboration networks." Proceedings of the ACM Web Conference 2023. 2023.

[4] Liu, Yang, et al. "FedBCD: A communication-efficient collaborative learning framework for distributed features." IEEE Transactions on Signal Processing 70 (2022): 4277-4290.