GPU-based Private Information Retrieval for On-Device Machine Learning Inference

Running a recommendation model on your phone keeps your data private until it needs a lookup in a table too big to fit on the phone — private information retrieval fetches the entry without the server learning which one, and this paper makes that fast enough to actually serve traffic.

Why this problem is worth solving

The privacy argument for on-device inference is straightforward: if the model runs on your phone, your features never reach anyone’s server. The obstacle is equally straightforward. Recommendation models — which are a large fraction of production machine learning — work by mapping categorical features like click and search history into dense vectors, and the

The paper studies a real production model where seven of the twenty-five most important features are ones that can only be used on the device, and removing them measurably degrades the model. So the choice is between a worse model and a leaky one — unless the lookup itself can be made private.

What people did before, and where it broke

• Store the table on the device. The pure on-device solution, and impossible by two orders of magnitude for these tables.
• Drop the private features. Works, and the paper’s production case shows it costs real accuracy, since those features rank among the most important the model has.
• Fetch the row in the clear. Fast and equivalent to disclosing the feature, since the index is the private value.
• Single-server private information retrieval. No non-collusion assumption needed, and substantially more expensive in both computation and communication than the two-server scheme.
• Oblivious RAM. Hides access patterns rather than single queries, and needs the client to maintain state across accesses — awkward when many clients query the same table.
• Run the CPU implementation faster. The baseline here is an optimised implementation from Google Research that uses the processor’s built-in AES instructions. On a million-entry table a single thread manages 1.3 queries per second; thirty-two threads manage 21.
• Batch private information retrieval. Standard techniques fetch several entries at once by splitting the table into bins and querying each once, which cuts computation roughly by the number of bins. The catch is that two wanted entries can land in the same bin, and only one comes back — queries get dropped probabilistically.
• The gap. The cryptography was sound and roughly a thousand times too slow, and the standard way of making it cheaper silently damaged the thing it was serving.

How the lookup gets fast enough to serve

The idea

Both natural parallelisations fail, in opposite directions. The expensive part of this protocol is expanding a key into a share of a one-hot vector, which is structured as a binary tree whose leaves are the table rows and where computing any node requires its parent. Two ways to parallelise that suggest themselves immediately. Assign each thread one root-to-leaf branch, and every thread recomputes the whole path above its leaf — the work explodes. Evaluate level by level instead, and the work is optimal but every intermediate node must be stored, so memory grows with the table size times the batch size, which caps the batch and leaves the GPU underused. The paper’s answer is a traversal that expands a bounded window of nodes at a time and discards what it no longer needs, bringing memory down to something logarithmic in the table size while keeping the work optimal — which permits batches large enough to actually saturate the hardware, worth up to an order of magnitude on its own. Finding that a cryptographic primitive’s obvious parallel decompositions both mismatch GPU hardware, and designing a third that respects the memory hierarchy, is the same exercise performed in accelerating zero-knowledge proof systems, covered elsewhere on this site.

Do not materialise what you are about to consume. Having reduced the expansion’s memory footprint, the design would immediately give it back: the next step multiplies the expanded vector by the table, and computing the vector in full before multiplying restores the memory cost the traversal just eliminated. So the two operations are fused. When a thread reaches a leaf it immediately dots that table row into an accumulator held in registers, and threads in a block combine their partial sums at the end. The memory bound holds, global memory traffic falls, and the two kinds of work interleave — worth more than 1.5 times in both throughput and latency. There is a related and pleasingly mundane observation alongside it: processors have dedicated AES instructions and GPUs do not, so the standard cipher choice is a CPU habit that costs dearly here. Swapping in a stream cipher from TLS nearly quadruples throughput at the same security parameter, and the paper is careful to note that the still-faster options have had less scrutiny. Finding that the real bottleneck in a private computation sits somewhere other than where the cryptography literature would point, and that it moves once the whole pipeline is measured on the hardware actually available, is the recurring lesson of end-to-end characterisations of private inference systems, covered elsewhere on this site.

The consumer is a machine learning model, and that changes what counts as correct. The final speedups come from refusing to treat the retrieval as an isolated cryptographic service. Batch retrieval saves computation by splitting the table into bins and querying each once, but two wanted rows in one bin means one gets dropped — a hard error for a database and a small accuracy cost for a model. Knowing the consumer, the paper shapes the bins: profile the training data, split off a hot table holding the most frequently requested rows, and co-locate rows that tend to be requested together so a single query returns several wanted entries rather than colliding on them. The right amount of co-location turns out to differ by domain — higher for language, where words have natural associations, lower for recommendation. Exploiting the strong skew and correlation in embedding table accesses to serve them from a much smaller working set is the same observation that drives caching designs for embedding storage, covered elsewhere on this site, applied here to decide what a private query should even ask for.

