[논문리뷰]REFRAG: Rethinking RAG based Decoding(Meta, 2025)

Xiaoqiang Lin, Aritra Ghosh, Bryan Kian Hsiang Low, Anshumali Shrivasava, and Vijai Mohan. 2025. REFRAG: Rethinking RAG based Decoding. arXiv:2509.01092 [cs.CL] https://arxiv.org/abs/2509.01092

Problem Statement

이 논문은 RAG(Retrieval-Augmented Generation) 기반 응용에서 효율적인 디코딩 구조를 제안하는 연구이다. 기존 LLM은 긴 컨텍스트 입력 시 높은 지연(latency)과 메모리 사용량 증가 문제로 인해 웹 스케일 검색, 멀티턴 대화, 문서 요약 등 실시간성이 중요한 응용에서 성능 저하를 겪는다. 특히 RAG에서는 입력의 대부분이 검색된 문서(passage)들이며, 그중 일부만 실제 질의에 유효하다. 저자들은 이러한 RAG 특유의 구조적 특성을 고려한 디코딩 최적화가 필요하다고 주장한다.

• [비효율적인 토큰 할당] RAG 컨텍스트에는 정보가 희소하고 많은 문서가 관련성이 낮거나나 반복적으로 등장하지만, 기존 연구들은 모든 토큰에 동일하게 메모리와 연산을 할당한다. 이로 인해 불필요한 자원 낭비가 발생한다.
• [검색 정보의 낭비] RAG에서 retrieval 과정은 이미 문서를 벡터화하고 쿼리와의 연관성을 계산하여 다양한 정보를 제공하지만, 디코딩 과정에서는 이러한 정보를 활용하지 않고 폐기한다. 이는 중복 계산을 초래한다.
• [비정형적이고 희소한 어텐션] RAG에서는 다양성과 중복 제거(deduplication)로 인해 문서 간 상호작용이 거의 없어 block-diagonal 형태의 희소 어텐션 패턴이 나타난다. 그러나 기존 연구들은 이러한 구조적 특성을 고려하지 않아 비효율이 발생한다.

REFRAG 저자들이 지적하는 가장 근본적인 문제의 핵심은 RAG 컨텍스트의 “희소한 어텐션(Sparse Attention)” 패턴이다. RAG 컨텍스트가 불필요한 계산을 유발하는 독특한 내부 구조를 가지고 있다는 것이다. 이를 실험적으로 보여주기위해, 여러 문서를 하나로 합쳐 LLM에 입력시켰을 경우 문서 쌍간의 어텐션 값을 시각화 한 것이 위의 그림이다.

LLaMA-2-7B 모델의 32개 레이어(layer 0 ~ layer 31) 각각에서 계산된 어텐션 값의 평균을 시각화한 것으로, 여기서 P0, P1, P2, P3, P4는 RAG 시스템이 사용자의 질문에 답변하기 위해 외부 데이터베이스(예: 벡터 DB)에서 검색해 온 5개의 서로 다른 문서(문서)를 의미한다.

• <span style=”color:coral”>대각 성분 (Red)</span>: 대각선의 붉은 값들은 각 문서별로 자기 자신과의 (e.g., P0-P0 & P1-P1) 어텐션 결과를 보여준다. 동일한 문서의 내부 토큰들 간의 어텐션은 당연히 매우 강할수 밖에 없으며 이는 너무 자명한 사실이다.
• <span style=”color:aqua”>비대각 성분 (Blue)</span>: 서로 다른 문서들 간의 어텐션 값이 낮아 대부분이 옅은 파란색을 띈다. 이는 서로 다른 문서 내의 토큰들 간의 연관성이 떨어짐을 의미한다.

다양성 확보나 중복 제거 과정을 거친 컨텍스트들은 실제 추론 단계에서 서로 의미적으로 연결될 필요가 거의 없기 때문에, 결과적으로 컨텍스트 간 어텐션이 희소하게 분포하는 블록-대각선(block-diagonal) 구조가 나타난다는 것이다. REFRAG는 이러한 구조적 특성에서 비롯되는 불필요한 계산을 제거함으로써 추론 속도를 크게 높이는 아이디어에서 출발한다.

Methodology

Overview

