Efficient Softmax Reformulation for Homomorphic Encryption via Moment Generating Function

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.

💡 Research Summary

The paper addresses a fundamental bottleneck in homomorphic encryption (HE)–based inference for transformer models: the evaluation of the softmax function. Softmax requires exponentiation of every input token, a shared normalization term (the sum of exponentials), and a subtraction of the maximum value for numerical stability. In the CKKS scheme, only addition and multiplication are natively supported, so exponentials and divisions must be approximated by polynomials, which dramatically increases multiplicative depth, consumes the limited level budget, and often forces costly bootstrapping. Existing approaches fall into two categories. “Softmax approximation” methods directly approximate the exponential and reciprocal with high‑degree polynomials, preserving the original softmax semantics but incurring high depth. “Softmax replacement” methods substitute a simpler surrogate (e.g., ReLU, Gaussian kernels) that avoids expensive operations but typically degrades accuracy, especially on large‑scale vision and language tasks.

The authors propose MGF‑softmax, a reformulation that eliminates both the max‑subtraction and the division. They reinterpret the denominator Σ exp(x_i) as n · M_X(1), where M_X(t)=E