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.
翻译:同态加密(HE)是隐私保护机器学习的重要框架,支持直接在加密数据上进行推理。然而,在HE中评估Transformer架构的核心组件——softmax函数——仍然极具挑战性,这源于其多元结构、指数函数引起的大动态范围以及归一化过程中对精确除法的需求。本文提出MGF-softmax,一种基于矩生成函数(MGF)的新型softmax重构方法,该方法将softmax分母替换为其基于矩的对应形式。该重构在保持softmax关键特性的同时,显著降低了乘法深度,并且随着输入标记数量的增加渐近收敛于精确softmax。在Vision Transformer和大型语言模型上的大量实验表明,MGF-softmax在加密推理中提供了高效且精确的softmax近似。特别地,该方法在实现接近高深度精确方法推理精度的同时,通过降低乘法深度显著减少了计算成本。