Technical Summary: PowerSoftmax: Towards Secure LLM Inference over Encrypted Data

Problem Statement

Modern privacy-preserving Large Language Models (LLMs) often rely on Homomorphic Encryption (HE), specifically the CKKS scheme, which allows computations on encrypted data without decryption. However, CKKS supports only polynomial operations. Standard Transformer architectures rely heavily on non-polynomial components, most notably the Softmax function in self-attention, which involves exponentials, division, and logarithmic operations (e.g., for numerical stability via the log-sum-exp trick).

Existing approaches to bridge this gap face significant limitations:

1. Post-Training Approximation (PTA): Directly approximating non-polynomial functions in pre-trained models using high-degree polynomials. This often results in inefficient inference due to high multiplicative depth and accumulated noise.
2. Approximation-Aware Training (AAT): Replacing non-polynomial primitives with easier-to-approximate alternatives (e.g., pointwise attention or scaled ReLU) before training. While effective for smaller models, these methods have struggled with stability and scalability, failing to support LLMs with billions of parameters or deep architectures (e.g., 32 layers).

The core challenge is designing a Transformer architecture that is inherently compatible with HE constraints (low-degree polynomials, stable numerical ranges) while maintaining the expressive power, reasoning capabilities, and in-context learning (ICL) of standard Softmax-based Transformers at the billion-parameter scale.

Methodology

The authors propose PowerSoftmax, a novel, HE-friendly variant of self-attention designed to replicate the core principles of Softmax while remaining amenable to polynomial approximation. The methodology involves three distinct stages:

1. HE-Friendly Attention Design

The authors distill the essential properties of Softmax: normalization to $[0, 1]$, amplification of score differences, and order preservation. They replace the exponential function $\exp(x)$ in Softmax with a power function $x^p$ (where $p$ is a positive even integer).
The proposed attention mechanism is defined as:
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
This formulation avoids exponentials and division by zero issues inherent in direct polynomial approximations of Softmax.

2. Training Stability and Numerical Robustness

To address training instability caused by the power function (where inputs $>1$ explode and inputs $<1$ vanish), the authors introduce two key modifications:

• Stable Variant: Inspired by the log-sum-exp trick, they normalize inputs by a constant $c = \|x\|_\infty + \delta$ before applying the power function. This ensures inputs remain within a manageable range $(0, 1)$, preventing overflow/underflow and preserving floating-point precision.
• $\epsilon$-Bounded Division: To handle the division operation, which is non-polynomial and sensitive near zero, they modify the denominator to $\epsilon + \sum x_i^p$. This creates a $1/\epsilon^2$-Lipschitz continuous function, making the subsequent polynomial approximation of the division significantly more tractable and reducing the required multiplication depth.

3. Length-Agnostic Approximation

Standard division approximation (e.g., Goldschmidt algorithm) struggles with the unbounded sum of attention scores, which scales linearly with sequence length $L$. The authors propose a Length-Agnostic PowerSoftmax:
$$\text{Length-Agnostic}(x)_j = \frac{x_j^p}{L \cdot \text{Mean}(x^p)}$$
Since $L$ (sequence length) is public, the term $1/L$ can be computed directly without approximation. The model approximates the division of the mean rather than the sum, ensuring the approximation domain remains bounded regardless of sequence length.

4. Construction Pipeline (Algorithm 1)

The full pipeline for creating a polynomial LLM involves:

1. Architectural Modification: Replacing Softmax with the Stable PowerSoftmax variant and training from scratch (or via continual pre-training) with range-minimization objectives.
2. Range Minimization: Applying a specialized loss function to constrain the input ranges of non-polynomial layers (PowerSoftmax, GELU, LayerNorm) to ensure low-degree polynomial approximations are accurate.
3. Polynomial Replacement: Post-training, replacing the remaining non-polynomial operations (division, inverse square root) with low-degree polynomial approximations (e.g., Goldschmidt method) to yield a fully polynomial model.

Key Contributions

1. PowerSoftmax Mechanism: A new self-attention variant that replaces exponential scaling with polynomial scaling, specifically tailored for HE environments while maintaining the core design principles of attention.
2. Scalable Polynomial LLMs: The development of the first polynomial LLMs exceeding one billion parameters (specifically 1.4B parameters with 32 layers). These models demonstrate reasoning and in-context learning (ICL) capabilities comparable to standard transformers of the same size.
3. Stable Training and Inference Strategies: Introduction of numerically stable training variants and length-agnostic inference strategies that overcome the instability and scalability issues of previous polynomial Transformer approaches.
4. Comprehensive Evaluation: A detailed latency breakdown for encrypted computations and ablation studies validating the design choices, including the impact of the power parameter $p$ and the $\epsilon$-bounded division.

Experimental Results

The authors evaluated their method on causal (GPT/Pythia) and bidirectional (RoBERTa) architectures:

• Performance: The 1.4B parameter polynomial model achieved performance comparable to the original Softmax-based Pythia models on zero-shot and few-shot benchmarks (e.g., Lambada, PIQA, WinoGrande, ARC). Notably, the model demonstrated competitive performance on reasoning benchmarks like ARC-Challenge, a capability previously unobserved in polynomial LLMs.
• Bidirectional Models: A polynomial RoBERTa-Base variant showed only a ~1% degradation compared to the original Softmax-based RoBERTa on GLUE tasks (SST-2, QNLI, MNLI), outperforming previous polynomial approaches like Nexus.
• HE Inference: Experiments on the Pythia 70M model using the HElayers framework confirmed that encrypted inference yields results nearly identical to plaintext inference (maximal MSE loss of 0.005).
• Efficiency: The proposed method significantly reduces multiplication depth compared to state-of-the-art alternatives (Nexus, Cho et al.). For a 128-token sequence, PowerSoftmax reduced latency by a factor of ~9.7 compared to Taylor/Goldschmidt approaches and ~20 compared to square-strategy approaches.

Significance and Claims

The paper claims a breakthrough in the field of privacy-preserving machine learning by demonstrating that fully polynomial LLMs can scale to the billion-parameter regime without sacrificing reasoning or in-context learning capabilities.

• Scalability: Unlike previous works limited to small models or simple classification tasks, this approach successfully trains deep (32-layer) transformers with 1.4 billion parameters.
• Non-Interactive Security: By utilizing fully homomorphic encryption (FHE) with a fully polynomial model, the solution enables non-interactive inference, avoiding the communication overhead and potential vulnerabilities associated with multi-party computation (MPC) or interactive protocols.
• Practicality: The detailed latency analysis and the successful deployment of a 1.4B parameter model suggest that secure, private LLM inference is becoming practically feasible, paving the way for applications in data-sensitive sectors like healthcare and finance.

The authors acknowledge that while the models perform well, a full evaluation of auto-regressive generative abilities in sequential decoding over encrypted data remains a direction for future work.