일반적인 RAG 모델은 사용자의 질문에 대해 임베딩 모델을 거쳐 유사한 문서를 미리 구축한 데이터베이스에서 \(k\)개 찾고, LLM에 “지시문 + 질문 + 문서1 + 문서2 + …” 형태로 입력하여 응답을 생성한다. 하지만 이는 몇 가지 고질적인 문제점을 야기한다.

i. Long Context: 검색하는 문서의 수가 많아지거나 각각의 문서의 길이가 길어질수록 LLM의 추론시 시간 지연(latency)가 증가한다. 특히 Transformer의 입력 컨텍스트 길이(\(L\))에 따른 지연은 \(O(L^2)\)으로 제곱에 비례하기 때문에 기하급수적으로 증가한다. 논문에서는 특히 두 가지 추론 지연을 문제삼는다.

• TTFT (Time-to-first-token): 사용자가 입력을 모델에 넣고 난 뒤, 첫 번째 토큰이 출력되는데 걸리는 시간
• TTIT (Time-to-iterative-token)첫 토큰이 생성된 이후 그 다음 토큰들이 연속적으로 생성될 때 걸리는 평균 시간

ii. Memory Inefficiency : Long Context로 인한 KV 캐싱 메모리 증가

iii. Information Inefficiency: 검색된 문서에서도 실제로 문제를 푸는데 도움이 되는 핵심 문장을 일부분이고, 대부분은 오히려 노이즈로 동작할 수 있다.

REFRAG의 전체 동작 과정은 다음과 같다. 사용자의 질문과 검색된 문서들이 입력되면, 문서들은 일정 길이 \(k\) 단위로 잘려 청크 단위로 분할되고, 각 청크는 경량화된 인코더를 통해 청크 임베딩으로 변환된다.

이렇게 얻어진 임베딩은 Projection Layer를 거쳐 디코더의 입력 차원과 정렬되며, RL 기반 정책 네트워크가 청크별 중요도를 판별해 중요한 것은 원래 토큰으로 확장하고, 덜 중요한 것은 압축된 embedding만 유지한다. 최종적으로 디코더는 질문 토큰, 확장된 청크 토큰, 그리고 압축된 청크 임베딩을 함께 받아 autoregressive하게 답변을 생성한다.

이 과정에서 REFRAG는 컨텍스트의 불필요한 계산을 줄여 KV-cache 메모리 사용량을 절감하고, time-to-first-token을 크게 단축하면서도 정확도를 유지한다.

Step 1. Query Processing and Context Retrieval

• 입력: 사용자의 질문 \(x_{1:q}\)
• 모듈: (a) Query Encoder(벡터 DB 검색용), (b) Tokenizer & Embedding(Decoder 입력용)
• 출력: (a) Retriever된 컨텍스트 \(x_{q+1:T}\), (b) 질문 토큰 임베딩 \(\{ e_1, \cdots, e_q \}\)

REFRAG의 입력은 사용자의 질문 \(x_{1:q}\)과 retrieval로 얻은 컨텍스트 \(x_{q+1:T}\)이다. 이 때, 아래첨자는 토큰의 인덱스를 의미한다. 먼저 사용자의 질문은 (a) Query Encoder (b) Decoder Toknizer로 각각 입력된다. (a)를 통해 질문의 임베딩으로부터 top-\(K\)개의 문서를 검색하고, (b)를 통해 최종 생성 단계에서 사용할 수 있도록 토크나이즈 후 질문 임베딩 \(\{e_i\}\)로 변환된다.

Step 2. Context Chunking and Compression

• 입력: 컨텍스트 토큰열 \(x_{q+1:T}\)
• 모듈: (a) 청크 분할(Chunking), (b) 경량 인코더 \(M_{enc}\), (c) Projection Layer \(\phi\)
• 출력: 청크 임베딩 \(\{c_i\}\), 프로젝션 후 디코더 차원 임베딩 \(\{e^{cnk}_i\}\)

검색된 컨텍스트는

검색된 컨텍스트 \(x_{q+1:T}\)는 길이 \(k\)로 나뉘어 \(L = \frac{s}{k}\)개의 청크 \(\{C_1,\dots,C_L\}\)로 분할된다. 각 청크는 경량 인코더를 통해 벡터 표현 \(c_i = M_{enc}(C_i)\)로 압축된다. 여기서 압축(compression)이란, 크기가 작은 인코더를 통해 컨텍스트의 토큰 임베딩을 mean pooling등의 방법으로 aggreagation하여 하나의 \(\mathbb{R}^{1 \times d}\)크기의 벡터로 표현하는 것을 의미한다.

