NUDGE: Lightweight Non-Parametric Fine-Tuning of Embeddings for Retrieval

Comment [W1]. "Fine-tuning the embedding directly makes the database difficult to update. Specifically, if new photos are added to the database, their embeddings must also be trained. This requires relabeling (or pseudo labeling) at least a few queries whose GT answers include the new DB images, which incurs additional cost (or the new DB embedding will not be updated). This weakness, however, does not appear in previous two lines of works, ie. fine-tuning the foundation model, or training a new lightweight model to refine the embeddings. Consequently, this limitation somewhat impacts the practicality of NUDGE."

Response. Thank you for your comment. We have provided a detailed discussion on generalization in our response to Q1. Here, we point out that there are many applications where the set of documents to be queried is relatively fixed, which creates a setting particularly well-suited for NUDGE. Some examples include a corpus of hospital or medical regulations, laws passed by the US Congress, the documentation for software, and textbooks on a specific topic. Journalists, legal experts, software users, and students frequently query such fixed corpora, and improving retrieval accuracy is beneficial in all such settings. Moreover, even if the corpus set changes, NUDGE is useful as long as the query distribution does not change. A journalist studying historical crime rates in a city still benefits from better retrieval of court cases in the past, even if new court cases are added to the corpus. In fact, NUDGE has already been incorporated into a popular open-source library for RAG applications, showing the interest to include NUDGE in real-world applications for better retrieval.

Comment [W2]. "NUDGE is unable to refine query embeddings. It has been well established that query embedding refinement can significantly enhance retrieval performance (as evidenced by query expansion and its variants). Both fine-tuning the foundation model and training an embedding update model are capable of updating query embeddings. However, NUDGE cannot accomplish this because the GT label for the input query is unknown, making it impossible to train. This limitation leaves considerable room for improvement."

Response. We thank the reviewer for pointing out query expansion techniques. Indeed, since NUDGE only modifies data embeddings, it is orthogonal to query rewriting and expansion techniques, and can be used together with such approaches out-of-the-box. To show this, we have added a new experiment, combining NUDGE with HyDE, a popular query rewriting technique. This experiment is included in the newly added Appendix E.8. The results in Table 16 show that incorporating HyDE within NUDGE improves the accuracy over using NUDGE alone (although NUDGE alone is more accurate compared with using HyDE without NUDGE). This shows that, as the reviewer suggests, NUDGE can be combined with existing query rewriting methods to improve accuracy.

Moreover, we also note that, as discussed in our response to Q1, NUDGE can be performed on top of embeddings created by PTFT and Adaptor, allowing NUDGE to combine query embedding refinements done by PTFT and Adaptor with improved data embeddings created by NUDGE.

Comment [W3]. "The generalizability of the proposed method is questionable. Intuitively, embedding update methods are likely to suffer from significant overfitting issues. The authors have invested considerable effort to address this, including constraining the Δ and normalizing the updated embeddings (either pre- or post-update). While NUDGE-N ultimately achieves acceptable results for OOD data, overall confidence in its generalizability remains limited. For OOD scenarios, NUDGE-N shows only minor improvements on the BGE-S benchmark, with many other benchmarks left untested. When the embedding is not normalized (i.e., NUDGE-M), the generalizability appears notably poor. While it's plausible that the transition from NUDGE-M to NUDGE-N could yield such significant improvements, additional experiments elaborating on generalizability would be beneficial, as this might constitute a major drawback of NUDGE. Nevertheless, if NUDGE can perform adequately across most OOD benchmarks, or even in its current form, it may still be considered acceptable given the other strengths of this work."

Response. Thank you for your comment. Our response to Q1 and Q2 provides a detailed discussion on generalization. Here, we additionally note that we have also included Table 18 in Appendix E.9, which presents generalization results using the GTE-L model. The table shows that similar observations on the BGE-S model also hold for GTE-L.

Comment [W4]. "Nit: The authors should follow the ICLR 2025 template guidelines by placing table captions above the tables."

Response. Thanks for pointing this out. We have moved the captions to above the tables.

