PUMA：LLaMA-7Bの安全な推論を5分で実現

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

July 24, 2023

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

cs.AI

要旨

ChatGPTを代表として、多くの企業が大規模なTransformerモデルに基づくサービスを提供し始めています。しかし、そのようなサービスを利用すると、ユーザーのプロンプトがモデル提供者に漏洩することを避けられません。これまでの研究では、安全なマルチパーティ計算（MPC）を用いてTransformerモデルの安全な推論を検討し、モデルパラメータとクライアントのプロンプトを秘匿する方法が研究されてきました。それにもかかわらず、これらのフレームワークは、モデルの性能、効率、および展開の面でまだ制限があります。これらの制限に対処するため、我々は高速かつ安全なTransformerモデル推論を可能にするフレームワークPUMAを提案します。我々のフレームワークは、GeLUやSoftmaxなどの高コストな関数に対して高品質な近似を設計し、モデルの性能を維持しながら安全な推論のコストを大幅に削減します。さらに、Transformerアーキテクチャを損なうことなく、所望の機能を忠実に実装する安全なEmbeddingおよびLayerNorm手順を設計しました。PUMAは、最先端のMPCフレームワークであるMPCFORMER（ICLR 2023）よりも約2倍高速であり、ファインチューニングなしで平文モデルと同等の精度を達成しています（これはこれまでの研究では達成できなかったことです）。 さらに、PUMAはLLaMA-7Bを約5分で評価し、1トークンを生成することができます。我々の知る限り、このようなパラメータサイズのモデルがMPC下で評価されたのは初めてのことです。PUMAは、SecretFlow-SPUのGithubリポジトリでオープンソース化されています。

English

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, which 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 2x faster than the state-of-the-art MPC framework MPCFORMER(ICLR 2023) and has similar accuracy as plaintext models without fine-tuning (which the previous works failed to achieve). One more thing, PUMA can 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.