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

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.

• The paper introduces TTT-Embed, which distills reranker or LLM ranking rewards into a reusable residual query vector while keeping the encoder and document index unchanged.
• The paper finds that optimal reward sharing shifts from global to task-wise to query-wise as budgets increase, delivering gains of up to 8.36 nDCG@10 over raw retrieval and outperforming cost-matched reranking.
• The paper shows that task-wise adaptation transfers to unrewarded queries and unseen tasks, while recovering 8.00 MTEB points lost through domain fine-tuning without changing the specialized model’s weights.

TTT-Embed (Test-Time Tuning of Embeddings) is a framework for adapting dense retrieval systems at inference time by distilling external ranking rewards into a lightweight, reusable vector in the output embedding space of a frozen encoder (2608.12569). The work addresses a structural inefficiency in deployed retrieval pipelines: reranker or LLM-judge feedback is typically consumed once to reorder a single candidate list and then discarded, while the alternative—updating retriever weights—requires parameter access that is unavailable for closed-source APIs and computationally prohibitive for foundation-model-scale encoders such as Qwen3-Embedding (0.6B–8B) or the Gemini embedding family. TTT-Embed instead learns a residual query-side vector $v_g$ and retrieves with $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$, requiring only vector outputs from the embedding system and leaving both the encoder and the ANN document index untouched.

Method

The framework operates on scalar relevance scores $r_{qi}$ assigned by an external reward model (a cross-encoder reranker or LLM judge) to the top-K candidates $C_q$ returned by the base retriever for a subset $F \subseteq E$ of rewarded queries. The reward budget is normalized as $b = B/N$, the average number of scored query–document pairs per query in the workload. Learning proceeds by listwise knowledge distillation: a teacher distribution is formed by softmaxing reward scores at temperature $\tau_T$, a student distribution is induced by the adapted query embedding over the same candidates at temperature $\tau_S$, and $v_g$ minimizes a weighted sum of KL divergences with ridge regularization. The learned vector is applied through a second search over the unchanged index; the paper proves a ranking-equivalence proposition showing that scoring with the normalized corrected query induces exactly the same pairwise ordering as the additive score $x_q^\top z_d + \alpha_g v_g^\top z_d$.

The central design degree of freedom is the sharing scope $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$0: a single learned vector may be global-wise (pooled across all tasks, reused by every served query), task-wise (pooled within one retrieval task), or query-wise (private to one query). The learning rule is identical across scopes; only the group assignment of rewarded queries changes. Query weights are scope-balanced—averaging within each active task first, then across tasks uniformly—for the global state, preventing large tasks from dominating the pooled vector.

Deployment magnitude is handled by a tuning-free, evidence-adaptive shrinkage rule $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$1, where $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$2 is the number of rewarded queries contributing to the state. The rule is justified by a Gaussian normal-means model under which this factor is the Bayes estimator under squared error, with the zero-centered prior carrying the weight of one pseudo-query. The authors are explicit that this model is a post-hoc justification for uncertainty calibration, not an assumption used during optimization of $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$3 itself.

Main results

Across five embedding models (Qwen3-Embedding-0.6B and 4B, EmbeddingGemma-300M, and two API-only Gemini embedding models) and 15 MTEB English retrieval tasks comprising 12,263 queries, the optimal sharing scope shifts systematically with budget. On the discrete grid, global sharing is best at very low budgets ($$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$4), task-wise sharing dominates from $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$5 through $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$6, and query-wise optimization becomes optimal at $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$7 and above. Representative scope-optimal operating points show gains over raw retrieval of +1.23 nDCG@10 at $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$8 (global), +6.36 at $$\widetilde{x}_q = \operatorname{unit}(x_q + \alpha_g v_g)$$9 (task-wise), and +8.36 at $r_{qi}$0 (query-wise), with the scope-optimal configuration exceeding cost-matched direct reranking at all three points (+1.23, +3.15, and +0.56 respectively). Task-wise sharing also achieves the highest area under the gain–budget curve (AUBC) for all five models, making it the most robust choice when performance is integrated over the full log-budget range.

