[논문 리뷰] $k$NN Prompting: Beyond-Context Learning with Calibration-Free Nearest Neighbor Inference

Benfeng Xu, Quan Wang외 4명|arXiv (Cornell University)|2023. 03. 24.

Domain Adaptation and Few-Shot Learning인용 수 12

한 줄 요약

본 논문은 kNN Prompting을 소개한다. 이는 보정 없이, 기울기 없이 작동하는 방법으로, LM 분포에 대한 최근접 이웃 접근법을 사용하여 컨텍스트 길이를 넘는 무제한 학습 데이터를 활용하고, 여러 LLM에서 강력한 few-shot 및 fully supervised 성능을 이끈다.

ABSTRACT

In-Context Learning (ICL), which formulates target tasks as prompt completion conditioned on in-context demonstrations, has become the prevailing utilization of LLMs. In this paper, we first disclose an actual predicament for this typical usage that it can not scale up with training data due to context length restriction. Besides, existing works have shown that ICL also suffers from various biases and requires delicate calibration treatment. To address both challenges, we advocate a simple and effective solution, $k$NN Prompting, which first queries LLM with training data for distributed representations, then predicts test instances by simply referring to nearest neighbors. We conduct comprehensive experiments to demonstrate its two-fold superiority: 1) Calibration-Free: $k$NN Prompting does not directly align LLM output distribution with task-specific label space, instead leverages such distribution to align test and training instances. It significantly outperforms state-of-the-art calibration-based methods under comparable few-shot scenario. 2) Beyond-Context: $k$NN Prompting can further scale up effectively with as many training data as are available, continually bringing substantial improvements. The scaling trend holds across 10 orders of magnitude ranging from 2 shots to 1024 shots as well as different LLMs scales ranging from 0.8B to 30B. It successfully bridges data scaling into model scaling, and brings new potentials for the gradient-free paradigm of LLM deployment. Code is publicly available.

연구 동기 및 목표

• 프롬프트 길이를 넘는 데이터를 재활용하여 in-context learning(ICL)의 컨텍스트 길이 한계를 해결한다.
• 직접 라벨-토큰 정렬이 아닌 LM 분포를 이용한 매칭으로 보정 기법에의 의존성을 제거한다.
• 프롬프트 길이와 학습 데이터 사용을 분리하여 규모가 큰 데이터 세트로 학습을 가능하게 함으로써 확장 가능한 데이터 활용을 가능하게 한다.
• 다양한 데이터 세트와 모델 규모에 걸친 anchor 기반 표현의 강건성, 분할 전략 및 효과를 조사한다.

제안 방법

• 학습 데이터를 시연 세트 D와 앵커 세트 A로 나눈다.
• D와 각 앵커 x^a를 포함하는 프롬프트로 LLM에 질의하여 전체 LM 분포 p(v|P,θ)를 얻고 이를 핵 k_i로 캐시하되 y^a를 값으로 데이터 스토어에 저장한다.
• 추론 시 테스트 인스턴스에 대해 p(v|P_test,θ)를 계산하고 데이터스토어에서 KL 발산 D_KL(p_test || k_i)을 이용해 k개의 최근접 이웃을 검색한다.
• y를 예측하기 위해 k개의 최근접 이웃의 y^a 라벨을 다수결 비슷한 방식으로 합산하여 결정한다: argmax_y sum_{i in NN^k(p_test, K)} 1(y_i^a = y).
• 이 프레임워크는 LM 출력과 고정된 작업-특정 라벨 공간을 정렬하지 않기 때문에 보정이 필요 없고, 앵커 세트를 사용하여 무제한 데이터를 활용하므로 컨텍스트를 넘어선 학습을 지원한다.
• 학습이 필요하지 않으며, 기존 LLM 및 학습 데이터로 작동하도록 설계되었다.

실험 결과

연구 질문

• RQ1 calibration-free 이웃 탐색추론이 few-shot 설정에서 보정 기반 ICL 방법에 비해 성능을 향상시킬 수 있는가?
• RQ2 컨텍스트 길이를 넘는 더 많은 학습 데이터가 가능해질 때 kNN Prompting은 다양한 모델 규모와 과제에서 어떻게 규모 확장을 할 수 있는가?
• RQ3 시연-앵커 분할이 성능에 미치는 영향은 무엇이며, 학습 분할의 변화에 대한 방법의 강건성은 어느 정도인가?
• RQ4 검색에 전체 LM 분포를 이용하는 것(라벨 단어만 사용하는 것과 대비하여)이 정확도와 강건성에 어떤 영향을 주는가?

주요 결과

• 엄격한 few-shot 설정에서 최첨단 보정 기반 방법을 능가하며(일부 과제에서 최대 +7.07).
• 컨텍스트를 넘어선 확장 시나리오에서 강한 증가를 보이며, 학습 데이터가 2에서 1024샷으로 증가하고 0.8B에서 30B 매개변수의 LLM들 간에도 점진적으로 개선된다.
• 32샷만으로도 피크에서 평균 이득 +13.58(0.8B 모델) 및 완전한 감독 설정에서 최대 +18.84까지 달성한다.
• vanilla ICL 대비 강건성 개선을 보여주며, 출력 불안정성과 시드 간 표준편차를 낮춘다.
• 다양한 데이터 규모에서 ICL Ensemble 및 표준 미세조정보다 우수한 성능을 보이고, 0.8B–30B를 포함한 다양한 모델 크기에서도 경쟁력을 유지한다.