On-device machine learning (ML) inference can enable the use of private user data on user devices without remote servers. However, a pure on-device solution to private ML inference is impractical for many applications that rely on embedding tables that are too large to be stored on-device. To overcome this barrier, we propose the use of private information retrieval (PIR) to efficiently and privately retrieve embeddings from servers without sharing any private information during on-device ML inference. As off-the-shelf PIR algorithms are usually too computationally intensive to directly use for latency-sensitive inference tasks, we 1) develop a novel algorithm for accelerating PIR on GPUs, and 2) co-design PIR with the downstream ML application to obtain further speedup. Our GPU acceleration strategy improves system throughput by more than $20 \times$ over an optimized CPU PIR implementation, and our co-design techniques obtain over $5 \times$ additional throughput improvement at fixed model quality. Together, on various on-device ML applications such as recommendation and language modeling, our system on a single V100 GPU can serve up to $100,000$ queries per second -- a $>100 \times$ throughput improvement over a naively implemented system -- while maintaining model accuracy, and limiting inference communication and response latency to within $300$KB and $<100$ms respectively.

翻译：设备端机器学习推理能够利用用户设备上的私有用户数据，而无需依赖远程服务器。然而，对于许多依赖嵌入表（这些表规模过大，无法存储在设备上）的应用而言，纯粹的设备端私有机器学习推理方案并不实用。为克服这一障碍，我们提出使用私有信息检索（PIR），在设备端机器学习推理过程中，高效且私密地从服务器检索嵌入信息，无需泄露任何私有数据。由于现成的PIR算法通常计算密集度太高，无法直接用于对延迟敏感的推理任务，我们：1) 开发了一种新颖的算法，用于在GPU上加速PIR；2) 将PIR与下游机器学习应用协同设计，以进一步获得加速效果。我们的GPU加速策略相较于优化的CPU PIR实现，系统吞吐量提升了20倍以上；而协同设计技术在固定模型质量下，额外实现了超过5倍的吞吐量提升。结合来看，在推荐系统和语言建模等多种设备端机器学习应用中，我们的系统在单个V100 GPU上每秒可处理高达10万次查询——相较于简单实现的系统，吞吐量提升了超过100倍——同时保持模型准确率，并将推理通信量和响应延迟分别限制在300KB以内和100毫秒以内。