Reproduce “A First Look At Efficient And Secure On-Device LLM Inference Against KV Leakage” on OpenResearch

About this paper

Huan Yang, Deyu Zhang, Yudong Zhao, Yuanchun Li, Yunxin Liu

This research identifies and demonstrates a critical vulnerability where Key-Value (KV) pairs can leak from mobile GPU memory during on-device LLM inference, enabling reconstruction of user conversations. It proposes KV-Shield, a scheme that leverages permutation and Trusted Execution Environments (TEEs) to protect these KV pairs, confirming its theoretical correctness while noting current performance overheads.

Key insights

• Existing FHE and TEE paradigms are not directly viable for securing on-device LLM inference due to their inherent performance and resource limitations in mobile environments.
• A hybrid security model, which delegates heavy computation to the untrusted GPU while confining lightweight, security-sensitive operations to the TEE, presents a more practical path for on-device LLM protection.
• Permuting the foundational linear layer weights within the self-attention mechanism can effectively scramble KV caches without altering the LLM's core architecture or output accuracy, if correctly inverse-permuted.

Summary by alphaXiv.