A Communication-Efficient Vertical Federated Learning via Neuron Selection
Cheng Wu Yang, Yiqiang Cheng, Xiaodong Yang, Yuting He · 2023
Vertical federated learning is a privacy-preserving machine learning approach in which different parties hold different features of a dataset. However, the communication cost during the training of vertical federated neural networks is high, due to the information exchange at each step. Prior researches has employed strategies such as multiple local iterations and compression to minimize communication cost, but these approaches often overlook the redundancy present in neural network structures. In this paper, we propose a method for selecting output neurons in local networks to decrease communication cost. First, we define an importance score based on the first-order Taylor expansion of loss increment for each neuron. To avoid the overhead of computing the total gradient in calculating first-order importance, we associate a trainable importance parameter with each neuron, ensuring that the order of these parameters reflects the order of the first-order importance score. These parameters are trained jointly with neural network models, introducing minimal additional communication and computation overhead. Subsequently, we employ a two-step selection procedure to remove unimportant neurons in each local network according to the importance order during the early stages of training, thereby reducing communication cost. Our method effectively decreases communication cost in both training and inference phases. Experimental results on the MNIST and CIFAR10 datasets demonstrate that our approach reduces communication cost to 77.78% and 32.29% respectively, while incurring only a negligible drop in model performance.