NSIO: Neural Search Indexing Optimization
NSIO (Neural Search Indexing Optimization) investigates how to make the Differentiable Search Index (DSI), a model that memorises a document corpus and retrieves documents by generating their IDs, more accurate, efficient, and memory-friendly.
Background: Differentiable Search Index
Traditional search systems use two separate components: an indexer (inverted index or dense vector store) and a retrieval model. DSI collapses both into a single sequence-to-sequence transformer: given a query, it generates the document ID of the most relevant document directly, no separate index needed. This approach (Tay et al., 2022) shows strong results on MS MARCO but is expensive to fine-tune at scale.
Optimisations
Data Augmentation
- Num2Word: convert numeric tokens to their word forms, reducing out-of-vocabulary issues.
- Stopwords Removal: reduce noise in document representations.
- POS-MLM: Part-of-Speech guided masked language modelling to generate diverse query variants.
Parameter-Efficient Fine-Tuning (PEFT)
Avoid full model fine-tuning, adapt only a small number of parameters:
| Method | Description |
|---|---|
| LoRA | Low-Rank Adapters injected into attention weight matrices |
| QLoRA | LoRA with 4-bit quantised base model, drastically reduces GPU memory |
| AdaLoRA | Adaptive rank allocation, higher rank for more important layers |
| ConvoLoRA | Novel convolutional LoRA variant for capturing local patterns |
Evaluation
All experiments evaluated on the MS MARCO document ranking benchmark. Metrics: MRR@10, Recall@100. QLoRA achieves near full-fine-tune accuracy at a fraction of the memory footprint.
Technology
Python, Hugging Face Transformers + PEFT library, PyTorch, Jupyter Notebooks.