Efficient Softmax Reformulation for Homomorphic Encryption via Moment Generating Function
Hanjun Park, Byeong-Seo Min, Jiheon Woo, Min-Wook Jeong, Jongho Shin, Yongwoo Lee, Young-Sik Kim, Yongjune Kim
TL;DR
This work tackles the challenge of evaluating softmax under homomorphic encryption by introducing MGF-softmax, a reformulation based on the moment generating function that replaces the traditional softmax denominator with a moment-based term. By expressing the denominator as $n M_X(1)$ and leveraging the cumulant generating function, the method eliminates the need for max-subtraction and division, substantially reducing multiplicative depth and the associated bootstrapping burden. The authors provide theoretical properties, including shift invariance and a data-adaptive domain scaling, and derive an error bound showing the relative error vanishes as the input size grows. Empirically, MGF-softmax achieves near-plaintext accuracy on large language models and Vision Transformers while requiring much lower depth (7–10) and runtime, outperforming softmax approximation baselines and softmax replacements on challenging benchmarks. This approach offers a scalable, privacy-preserving pathway for efficient encrypted inference in transformer-based architectures.
Abstract
Homomorphic encryption (HE) is a prominent framework for privacy-preserving machine learning, enabling inference directly on encrypted data. However, evaluating softmax, a core component of transformer architectures, remains particularly challenging in HE due to its multivariate structure, the large dynamic range induced by exponential functions, and the need for accurate division during normalization. In this paper, we propose MGF-softmax, a novel softmax reformulation based on the moment generating function (MGF) that replaces the softmax denominator with its moment-based counterpart. This reformulation substantially reduces multiplicative depth while preserving key properties of softmax and asymptotically converging to the exact softmax as the number of input tokens increases. Extensive experiments on Vision Transformers and large language models show that MGF-softmax provides an efficient and accurate approximation of softmax in encrypted inference. In particular, it achieves inference accuracy close to that of high-depth exact methods, while requiring substantially lower computational cost through reduced multiplicative depth.