PaperlayerView exact revision on arXiv

Source-linked AI summary

PUMA: Secure Inference of LLaMA-7B in Five Minutes

Ye Dong, Wen-jie Lu, Yancheng Zheng, Haoqi Wu, Derun Zhao, Jin Tan, Zhicong Huang, Cheng Hong, Tao Wei, Wenguang Chen

arXiv:2307.12533v3cs.CR

TL;DR

Transformer services raise privacy concerns because users may have to reveal prompts or providers may have to release proprietary weights. PUMA uses specialized nonlinear approximations and secure Embedding and LayerNorm procedures for end-to-end inference. It reports plaintext-like accuracy, about 2× faster inference than MPCFORMER, and LLaMA-7B evaluation in around five minutes per generated word.

  • Problem

    Secure Transformer inference aims to protect prompts and model weights, but existing approaches remain limited in performance, efficiency, or deployment.

    Source: §1 · pp. 1–2

  • Method

    PUMA combines accurate polynomial approximations for expensive functions with secure Embedding and LayerNorm protocols for end-to-end Transformer inference.

    Source: §6 · p. 9

  • Results

    PUMA’s precision is similar to plaintext models and it is about 2× faster than MPCFORMER; LLaMA-7B evaluation takes around 5 minutes to generate one word.

    Source: §1 · p. 2

  • Takeaways & Limitations

    PUMA moves secure inference closer to addressing privacy concerns in Transformer-based DLaaS, although inference cost remains high.

    Source: §6 · p. 9

    Prior cheaper approximations of GeLU and softmax introduce utility loss and require fine-tuning, which is unfriendly to data-limited participants.

    Source: §1 · p. 2

Abstract

from arXiv · show

With ChatGPT as a representative, tons of companies have began to provide services based on large Transformers models. However, using such a service inevitably leak users' prompts to the model provider. Previous studies have studied secure inference for Transformer models using secure multiparty computation (MPC), where model parameters and clients' prompts are kept secret. Despite this, these frameworks are still limited in terms of model performance, efficiency, and deployment. To address these limitations, we propose framework PUMA to enable fast and secure Transformer model inference. Our framework designs high quality approximations for expensive functions such as GeLU and softmax, and significantly reduce the cost of secure inference while preserving the model performance. Additionally, we design secure Embedding and LayerNorm procedures that faithfully implement the desired functionality without undermining the Transformer architecture. PUMA is about $2\times$ faster than the state-of-the-art framework MPCFORMER(ICLR 2023) and has similar accuracy as plaintext models without fine-tuning (which the previous works failed to achieve). PUMA can even evaluate LLaMA-7B in around 5 minutes to generate 1 token. To our best knowledge, this is the first time that a model with such a parameter size is able to be evaluated under MPC. PUMA has been open-sourced in the Github repository of SecretFlow-SPU.

1 Introduction

Secure Transformer inference must protect prompts and model weights while addressing gaps in efficiency, accuracy, and deployment. PUMA targets these gaps with specialized nonlinear approximations and secure Embedding and LayerNorm procedures, reporting plaintext-like precision and faster inference than MPCFORMER.

1 Introduction

  • Transformer services can expose users’ private prompts or require providers to release proprietary trained weights.MPC is presented as a way to keep data and model weights private during inference.Source: §1 · pp. 1–2
  • Prior secure-inference approaches face high nonlinear-function costs, retraining requirements, or architecture changes that impede deployment.Faithful GeLU and softmax methods can be costly, while cheaper approximations can reduce utility and require fine-tuning; some architecture changes conflict with plaintext systems.Source: §1 · p. 2
  • The paper asks whether pretrained large Transformers can be evaluated securely and efficiently with plaintext-like accuracy and no further retraining.Source: §1 · p. 2
  • PUMA is proposed as a fast, accurate, end-to-end secure Transformer inference framework.Its nonlinear approximations are designed around the specialized properties of functions such as GeLU and softmax to pursue both accuracy and efficiency.Source: §1 · p. 2
  • PUMA’s precision is similar to plaintext models and it is about 2× faster than MPCFORMER; it evaluates LLaMA-7B in around 5 minutes to generate one word.The authors report this as the first evaluation under MPC of a language model of this size.Source: §1 · p. 2
  • PUMA implements Transformer layers including Embedding and LayerNorm in MPC, enabling secure evaluation of pretrained plaintext models without retraining or other modifications.The authors describe it as the first open-sourced MPC solution to support this use case accurately.Source: §1 · p. 2

