Power-Softmax: Towards Secure LLM Inference over Encrypted Data

TL;DR

This paper introduces a new polynomial-friendly self-attention mechanism enabling secure inference of large language models over encrypted data, achieving comparable reasoning abilities to standard models.

Contribution

It presents the first polynomial LLMs with over a billion parameters, surpassing previous models in size and maintaining strong reasoning and ICL capabilities.

Findings

First polynomial LLMs over a billion parameters.

Models demonstrate reasoning and ICL comparable to standard transformers.

Provides latency analysis for encrypted inference.

Abstract

Modern cryptographic methods for implementing privacy-preserving LLMs such as \gls{HE} require the LLMs to have a polynomial form. Forming such a representation is challenging because transformers include non-polynomial components, such as \Softmax and layer normalization. Previous approaches have either directly approximated pre-trained models with large-degree polynomials, which are less efficient over HE, or replaced non-polynomial components with easier-to-approximate primitives before training, e.g., \Softmax with pointwise attention. The latter approach might introduce scalability challenges. We present a new HE-friendly variant of self-attention that offers a stable form for training and is easy to approximate with polynomials for secure inference. Our work introduces the first polynomial LLMs over a billion parameters, exceeding the size of previous models by more than tenfold. The resulting models demonstrate reasoning and in-context learning (ICL) capabilities comparable to standard transformers of the same size, representing a breakthrough in the field. Finally, we provide a detailed latency breakdown for each computation over encrypted data, paving the way for further optimization, and explore the differences in inductive bias between models relying on our HE-friendly variant and standard transformers.