Following one private lookup

Neither key alone reveals anything about the index. The privacy guarantee rests entirely on the two servers not comparing notes — which is the deployment assumption the whole scheme carries.

Why it works

• Bounded-window traversal escapes the work-versus-memory dilemma. Branch-parallel wastes work, level-by-level wastes memory → expanding a fixed window and discarding it gives optimal work at logarithmic memory → so the batch size can grow until the GPU is saturated, worth up to ten times.
• Fusion protects the memory saving. Materialising the expanded vector before multiplying would restore the original memory cost → accumulating each leaf into registers as it is produced keeps the bound → and cuts global memory traffic, giving over 1.5 times in throughput and latency together.
• The cipher choice was inherited from the wrong machine. Processors accelerate AES in hardware and GPUs do not → a stream cipher at the same security level runs nearly four times faster on this hardware → which costs nothing in security because it is the cipher TLS already uses.
• Large tables need less batching, not more. Beyond about four million entries a single query already fills the GPU → dedicating the whole device to one query at a time cuts latency without hurting throughput → so the scheduling adapts to table size rather than fixing a batch.
• Collisions are cheaper when the bins are chosen with knowledge of the workload. Batch retrieval drops entries that collide in a bin → putting frequently co-requested rows in the same bin turns collisions into hits → reducing computation by 1.9 to 7.4 times at fixed accuracy.

What it costs (trade-offs and scope)

• It requires two servers that do not collude. The entire privacy guarantee rests on this. It is a standard assumption in this literature and it is an assumption about organisations, not about mathematics.
• The headline combines three things, one of which is accuracy. More than 100 times over the CPU baseline reflects GPU acceleration plus co-design plus a relaxed accuracy target. At preserved model quality the improvement averages 5 to 39 times; allowing degradation raises it to 40 to 124 times.
• Batch retrieval drops queries by design. Two wanted entries in one bin means one is lost. The co-design reduces how often this hurts; it does not eliminate it, and the throughput gains are partly bought with it.
• The GPU does not always win on latency. On a small 16-thousand-entry table, thirty-two CPU threads answer in 0.71 milliseconds against the GPU’s 3.2. The GPU wins on throughput there by more than twenty times, but a latency-bound deployment with small tables would not benefit.
• Faster ciphers come with weaker assurance. The fastest option measured is roughly seven times faster than AES, and the authors state plainly that it has not been analysed as thoroughly, leaving the choice as an explicit security trade-off.
• The co-design parameters need a grid search per application. Hot table size, bin count and co-location factor are all swept on a validation set, and the best co-location differs between language and recommendation workloads.
• Multi-GPU scaling is argued, not measured. Splitting a table across devices and summing at the end should scale linearly, but all reported numbers come from a single V100.
• Language model coverage is narrow. The technique protects word embeddings for a small recurrent model; the authors note that private inference for large language models needs much more than embedding access moved off the device.

The evaluation

Check yourself

Sources

The paper is “GPU-based Private Information Retrieval for On-Device Machine Learning Inference” by Maximilian Lam, Jeff Johnson, Wenjie Xiong, Kiwan Maeng, Udit Gupta, Yang Li, Liangzhen Lai, Ilias Leontiadis, Minsoo Rhu, Hsien-Hsin S. Lee, Vijay Janapa Reddi, Gu-Yeon Wei, David Brooks and G. Edward Suh, published at ASPLOS 2024 Volume 1, pages 197-214, DOI 10.1145/3617232.3624855. The authors span Harvard, Meta, Virginia Tech, Penn State, Cornell and KAIST.

This explainer is written from the full text, read via the arXiv version at arxiv.org/abs/2301.10904. All figures above come from the paper’s own sections 2 through 5: the embedding table sizes and the 34-megabyte average app size, the production model where seven of the top twenty-five features are device-only, the memory-bounded traversal with its window parameter of 128 and the resulting logarithmic memory bound, the over 1.5 times gain from operator fusion, the throughput table giving 1,358 queries per second on a million-entry table against 1.3 and 21.2 for one and thirty-two CPU threads, the 4.18 millisecond latency on four million entries against 160 milliseconds, the cipher comparison of 965 queries per second for AES-128 against 3,640 for ChaCha20 and 7,447 for SipHash, the average improvements of 5 to 39 times at preserved quality and 40 to 124 times with relaxed accuracy, the co-design savings of 1.9 to 7.4 times in computation and 1.7 to 2.6 times in communication, and the 300 kilobyte and 300 millisecond budgets. Multi-GPU scaling is described in section 3.2.7 as following from the decomposition; all measurements use a single V100.