2 Related Work

MPC methods for secure inference have mostly targeted convolutional or deep neural networks, with recent Transformer approaches still limited in performance, efficiency, or deployment. The paper compares PUMA chiefly with MPCFORMER under the same three-party setting.

2 Related Work

  • Most earlier MPC inference approaches target convolutional or deep neural networks and cannot be directly extended to Transformers.Source: §2 · p. 2
  • Recent Transformer MPC approaches remain limited in model performance, efficiency, or deployment.Source: §2 · p. 3
  • MPCFORMER is the only open-sourced approach among those works and uses CrypTen’s three-party setting with a non-colluding third party, which the authors use for comparison with PUMA.They compare under the same three-party setting.Source: §2 · p. 3

3 Background

This section establishes Transformer components and the notation and MPC machinery used to describe their secure evaluation. It also specifies the sharing model and adversary scope.

3.2 Transformer Model

  • A Transformer comprises Embedding, Attention, Feed-Forward Network, and LayerNorm sub-layers; attention uses softmax, while FFN commonly uses GeLU.LayerNorm is defined over a vector using its mean and standard deviation.Source: §3.2 · p. 3

3.3 2-out-of-3 Replicated Secret Sharing

  • J·K denotes 2-out-of-3 replicated secret sharing; arithmetic sharing uses a ring, while Boolean sharing uses bitwise operations.Each party holds two of three random shares whose sum reconstructs the secret modulo 2^ℓ.Source: §3.3 · p. 3
  • The MPC toolkit combines local addition with secure multiplication and black-box protocols for equality, comparisons, reciprocals, maxima, and other operations.Fixed-point multiplication requires secure truncation to avoid overflow.Source: §3.3 · pp. 3–4
  • PUMA assumes a semi-honest adversary corrupting at most one of three computing parties; inference-result attacks are outside its scope.Vector and matrix operations extend the described operations.Source: §3.3 · p. 4

4 Secure Design of PUMA

PUMA describes secure inference as sequential MPC over secret-shared model weights and layer outputs, with protocols for Transformer nonlinear and normalization layers. Its designs target efficient evaluation while retaining plaintext-model workflows.

4.1 Overview of PUMA

  • Each layer starts and ends with 2-out-of-3 replicated shares of weights and activations, allowing layers to be combined across Transformer depths.The framework describes a model owner, client, and three computing parties that execute MPC and return results to the client.Source: §4.1 · p. 4

4.2 Protocol for Secure GeLU

  • The piecewise GeLU fit reports max error < 0.01403, median error < 4.41e−05, and mean error < 0.00168.It exploits GeLU's near-linear tails and uses low-degree polynomials over the central interval.Source: §4.2 · pp. 4–5

4.3 Protocol for Secure Softmax

  • Secure softmax uses a clipped Taylor-series exponential on negative inputs and replaces repeated divisions with one reciprocal and multiplications.The protocol subtracts the input maximum for numerical stability; its division optimization is stated to reduce computation and communication while maintaining accuracy.Source: §4.3 · pp. 5–6

4.4 Protocol for Secure Embedding

  • Secure embedding shares the token id and embedding table, computes the one-hot vector at the computing parties, and obtains the selected embedding without changing the plaintext workflow.This design uses additional equality and Boolean-arithmetic multiplication operations.Source: §4.4 · p. 6

4.5 Protocol for Secure LayerNorm

  • The LayerNorm design computes the inverse square root of the variance and broadcasts it for normalization, rather than sequentially computing square root, reciprocal, and multiplication.The authors state that inverse-square-root computation has cost similar to square root in MPC.Source: §4.5 · p. 6

