Maximally-Informative Retrieval for State Space Model Generation

• URL: http://arxiv.org/abs/2506.12149v1
• Date: Fri, 13 Jun 2025 18:08:54 GMT
• Title: Maximally-Informative Retrieval for State Space Model Generation
• Authors: Evan Becker, Benjamin Bowman, Matthew Trager, Tian Yu Liu, Luca Zancato, Wei Xia, Stefano Soatto,
• Abstract summary: We introduce Retrieval In-Context Optimization (RICO) to minimize model uncertainty for a particular query at test-time.<n>Unlike traditional retrieval-augmented generation (RAG), which relies on externals for document retrieval, our approach leverages direct feedback from the model.<n>We show that standard top-$k$ retrieval with model gradients can approximate our optimization procedure, and provide connections to the leave-one-out loss.
• Score: 59.954191072042526
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Given a query and dataset, the optimal way of answering the query is to make use all the information available. Modern LLMs exhibit impressive ability to memorize training data, but data not deemed important during training is forgotten, and information outside that training set cannot be made use of. Processing an entire dataset at inference time is infeasible due to the bounded nature of model resources (e.g. context size in transformers or states in state space models), meaning we must resort to external memory. This constraint naturally leads to the following problem: How can we decide based on the present query and model, what among a virtually unbounded set of known data matters for inference? To minimize model uncertainty for a particular query at test-time, we introduce Retrieval In-Context Optimization (RICO), a retrieval method that uses gradients from the LLM itself to learn the optimal mixture of documents for answer generation. Unlike traditional retrieval-augmented generation (RAG), which relies on external heuristics for document retrieval, our approach leverages direct feedback from the model. Theoretically, we show that standard top-$k$ retrieval with model gradients can approximate our optimization procedure, and provide connections to the leave-one-out loss. We demonstrate empirically that by minimizing an unsupervised loss objective in the form of question perplexity, we can achieve comparable retriever metric performance to BM25 with \emph{no finetuning}. Furthermore, when evaluated on quality of the final prediction, our method often outperforms fine-tuned dense retrievers such as E5.

Related papers

• SPaRFT: Self-Paced Reinforcement Fine-Tuning for Large Language Models [51.74498855100541]
  Large language models (LLMs) have shown strong reasoning capabilities when fine-tuned with reinforcement learning (RL)<n>We propose textbfSPaRFT, a self-paced learning framework that enables efficient learning based on the capability of the model being trained.
  arXiv  Detail & Related papers  (2025-08-07T03:50:48Z)
• Causal LLM Routing: End-to-End Regret Minimization from Observational Data [3.3580884064577616]
  LLM routing aims to select the most appropriate model for each query.<n>Prior approaches typically adopt a decoupled strategy, where the metrics are first predicted and the model is then selected based on these estimates.<n>We propose a causal end-to-end framework that learns routing policies by minimizing decision-making regret from observational data.
  arXiv  Detail & Related papers  (2025-05-21T21:34:18Z)
• Teaching Dense Retrieval Models to Specialize with Listwise Distillation and LLM Data Augmentation [43.81779293196647]
  We show that standard fine-tuning methods can unexpectedly degrade effectiveness rather than improve it, even for domain-specific scenarios.<n>We explore a training strategy that uses listwise distillation from a teacher cross-encoder, leveraging rich relevance signals to fine-tune the retriever.<n>Our results also reveal that synthetic queries can rival human-written queries in training utility.
  arXiv  Detail & Related papers  (2025-02-27T03:07:49Z)
• DUET: Optimizing Training Data Mixtures via Feedback from Unseen Evaluation Tasks [40.91931801667421]
  Our paper presents a novel global-to-local algorithm that interleaves influence function as a data selection method with Bayesian optimization to optimize data mixture via feedback from a specific unseen evaluation task.<n>By analyzing DUET's cumulative regret, we theoretically show that DUET converges to the optimal training data mixture for an unseen task even without any data knowledge of the task.
  arXiv  Detail & Related papers  (2025-02-01T01:52:32Z)
• What Do Learning Dynamics Reveal About Generalization in LLM Reasoning? [83.83230167222852]
  We find that a model's generalization behavior can be effectively characterized by a training metric we call pre-memorization train accuracy. By connecting a model's learning behavior to its generalization, pre-memorization train accuracy can guide targeted improvements to training strategies.
  arXiv  Detail & Related papers  (2024-11-12T09:52:40Z)
• GLoRe: When, Where, and How to Improve LLM Reasoning via Global and Local Refinements [18.939302106382094]
  State-of-the-art language models can exhibit impressive reasoning refinement capabilities on math, science or coding tasks. But even the best models struggle to identify textitwhen and where to refine without access to external feedback. We propose Stepwise ORMs which are trained, only on synthetic data, to approximate the expected future reward of the optimal policy.
  arXiv  Detail & Related papers  (2024-02-13T20:16:29Z)
• An Information Theoretic Approach to Machine Unlearning [43.423418819707784]
  To comply with AI and data regulations, the need to forget private or copyrighted information from trained machine learning models is increasingly important.<n>In this work, we address the zero-shot unlearning scenario, whereby an unlearning algorithm must be able to remove data given only a trained model and the data to be forgotten.<n>We derive a simple but principled zero-shot unlearning method based on the geometry of the model.
  arXiv  Detail & Related papers  (2024-02-02T13:33:30Z)
• Zero-shot Retrieval: Augmenting Pre-trained Models with Search Engines [83.65380507372483]
  Large pre-trained models can dramatically reduce the amount of task-specific data required to solve a problem, but they often fail to capture domain-specific nuances out of the box. This paper shows how to leverage recent advances in NLP and multi-modal learning to augment a pre-trained model with search engine retrieval.
  arXiv  Detail & Related papers  (2023-11-29T05:33:28Z)
• Model-based Offline Imitation Learning with Non-expert Data [7.615595533111191]
  We propose a scalable model-based offline imitation learning algorithmic framework that leverages datasets collected by both suboptimal and optimal policies. We show that the proposed method textitalways outperforms Behavioral Cloning in the low data regime on simulated continuous control domains.
  arXiv  Detail & Related papers  (2022-06-11T13:08:08Z)
• Imputation-Free Learning from Incomplete Observations [73.15386629370111]
  We introduce the importance of guided gradient descent (IGSGD) method to train inference from inputs containing missing values without imputation. We employ reinforcement learning (RL) to adjust the gradients used to train the models via back-propagation. Our imputation-free predictions outperform the traditional two-step imputation-based predictions using state-of-the-art imputation methods.
  arXiv  Detail & Related papers  (2021-07-05T12:44:39Z)
• S^3-Rec: Self-Supervised Learning for Sequential Recommendation with Mutual Information Maximization [104.87483578308526]
  We propose the model S3-Rec, which stands for Self-Supervised learning for Sequential Recommendation. For our task, we devise four auxiliary self-supervised objectives to learn the correlations among attribute, item, subsequence, and sequence. Extensive experiments conducted on six real-world datasets demonstrate the superiority of our proposed method over existing state-of-the-art methods.
  arXiv  Detail & Related papers  (2020-08-18T11:44:10Z)

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.