Test-Time Optimization of Query Embeddings with Ranking Aware Reward Maximization

2026-08-12

AI Summary

This work addresses the challenge of improving retrieval performance without modifying frozen dense retriever weights, access to ground-truth labels, or index reconstruction. The authors propose TTT-Embed, a framework that learns lightweight vectors within the frozen embedding space by optimizing query embeddings at test time using only scalar ranking rewards. It introduces a multi-granularity reward reuse mechanism—spanning global, task-level, and query-level—to dynamically balance reward generalization and specificity. Notably, this approach is the first to transform ranking rewards into reusable test-time states, enabling efficient adaptation of closed-source models. Evaluated across five embedding models and fifteen MTEB tasks, TTT-Embed achieves up to +8.36 nDCG@10 improvement, with gains of +8.57 and +4.71 nDCG@10 on unseen queries and tasks, respectively, while effectively mitigating catastrophic forgetting and even recovering or surpassing original model performance.

Abstract

Dense retrievers rank documents using vector similarity between a frozen encoder and a precomputed index. While test-time ranking rewards from a reranker or LLM judge can improve results, existing methods discard this signal after a single query. Updating the retriever's weights makes rewards reusable, but this requires parameter access, which is unavailable for closed-source models, and is computationally prohibitive. We propose TTT-Embed (Test-Time Tuning of Embeddings), a framework that distills ranking rewards into a lightweight, learned vector within the output embedding space of a frozen model. This vector is optimized purely from scalar ranking scores assigned to the retriever's own candidate documents, requiring no access to model weights, ground-truth labels, or modifications to index. A single scope parameter controls rewards reuse (global, task, or query), enabling a principled trade-off between reusability and specificity under a fixed reward computation budget. We demonstrate that as the available reward budget scales, the optimal sharing scope shifts dynamically from global-wise to task-wise and finally to query-wise. Evaluated across five embedding models and 15 MTEB retrieval tasks, TTT-Embed improves test-time retrieval by up to +8.36 nDCG@10. Crucially, the learned states generalize effectively to unseen queries (up to +8.57 nDCG@10) and unseen tasks (up to +4.71 nDCG@10). Furthermore, TTT-Embed successfully resolves catastrophic forgetting: by leaving base weights entirely frozen, it recovers degraded general capabilities (up to +8.00 nDCG@10, even surpassing the original base model) while preserving in-domain specialization. These results establish ranking rewards as a reusable test-time state, enabling budget-efficient adaptation for any embedding model, including closed-source APIs.

Problem

• dense retrieval
• test-time adaptation
• ranking rewards
• embedding models
• closed-source models

Innovation

• test-time optimization
• query embedding
• ranking reward
• frozen retriever
• reusable adaptation