다음으로 Projection layer가 이를 디코더 입력 공간에 맞게 변환하여 \(e^{cnk}_i = \phi(c_i)\)를 생성한다. 이 결과물은 질문 토큰 임베딩과 함께 디코더로 전달되며, 원래 컨텍스트 길이 \(s\)를 약 \(k\)배 줄이는 효과를 낸다. 또한 청크 임베딩은 사전 계산 가능(precomputable)하여 여러 요청에 재사용 가능하므로, 효율성과 속도를 모두 확보할 수 있다.

Step 3. Selective Chunk Expansion Policy

• 입력: 청크 임베딩 \(\{c_i\}\), 원본 토큰 임베딩 시퀀스 \(\{e_{i,1},\dots,e_{i,k}\}\)
• 모듈: (a) Projection \(\phi\) (압축 경로), (b) RL 기반 정책 네트워크(Selective Compression Policy)
• 출력: 압축된 청크 임베딩 \(\{e^{cnk}_i\}\), 또는 확장된 원본 토큰 임베딩 \(\{e_{i,1},\dots,e_{i,k}\}\)

모든 청크를 단일 벡터로 대체하면 정보 손실이 발생할 수 있으므로, REFRAG는 RL 기반 정책(Policy)으로 어떤 청크를 압축할지, 혹은 원본 토큰 단위로 확장할지를 결정한다.

압축을 선택한 청크는 \(e^{cnk}_i=\phi(c_i)\) 형태로, 확장을 선택한 청크는 \(\{e_{i,1},\dots,e_{i,k}\}\) 형태로 디코더에 전달된다. 정책은 next-paragraph prediction perplexity의 음수값을 보상으로 삼아 학습되며, 이는 곧 디코더의 다음 토큰 예측 정확도가 높아질수록 높은 보상을 받도록 설계되었음을 의미한다. 이 방식은 대부분의 청크를 압축 상태로 유지하면서도, 답변 생성에 중요한 일부 청크는 확장해 정확도를 유지할 수 있도록 한다.

Final Processing by Decoder-only Foundation Model

