Towards Efficient and Secure Distributed Machine Learning Systems on Edge Computing
COM3 Level 2
MR21, COM3 02-61
Abstract:
Distributed machine learning (DML) has demonstrated tremendous potential in real-world applications, enabling the utilization of geographically distributed data and computational resources. With the growing development of edge computing, an increasing number of researchers are attempting to deploy machine learning solutions on resource-constrained edge devices for real-time data analysis. However, distributed machine learning systems on edge devices still face significant challenges. (1) To train a global model on the server, local clients on edge devices are required to perform intensive computation and communication, which is often impractical in real-world scenarios. (2) The distributed nature of such systems exposes them to adversarial manipulations at both the client and aggregation levels. (3) Resource-constrained edge devices are limited in computational capability, making it difficult to deploy and infer large-scale models efficiently.
Firstly, since federated learning (FL) is a general paradigm of distributed machine learning, we address its inherent communication and privacy-security challenges from the perspective of both data and model heterogeneity. To this end, we propose a novel taxonomy of existing OFL methods and also propose a novel one-shot federated learning (OFL) method termed FedLPA (Layer-wise Posterior Aggregation). FedLPA infers the posterior distribution of each layer within local models using layer-wise Laplace approximation and aggregates them to update the global parameters. Furthermore, balancing model accuracy against computational overhead, and to better enable personalized models tailored to individual devices, we propose MH-pFedHN. This framework quantifies clients possessing diverse architectures via customized embedding vectors, subsequently generating client-specific model parameters through a server-side hypernetwork.
Secondly, we observe that current FL systems suffer from bi-directional untrustworthiness: the server is vulnerable to Byzantine attacks from malicious clients, while clients are concerned about privacy breaches from a curious server. To address this, we propose the first Collusion-resilient Byzantine-Robust Federated Learning framework, CBRFL, which counteracts Byzantine attacks while preserving client data privacy. CBRFL leverages multi-party fully homomorphic encryption (FHE), allowing each client to protect its data privacy using a private key—even when Byzantine clients collude with the server. Moreover, it achieves Byzantine-robust aggregation directly in ciphertext via a secure cosine similarity mechanism, which measures distances and identifies malicious gradients through Byzantine consensus.
Thirdly, we explore deploying the popular Vision Transformer (ViT) and Vision Mamba (ViM) on real-world, resource-constrained edge devices. We design ED-ViT and ED-ViM, novel frameworks capable of efficiently splitting and executing complex Transformer and Mamba across multiple edge devices. Both approaches partition the large-scale models into several sub-models using methods tailored to their respective architectures, with each sub-model responsible for processing a specific subset of data classes. To further reduce computational overhead and inference latency, we introduce a class-wise pruning technique that reduces the size of each sub-model while maintaining accuracy.
Collectively, our studies enhance both the efficiency and security of distributed machine learning systems and demonstrate strong potential for deployment in edge computing environments.