Few-shot Learning with Retrieval Augmented Language Models

Essence

Figure 1: Atlas는 사전학습과 미세조정 단계 모두에서 검색을 활용하는 검색 증강 언어 모델로, 지식 기반 작업에서 강력한 few-shot 성능을 보임

본 논문은 매개변수 메모리에 의존하지 않고 외부 지식 소스를 활용하는 검색 증강 언어 모델(Atlas)을 제시하여, 550억 개 매개변수 모델보다 50배 적은 매개변수(110억)로 우수한 few-shot 학습 성능을 달성한다.

Motivation

Known: 대규모 언어 모델(LLM)은 few-shot 학습에서 우수한 성능을 보이지만, 질의응답이나 사실 검증 같은 지식 집약적 작업에서는 막대한 매개변수가 필요하다.
Gap: 검색 증강 모델(retrieval-augmented models)이 지식 집약적 작업에서 효과적임이 알려져 있으나, few-shot 설정에서의 성능은 불명확하다. 또한 few-shot 학습이 모델의 매개변수에 방대한 정보 저장을 필수로 하는지 검증되지 않았다.
Why: 지식은 외부 비파라미터 메모리(external non-parametric memory)로 아웃소싱할 수 있으며, 이를 통해 모델의 적응성, 해석가능성, 효율성을 개선할 수 있다.
Approach: 검색 증강 아키텍처를 활용하되, 사전학습 단계에서 retriever와 언어 모델을 함께 학습하고, 네 가지 novel한 retriever 학습 목적 함수(ADist, EMDR2, PDist, Margin-MSE)를 제시하여 few-shot 성능을 최적화한다.

Achievement

Few-shot 성능 우수성: NaturalQuestions에서 64개 예시만으로 42.4% 정확도 달성 (PaLM 540B 대비 3% 우수), Wikipedia 인덱스만 사용 시 45.1% 달성
Full-dataset 성능 최고 성과: NaturalQuestions (+8.1%), TriviaQA (+9.3%), FEVER, KILT의 5개 작업에서 새로운 최고 성능 기록
MMLU 성능: 매개변수 15배 많은 모델과 동등하거나 우수한 성능 달성
인덱스 압축: Product quantization을 이용한 압축 인덱스가 비압축 인덱스와 유사한 성능을 유지하면서 5배 메모리 감소 달성

How

아키텍처 설계

Retriever: Contriever 기반의 dual-encoder 구조로, 쿼리와 문서를 독립적으로 임베딩하여 dot-product로 유사도 계산. MoCo 대조 손실로 사전학습.
Language Model: T5 sequence-to-sequence 아키텍처에 Fusion-in-Decoder(FiD) 적용. 각 문서를 인코더에서 독립적으로 처리한 후 디코더에서 결합된 출력에 대해 교차-주의(cross-attention) 수행.

Retriever 학습 목적 함수

Attention Distillation (ADist): 언어 모델의 교차-주의 점수(α_n ∥v_n∥²)를 retriever 확률 분포로 KL-divergence 최소화하여 증류
EMDR² (End-to-end Multi-Document Reader and Retriever): EM 알고리즘 영감, 검색된 문서를 잠재 변수로 취급하여 $\log[\sum_{k=1}^K p_{lm}(a|q,d_k)p_{retr}(d_k|q)]$ 최대화
Perplexity Distillation (PDist): 언어 모델이 각 문서로부터 받는 perplexity 개선을 학습 신호로 활용하는 단순한 손실 함수
Margin-MSE: 문서 쌍 간의 상대적 순위를 고려한 마진 기반 평균제곱오차(MSE) 손실

사전학습 전략

다양한 자가지도학습(self-supervised) 작업 활용: Masked Language Modeling(MLM), Span Corruption, 대조 학습
Retriever와 언어 모델의 공동 사전학습이 few-shot 성능의 핵심

미세조정 전략

Retriever와 언어 모델 양쪽을 효율적으로 적응시키기 위한 다양한 파라미터 업데이트 전략 제시
Task별로 선택적 업데이트 가능

Originality

Novel 학습 목적 함수: Attention 점수에서 value norm 고려(α_n ∥v_n∥²)라는 개선된 attention distillation 제시
Systematic 설계 연구: Few-shot 학습을 위한 검색 증강 모델의 사전학습과 미세조정 전략을 체계적으로 연구한 첫 수준의 연구
Scalable 아키텍처: Fusion-in-Decoder를 통해 검색 문서 수에 따른 quadratic 복잡도 문제 해결
Index 업데이트 가능성: 외부 문서 인덱스의 동적 업데이트가 가능하여 지식의 시간적 변화에 대응 가능
Interpretability: 검색된 문서를 명시적으로 확인 가능하여 모델 결정에 대한 설명가능성 제공

Limitation & Further Study

Index 크기 제약: 대규모 인덱스 구축 및 관리의 계산 비용이 여전히 존재하며, 인덱스 품질의 편향 문제 검토 필요
Retriever 성능 의존성: 전체 시스템의 성능이 retriever의 성능에 크게 의존하므로, retriever 오류 분석 및 개선 필요
Generalization across domains: 특정 도메인에서 학습한 모델이 다른 도메인으로의 전이 가능성에 대한 추가 연구 필요
Real-time 검색: 매우 대규모 인덱스에서의 실시간 검색 지연 시간 개선 필요
Cross-lingual few-shot learning: 다국어 few-shot 학습 설정에서의 성능 평가 필요

Evaluation

Novelty: 4.5/5 Technical Soundness: 4.5/5 Significance: 4.5/5 Clarity: 4.5/5 Overall: 4.5/5

총평: 본 논문은 검색 증강 언어 모델의 few-shot 학습 능력을 체계적으로 탐구하여, 매개변수 효율성과 성능 간의 새로운 균형점을 제시한 고품질 연구다. 특히 실무 적용 가능성과 지식 업데이트 용이성 측면에서 학계와 산업계에 모두 기여할 수 있는 중요한 작업이다.

같이 보면 좋은 논문

A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models

기반 연구검색 증강 방식의 언어 모델 설계에 이론적 기반을 제공하는 연구이다.

다른 접근검색 증강 생성에 대한 다른 기술적 접근법을 제시

REALM: Retrieval-Augmented Language Model Pre-Training

다른 접근검색 증강 언어 모델의 few-shot 학습 성능을 향상시키는 다른 접근법을 제시한다.

Splade v2: Sparse lexical and expansion model for information retrieval

다른 접근정보검색을 위한 다른 신경망 기반 표현 학습 방법

Improving RAG for LLM-based Mental Health Risk Detection through Contrastive Embeddings

다른 접근매개변수 효율적 검색 증강 접근의 대안적 구현을 제시한다.

후속 연구SPECTER2 유사도 0.93로 Multimodal Biomedical Data Fusion와 Scientific Information Extraction and QA가 맞닿아, 'Few-shot Learning with Retrieval Augmented Language Models'가 이 ICML 2026 논문의 배경·대안·응용 맥락을 보완한다.

Assisting in writing wikipedia-like articles from scratch with large language models

후속 연구LLM 기반 장문 생성의 이론적 기반을 제공하는 연구로 보임

PaSa: An LLM Agent for Comprehensive Academic Paper Search

응용 사례검색 증강 언어 모델을 실제 few-shot 태스크에 적용한 사례를 다룬다.