Post

REFRAG: Rethinking RAG based Decoding

Posted Updated
By Jeong Mo, Hong
13 min read
REFRAG: Rethinking RAG based Decoding

논문 링크: https://arxiv.org/pdf/2509.01092

  • 개발: Meta Superintelligence
  • 핵심 키워드: RAG, Chunk Embedding, Selective Compression, RL Policy, TTFT

핵심 요약

REFRAG는 Meta Superintelligence가 개발한 차세대 RAG 기술로, 기존 RAG의 근본적인 비효율성을 해결한다.

핵심 아이디어: “토큰 개념 자체를 바꾼다”

구분기존 RAGREFRAG
전달 형태검색된 문서 청크를 모든 토큰 그대로 LLM에 전달청크를 압축된 임베딩 벡터로 변환 후 전달
문제점긴 컨텍스트 → 느린 응답, 높은 비용-
성과-30배 빠른 응답 속도

5단계 동작 프로세스

1

[문서] → [청크 분할] → [임베딩 변환] → [캐싱] → [검색] → [선택적 복원] → [LLM 생성]
  1. 사전 변환 (Encoding)
    • 약 128 토큰 크기의 청크를 경량 인코더로 임베딩 벡터로 압축
  2. 캐싱 (Pre-computation)
    • 변환된 임베딩을 미리 계산해서 저장
    • 실시간 연산 부담 제거
  3. 검색 (Retrieval)
    • 사용자 질의와 관련된 청크를 검색
    • 여기까진 기존 RAG와 동일
  4. 선택적 복원 (Selective Expansion) — 핵심
    • RL 기반 정책 네트워크가 “정말 중요한 청크”만 선별
    • 선별된 것만 전체 토큰으로 확장
    • 나머지는 압축된 임베딩 상태 유지
  5. 생성 (Generation)
    • LLM 입력: 질의 + 복원된 토큰 + 압축 임베딩들
    • LLM이 임베딩을 직접 이해하며 응답 생성

핵심 기술 요소

  • 정책 네트워크 (RL): Perplexity를 줄이면서 품질을 유지하는 보상 함수로 학습. 어떤 청크를 복원할지 결정
  • Chunk Embedding: LLM이 직접 이해 가능한 형태의 압축 표현
  • KV Cache/Attention 비용 절감: 토큰 수가 줄어 연산량 대폭 감소

왜 빠른가?

기존 RAG는 검색된 모든 문서를 “풀 텍스트”로 LLM에 넣어 컨텍스트가 길어지고, 이는 곧 느린 응답과 높은 비용으로 이어진다. REFRAG는 대부분의 정보를 압축된 임베딩으로 전달하고, 꼭 필요한 부분만 토큰으로 복원하므로 LLM이 처리해야 할 실제 토큰 수가 크게 줄어든다.

1. 서론 (Introduction)

REFRAG 논문은 RAG(Retrieval-Augmented Generation) 시스템의 추론 효율성 문제를 다룬다. LLM이 긴 컨텍스트를 처리할 때 발생하는 주요 문제점:

  • 비효율적인 토큰 할당: RAG 컨텍스트에는 정보가 희소하게 분포되어 있어, 검색된 많은 패시지가 실제로는 유용하지 않음
  • 낭비되는 인코딩 정보: 검색 과정에서 이미 처리된 청크 인코딩과 쿼리와의 상관관계 정보가 디코딩 시 버려짐
  • 희소하고 구조화된 어텐션: 다양성 확보나 중복 제거로 인해 대부분의 컨텍스트 청크들이 서로 관련이 없어, 청크 간 크로스-어텐션이 대부분 0에 가까움

2. 주요 기여 (Our Contributions)

REFRAG(REpresentation For RAG)는 RAG 컨텍스트의 효율적인 디코딩을 위한 새로운 메커니즘을 제안한다. LLM 아키텍처 수정이나 새로운 디코더 파라미터 추가 없이 지연시간, TTFT, 메모리 사용량을 크게 줄인다.

핵심 기여:

  1. 사전 계산된 청크 임베딩 활용: 검색된 패시지의 토큰 대신 압축된 청크 임베딩을 디코더에 직접 입력
  2. “어디서든 압축” 기능: 임의의 위치에서 토큰 청크 압축을 지원하면서 오토레그레시브 특성 유지
  3. RL 기반 선택적 압축 정책: 전체 토큰 입력이 필요한 경우와 근사 청크 임베딩으로 충분한 경우를 동적으로 결정

3. 모델 아키텍처 (Model Architecture)

REFRAG는 디코더 전용 기반 모델(예: LLaMA)과 경량 인코더 모델(예: RoBERTa)로 구성된다.

작동 방식:

  1. 질문과 컨텍스트가 입력되면, 컨텍스트를 k 크기의 L개 청크로 분할
  2. 인코더가 각 청크를 처리하여 청크 임베딩 생성
  3. 프로젝션 레이어로 임베딩을 디코더 토큰 임베딩 크기에 맞게 변환
  4. 변환된 청크 임베딩과 질문의 토큰 임베딩을 함께 디코더에 입력

세 가지 주요 장점:

  • 디코더 입력 길이 단축으로 토큰 할당 효율성 향상
  • 검색에서 사전 계산된 청크 임베딩 재사용으로 중복 계산 제거
  • 어텐션 계산 복잡도가 토큰 수 대신 청크 수에 대해 이차적으로 증가

4. 지연시간 및 처리량 개선 (Latency and Throughput Improvement)