Question Q2. "For out-of-sample results, it would be interesting to see what would happen if an actual feature transformation model was trained on top of the embeddings found by this method."

Response. Based on the reviewer’s suggestion, and as discussed in our response to Q1, we have now included in Tables 6 and 7 results that show how NUDGE can be combined with both Adaptor (which learns a linear feature transformation) and PTFT. Please see our response to Q1 for a discussion of these results, which shows that combining NUDGE with Adaptor and PTFT allows for improving accuracy both on in-distribution samples through NUDGE, and on out-of-distribution samples by inheriting the generalization ability of Adaptor and PTFT.

Here, we note that the reviewer seems to suggest performing Adaptor after using NUDGE on the embeddings, referred to NUDGE-N+A (instead of A+NUDGE-N which performs NUDGE after transforming embeddings by Adaptor). In fact, our experiment showed that NUDGE-N+A performs identical to NUDGE-N (we have omitted this result for brevity). We suspect this is because NUDGE already changes the embeddings until a point where further modifications worsen accuracy. In contrast, in A+NUDGE-N, there is still room for improvement after applying Adaptor, so performing NUDGE on top of Adaptor is beneficial.

Comment W1. "The "elephant in the room" for this method is that it does not work on out-of-sample documents. Since the embeddings of the database are adjusted at fine-tuning time, new documents cannot benefit from it, unless new training data is added for them as well, which becomes a really narrow application setting. This is all the more apparent since out-of-distribution on the queries side is considered in this work."

Response. Thank you for the comment, please see our response to Q1 where we discuss in detail the points raised.

Comment W2. "The method works only on densely annotated training data, since a document that is never returned to a query will not be adapted. This is apparent in the KILT dataset where documents that do not appear in any result are simply removed."

Response. Thank you for the comment, please see our response to Q1 where we discuss in detail the points raised, especially with the added experiment evaluating the approach’s accuracy on unlabeled documents, as well as new experiments showing that NUDGE can be combined with parametric methods if generalization is needed in an application.

Here, we additionally note that although it is correct that NUDGE does not modify embeddings for which there is no ground-truth in the training set, it is not necessary that all of the corpus be annotated to improve accuracy. This is because, in practice, query labels are often not uniformly distributed, but instead are concentrated on a smaller subset of the dataset. For instance, consider the benchmark datasets NQ, NFCorpus and SciFact that are created from real-world queries. As the newly added Fig. 7 shows, a significant portion (around 40% or more) of queries on these datasets ask for data records accessed by other queries. Improving retrieval accuracy for those embeddings alone is sufficient to increase retrieval accuracy in general, as evidenced by the benefits of NUDGE on these datasets.

Comment W3. "Because of this particular setting, the queries in experiments must be split in train+val+test, which is unusual. Therefore there are no comparisons with state-of-the-art methods."

Response. We are unsure what state-of-the-art methods the reviewer is referring to, but we would be happy to compare with them given the chance. Our method is a fine-tuning approach for which train+val+test split is standard (where validation set is used for hyperparameter tuning). For example, many of the BEIR and KILT datasets already come with a train+val+test split. In case the reviewer is referring to recent query rewriting methods, we note that NUDGE is indeed orthogonal to such methods, and can be used together with them. We have added a new experiment showing how NUDGE can be combined with HyDE, a popular query rewriting technique that rewrites the queries in terms of hypothetical documents that answer the queries. The new experiment is presented in Appendix E.8, and shows that combining NUDGE with HyDE performs better than using NUDGE alone (which still performs significantly better than HyDE alone).

Comment W2. Lack a detailed ablation study that examines how skewed distributions impact both in-distribution and out-of-distribution performance could reveal if other normalization techniques or dataset-specific adjustments would better manage this overfitting.

Response. We assume that by skewed distribution the reviewer is referring to the skewness in the distribution of data labels. Indeed, our experiments already consider various label distributions with different skewness. We have now added Fig. 7 which depicts the label distribution in the datasets used in our experiments, as well as a discussion regarding the impact of skewness in label distribution on our results in Sec. 4.1 with details in Appendix E.12. Our newly added discussion reveals an interesting relationship between the impact of normalization and skewness in label distribution. We also present the discussion below.

