Power-Softmax: Towards Secure LLM Inference over Encrypted Data

📄 arXiv: 2410.09457v1 📥 PDF

作者: Itamar Zimerman, Allon Adir, Ehud Aharoni, Matan Avitan, Moran Baruch, Nir Drucker, Jenny Lerner, Ramy Masalha, Reut Meiri, Omri Soceanu

分类: cs.LG, cs.CR

发布日期: 2024-10-12

💡 一句话要点

提出Power-Softmax，实现加密数据上安全LLM推理，模型参数超十亿

🎯 匹配领域: 支柱五：交互与反应 (Interaction & Reaction)

关键词: 同态加密 安全推理 大语言模型 Power-Softmax 自注意力机制 隐私保护 多项式近似

📋 核心要点

  1. 现有同态加密LLM方法依赖多项式形式，但Transformer中的Softmax等非多项式组件难以高效近似。
  2. 提出Power-Softmax，一种HE友好的自注意力变体，易于多项式近似，并保持训练稳定性。
  3. 构建了首个32层、超10亿参数的多项式LLM，性能与同等规模的标准Transformer相当。

📝 摘要（中文）

本文提出了一种新的同态加密（HE）友好的自注意力机制变体，该变体具有稳定的训练形式，并且易于用多项式近似，从而实现安全推理。Transformer包含非多项式组件，如Softmax和层归一化，这给构建多项式形式的LLM带来了挑战。先前的方法要么直接用大度多项式近似预训练模型（效率较低），要么在训练前用更容易近似的基元替换非多项式组件（可能引入可扩展性问题）。本文提出的方法构建了首个具有32层和超过10亿参数的多项式LLM，超过了先前模型的规模十倍以上。实验表明，生成的模型具有与相同大小的标准Transformer相当的推理和上下文学习（ICL）能力，代表了该领域的突破。最后，本文提供了加密数据上每次计算的详细延迟分解，为进一步优化铺平了道路，并探讨了依赖于HE友好变体的Transformer与标准Transformer之间的归纳偏差差异。

🔬 方法详解

问题定义：论文旨在解决在加密数据上安全地进行大规模语言模型（LLM）推理的问题。现有的方法要么直接用高阶多项式近似预训练模型，导致效率低下；要么在训练前替换非多项式组件，可能影响模型的可扩展性和性能。因此，需要一种既能保持模型性能，又能方便进行同态加密推理的方法。

核心思路：论文的核心思路是设计一种新的、同态加密友好的自注意力机制，即Power-Softmax。这种机制在训练时具有稳定性，并且可以很容易地用低阶多项式进行近似，从而在加密数据上实现高效的推理。通过替换标准Softmax，使得整个模型更容易转换为多项式形式，从而兼容同态加密。

技术框架：整体框架与标准的Transformer类似，但关键在于将标准的Softmax替换为Power-Softmax。模型包括嵌入层、多层Transformer块（每个块包含自注意力层和前馈网络）以及输出层。自注意力层是核心，其中Power-Softmax取代了传统的Softmax。整个模型在明文数据上进行训练，然后将其转换为多项式形式，以便在加密数据上进行推理。

关键创新：最重要的技术创新点在于Power-Softmax的设计。与直接近似Softmax不同，Power-Softmax本身就是一个设计好的、更易于多项式近似的函数。这避免了高阶多项式近似带来的计算复杂性，并提高了同态加密推理的效率。此外，该方法构建了当时最大的多项式LLM，证明了其可扩展性。

关键设计：Power-Softmax的具体形式未知，但可以推断其设计目标是：1）易于用低阶多项式近似；2）在训练过程中保持数值稳定性；3）尽可能接近标准Softmax的功能，以减少对模型性能的影响。论文中可能包含关于Power-Softmax的具体数学表达式、近似多项式的阶数选择以及训练过程中的超参数设置等细节。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，使用Power-Softmax构建的32层、超过10亿参数的LLM，在推理和上下文学习能力上与同等规模的标准Transformer相当。这是首个如此大规模的多项式LLM，代表了在加密数据上进行LLM推理的重大突破。论文还提供了加密数据上计算的详细延迟分解，为后续优化提供了依据。

🎯 应用场景

该研究成果可应用于需要保护用户隐私的场景，例如金融、医疗等领域。用户可以在不泄露数据内容的前提下，利用LLM进行智能分析和决策。未来，该技术有望推动隐私计算的发展，促进安全多方计算和联邦学习等技术的应用。

📄 摘要（原文）

Modern cryptographic methods for implementing privacy-preserving LLMs such as Homomorphic Encryption (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 with 32 layers and 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 transformers relying on our HE-friendly variant and standard transformers. Our code is attached as a supplement.