Efficient and performant Transformer private inference with heterogeneous attention mechanisms

With the development of large-scale models, Transformer architectures have gained widespread adoption. However, privacy concerns become critical when model inference involves separate ownership of data and model parameters. Existing MPC-based methods for private inference suffer from significant overhead and high latency, where replacing traditional Softmax attention mechanisms with faster alternatives serves as a promising research direction. To achieve a better balance between Transformer model performance and inference speed, we explore the impact of attention mechanisms and attention heads on model performance. First, we found that the performance of the attention mechanism is closely related to the downstream task dataset, and the attention mechanism that is faster on specific datasets can actually achieve better model performance. Additionally, we discovered that for attention mechanisms that experience a performance decline, appropriately restoring the attention heads of the Softmax mechanism can significantly enhance performance. We further observed that the selection of key attention heads under different mechanisms is consistent, providing a basis for designing search strategies adapted to different scenarios. Based on these findings, we propose an MPC-friendly attention mechanism replacement method that enables Transformer private inference to be more efficient and performant. This method incorporates two strategies for selecting and replacing attention mechanisms to address diverse scenario requirements, and the resulting heterogeneous attention mechanism significantly improves the speed of private inference while maximizing model performance. With experiments on different downstream tasks, we demonstrated that our method improves average model performance by 1.94 % compared to standard pre-training models, with an inference speed increase of approximately 3 × . Compared to the state-of-the-art methods, our approach enhances model performance by 1.01–8.03 %, with faster inference speeds. Additionally, when evaluated using comprehensive metrics, our method shows improvements of 4.15 × to 8.97 × compared to other approaches.