Summary of added discussion in Sec 4.1 and Appendix E.12. We consider the relationship between skewness in label distribution and normalization to understand the difference in accuracy between NUDGE-M and NUDGE-N and the need for normalization. Comparing accuracy results in Table 4-5 and label distribution in Fig. 7, we see that the impact of normalization depends on whether label distribution is heavy-tailed (leading to many test labels not being present in the training set) or if it is skewed towards a fixed set of labels (in which test labels are likely seen during training). To understand why, first, note that if a data record is not seen during training (i.e., it's not a ground truth to any training query), its embedding will not be modified by NUDGE-M. Meanwhile, NUDGE-M will likely increase the magnitude of the data embeddings that are represented in the training set. As a result, the data embeddings seen during training will have a larger magnitude than data embeddings not seen during training. Smaller magnitude causes the inner product similarity of such unseen data embeddings with test queries to be generally lower, which is detrimental if many test queries ask for data embeddings whose ground truth was not seen during training. Thus, heavy-tailed distributions benefit from normalization, while distributions skewed towards a fixed set of labels are less impacted by it.

Using the above argument, the label distribution can provide an explanation for the gap between NUDGE-M and NUDGE-N in SciFact, TriviaQA, NQ, and HotpotQA, where, as Fig. 7 shows, the ground truth for between 40-60% of test queries was unseen during training. It can also provide an explanation for why NUDGE-M performs similarly to NUDGE-N on CoCo, Flicker, and Fever datasets, where, as shown in Fig. 7, almost no test query asks for data records not seen during training. Moreover, our Out-of-Distribution experimental results (Sec. 4.2) can also be explained using the label distribution, where out-of-distribution test queries ask for labels unseen during training, leading to poor accuracy for unnormalized embeddings. We also see that NUDGE-M performs worse than NUDGE-N in that setting. Overall, this suggests that NUDGE-M (with unnormalized embeddings) has poor generalization outside of the training set, and normalization, as done in NUDGE-N, helps solve this problem.

Comment W5. "Table 3’s results are overly generalized under “across text datasets,” obscuring which datasets contributed to the reported metrics."

Response. The paper already contains the detailed per dataset results for all the datasets in Table 3; specifically, Table 5 shows per dataset results for BGE-S, Table 9 shows per dataset results for GTE-L and Table 10 shows per dataset result for TE3-L. In fact, the reviewer already refers to such per dataset results in their comment below, indicating that per dataset results are clearly presented in the paper.

Comment W6. "In BEIR benchmark results (Table 5 and Table 9), the first column shows identical performance across all methods, which is counterintuitive and suggests either data or methodology issues."

Response. First, we note that due to the use of a validation set, a method that performs worse than No Fine-Tuning on a dataset will default to No Fine-Tuning. For example, as Fig. 2 shows, for ArguAna, NUDGE-M does worse than No Fine-Tuning for all values of $\gamma>0$. Thus, the approach sets $\gamma=0$, resulting in a method identical to No Fine-Tuning. This means that different approaches do in fact perform differently, but since we use a validation set, a performance worse than No Fine-Tuning will not be reported for any approach. This is a fair comparison, as in practice, users can default to No Fine-Tuning, if a fine-tuning method does not improve accuracy.

Consequently, the results in Table 5 and Table 9 show that no approach among NUDGE-M, NUDGE-N, PTFT and Adaptor improves over No Fine-Tuning. Interestingly, based on the Reviewer j9KJ’s suggestion, we have included the method NUDGE-S (see Appendix C.1 for details of the approach and Appendix E.6 for experiments), which performs sparse data updates for fine-tuning, and indeed performs better than No Fine-Tuning on ArguAna. This suggests that sparse updates are more suitable for that dataset, while NUDGE-M, NUDGE-N, PTFT and Adaptor perform dense updates. Overall, we have included ArguAna as an example of a dataset where no standard approach performs well, presenting a diverse set of datasets that thoroughly evaluate different approaches.

Comment W7. "The GPU speedup in Table 6 is marginal compared to the CPU, which raises efficiency concerns, as an A100 GPU should provide significantly higher speedup for such operations relative to CPU processing."

Response. First, we emphasize that the goal of Table 6 (Table 8 in the revised paper) is to show that NUDGE runs within a few minutes on both GPU and CPU, while Adaptor, although fast on GPU, takes more than an hour on CPU, and PTFT takes hours on GPU. The focus of our paper has not been on providing a GPU-optimized implementation of any of the algorithms, and the runtime analysis is for a standard implementation of all algorithms. For PTFT and Adaptors, the implementation closely follows that of the SentenceTransformers and LlamaIndex libraries, respectively.

Having said that, we disagree with the reviewer’s suggestion that just because an A100 GPU has significantly more compute capacity, it can provide significantly higher speed-ups. In fact, the speed-ups the GPU can provide does not depend only on its capacity, but on how much of the capacity can be used to perform the operations (which depends on how much parallelism is possible for the specific operations). For the experiments in the paper where runtime is presented, the embedding dimension is 384, so most operations involve matrices where each row is 384 dimensional. A dimensionality of 384 is much smaller than the maximum capacity of A100 GPUs and processing such matrices does not utilize the full capacity of the GPU. This leads to lower speed-ups than using A100 GPUs on larger matrices. To show this, we performed a micro-benchmark to study the effect of matrix dimension on runtime difference between GPU and CPU. The newly added Fig. 6 in Appendix E.10 in the paper shows the runtime of GPU and CPU when performing 13 consecutive matrix addition and multiplication operations on matrices of dimensions $d\times d$, varying $d$. We see that at $d=384$, GPU is 2.3x faster, but at $d=3072$ the GPU is 1600x faster. To summarize, for lightweight methods such as NUDGE and Adaptor which involve small matrix operations, speed-ups of using an A100 GPU are not proportional to its capacity.

Thanks for the response,

For W1 and W2, I am questioning the necessity of NP-hard modeling.

Inner product similarity-based search (MIPS) does not inherently require NP-hardness, as gradient-based approaches are sufficient to optimize for inner product similarity. Could you clarify why NP-hard modeling is necessary for addressing MaxA-EFT?

For W4, I am asking clarification of the notation. If $\Delta_i$ represents the $i$-th row of $\Delta$, what does $\Delta_{Y_i}$ in Equations (2) and (32) refer to, given that $Y$ is the set of ground-truth labels? The authors could have resolved this ambiguity with just a few wording adjustments, but they chose not to do so.

Comment "Thank you for reiterating the novelty … that is as good as the optimal one."

Response. The reviewer’s comment conflates the multiple optimization problems studied in the paper. Specifically, in the paper, we study (1) MaxA-EFT optimization problem, which is the unconstrained optimization problem of fine-tuning embeddings while maximizing accuracy and (2) the constrained variations of the optimization problem, BiMax-M and BiMax-N.

(1) MaxA-EFT is the unconstrained optimization problem of fine-tuninig embeddings to maximize retrieval accuracy. It is a natural formulation of the problem of fine-tuning embeddings to maximize accuracy. Given that fine-tuning embeddings to maximize retrieval accuracy is the focus of the paper, MaxA-EFT is a problem that should be solved in the paper. However, because it is NP-hard, it cannot be solved optimally. Thus, instead of solving MaxA-EFT directly, we solve tractable variations of the problem. In other words, because MaxA-EFT is NP-hard, we don’t solve it directly but instead formulate tractable variants of the problem. Clearly, NP-hardness is necessary to motivate this discussion in the paper.

(2) The paper specifically formulates two tractable variants of the optimization problem BiMax-M and BiMax-N, both of which are constrained optimization problems. We provide a thorough theoretical and empirical analysis of how these problems can be solved, presenting both closed-form and iterative gradient-based solutions. As the reviewer points out and already discussed in Appendix C.2, the gradient-based solutions can be used to solve BiMax-M and BiMax-N. Appendix C.2 specifically presents NUDGE-IM and NUDGE-IN, two variants that use iterative gradient-based optimization, and we theoretically show that NUDGE-IM is equivalent to NUDGE-M under the correct hyper-parameter setting (although suboptimal hyper-parameter setting leads to suboptimally). Nonetheless, iterative variants are less efficient compared with closed-form solutions, since their iterative nature means they need to go over the training/validation set multiple times. Comparing iterative variants with closed-form solutions, Appendix C2 already states, “These approaches can be less efficient and suboptimal, but can still provide accurate solutions, and are simple and flexible.” Overall, we have provided a thorough study of how the constrained optimization problems can be solved, and we leave it to the practitioners to decide if simplicity of solution is preferred in practice, or optimality and efficiency.

Comment "The authors should also avoid … is defined as the label set."

Response. The reviewer seems to have misunderstood the notation. In our notation, $Y$ is the label set, and $Y_i$ is the $i$-th element of the label set. In a retrieval setting, the ground-truth answer (i.e., the label) for a query is the index of the data record that matches the query.

In case the reviewer is implying that using $Y$ as a label set is uncommon, we note that the contrary is true, and using $Y$ to refer to the label set is very common in machine learning, please see the Deep Learning book by Goodfellow et al.

First, if by “necessary” the reviewer means the experimental results presented in the paper do not depend on NP-hardness of MaxA-EFT, then we agree with the reviewer that NP-hardness is not “necessary”. If a reader is only interested in our proposed algorithms and their practical utility, then they can safely ignore our NP-hardness result and only focus on other relevant sections. We believe our experimental results are impressive as they are and our methodology is novel and interesting; so the paper already makes significant contributions to improve retrieval accuracy through fine-tuning for practical applications.

However, the reviewer seems to question the need for NP-hardness altogether and suggests studying NP-hardness is not relevant to machine learning problems. We note that the contrary is true, and studying NP-hardness in various settings is, and has been for decades, an active area of research (e.g., Brand et al 2024 NeurIPS, Froese & Hertrich 2023 NuerIPS, Shamir 2018 JMLR and Blum & Rivest 1988 NeurIPS below and the references therein). In fact, training neural networks, in general, is known to be NP-hard (Blum & Rivest 1988), meaning that when using parametric neural network-based methods to train classification or retrieval models, NP-hardness is already known. This explains why no new NP-hardness result is presented in papers when using parametric methods.

Moreover, we also disagree with the reviewer's suggestion that just because one problem that maximizes validation accuracy is NP-hard, other related problems also become NP-hard. This is already apparent from the paper, since all three of MaxA-EFT, BiMax-M and BiMax-N maximize validation accuracy, but the latter two are not NP-hard and can be solved optimally in polynomial time. Consequently, NP-hardness must be studied specifically for the problem at hand, and unfounded extrapolations can lead to incorrect conclusions.

Overall, MaxA-EFT is a novel problem (thus unknown if it is NP-hard) and solving it is the central focus of this paper. We believe the NP-Hardness result is relevant and useful because it allows us to know that an exact solution to MaxA-EFT cannot be obtained and justifies the use of our alternative problem formulations to solve the fine-tuning problem. Moreover, the NP-Hardness result also contributes to our general understanding of the non-parametric fine-tuning problem (e.g., future work now knows the solution cannot be improved by solving MaxA-EFT optimally). Nonetheless, if a reader is only interested in practical utility and not NP-Hardness of the problem, they are free to ignore that portion of the paper.

Brand, Cornelius, Robert Ganian, and Mathis Rocton. "New complexity-theoretic frontiers of tractability for neural network training." Advances in Neural Information Processing Systems 36 (2024).

Froese, Vincent, and Christoph Hertrich. "Training neural networks is np-hard in fixed dimension." Advances in Neural Information Processing Systems 36 (2024).

Shamir, Ohad. "Distribution-specific hardness of learning neural networks." Journal of Machine Learning Research 19.32 (2018): 1-29.

Blum, Avrim, and Ronald Rivest. "Training a 3-node neural network is NP-complete." Advances in neural information processing systems 1 (1988).