• 입력: 질문 임베딩 \(\{e_1,\dots,e_q\}\), 압축 청크 임베딩 \(\{e^{cnk}_1,\dots,e^{cnk}_{L'}\}\), 확장 청크 토큰 임베딩 \(\{e_{j,1},\dots,e_{j,k}\}\)
• 모듈: Decoder-only Foundation Model \(M_{dec}\)
• 출력: 최종 답변 시퀀스 \(y\)

최종 단계에서 디코더는 질문 임베딩과 선택적으로 압축/확장된 컨텍스트 임베딩을 모두 받아, 오토리그레시브 방식으로 답변 \(y\)를 생성한다.

$$y \sim M_{dec}\big(\{e_1,\dots,e_q,\, e^{cnk}_{1:L'},\, e_{j,1},\dots,e_{j,k}\}\big)$$

압축률이 높아질수록 KV-cache 크기와 TTFT(latency to first token)가 크게 감소하며, selective expansion이 더해지면 성능 저하 없이 효율성을 극대화할 수 있다. 이 단계는 REFRAG의 설계 목표인 속도·메모리 최적화와 정확도 유지를 동시에 달성하는 핵심이다.

Experiments

Verification of Inference Accelearation

이 실험 REFRAG가 기존 방법(특히 CEPE 및 LLaMA)에 비해 추론 속도를 얼마나 가속할 수 있는지를 보여준다. 평가 지표는 TTFT(latency to first token), TTIT(time to iterative token), Throughput(tokens/sec) 세 가지이다.

실험 결과, 컨텍스트 길이가 16,38416{,}38416,384일 때 REFRAG는 TTFT에서 캐시 사용 시 16.53배, 비캐시 시 8.59배의 가속을 달성하였다. 이는 CEPE의 2.01배(캐시) 및 1.04배(비캐시)를 크게 상회한다. Throughput 역시 최대 6.78배 향상되어, 단순히 첫 토큰만 빠른 것이 아니라 전체 토큰 생성 속도에서도 우월함을 입증했다. 이 결과는 REFRAG의 핵심 설계(경량 인코더와 병렬 처리, selective expansion)가 실제 추론 환경에서 강력한 가속 효과를 제공함을 뒷받침한다.

Perplexity under varying compression rates

이 실험은 REFRAG의 Selective Compression Policy 효과를 검증한 것이다. 다양한 압축률 \(k\)에서, 전체 청크 중 일부만 압축하지 않고 확장할 때의 성능(perplexity)을 비교하였다. 정책은 4가지 기준으로 나뉜다:

• RL Policy: 강화학습 기반 학습된 정책
• Perplexity-desc: perplexity가 낮은 청크부터 압축
• Perplexity-asc: perplexity가 높은 청크부터 압축
• Random: 임의 선택

결과적으로, RL 정책이 다른 휴리스틱이나 랜덤 방식보다 일관되게 낮은 perplexity를 보여주었다. 이는 RL 정책이 답변에 중요한 청크를 효과적으로 식별하여 원본 형태로 유지함으로써, 압축률이 높아져도 성능 저하를 최소화할 수 있음을 시사한다.

RAG performance comparison under retriever scenarios

이 실험은 REFRAG를 실제 RAG 시스템에서 적용했을 때의 성능을 보여준다. 두 가지 상황이 비교된다:

1. Strong Retriever: 검색기로부터 높은 품질의 문서가 주어지는 경우
2. Weak Retriever: 검색 품질이 낮아 노이즈가 많은 경우

결과적으로, strong retriever 상황에서는 REFRAG가 LLaMA와 거의 동등한 성능을 유지하면서도 훨씬 낮은 추론 지연(latency)으로 동작하였다. 반면 weak retriever 상황에서는 REFRAG가 오히려 LLaMA보다 더 나은 성능을 보였다. 이는 REFRAG의 selective compression이 노이즈 문서의 불필요한 계산을 억제하는 역할을 하기 때문으로, 효율성과 견고함을 동시에 달성했음을 보여준다.

Perplexity on output tokens given fixed-length context

이 실험은 출력 토큰 \(x_{s+1:s+o}\)의 perplexity를 베이스라인들과 비교 평가한 실험이다. 입력 컨텍스트 길이는 \(s=2048\), 출력 길이 \(o\)는 \(\{512, 1024, 2048\}\)로 다양하게 설정되었다. 결과적으로, REFRAG-8과 REFRAG-16이 CEPE 및 다른 축약 모델들(LLaMA-256 등)보다 낮은 perplexity를 보였다. 특히 LLaMA-256은 REFRAG-8과 같은 입력 크기를 사용하지만, REFRAG-8이 일관되게 더 좋은 성능을 기록했다. 이는 청크 임베딩 기반 압축이 단순한 입력 절단보다 훨씬 효과적임을 보여준다.

Perplexity on output tokens given extended context lengths

이 실험은 출력 길이를 \(o=2048\)로 고정하고, 컨텍스트 길이 \(s\)를 \(s \in \{ 4096, 8192, 16384\}\)로 늘려가며 perplexity를 측정한 실험이다. REFRAG는 훈련 시 \(s+o=6144\) 토큰만 보았음에도, 긴 문맥 길이에서도 성능 저하 없이 확장성을 보였다. 예를 들어, \(k=16\) 압축률에서는 CEPE보다 평균적으로 9.3% 낮은 perplexity를 기록했고, 동시에 TTFT는 LLaMA 대비 16.53배, CEPE 대비 2.01배 빠르게 동작하였다. 압축률을 32로 올렸을 때도 CEPE와 동등한 perplexity를 유지하면서 TTFT 가속은 30배 이상으로 증가했다.

RAG performance under strong retriever

이 실험은 검색된 passage 수(\(N\))에 따라 RAG 성능이 어떻게 변하는지를 비교한 것이다. REFRAG는 LLaMA와 동등하거나 더 나은 성능을 유지했으며, 특히 passage 수가 많아질수록 계산량과 메모리 사용 측면에서 훨씬 효율적이었다. CEPE와 비교했을 때도, REFRAG는 더 낮은 latency로 더 좋은 답변 품질을 제공했다. 이 실험은 REFRAG가 RAG 환경에서 실제로 확장 가능한 대안임을 보여주는 핵심 근거이다.

NR 카테고리 내 다른 글 보러가기