NUDGE: Lightweight Non-Parametric Fine-Tuning of Embeddings for Retrieval (2409.02343v1)

Abstract: $k$-Nearest Neighbor search on dense vector embeddings ($k$-NN retrieval) from pre-trained embedding models is the predominant retrieval method for text and images, as well as Retrieval-Augmented Generation (RAG) pipelines. In practice, application developers often fine-tune the embeddings to improve their accuracy on the dataset and query workload in hand. Existing approaches either fine-tune the pre-trained model itself or, more efficiently, but at the cost of accuracy, train adaptor models to transform the output of the pre-trained model. We present NUDGE, a family of novel non-parametric embedding fine-tuning approaches that are significantly more accurate and efficient than both sets of existing approaches. NUDGE directly modifies the embeddings of data records to maximize the accuracy of $k$-NN retrieval. We present a thorough theoretical and experimental study of NUDGE's non-parametric approach. We show that even though the underlying problem is NP-Hard, constrained variations can be solved efficiently. These constraints additionally ensure that the changes to the embeddings are modest, avoiding large distortions to the semantics learned during pre-training. In experiments across five pre-trained models and nine standard text and image retrieval datasets, NUDGE runs in minutes and often improves NDCG@10 by more than 10% over existing fine-tuning methods. On average, NUDGE provides 3.3x and 4.3x higher increase in accuracy and runs 200x and 3x faster, respectively, over fine-tuning the pre-trained model and training adaptors.

• The paper introduces NUDGE, a non-parametric approach that directly optimizes pre-trained embeddings to significantly improve k-NN retrieval.
• It presents two methods, NUDGE-M and NUDGE-N, with closed-form solutions that efficiently address NP-Hard challenges while preserving embedding semantics.
• Empirical results across text and image datasets demonstrate up to a 14.3% boost in NDCG@10 and dramatically faster fine-tuning compared to traditional methods.

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

The paper "NUDGE: Lightweight Non-Parametric Fine-Tuning of Embeddings for Retrieval" introduces a novel method for fine-tuning pre-trained embeddings, specifically to enhance $k$-Nearest Neighbor ($k$-NN) retrieval. The presented approach, NUDGE (Non-parametric Updater for Dense Graph Embeddings), offers a non-parametric method that optimizes embeddings directly, rather than relying on traditional parametric methods like fine-tuning model parameters or employing adaptor models. This essay will summarize the paper's contributions, methodology, results, and implications for future AI developments.

Contributions and Methodology

The paper's main contributions are:

1. Formalization and Theoretical Analysis:
   • The authors formalize the problem of non-parametric embedding fine-tuning.
   • They establish that the maximum accuracy embedding fine-tuning problem (MaxA-EFT) is NP-Hard.
2. Introduction of NUDGE Methods:
   • NUDGE-M and NUDGE-N methods are introduced, providing closed-form, efficient solutions to constrained optimization problems that are variations of MaxA-EFT.
3. Extensive Empirical Evaluation:
   • The experiments conducted span five pre-trained models and nine standardized retrieval datasets, showcasing significant improvements in retrieval metrics.

Optimization Formulations and Methodologies

MaxA-EFT and MaxS-EFT:

• MaxA-EFT directly aims to maximize the number of correctly answered queries but is proven to be NP-Hard.
• MaxS-EFT is an alternative surrogate optimization problem focused on maximizing the similarity between query and data embeddings, but it is unbounded without proper constraints.

NUDGE-M and NUDGE-N:

• NUDGE-M:
  • Solves a variation of MaxS-EFT with constraints on the magnitude of changes made to embeddings.
  • Utilizes a bounded optimization approach to avoid overfitting, ensuring efficient updates via closed-form solutions.
• NUDGE-N:
  • Further constrains embeddings to remain normalized.
  • These constraints help maintain the semantic integrity of embeddings during fine-tuning, providing robustness against overfitting and enhancing out-of-distribution generalization.

Experimental Results

NUDGE methods significantly outperform existing parametric methods like fine-tuning pre-trained models (PTFT) and training adaptors:

• Text Retrieval:
  • Across models like BGE-S, GTE-L, and TE3-L, NUDGE-M and NUDGE-N provide significant improvements in NDCG@10, sometimes boosting accuracy by up to 14.3% over no fine-tuning and outstripping PTFT and adaptors by up to 10%.
  • For instance, BGE-S saw an average NDCG@10 improvement of 8.4% for NUDGE-M compared to minor improvements for PTFT (3.8%) and adaptors (2.9%).
• Image Retrieval:
  • With CLIP-B and CLIP-L embeddings, NUDGE methods yield up to 14.3% improvement in NDCG@10 over no fine-tuning.
  • Especially notable was NUDGE-N, which consistently outperformed both the non-fine-tuned baseline and adaptors, demonstrating robustness across varying query workloads.

Efficiency

NUDGE's computational efficiency was highlighted as a significant advantage:

• Fine-tuning times are dramatically lower compared to PTFT and adaptors. NUDGE methods generally execute within minutes, while PTFT may require hours.
• On a practical dataset, BGE-S's fine-tuning using NUDGE took approximately 1-2 minutes on GPU, compared to 447 minutes for PTFT.

Implications and Future Directions

Practical Implications:

• NUDGE presents a valuable approach for applications requiring rapid and resource-light fine-tuning. By avoiding extensive model re-training and deployment costs associated with traditional methods, NUDGE offers a scalable solution suitable for environments with limited computational resources but stringent accuracy requirements.

Theoretical Implications:

• The exploration of non-parametric methods paves the way for future research into direct optimization techniques. This shift from adjusting model parameters to directly modifying data embeddings could be beneficial in various domains where pre-trained models are leveraged for specific tasks.

Future Developments:

• Integration into vector databases could provide seamless deployment of NUDGE, offering near real-time accuracy improvements for retrieval systems.
• Extending the framework to other AI subsystems, such as recommender systems, may be a logical progression, given the shared reliance on embedding-based representations.

Speculation on AI's Future Developments:

• Non-parametric fine-tuning may align well with the trend towards modular and component-based AI systems. The ability to fine-tune components without extensive re-training could facilitate more efficient and dynamic AI architectures.

In conclusion, NUDGE represents a significant advancement in the field of retrieval-oriented fine-tuning techniques. Its non-parametric nature, alongside impressive efficiency and accuracy gains, positions it as a compelling alternative to traditional fine-tuning methods, with broad future applications across AI-driven systems.