Speculative RAG: Enhancing Retrieval Augmented Generation through Drafting
Retrieval augmented generation (RAG) combines the generative abilities of large language models (LLMs) with external knowledge sources to provide more accurate and up-to-date responses. Recent RAG advancements focus on improving retrieval outcomes through
arxiv.org
0. Abstract
RAG는 LLM의 생성 기능과 외부 지식을 결합하여 답변을 제공한다. 최근 RAG는 반복적인 LLM의 개선 또는 self-critique을 활용해 검색 결과는 개선하는 데 중점을 두고 있다. Speculative RAG는, 큰 규모의 생성형 LM이 작은 규모의 distill된 RAG 모델들의 초안을 병렬로 효과적으로 검증한다. 각각의 초안은 별개의 하위 집합에서 생성되므로, 초안 당 입력 토큰 수를 줄이면서 증거에 대한 다양한 관점을 제공 한다. 이는 각 하위 집합에 대한 이해도를 높이고 긴 문맥에 대한 잠재적인 입장 편향을 완화한다. 이 방식은 소규모 specialist LM에게 초안을 작성하도록 하고, 대규모 generalist LM이 생성된 초안들에 대해 검증함으로써 가속화한다. Speculative RAG는 TriviaQA, MuSiQue, PubHealth 및 ARC-Challenge 벤치마크에서 지연 시간을 줄이면서 SOTA를 달성했다. 특히 PubHealth에서 기존 RAG 시스템에 비해 지연 시간을 51% 줄이면서 정확도를 최대 12.97%까지 향상시켰다.
1. Introduction
LLM은 질문 답변 작업에서 괄목할 만한 성공을 거두었다. 하지만 최신 정보나 모호한 사실을 요구하는 knowledge-intensive에는 사실관계의 부정확성으로 인해 어려움을 겪고 hallucination을 생성할 수 있다. RAG는 방대한 데이터베이스에 쉽고 효율적으로 액세스할 수 있을 뿐만 아니라 지식 통합을 통합하여 문제를 해결한다.
이로 인해 LLM에 대한 입력 길이를 증가되며, 특히 검색된 긴 문서를 인코딩할 때 추가적인 지연 시간이 발생하고 더 복잡한 추론이 필요하다. 최근의 연구들은 LLM의 컨텍스트 길이 제한을 확장하는 방법을 모색했지만, 확장된 컨텍스트에 대해 근거에 기반한 추론을 달성하는 것은 어렵다. 따라서 RAG에서 효율성과 효과성 사이의 균형을 맞추는 것이 주요 연구 주제였다. RAG 시스템에 대한 기존 연구는 주로 검색 결과의 문맥 정보 품질을 개선하는 데 집중하지만, 이러한 시스템과 관련된 지연 문제를 무시하는 경우가 많다. 이러한 방법은 여러 번의 개선 반복 또는 자체 비평 능력을 위한 instruction-tuining에 의존한다. 이러한 개선 사항을 일반 LM에 통합하려면 추가 교육이 필요하거나 지연 시간이 증가하여 실제 활용에서 어려움을 겪을 수 있다.
더 작은 모델로 여러 초안 토큰을 동시에 생성하고 기본 모델과 병렬로 검증함으로써 자동 회귀적 LM 추론을 가속화하는 Speculative Decoding(Leviathan 외, 2023; Chen 외, 2023a; Xia 외, 2024)에서 영감을 받았다. Speculative RAG는 검색된 문서를 하위 집합으로 분할하여 답변 후보 초안을 작성합니다. (Step 1)검색된 문서를 콘텐츠 유사도별로 클러스터링하고 각 클러스터에서 하나의 문서를 샘플링하여 하위 집합을 형성함으로써 중복을 최소화하고 다양성을 최대화한다. 그런 다음 이러한 (Step 2)문서 하위 집합은 작고 전문화된 RAG 모듈들에 공급되어 해당 근거를 가진 답변 초안을 신속하게 병렬로 생성한다. 그 후 (Step 3)생성형 LM은 반복될 가능성이 있는 문서의 세부 검토를 생략하고 대신 가장 정확한 답변을 결정하기 위해 근거에 대한 초안을 검증( 답변 초안과 근거의 조건부 생성 확률을 신뢰 점수로 계산)하는 데 집중한다.
- 소규모의 전문 RAG 초안 작성자를 활용해 고품질의 초안 답변을 생성하는 새로운 RAG 프레임워크를 도입했습니다. 각 초안은 검색된 문서의 고유한 하위 집합에서 파생되므로 초안당 입력 토큰 수를 줄이면서 다양한 관점을 제공합니다.
- RAG 초안 작성기와 함께 작동하는 제너럴리스트 LM은 추가적인 조정이 필요하지 않습니다. 가장 유력한 초안을 검증하고 최종 답변에 통합하기만 하면 됩니다. 이 접근 방식은 각 하위 집합에 대한 이해도를 높이고 잠재적인 중간 손실(Liu et al., 2024) 현상을 완화합니다.
- 이 방법은 소규모 전문가 LM에게 초안 작성을 위임하고 대규모 제너럴리스트 LM이 초안에 대해 편향되지 않은 단일 검증 패스를 병렬로 수행함으로써 RAG를 크게 가속화합니다. 4개의 자유형 질문-답변 및 폐쇄형 세트 생성 벤치마크에 대한 광범위한 실험을 통해 이 방법의 탁월한 효과와 효율성이 입증되었습니다.
2. Related Works
RAG
외부 데이터베이스에서 관련 문서를 검색하여 이를 생성 과정에 통합함으로써 LLM의 성능을 향상시킨다.
- SAIL: 검색 데이터를 기반으로 사전 학습된 LLM을 미세 조정하여 관련 없는 내용을 필터링
- SelfReflective RAG: 검색 및 주석 작업을 개선하기 위해 refleciton token 도입
→ 한계: 추가적인 미세 조정은 리소스 소모가 크며, 모델의 망각 또는 과적합 문제를 야기 - Corrective RAG: 가벼운 검색 평가자를 제안했지만 고차원 추론 능력 부족
Speculative Decoding
자동 회귀 디코딩의 지연 시간을 줄이기 위해 "초안 작성 후 검증(draft-then-verify)" 방식을 사용한다.
- 기존
- 초안 작성: 작은 모델이 미래 토큰을 초안으로 작성.
- 검증: 목표 모델이 병렬적으로 초안을 확인.
- Speculative RAG
- 토큰 수준 → 답변 수준 초안 작성: 답변의 문맥적 통합을 강화
- 검증 기준 차별화: 전통적으로는 생성 확률에 기반하여 토큰을 수용 또는 거절 → 새로운 방식으로 보다 효율적이고 맥락에 맞는 답변 생성.
3. Speculative Retrieval Augmented Generation through Drafting
Knoweldge intensive task(지식 집약적인 업무)는 Q, D, A가 필요하다.
- Q: 추가적인 지식이 필요한 question or statement
- D: ${d_1, ..., d_n}$으로, DB에서 검색한 n개 문서들
- A: answer (T/F, 객관식, etc.)
3.1 Overview
- 문서 클러스터링:
- 검색된 문서를 질문(Q)과의 연관성에 따라 클러스터로 나눔(Line 2)
- 각 클러스터는 검색 결과에서 하나의 관점을 대표
- 다양성 확보:
- 각 클러스터에서 하나의 문서를 샘플링하여 서브셋(δ ⊂ D)을 구성(Line 5~8)
- 서브셋은 중복을 최소화하고 관점의 다양성을 극대화하도록 설계
- RAG 드래프터를 통한 초안 생성:
- 각 서브셋(δ)과 질문(Q)을 RAG 드래프터(MDrafter)에 전달(Line 12)
- MDrafter는 문서를 분석해 답변 초안(α)과 근거(β)를 병렬적으로 생성
- MDrafter는 검색된 문서를 이해하고 근거를 생성하도록 특화된 소형 모델로, 병렬 처리를 통해 높은 효율성을 제공
- 초안 검증:
- 생성된 답변 초안(α)과 근거(β)에 대해, 일반 LLM(MVerifier)을 사용해 신뢰 점수를 계산(Line 13)
- MVerifier는 사전 학습된 언어 모델 능력을 활용하며, 추가적인 조정 없이 근거를 기반으로 초안을 검증
- 최종 답변 선택:
- 신뢰 점수가 가장 높은 초안을 최종 답변으로 선택(Line 15)
- 선택된 답변은 일반 LLM의 생성 결과에 통합
3.2 Specialist RAG Drafter
일반적인 LLM을 사용하는 대신, RAG 시나리오에서는 전문화된 소형 LLM Drafter를 활용한다. Drafter는 검색된 문서를 이해하고, 이를 기반으로 답변 초안(α)과 근거(β)를 생성하는 데 특화되었다.
- Instruction Tuning
- 질문(Q), 답변(A), 문서(D)에 대해, 문서(D)의 핵심 정보를 추출하여 합리적 근거(E)를 생성
- E는 답변이 질문에 적합하다는 것을 설명
- Multi-Perspective Sampling
- 질문(Q)에 대해 검색된 문서들을 임베딩(emb(d))으로 변환한 뒤, K-Means 클러스터링을 적용
- 각 클러스터에서 하나의 문서를 샘플링하여 중복을 최소화하고 다양성을 극대화한 문서 서브셋(δ) 생성
- RAG Drafting
- 생성된 서브셋(δ)과 질문(Q)을 기반으로, Drafter가 답변 초안(α)과 근거(β)를 병렬적으로 생성
- 이 초안은 질문과 다중 관점의 문서를 연결하여 논리적이고 신뢰성 있는 답변을 제공
3.3 Generalist RAG Verifier
- Evaluation Scores
- Verifier는 Drafter가 생성한 초안(α)과 근거(β)를 평가하여 신뢰 점수 계산
- Self-Consistency Score (자기 일관성 점수): 초안과 근거가 질문과 일관성을 가지는지 평가
- Self-Reflection Score (자기 반성 점수): 초안이 답변으로서 신뢰할 수 있는지를 "예/아니오"로 평가
- Computation Method
- Verifier는 사전 학습된 언어 모델 능력을 활용하며, 초안을 검증하기 위해 추가 조정이 필요 없음
- 자기 회귀(auto-regressive) 기능을 활용하여 초안에 대한 신뢰 점수를 효율적으로 계산
- 계산된 점수(자기 일관성, 자기 반성 점수)를 종합하여 가장 높은 점수를 가진 초안을 최종 답변으로 선택
4. Experiments
4.1 Baselines
- TriviaQA: 사실 기반 질문에 대한 정답 생성
- MusiQue: 다중 선택형 질문을 포함하여 과학적 데이터 기반
- PubHealth: 의료적 주장을 포함한 답변 생성
- ARC-Challenge: 초등학교 과학 문제로 구성된 다중 선택형 질문
4.2 Experiment Settings
- RAG drafter: Mistral-7B(v0.1)
- RAG verifier: Fine-tuning 없이 Mistral7B(v0.1) 또는 Mixtral8x7B(v0.1)
- Retrival Process: light-weight 명령어 인식 임베딩 모델인 InBedderRoberta를 사용해 검색된 문서의 임베딩을 미리 계산
- Inference: Greedy decoding(temperature = 0)을 사용하는 vLLM 프레임워크
Self-RAG와 동일한 실험 설정을 채택하고 더 까다로운 벤치마크인 MuSiQue를 포함한다. 증거 인용보다는 RAG 추론에 초점을 맞추기 때문에 다른 두 개의 긴 형식의 생성 벤치마크인 Biography와 ALCE-ASQA는 생략한다. TriviaQA, PubHealth, ARC-Challenge에서는 상위 10개 문서를 검색하고 쿼리당 5개 초안을 생성하며(m=5), 각 초안은 2개 문서의 하위 집합을 기반으로 한다(k=2). MuSiQue 데이터 세트의 경우, 더 복잡한 추론으로 인해 상위 15개 문서를 검색하고 각 쿼리에 대해 10개 초안을 생성하며(m=10), 각 초안은 6개 문서의 하위 집합을 사용한다.
4.3 Main Results
Superior Performance over Baselines
Speculative RAG는 모든 데이터셋에서 기준 모델과 고급 모델보다 더 높은 정확도를 달성했고, 특히 Verifier-8x7B + Drafter-7B 조합이 가장 경쟁력 있는 Standard RAG 모델(Mixtral-Instruct-8x7B)을 다음과 같이 능가했다.
- TriviaQA: +0.33%
- MuSiQue: +2.15%
- PubHealth: +12.97%
- ARC-Challenge: +2.14%
Drafter-7B 단독으로도 대부분의 환경에서 Self-Reflective RAG 및 Corrective RAG를 능가했다.
Effective Instruction Tuning for RAG Drafter
Instruction Tuning은 drafter 모델의 추론 능력을 크게 향상.
Mistral-7B와 Drafter-7B를 비교했을 때, Drafter-7B의 성능이 눈에 띄게 향상됨.
Mixtral-8x7B와 Drafter-7B의 조합은 다음과 같은 정확도 향상을 보여줌
- TriviaQA: +14.39%
- MuSiQue: +12.41%
- PubHealth: +39.52%
- ARC-Challenge: +31.83%
Mistral-7B와의 조합에서도 비슷한 개선이 나타났다.
- TriviaQA: +19.76%
- MuSiQue: +14.32%
- PubHealth: +40.94%
- ARC-Challenge: +33.44%
RAG 드래프터가 검색 문서에서 불필요한 중복을 최소화하고, 다양한 관점을 반영한 높은 품질의 답변 초안을 생성했기 때문이다.
Reliable Scoring by RAG Verifier
일반 LLM(Verifier)의 검증 능력이 성능 향상에 크게 기여했다. Drafter-7B 단독과 Verifier-7B + Drafter-7B 조합을 비교하면 성능이 크게 향상되었다.
- RAG 드래프터는 검색 문서에 기반한 답변 초안을 생성
- 일반 LLM은 이 초안을 논리적 근거(β)에 비추어 검증
검증 접근 방식은 간단하면서도 효과적이며, SPECULATIVE RAG의 성능을 입증했다.
4.4 Latency Analysis
실험 설정
- 각 데이터셋에서 100개의 샘플을 무작위로 추출
- 실세계 시나리오를 시뮬레이션하기 위해 배치(batch) 없이 개별적으로 처리
- 병렬 설정
- Standard RAG
- Tensor 병렬화를 사용하여 Mixtral-Instruct8x7B를 GPU 메모리에 적합하게 실행
- Tensor 병렬화 크기: 4, 8, 16으로 실험
- 그러나, tensor 병렬화 크기를 늘리면 통신 및 집계 오버헤드로 인해 효율성 증가가 제한적
- SPECULATIVE RAG
- Drafter-7B를 5개의 엔드포인트에서 병렬 초안 생성(TriviaQA, PubHealth, ARC-Challenge)
- MuSiQue의 경우 초안 수가 더 많아 10개의 엔드포인트 실행
- Standard RAG
실험 결과
Tensor 병렬화는 Standard RAG에서 Tensor 병렬화를 통해 GPU 메모리 적합성을 확보했으나, 통신 및 집계 오버헤드로 인해 효율성이 제한적이고, Mixtral-Instruct8x7B의 latency는 데이터셋에 따라 검색된 문서 길이의 차이로 변동된다. Speculative RAG는 검색된 문서를 병렬적으로 처리하고 작은 drafter 모델을 활용하여 답변 초안을 생성하며, 레이턴시를 대폭 줄이는 동시에 높은 성능 유지한다. Standard RAG 대비 최대 51.25%의 latency 절감 효과를 보여주었다. 효율성과 성능 모두에서 우수한 결과를 달성하며, 실세계 애플리케이션에서 시간과 자원을 절약할 수 있는 유망한 접근 방식이다.
4.5 Ablation Studies
Diversity and reduced redundancy in retrieval improves draft quality significantly
Speculative RAG는 검색된 문서를 다양한 관점으로 클러스터링하고, 각 클러스터에서 하나의 문서를 샘플링하여 초안을 생성한다.
- Clustering 없이 랜덤 샘플링: 문서 서브셋을 무작위로 선택
- 단일 cluster에서 샘플링: 한 클러스터에서 모든 문서를 선택
Clustering 기반 샘플링이 최고의 성능을 달성:
- TriviaQA: +1.88% (73.02 → 74.24)
- PubHealth: +2.23% (74.37 → 76.60)
랜덤 샘플링은 다양성을 제공하지만 중복 문서 포함 가능성이 높아 성능 저하된다. 또 단일 클러스터 샘플링은 관점의 다양성이 부족하여 성능이 가장 낮다.
Scoring method on self-consistency and self-reflection refines draft quality effectively
검증 단계에서 사용되는 세 가지 신뢰 점수는 아래와 같고, 구성 요소를 하나씩 제거한 후 성능 변화를 평가했다.
- ρDraft: 초안 자체의 품질 평가
- ρSelf-contain: 초안의 자기 포함성 평가
- ρSelf-reflect: 초안의 자기 반성 점수.
$ρ_{Draft}$를 제거하면,
- TriviaQA: -0.19%
- PubHealth: -1.12%
성능 감소는 적지만, 작은 RAG drafter의 검증 능력이 제한적임을 반영했다.
$ρ_{Self-contain}$를 제거하면,
- TriviaQA: -1.88%
- PubHealth: -0.92%
$ρ_{Self-reflect}를 제거하면,
- TriviaQA: -2.20%
- PubHealth: -0.71%
Self-contain(자기 포함성)과 Self-reflect(자기 반성) 점수는 각각 다른 논리적 요소를 평가하며, 모두 중요하다.
검증 없이 랜덤 선택하면,
- TriviaQA: -5.69% (68.55 → 74.24)
- PubHealth: -5.37% (71.23 → 76.60)
성능 대폭 감소했다.
4.6 Effects of Generated Rationale for Verification
Speculative RAG는 RAG drafter가 생성한 근거(β)를 통해 답변 초안(α)의 신뢰도를 평가한다.
근거(β)의 역할 및 평가 방법
근거는 관련된 핵심 정보를 강조하고, 중복된 정보를 생략한다. 또한 초안과 이를 뒷받침하는 문서 간의 논리적 연결을 강화한다.
- (a) 검색된 문서만 사용: ρ = Score(α|Q, δ): 생성된 근거(β)를 제거하고, 검색된 문서(δ)만 사용
- (b) 근거 + 검색된 문서 사용: ρ = Score(α|Q, β, δ): 생성된 근거(β)에 검색된 문서(δ)를 추가
- ρ = Score(α|Q, β): 생성된 근거(β)만 사용하여 답변 초안을 검증
결과
검색된 문서만 사용 (ρ = Score(α|Q, δ))
긴 검색 문서 사용은 성능 향상에 기여하지 않으며, 검증 효율성 저하 및 latency가 증가했다.
- TriviaQA: 성능 감소 (74.24 → 70.05)
- PubHealth: 성능 감소 (76.60 → 75.89)
근거 + 검색된 문서 사용 (ρ = Score(α|Q, β, δ))
근거(β)에 검색 문서를 추가하면 성능 개선 없이 latency가 증가했다.
- TriviaQA: 성능 감소 (74.24 → 72.04)
- PubHealth: 성능 감소 (76.60 → 75.63)
근거만 사용 (ρ = Score(α|Q, β))
가장 높은 정확도를 기록했기 떄문에, 생성된 근거(β)는 검색 문서의 요약 및 논리적 연결로서 충분히 높은 품질을 제공한다.
- TriviaQA: 74.24
- PubHealth: 76.60
4.7 Effects of Draft Number and Document Subset Size
Draft 수의 변화에 따른 성능
- MVerifier-7B + MDrafter-7B 조합을 사용
- Draft 수를 5, 10, 15, 20으로 조정하여 TriviaQA와 PubHealth 데이터셋에서 실험
- 각 Draft는 두 개의 문서를 문맥으로 사용
- 총 10개의 검색된 문서를 기반으로 최대 20개의 Draft 생성
- 더 많은 Draft가 다양한 관점을 더 포괄적으로 반영할 수 있기 때문에, Draft 수가 증가함에 따라 성능이 지속적으로 개선
- Speculative RAG는 병렬 처리를 통해 추가 latency 없이 여러 Draft를 동시에 생성할 수 있어 효율적
문서 서브셋 크기의 변화에 따른 성능
- Draft 생성 시 문서 서브셋 크기를 1, 2, 4, 6으로 조정하여 TriviaQA와 PubHealth에서 실험
- TriviaQA: 문서가 많아질수록 질문의 복잡성을 해결하는 데 유리
- PubHealth: 추가 문서가 항상 성능 개선으로 이어지지는 않음
- 문서 하나만 포함한 경우에도 Mistral-Instruct7B를 능가
- 문서 2개 이상 포함하면 Mixtral-Instruct8x7B를 능가
- 이는 SPECULATIVE RAG의 초안 설계가 매우 효과적임을 입증
5. Conclusion
Speculative RAG는 RAG 작업을 drafting(작은 크기의 전문화된 RAG Drafter가 수행)과 verification(일반적인 LLM) 두 단계로 분해했다. 이를 이용해 병렬적으로 다양한 문서 서브셋에서 초안을 생성하여 고품질 답변 후보를 제공하고, 긴 문맥에서 발생하는 위치 편향(position bias)을 줄이고, 입력 토큰 수를 감소시킨다. 결과적으로, 정확도와 속도 모두에서 큰 개선을 달성했다. 하지만 추가적인 drafter 학습과 복잡성 증가는 적용 시 고려해야 할 요소이다.
