提出KV-Shield防护设备端大模型推理中的密钥泄漏问题
A First Look At Efficient And Secure On-Device LLM Inference Against KV Leakage
- 通过初始化阶段打乱权重矩阵，使键值对混淆
- 运行时逆向置换注意力向量，保证输出正确性
- 在可信执行环境中操作，防止攻击者重建对话
在终端设备上运行大语言模型因隐私保护优势受到广泛关注。随着轻量级模型和专用GPU的发展，设备端推理已达到所需的准确率和性能。然而我们发现，基于GPU的推理会泄露敏感的中间信息——键值对（KV pairs），攻击者可利用这些信息重构用户完整对话，造成严重风险。现有方案如全同态加密（FHE）和可信执行环境（TEE）要么计算开销过大，要么资源受限。为此，我们设计了KV-Shield，包含两个阶段：初始化阶段对权重矩阵进行置换，使所有KV对相应混淆；运行阶段对注意力向量逆向置换，确保层输出正确。所有置换操作均在TEE内完成，确保不安全的GPU无法获取原始KV对，从而防止对话重建。最后，我们理论分析了KV-Shield的正确性及其优劣与开销。
原文摘要 · Abstract (English)
Running LLMs on end devices has garnered significant attention recently due to their advantages in privacy preservation. With the advent of lightweight LLM models and specially designed GPUs, on-device LLM inference has achieved the necessary accuracy and performance metrics. However, we have identified that LLM inference on GPUs can leak privacy-sensitive intermediate information, specifically the KV pairs. An attacker could exploit these KV pairs to reconstruct the entire user conversation, leading to significant vulnerabilities. Existing solutions, such as Fully Homomorphic Encryption (FHE) and Trusted Execution Environments (TEE), are either too computation-intensive or resource-limited. To address these issues, we designed KV-Shield, which operates in two phases. In the initialization phase, it permutes the weight matrices so that all KV pairs are correspondingly permuted. During the runtime phase, the attention vector is inversely permuted to ensure the correctness of the layer output. All permutation-related operations are executed within the TEE, ensuring that insecure GPUs cannot access the original KV pairs, thus preventing conversation reconstruction. Finally, we theoretically analyze the correctness of KV-Shield, along with its advantages and overhead.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。