Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Comet: Accelerating Private Inference for Large Language Model by Predicting Activation Sparsity

Guang Yan, Yuhui Zhang, Zimu Guo, Lutan Zhao, Xiaojun Chen, Chen Wang, Wenhao Wang, Dan Meng, Rui Hou

TL;DR

Comet tackles privacy-preserving inference for large language models in MPC settings by predicting activation sparsity and skipping zero-valued activations in both preceding and following linear layers. It employs shallow, low-rank predictors with an oblivious shuffle to hide sparsity patterns, enabling efficient plaintext indexing, and introduces SOMM and SIMM protocols that exploit spatial locality to reduce inter-server communication. A KV cache manager further preserves cache continuity by merging misses and prefetching for not-activated heads, maintaining model accuracy while lowering latency. Across four LLMs, Comet yields end-to-end speedups of 1.87x–2.63x and communication reductions of 1.94x–2.64x compared with six baselines, with larger models benefiting most, indicating practical viability for privacy-preserving private inference in cloud settings.

Abstract

With the growing use of large language models (LLMs) hosted on cloud platforms to offer inference services, privacy concerns about the potential leakage of sensitive information are escalating. Secure multi-party computation (MPC) is a promising solution to protect the privacy in LLM inference. However, MPC requires frequent inter-server communication, causing high performance overhead. Inspired by the prevalent activation sparsity of LLMs, where most neuron are not activated after non-linear activation functions, we propose an efficient private inference system, Comet. This system employs an accurate and fast predictor to predict the sparsity distribution of activation function output. Additionally, we introduce a new private inference protocol. It efficiently and securely avoids computations involving zero values by exploiting the spatial locality of the predicted sparse distribution. While this computation-avoidance approach impacts the spatiotemporal continuity of KV cache entries, we address this challenge with a low-communication overhead cache refilling strategy that merges miss requests and incorporates a prefetching mechanism. Finally, we evaluate Comet on four common LLMs and compare it with six state-of-the-art private inference systems. Comet achieves a 1.87x-2.63x speedup and a 1.94x-2.64x communication reduction.

Comet: Accelerating Private Inference for Large Language Model by Predicting Activation Sparsity

TL;DR

Comet tackles privacy-preserving inference for large language models in MPC settings by predicting activation sparsity and skipping zero-valued activations in both preceding and following linear layers. It employs shallow, low-rank predictors with an oblivious shuffle to hide sparsity patterns, enabling efficient plaintext indexing, and introduces SOMM and SIMM protocols that exploit spatial locality to reduce inter-server communication. A KV cache manager further preserves cache continuity by merging misses and prefetching for not-activated heads, maintaining model accuracy while lowering latency. Across four LLMs, Comet yields end-to-end speedups of 1.87x–2.63x and communication reductions of 1.94x–2.64x compared with six baselines, with larger models benefiting most, indicating practical viability for privacy-preserving private inference in cloud settings.

Abstract

With the growing use of large language models (LLMs) hosted on cloud platforms to offer inference services, privacy concerns about the potential leakage of sensitive information are escalating. Secure multi-party computation (MPC) is a promising solution to protect the privacy in LLM inference. However, MPC requires frequent inter-server communication, causing high performance overhead. Inspired by the prevalent activation sparsity of LLMs, where most neuron are not activated after non-linear activation functions, we propose an efficient private inference system, Comet. This system employs an accurate and fast predictor to predict the sparsity distribution of activation function output. Additionally, we introduce a new private inference protocol. It efficiently and securely avoids computations involving zero values by exploiting the spatial locality of the predicted sparse distribution. While this computation-avoidance approach impacts the spatiotemporal continuity of KV cache entries, we address this challenge with a low-communication overhead cache refilling strategy that merges miss requests and incorporates a prefetching mechanism. Finally, we evaluate Comet on four common LLMs and compare it with six state-of-the-art private inference systems. Comet achieves a 1.87x-2.63x speedup and a 1.94x-2.64x communication reduction.
Paper Structure (25 sections, 2 theorems, 16 figures, 3 tables, 3 algorithms)

This paper contains 25 sections, 2 theorems, 16 figures, 3 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Understanding of Private Inference
  3. MPC-based private inference
  4. LLM private inference has high communication overhead
  5. Prevalent activation sparsity enables significant reductions in communication for private inference
  6. Comet System Overview
  7. Design motivation
  8. Threat model
  9. Design overview
  10. Activation Sparsity Predictor with Oblivious Shuffle
  11. Basic design
  12. Enhancing indexing efficiency with oblivious shuffle
  13. Private Inference Protocol Exploiting Spatial Locality of Sparsity Distribution
  14. Sparse output matrix multiplication protocol for preceding linear layer
  15. Sparse input matrix multiplication protocol for subsequent linear layer
  16. ...and 10 more sections

Key Result

Theorem 1

Given $[\![ \mathbf{X} ]\!]\xspace$, $[\![ \mathbf{Y} ]\!]\xspace$, and the sparsity distribution $\mathcal{S}$ of $\mathbf{XY}$, Protocol protocol:SOMM achieves the minimal communication cost for computing $[\![ \mathbf{XY} ]\!]\xspace$. Furthermore, under this minimal communication cost, the proto

Figures (16)

  • Figure 1: Activation sparsity of LLMs.
  • Figure 2: MPC-based private inference time breakdown with an input length of 512, output length of 16, and bandwidth of 5Gbps. The numbers on the bars is the proportion of communication time.
  • Figure 3: Activation sparsity of different LLMs, evaluated on the same dataset used in Section \ref{['sec:evalution']}.
  • Figure 4: Communication cost breakdown for OPT-6.7B with an input length of 512 and output length of 16.
  • Figure 5: System overview of Comet.
  • ...and 11 more figures

Theorems & Definitions (2)

  • Theorem 1
  • Theorem 2