5 Experimental Evaluations

Across seven NLP models, the evaluations compare PUMA’s precision with plaintext and its inference costs with MPCFORMER. They also test scalability and secure LLaMA-7B inference.

5 Experimental Evaluations

  • The evaluation covers seven NLP models, including BERT, GPT2, and LLaMA-7B, on GLUE tasks and Wikitext-103 V1.Experiments used three Alibaba Cloud servers; the reported BERT and GPT2 cost comparisons use different input-length settings.Source: §5 · pp. 6–7

5.1 Precision

  • PUMA’s accuracy differs from plaintext by at most 0.011 across the reported BERT datasets, while GPT2 perplexity differences are at most 0.02.For CoLA, replacing LayerNorm with BatchNorm drops the plaintext BERT-Base MCC from 0.616 to −0.020; PUMA’s MCC is 0.613.Source: §5.1 · p. 7, Table 1 · p. 7, Table 2 · p. 7

5.2 Inference Costs

  • 1.375∼1.916× faster on BERT models and 2.250∼2.414× faster on GPT2 models, PUMA also uses 1.079∼1.195× and 1.325∼1.884× less communication, respectively, than MPCFORMER.The BERT costs are for one length-128 sentence; GPT2 costs are for a length-32 input generating one token.Source: §5.2 · p. 8, Table 3 · p. 8, Table 4 · p. 8

5.3 Scalability

  • 1.631∼1.837× faster for BERT-Base and 1.744∼2.686× faster for GPT2-Base across tested input lengths, PUMA’s output-length gains range from 1.279∼2.700×.For GPT2, efficiency gains decrease as input or output token counts rise; the authors attribute this to extra one-hot embedding costs.Source: §5.3 · p. 8, Table 5 · p. 8, Figure 1 · p. 9

5.4 Evaluating LLaMA-7B in Five Minutes.

  • With an 8-token input, secure LLaMA-7B inference generates one token in around 200 seconds with 1.794 GB communication; it produces the same tokens as plaintext.The authors identify this as the first evaluation of LLaMA-7B using MPC.Source: §5.4 · p. 9, Table 6 · p. 9

6 Conclusion

PUMA is an MPC framework for secure Transformer inference, using replicated secret sharing, accurate polynomial approximations, and secure Embedding and LayerNorm protocols. Its inference cost remains high, but the authors identify quantization and hardware acceleration as possible future directions.

6 Conclusion

  • PUMA combines replicated secret sharing with accurate polynomial approximations and secure Embedding and LayerNorm protocols for end-to-end Transformer inference.The approximations target expensive functions.Source: §6 · p. 9
  • The authors say PUMA brings secure inference one step closer to addressing privacy concerns in Transformer-based DLaaS, although its inference cost remains high.Source: §6 · p. 9
  • The authors suggest combining PUMA with quantization methods and hardware acceleration to make secure inference of large Transformer models in seconds possible.Source: §6 · p. 9

A Details of Experimental Models

The experimental models span BERT, RoBERTa, and GPT-2 variants, with differing layer counts, hidden sizes, head counts, and parameter totals.

A Details of Experimental Models

B PUMA for LLaMA-7B

For LLaMA-7B, PUMA approximates SiLU with piece-wise low-degree polynomials using different coefficients. With fixed randomness, its generated tokens match plaintext LLaMA-7B for more than 20 tokens in the illustrated example.

B PUMA for LLaMA-7B

  • LLaMA uses SiLU rather than GeLU, and PUMA approximates SiLU with similar piece-wise low-degree polynomials using different coefficients.Source: B PUMA for LLaMA-7B · p. 12
  • In Figure 2, PUMA and plaintext LLaMA-7B produce the same tokens for more than 20 generated tokens with fixed randomness.The example uses the prompt “Q: What is the largest animal?”Source: B PUMA for LLaMA-7B · p. 12, Figure 2 · p. 13
  • Figure 2 presents outputs from LLaMA-7B in plaintext and PUMA.Source: Figure 2 · p. 13

Related papers

Loading 2307.12533v3…