HawkEye: Statically and Accurately Profiling the Communication Cost of Models in Multi-party Learning

TL;DR

HawkEye introduces a static, prefix-structure based framework to accurately profile model communication costs in multi-party learning without dynamic profiling. It pairs a static profiling method with an Autograd-compatible PyTorch-like library, including an optimized broadcast operator, to enable end-to-end profiling of forward and backward passes. Across multiple MPL frameworks and models, HawkEye achieves close agreement with dynamic profiling while dramatically reducing profiling time, and it supports practical case studies in security models, optimizer choices, and computational-graph optimization. This work promises faster, framework-agnostic MPC-friendly model design, reducing the iteration burden for researchers and practitioners designing privacy-preserving ML systems.

Abstract

Multi-party computation (MPC) based machine learning, referred to as multi-party learning (MPL), has become an important technology for utilizing data from multiple parties with privacy preservation. In recent years, in order to apply MPL in more practical scenarios, various MPC-friendly models have been proposedto reduce the extraordinary communication overhead of MPL. Within the optimization of MPC-friendly models, a critical element to tackle the challenge is profiling the communication cost of models. However, the current solutions mainly depend on manually establishing the profiles to identify communication bottlenecks of models, often involving burdensome human efforts in a monotonous procedure. In this paper, we propose HawkEye, a static model communication cost profiling framework, which enables model designers to get the accurate communication cost of models in MPL frameworks without dynamically running the secure model training or inference processes on a specific MPL framework. Firstly, to profile the communication cost of models with complex structures, we propose a static communication cost profiling method based on a prefix structure that records the function calling chain during the static analysis. Secondly, HawkEye employs an automatic differentiation library to assist model designers in profiling the communication cost of models in PyTorch. Finally, we compare the static profiling results of HawkEye against the profiling results obtained through dynamically running secure model training and inference processes on five popular MPL frameworks, CryptFlow2, CrypTen, Delphi, Cheetah, and SecretFlow-SEMI2K. The experimental results show that HawkEye can accurately profile the model communication cost without dynamic profiling.

Figures (5)

• Figure 1: The architecture of the HawkEye Autograd library.
• Figure 2: The strawman implementation of broadcast multiplication and optimized broadcast multiplication.
• Figure 3: The online communication round profiling results of four secure CNN model inference processes outputted by CryptFlow2rathee2020cryptflow2 and HawkEye.
• Figure 4: The online communication round profiling results of secure BERT$_{\textsc{base}}$ model inference process outputted by CrypTencrypten2020 and HawkEye.
• Figure 5: The proportion of each operator's online communication size to the total online communication size on four MPL frameworks with different security models.