[Paper Review] $k$NN-Adapter: Efficient Domain Adaptation for Black-Box Language Models

Yangsibo Huang, Daogao Liu+3 more|arXiv (Cornell University)|Feb 21, 2023

Topic Modeling4 citations

TL;DR

kNN-Adapter is a lightweight, parameter-efficient method for adapting black-box large language models to new domains without fine-tuning. It adaptively learns token-wise interpolation weights between a pretrained language model and a retrieval-augmented datastore, achieving significant improvements in perplexity—especially in low-data and limited-access settings—outperforming both vanilla kNN-LM and fine-tuning in few-shot scenarios.

ABSTRACT

Fine-tuning a language model on a new domain is standard practice for domain adaptation. However, it can be infeasible when it comes to modern large-scale language models such as GPT-3, which can only be accessed through APIs, making it difficult to access the internal parameters of the model. In this paper, we propose $k$NN-Adapter, a method to effectively adapt these black-box large language models (LLMs) to a new domain. The $k$NN-Adapter builds on top of the retrieval-augmented language model, and adaptively learns to interpolate the output of the language model with retrieval results from a datastore consisting of the target domain data. Our experiments on four different domains demonstrate that $k$NN-Adapter significantly improves perplexity, and works particularly well in settings with limited access to LLMs. Additionally, we show that $k$NN-Adapter is more effective than fine-tuning when the amount of training data is limited. We also release a dataset to encourage further study.

Motivation & Objective

• To address the challenge of adapting large language models (LLMs) to new domains when fine-tuning is infeasible due to black-box API access.
• To improve the performance of retrieval-augmented language models like kNN-LM by learning dynamic, context- and token-specific interpolation coefficients.
• To enable effective domain adaptation in low-data regimes and under restricted model access, such as when only top-k token probabilities are available.
• To provide a more efficient alternative to full fine-tuning, especially when model parameters are inaccessible and computational resources are limited.
• To gain insight into the mechanisms that make retrieval-augmented models effective through analysis of learned interpolation weights.

Proposed method

• kNN-Adapter builds on kNN-LM by introducing a learnable, token-wise interpolation coefficient λ that dynamically combines the language model’s output distribution p_LM with the k-nearest neighbor retrieval distribution p_kNN.
• The interpolation coefficient λ is predicted by a small, trainable adapter network that takes as input the current context, the token to be predicted, and the retrieved k-nearest neighbors from a domain-specific datastore.
• The method uses a learnable temperature parameter for the retrieval distribution to further refine the interpolation, improving calibration and performance.
• The adapter is trained end-to-end using cross-entropy loss on a target domain dataset, with no updates to the frozen LLM weights.
• The approach maintains the efficiency of kNN-LM while enabling adaptive, context-aware fusion of retrieval-based knowledge with the LLM’s generalization capacity.
• The method is compatible with any black-box LLM that supports token probability queries and can be applied without architectural changes.

Experimental results

Research questions

• RQ1 Can a learnable, token-wise interpolation mechanism improve the performance of retrieval-augmented language models in domain adaptation tasks?
• RQ2 How does kNN-Adapter compare to fine-tuning when only limited data is available for the target domain?
• RQ3 Does kNN-Adapter maintain high performance under restricted access to the LLM, such as when only top-k token probabilities are returned?
• RQ4 What insights can be gained from analyzing the adaptive interpolation weights λ across different token types and domains?
• RQ5 Can kNN-Adapter outperform both vanilla kNN-LM and fine-tuned models in low-resource settings?

Key findings

• kNN-Adapter reduces perplexity by approximately 2 points across four different domains compared to the baseline kNN-LM, with only a marginal increase in inference cost.
• The method significantly outperforms fine-tuning when fewer than 10^4 training tokens are available, demonstrating superior sample efficiency in few-shot settings.
• kNN-Adapter maintains strong performance even when the LLM API only returns probabilities for the top-k tokens, making it suitable for real-world black-box deployment.
• The learned interpolation weights λ vary significantly across domains and token types; for example, Amazon and Enron show higher λ values for pronouns, while Wikitext-103 shows lower values due to fewer pronoun tokens in its datastore.
• The adaptive λ mechanism is crucial—static or group-wise coefficients yield suboptimal performance, confirming the importance of token-level personalization in retrieval-augmented inference.
• The study reveals that kNN-Adapter’s success stems from its ability to dynamically balance general language knowledge and domain-specific patterns, especially in low-resource scenarios.