Two aspects of these results deserve emphasis. First, the advantage over direct reranking at matched reward budgets indicates that the gains are not merely a restatement of reranker quality: distilling reward scores into the query embedding enables retrieval of documents outside the originally scored candidate set. Second, the learned state demonstrably generalizes. With 80% of queries rewarded at depth 10, the task-wise state improves the untouched 20% of queries by 3.17–8.57 points (5.57 average); a global state learned on 14 tasks improves the held-out 15th task (leave-one-task-out) by 0.85–4.71 points (3.09 average), with positive LOTO gains for both closed Gemini models. This establishes that the vector encodes transferable relevance information rather than query-specific overfitting.

Recovery of degraded general capability

A notable application addresses the specialization–generalization trade-off. Fine-tuning Qwen3-Embedding-0.6B on SKILLRET (an agent-skill retrieval benchmark released after the training cutoffs of all models used, mitigating contamination risk) raises in-domain nDCG@10 from 52.24 to 66.21 but degrades the 15-task MTEB average from 49.90 to 45.42. With task-wise TTT-Embed at $r_{qi}$1, the frozen specialized model's MTEB score rises to 53.42—an 8.00-point recovery that exceeds the original unspecialized base model by 3.51 points and cost-matched direct reranking by 1.84 points—while the in-domain SKILLRET specialization is preserved exactly, since no weights change. This is a strong claim: test-time embedding correction can fully reverse, and surpass, the catastrophic forgetting induced by domain fine-tuning.

Deployment analysis

Three analyses bear on practical use. Breadth versus depth: at a small budget ($r_{qi}$2), depth-primary allocation (more judged documents per rewarded query) outperforms breadth-primary allocation across all scopes, while at $r_{qi}$3 breadth becomes more valuable, indicating diminishing returns to per-query reward depth once a critical mass of judged documents is reached. Reward-model selection: the retrieval-specialized Qwen3-Reranker-4B is the strongest reward source (+3.25 over raw retrieval when used to reorder the same candidate sets), outperforming Gemini 3.1 Pro (+2.73) and Gemini 3.5 Flash (+2.10); notably, two general-purpose Qwen instruct models underperform raw retrieval (−3.32 and −4.87), so not every capable LLM yields useful ranking rewards. Objective compatibility: the knowledge-distillation objective leads at both shared scopes (+4.00 global, +6.38 task-wise at $r_{qi}$4) against a reverse-KL variant and a Softmax-GRPO adaptation, though reverse KL is marginally better query-wise; the appendix shows reverse KL is exactly entropy-regularized reward maximization and that Softmax-GRPO coincides with distillation only in a local uniform-policy, small-reward-spread regime.

Limitations and open questions

The paper concedes several boundaries. The scope-transition thresholds (e.g., task-wise optimal through $r_{qi}$5, query-wise from $r_{qi}$6) are grid-specific and explicitly "should not be interpreted as universal thresholds," and model-level curves differ—Gemini Embedding 2 shows a negative global-wise gain (−1.54) at the low-budget operating point. The Bayesian shrinkage rule rests on a normal-means surrogate whose equal prior and noise scales ($r_{qi}$7) are assumed rather than estimated. Reward allocation is handled by an externally imposed task-balanced water-filling protocol; the framework itself offers no principled query- or document-selection policy, and scope selection in deployment is left manual, with the authors identifying strategic scope-selection mechanisms and complexity-based dynamic adaptation as open problems. Evaluation is limited to English MTEB retrieval tasks and one specialization benchmark, and the method presumes the deployer controls downstream vector search over returned embeddings—an assumption satisfied by API-based embedding systems but not by fully opaque retrieval services.

Conclusion

TTT-Embed reframes ranking rewards as persistent test-time state rather than transient per-query signal, realized as a residual vector in the output embedding space of a frozen encoder. Its principal empirical findings are a budget-dependent optimal sharing scope (global → task → query as rewards scale), consistent superiority over cost-matched direct reranking at the scope-optimal operating point, positive transfer to unseen queries and tasks, and complete test-time recovery of catastrophic forgetting without weight modification. The framework is compatible with closed-source embedding APIs and leaves the ANN index untouched, making budgeted relevance feedback deployable in systems where conventional test-time training is impossible.