NUDGE: Lightweight Non-Parametric Embedding Fine-Tuning for Retrieval

Introduction

NUDGE is a lightweight, non-parametric tool designed for fine-tuning pre-trained embeddings, specifically for enhancing retrieval and RAG (Retrieval Augmented Generation) pipelines. Presented in the ICLR'25 paper "NUDGE: Lightweight Non-Parametric Embedding Fine-Tuning", this tool can significantly improve retrieval accuracy, often by over 10%, and runs in minutes.

Unlike traditional methods that modify model parameters, NUDGE operates by non-parametrically adjusting the data embeddings themselves. It solves a constrained optimization problem, moving data embeddings closer to the embeddings of training queries for which they are the ground-truth answers. The repository offers two variants, NUDGE-M and NUDGE-N, each with distinct optimization constraints.

NUDGE Overview

As illustrated above, NUDGE modifies data embeddings within a defined region to maximize similarity with training queries.

Installation

To get started with NUDGE, simply install it using pip:

pip install nudge-ft

Examples

NUDGE operates directly on pre-computed embeddings. You need to have your documents and training/validation queries already embedded, along with ground-truth answers for the queries.

Here's a basic workflow:

from nudge import NUDGEN, NUDGEM

train_set = {'q_embs':train_q_embs, 'q_ans_indx':train_q_ans_indx}
val_set = {'q_embs':val_q_embs, 'q_ans_indx':val_q_ans_indx}

finetuned_embs_nudge_n = NUDGEN().finetune_embeddings(data_embs, train_set, val_set)
finetuned_embs_nudge_m = NUDGEM().finetune_embeddings(data_embs, train_set, val_set)

For a complete end-to-end example, you can fine-tune embeddings on the nfcorpus dataset. This involves embedding data and queries using sentence_transformers and then applying NUDGE. A detailed example is available in the repository's notebook or can be run via python example.py.

# Install dependencies
pip install sentence_transformers datasets

# Load dataset and embed
from util.utils import load_hf_datasets, embed_data_and_query_sets
dataset_name = 'nfcorpus'
dataset, query_sets = load_hf_datasets(dataset_name)
data_emb, query_sets = embed_data_and_query_sets(dataset, query_sets, "BAAI/bge-small-en-v1.5")

# Fine-tune Embeddings
from nudge import NUDGEN
finetuned_embs_nudge_n = NUDGEN().finetune_embeddings(data_emb, query_sets['train'], query_sets['dev'])

# Use fine-tuned embeddings for retrieval
from util.knnretriever import kNNRetriever
nudge_n_res = kNNRetriever(finetuned_embs_nudge_n).retrieve_topk_from_emb_batch(k=10, q_embeds=query_sets['test']['q_embs'])

# Use non-fine-tuned embeddings to answer queries (for comparison)
no_ft_res = kNNRetriever(data_emb).retrieve_topk_from_emb_batch(k=10, q_embeds=query_sets['test']['q_embs'])

# Compare accuracy
from util.utils import calc_metrics_batch
metrics = [('recall',10), ('ndcg',10)]
no_ft_accs = calc_metrics_batch(metrics, no_ft_res, query_sets['test']['q_ans_indx'], query_sets['test']['q_ans_indx_rel'])
nudgen_accs = calc_metrics_batch(metrics, nudge_n_res, query_sets['test']['q_ans_indx'], query_sets['test']['q_ans_indx_rel'])
print(f"No Fine-Tuning {metrics[0][0]}@{metrics[0][1]}: {no_ft_accs[0]*100:.1f}, {metrics[1][0]}@{metrics[1][1]}: {no_ft_accs[1]*100:.1f}")
print(f"NUDGE-N {metrics[0][0]}@{metrics[0][1]}: {nudgen_accs[0]*100:.1f}, {metrics[1][0]}@{metrics[1][1]}: {nudgen_accs[1]*100:.1f}")

This example typically shows a significant improvement in metrics like recall and nDCG. For larger datasets, NUDGE also provides an optimization to reduce memory usage by filtering out data records not relevant to training or validation queries. An example for larger datasets is available in this notebook.

Why use NUDGE?

Efficiency: It's a lightweight tool that runs in minutes, making it highly efficient for fine-tuning.
Performance Boost: It consistently improves retrieval accuracy, often by over 10%, without altering the underlying embedding model.
Non-Parametric Approach: By directly adjusting data embeddings, NUDGE avoids the complexities and computational costs associated with fine-tuning large language models.
Ease of Integration: NUDGE seamlessly integrates into existing embedding-based retrieval pipelines, requiring only pre-computed embeddings and ground-truth answers.
Scalability: Includes optimizations for handling larger datasets, ensuring efficient memory usage.