GRATING: Low-Latency and Memory-Efficient Semantic Selection on Device

TL;DR

GRATING tackles the on-device bottleneck of semantic top-K selection by reframing the task to rely on relative rankings and leveraging sequence-level sparsity. It introduces monolithic forwarding with progressive cluster pruning, chunked execution, a dual-layer sliding window, and embedding table caching to achieve large latency and memory gains without retraining. Across microbenchmarks and three real-world on-device tasks, GRATING delivers substantial reductions in latency (up to 89%) and peak memory (up to 11.45x) while maintaining precision, and demonstrates strong deployability on edge platforms. The approach is training-free, orthogonal to existing compression methods, and broadly applicable to cross-encoder rerankers beyond the evaluated models. This significantly enhances the practicality of on-device AI services such as RAG, agent memory, and long-context selection.

Abstract

Semantic top-K selection with cross-encoder rerankers underpins of on-device AI services, such as retrieval-augmented generation, agent memory, and personalized recommendation. However, its latency and memory demands dominate end-to-end budgets on edge hardware. Revisiting the objective of top-K selection, we reveal that only relative rankings matter, not exact per-candidate scores. We further observe sequence-level sparsity: relative rankings stabilize early in intermediate layers, allowing pruning opportunities prior to completing full inference. Building on this insight, we propose monolithic forwarding and develop a training-free inference system, GRATING. By maintaining a global view of all candidates, it reduces latency through progressive cluster pruning. It also bounds peak memory usage by strategically overlapping I/O with computation via dual-layer sliding window and chunked execution. We evaluate GRATING against state-of-the-art baselines on rerankers from 0.6B to 8B parameters across Apple M2 and RTX 5070. GRATING consistently reduces latency by up to 89.0% and peak memory by up to 94.9% in microbenchmarks, without any loss in precision. Across three real-world on-device AI applications, GRATING lowers latency by 11.6%-51.0% and peak memory by 18.6%-77.8%, demonstrating substantial improvements in efficiency and deployability.

Figures (16)

• Figure 1: Typical on-device top-$K$ selection pipeline and per-stage cost. Per-stage latency and peak memory under a representative on-device semantic file search is reported.
• Figure 2: Score evolution across transformer layers reveals sequence-level sparsity. Each line presents the score evolution of a candidate. Candidates progressively diverge into distinct clusters as they pass through layers, and relative rankings stabilize early in intermediate layers (e.g., layers 7 and 13).
• Figure 3: GRATING overview with a working example.
• Figure 4: Progressive cluster pruning. Once the score dispersion exceeds the threshold during the layer progressing, we partition candidates into selected, deferred, and dropped clusters. Only the deferred candidates are retained and others are pruned.
• Figure 5: Dual-layer sliding window. Throughout the inference, we only reserve two memory buffers for model weights. When layer $i$ resident in one buffer is computing, the next layer $i{+}1$ is prefetched from storage into the other buffer to overlap I/O. When layer $i$ finishes, its buffer is released and reused to prefetch layer $i{+}2$, advancing a dual-layer sliding window.