Technical Summary: CGF-Softmax

Problem Statement

Homomorphic Encryption (HE) enables privacy-preserving machine learning by allowing computation on encrypted data without decryption. However, implementing non-polynomial operations, particularly the softmax function central to Transformer architectures, remains a significant bottleneck in strictly non-interactive HE settings.

The standard softmax formulation, $\text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}$, presents three primary challenges in HE:

1. Numerical Instability: The exponential function causes rapid growth, leading to overflow. Standard plaintext inference mitigates this by subtracting the maximum input value ($x_{\max}$), but the maximum operation requires complex comparison logic that is incompatible with HE and necessitates expensive polynomial approximations.
2. Division Operations: The normalization step requires division, which HE schemes do not natively support. Prior work relies on iterative polynomial approximations (e.g., Goldschmidt's algorithm or Chebyshev polynomials) for reciprocal functions, drastically increasing multiplicative depth.
3. Computational Cost: High-degree polynomial approximations required for both the exponential and division functions consume the limited level budget of leveled HE schemes (like CKKS), often necessitating costly bootstrapping operations to refresh ciphertexts.

Existing approaches fall into two categories: softmax approximation (high accuracy but high depth due to division/max) and softmax replacement (low depth but significant accuracy degradation, especially on large-scale models and complex benchmarks).

Methodology: CGF-Softmax

The authors propose CGF-Softmax, a reformulation of the softmax function that eliminates both the homomorphic division and the explicit maximum subtraction operations. The core innovation lies in interpreting the softmax denominator through the Cumulant Generating Function (CGF).

Theoretical Reformulation

Given an input vector $x$, the authors interpret the denominator $\sum \exp(x_i)$ as a scaled sample mean of exponential terms. By substituting the sample mean with the ensemble mean (expectation), the denominator is replaced by $n M_X(1)$, where $M_X(t)$ is the Moment Generating Function (MGF).

Using the CGF, defined as $K_X(t) = \ln M_X(t)$, the reformulated softmax is expressed as:
$$\text{softmax}_{\text{CGF}}(x)_i = \exp(x_i - K_X(1) - \ln n)$$

This formulation offers three critical advantages:

1. Elimination of Max Operation: Instead of subtracting the maximum value, the input is shifted by $K_X(1) + \ln n$. Since $K_X(1)$ can be approximated using low-order cumulants (mean $\mu$ and variance $\sigma^2$), this shift is data-adaptive and computable via basic HE operations (addition and multiplication) without comparisons.
2. Elimination of Division: The normalization term is absorbed into the exponent. The result is a single exponential function, removing the need for homomorphic division entirely.
3. Domain Scaling: The method naturally supports domain scaling. By scaling the exponent by $1/2^k$, the exponential function is evaluated over a smaller interval, and the original scale is recovered by squaring $k$ times. This avoids the iterative "normalize-and-square" process required by standard softmax, which involves repeated divisions.

Practical Instantiation

For practical HE evaluation, the authors adopt a second-order cumulant approximation:
$$K_X(1) \approx \mu + \frac{\sigma^2}{2}$$
Higher-order terms ($\kappa_j/j!$ for $j \ge 3$) are omitted. The authors argue that the factorial denominator intrinsically suppresses these terms, and empirical fine-tuning further reduces the magnitude of higher-order cumulants, making the approximation accurate enough for inference.

Key Contributions

1. CGF-Softmax Formulation: A novel reformulation that removes the need for homomorphic division and maximum subtraction, substantially reducing multiplicative depth.
2. Theoretical Analysis: The paper establishes key properties of CGF-softmax, including:
   • Positivity: Strictly positive outputs.
   • Order and Pairwise-Ratio Preservation: It preserves the relative ordering of logits and the exact pairwise output ratios ($\exp(x_i - x_j)$) of standard softmax.
   • Shift Invariance: Like standard softmax, it is invariant to constant shifts.
   • Asymptotic Error Bound: Using the Berry–Esseen theorem, the authors prove that the relative error converges to zero in probability as the input dimension $n$ increases.
3. Experimental Validation: Extensive evaluation on Vision Transformers (ViT/DeiT) and Large Language Models (LLaMA-3.2-1B) demonstrating high accuracy with reduced computational cost.

Experimental Results

The authors evaluated CGF-softmax on ImageNet-1k (Vision Transformers) and multiple NLP benchmarks (Clinc150, Banking77, SST-2) using the LLaMA-3.2-1B model.

• Multiplicative Depth: CGF-softmax achieves a multiplicative depth of 7–10 for large-scale models (e.g., LLaMA-3.2-1B, ViT-Base). In contrast, prior approximation methods (e.g., Cho et al., EncryptedLLM) require depths of 14–52 for smaller models or tasks.
• Runtime: On a 256×256 input matrix, CGF-softmax reduced total latency from 107.65s (Cho et al.) to 3.25s, primarily by eliminating bootstrapping and reducing ciphertext multiplications (CMult).
• Accuracy:
  • CGF-softmax maintains inference accuracy within <1% of the exact plaintext softmax baseline across all benchmarks.
  • It significantly outperforms softmax replacement methods (e.g., BPMax), which suffer from substantial accuracy drops on complex tasks.
  • It achieves near-plaintext accuracy while requiring 4–6× less depth than high-precision approximation baselines.
• Noise Behavior: By eliminating division, CGF-softmax reduces accumulated noise. At depth 7, it achieves a total noise of $3.05 \times 10^{-10}$, which is lower than standard softmax formulations even at greater depths.

Significance and Claims

The paper claims that CGF-softmax offers a scalable and efficient solution for secure inference in Transformer-based models. By reconciling high accuracy with substantially reduced multiplicative depth, it addresses the primary bottleneck in HE-based PPML.

The authors emphasize that while softmax replacement methods may suffice for small models on simple tasks, CGF-softmax is necessary for large-scale architectures and complex benchmarks where existing replacements fail to maintain accuracy. The method enables fully outsourced, non-interactive inference on large models (like LLaMA-3.2-1B) with computational costs that are previously unattainable for high-accuracy HE inference.

The paper modestly notes that for relatively small models (e.g., BERT) on simple tasks, existing low-depth replacement methods might still be sufficient, but CGF-softmax provides a superior trade-off for more demanding scenarios.