이론적 분석에 따르면, 짧은 컨텍스트 길이에서는 TTFT와 처리량에서 최대 k배 가속을, 긴 컨텍스트에서는 최대 k² 배 가속을 달성한다.

실험 결과 (컨텍스트 길이 16,384, k=16 기준):

  • TTFT 가속: 캐시 사용 시 16.53배, 캐시 미사용 시 8.59배 (CEPE 대비 각각 2.01배, 1.04배)
  • 처리량: LLaMA 대비 최대 6.78배 가속
  • k=32에서는 TTFT 가속이 32.99배(LLaMA 대비), 3.75배(CEPE 대비)까지 도달

5. 방법론 (Methodology)

5.1 지속적 사전학습 레시피 (Continual Pre-training Recipe)

인코더와 디코더를 정렬하기 위해 다음 단락 예측(next paragraph prediction) 태스크를 사용한다.

재구성 태스크 (Reconstruction Task):

  • 인코더에 첫 s개 토큰을 입력하고 디코더에서 이를 재구성
  • 디코더는 동결하고 인코더와 프로젝션 레이어만 학습
  • 목표: 인코더가 k개 토큰을 최소 정보 손실로 압축하고, 프로젝션 레이어가 디코더 토큰 공간에 효과적으로 매핑

커리큘럼 학습 (Curriculum Learning):

  • 청크 길이 k가 증가하면 가능한 토큰 조합이 V^k로 기하급수적 증가
  • 단일 청크 재구성부터 시작하여 점진적으로 난이도 증가
  • 쉬운 태스크(단일 청크)에서 어려운 태스크(L개 청크)로 데이터 혼합 비율 조정

5.2 선택적 압축 (Selective Compression)

RL 정책이 다음 단락 예측 퍼플렉시티를 음의 보상으로 사용하여 어떤 청크를 원본 형태로 유지할지 결정한다. 인코더와 디코더는 압축된 청크와 비압축 청크의 혼합 입력을 처리하도록 파인튜닝된다.

6. 실험 결과 (Experimental Results)

6.1 학습 및 평가 데이터셋

SlimPajama 데이터셋(Book, ArXiv 도메인)에서 200억 토큰으로 학습하고, PG19, ProofPile 데이터셋에서도 일반화 성능을 평가했다.

6.2 주요 결과

퍼플렉시티 성능:

  • 압축률 16에서 CEPE 대비 평균 9.3% 퍼플렉시티 개선
  • REFRAG₈과 REFRAG₁₆이 거의 모든 설정에서 다른 베이스라인을 일관되게 능가
  • LLaMA-2-7B의 원래 4k 컨텍스트 윈도우를 청크 임베딩을 통해 16배까지 확장 가능

6.3 Ablation Study

  • 커리큘럼 학습: 재구성 태스크 성공에 필수적
  • 재구성 태스크: 지속적 사전학습 성공의 핵심
  • RL 기반 선택적 압축: 다양한 압축률에서 랜덤 선택이나 휴리스틱 기반 선택보다 우수
  • 압축률 한계: 32까지는 경쟁력 유지, 64는 성능 저하 발생

7. 컨텍스트 학습 응용 (Contextual Learning Applications)

7.1 RAG (Retrieval Augmented Generation)

5개 도메인의 QA 데이터셋 110만 개로 파인튜닝하고, MMLU, BoolQ, SIQA 등 16개 RAG 벤치마크에서 평가했다.

결과:

  • 동일 패시지 수에서 강한 리트리버 설정: LLaMA와 동등한 성능 + 5.26배 TTFT 가속
  • 약한 리트리버 설정: LLaMA보다 0.71% 성능 향상 + 5.26배 TTFT 가속
  • 동일 지연시간에서(REFRAG 8패시지 vs LLaMA 1패시지): 평균 1.22~1.93% 성능 향상

7.2 멀티턴 대화 (Multi-Turn Conversation)

TopiOCQA, ORConvQA, QReCC 데이터셋에서 평가했다. LLaMA의 4k 토큰 컨텍스트 제한으로 긴 대화 히스토리가 잘리는 반면, REFRAG는 압축을 통해 많은 패시지에서도 강건한 성능을 유지한다.

논문은 네 가지 관련 연구 분야를 검토한다:

  1. 검색 증강 언어 모델링: RETRO, REPLUG, Fusion-in-Decoder 등
  2. 효율적인 긴 컨텍스트 LLM: StreamingLLM, CEPE 등
  3. 압축 트랜스포머: 재귀적 컨텍스트 압축 방법들
  4. 프롬프트 압축: LLMLingua, LongLLMLingua 등 (REFRAG와 상호보완적)

9. 결론 (Conclusion)

REFRAG는 RAG 컨텍스트에 내재된 희소성과 블록 대각선 어텐션 패턴을 활용하여 컨텍스트 표현을 압축, 감지, 확장하는 효율적인 디코딩 프레임워크이다. 다양한 긴 컨텍스트 응용(RAG, 멀티턴 대화, 긴 문서 요약)에서 퍼플렉시티나 다운스트림 정확도 손실 없이 최대 30.85배 TTFT 가속(기존 SOTA 대비 3.75배 개선)을 달성했다.

핵심 의의:

  • RAG 기반 시스템에 대한 특화된 처리의 중요성 입증
  • 지연시간에 민감하고 지식 집약적인 응용 프로그램에 실용적이고 확장 가능한 솔루션 제공
AI, Paper
rag refrag chunk-embedding selective-compression
This post is licensed under CC BY 4.0 by the author.