Amortizing Maximum Inner Product Search with Learned Support Functions

• URL: http://arxiv.org/abs/2603.08001v1
• Date: Mon, 09 Mar 2026 06:09:20 GMT
• Title: Amortizing Maximum Inner Product Search with Learned Support Functions
• Abstract summary: inner product search (MIPS) is a crucial computation in machine learning.<n>We propose amortized MIPS: a learning-based approach that trains neural networks to directly predict MIPS solutions.
• Score: 26.9616175011391
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Maximum inner product search (MIPS) is a crucial subroutine in machine learning, requiring the identification of key vectors that align best with a given query. We propose amortized MIPS: a learning-based approach that trains neural networks to directly predict MIPS solutions, amortizing the computational cost of matching queries (drawn from a fixed distribution) to a fixed set of keys. Our key insight is that the MIPS value function, the maximal inner product between a query and keys, is also known as the support function of the set of keys. Support functions are convex, 1-homogeneous and their gradient w.r.t. the query is exactly the optimal key in the database. We approximate the support function using two complementary approaches: (1) we train an input-convex neural network (SupportNet) to model the support function directly; the optimal key can be recovered via (autodiff) gradient computation, and (2) we regress directly the optimal key from the query using a vector valued network (KeyNet), bypassing gradient computation entirely at inference time. To learn a SupportNet, we combine score regression with gradient matching losses, and propose homogenization wrappers that enforce the positive 1-homogeneity of a neural network, theoretically linking function values to gradients. To train a KeyNet, we introduce a score consistency loss derived from the Euler theorem for homogeneous functions. Our experiments show that learned SupportNet or KeyNet achieve high match rates and open up new directions to compress databases with a specific query distribution in mind.

Related papers

• Outcome-Based Online Reinforcement Learning: Algorithms and Fundamental Limits [58.63897489864948]
  Reinforcement learning with outcome-based feedback faces a fundamental challenge.<n>How do we assign credit to the right actions?<n>This paper provides the first comprehensive analysis of this problem in online RL with general function approximation.
  arXiv  Detail & Related papers  (2025-05-26T17:44:08Z)
• A Refreshed Similarity-based Upsampler for Direct High-Ratio Feature Upsampling [54.05517338122698]
  A popular similarity-based feature upsampling pipeline has been proposed, which utilizes a high-resolution feature as guidance.<n>We propose an explicitly controllable query-key feature alignment from both semantic-aware and detail-aware perspectives.<n>We develop a fine-grained neighbor selection strategy on HR features, which is simple yet effective for alleviating mosaic artifacts.
  arXiv  Detail & Related papers  (2024-07-02T14:12:21Z)
• Towards Model-Size Agnostic, Compute-Free, Memorization-based Inference of Deep Learning [5.41530201129053]
  This paper proposes a novel memorization-based inference (MBI) that is compute free and only requires lookups. Specifically, our work capitalizes on the inference mechanism of the recurrent attention model (RAM) By leveraging the low-dimensionality of glimpse, our inference procedure stores key value pairs comprising of glimpse location, patch vector, etc. in a table. The computations are obviated during inference by utilizing the table to read out key-value pairs and performing compute-free inference by memorization.
  arXiv  Detail & Related papers  (2023-07-14T21:01:59Z)
• Learning Feature Matching via Matchable Keypoint-Assisted Graph Neural Network [52.29330138835208]
  Accurately matching local features between a pair of images is a challenging computer vision task. Previous studies typically use attention based graph neural networks (GNNs) with fully-connected graphs over keypoints within/across images. We propose MaKeGNN, a sparse attention-based GNN architecture which bypasses non-repeatable keypoints and leverages matchable ones to guide message passing.
  arXiv  Detail & Related papers  (2023-07-04T02:50:44Z)
• Provable Data Subset Selection For Efficient Neural Network Training [73.34254513162898]
  We introduce the first algorithm to construct coresets for emphRBFNNs, i.e., small weighted subsets that approximate the loss of the input data on any radial basis function network. We then perform empirical evaluations on function approximation and dataset subset selection on popular network architectures and data sets.
  arXiv  Detail & Related papers  (2023-03-09T10:08:34Z)
• On the Convergence of Distributed Stochastic Bilevel Optimization Algorithms over a Network [55.56019538079826]
  Bilevel optimization has been applied to a wide variety of machine learning models. Most existing algorithms restrict their single-machine setting so that they are incapable of handling distributed data. We develop novel decentralized bilevel optimization algorithms based on a gradient tracking communication mechanism and two different gradients.
  arXiv  Detail & Related papers  (2022-06-30T05:29:52Z)
• NFL: Robust Learned Index via Distribution Transformation [14.812854942243503]
  This paper tackles the approximation problem by applying a textit distribution transformation to the keys before constructing the learned index. A two-stage Normalizing-Flow-based Learned index framework (NFL) is proposed, which first transforms the original complex key distribution into a near-uniform distribution, then builds a learned index leveraging the transformed keys. Based on the characteristics of the transformed keys, we propose a robust After-Flow Learned Index (AFLI)
  arXiv  Detail & Related papers  (2022-05-24T06:03:19Z)
• Combinatorial optimization and reasoning with graph neural networks [7.8107109904672365]
  Combinatorial optimization is a well-established area in operations research and computer science. Recent years have seen a surge of interest in using machine learning, especially graph neural networks (GNNs) as a key building block for tasks.
  arXiv  Detail & Related papers  (2021-02-18T18:47:20Z)
• Fast Few-Shot Classification by Few-Iteration Meta-Learning [173.32497326674775]
  We introduce a fast optimization-based meta-learning method for few-shot classification. Our strategy enables important aspects of the base learner objective to be learned during meta-training. We perform a comprehensive experimental analysis, demonstrating the speed and effectiveness of our approach.
  arXiv  Detail & Related papers  (2020-10-01T15:59:31Z)
• Coresets for Near-Convex Functions [25.922075279588757]
  Coreset is usually a small weighted subset of $n$ input points in $mathbbRd$, that provably approximates their loss function for a given set of queries. We suggest a generic framework for computing sensitivities (and thus coresets) for wide family of loss functions. Examples include SVM, Logistic M-estimators, and $ell_z$-regression.
  arXiv  Detail & Related papers  (2020-06-09T19:49:19Z)
• Communication-Efficient Distributed Stochastic AUC Maximization with Deep Neural Networks [50.42141893913188]
  We study a distributed variable for large-scale AUC for a neural network as with a deep neural network. Our model requires a much less number of communication rounds and still a number of communication rounds in theory. Our experiments on several datasets show the effectiveness of our theory and also confirm our theory.
  arXiv  Detail & Related papers  (2020-05-05T18:08:23Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.