Amortizing Maximum Inner Product Search with Learned Support Functions
This paper proposes amortized MIPS, a regression-based approach that trains neural networks to directly predict MIPS solutions. The key insight is that the MIPS value function is the support function of the key set, whose gradient yields the optimal key. Two models are introduced: SupportNet, an input-convex network for regressing the support function, and KeyNet, a vector-valued network for directly regressing the optimal key. Experiments on BEIR show improved IVF hit rates under equal compute budgets.
content type paperpublished July 2026
Amortizing Maximum Inner Product Search with Learned Support Functions
AuthorsTheo X. Olausson†**, João Monteiro, Michal Klein, Marco Cuturi
View publication
Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring the identification of a vector taken within a database (the keys) that best aligns with a given query. We propose amortized MIPS: a regression-based approach that trains neural networks to directly predict MIPS solutions, amortizing the cost of repeatedly solving MIPS for queries drawn from a known distribution over a fixed key database. Our key insight is that the MIPS value function is the support function of the set of keys, a well-studied convex function whose gradient yields the optimal key. This motivates two complementary amortized models: SupportNet, an input-convex neural network trained to regress the support function, and KeyNet, a vector-valued network that directly regresses the optimal key. SupportNet can serve as a cluster router, steering queries toward relevant database partitions, while KeyNet can be used as a drop-in replacement for the original query, fed directly to off-the-shelf indexing pipelines. Our experiments on the BEIR benchmark show that, for document embeddings, learned SupportNets and KeyNets significantly improve IVF match rates when accounting for compute effort, whether measured in FLOPs, number of probes, or wall-clock time. Our code is available at: https://github.com/apple/ml-amips.
† MIT
** Work done while at Apple
Scalable Private Search with Wally
October 16, 2024research area Privacy
This paper presents Wally, a private search system that supports efficient semantic and keyword search queries against large databases. When sufficiently many clients are making queries, Wally’s performance is significantly better than previous systems. In previous private search systems, for each client query, the server must perform at least one expensive cryptographic operation per database entry. As a result, performance degraded…
Read more
Syntactic Code Search with Sequence-to-Tree Matching: Supporting Syntactic Search with Incomplete Code Fragments
June 20, 2024research area Human-Computer Interaction, research area Tools, Platforms, Frameworksconference Programming Language Design and Implementation (PLDI)
Lightweight syntactic analysis tools like Semgrep and Comby leverage the tree structure of code, making them more expressive than string and regex search. Unlike traditional language frameworks (e.g., ESLint) that analyze codebases via explicit syntax tree manipulations, these tools use query languages that closely resemble the source language. However, state-of-the-art matching techniques for these tools require queries to be complete and…
Read more