Don't do RAG | Believe I.Y.

Quick Summary
Implementation
Experiment
느낀점
Reference

Quick Summary

RAG는 검색 지연, 문서 선택 오류, 시스템 복잡성 증가 같은 문제를 일으킨다. 최근 LLM의 컨텍스트 길이가 크게 늘어나면서, 실시간 검색 전략인 캐시 증강 생성(CAG)이 제안된다.

CAG는 모델의 컨텍스트 윈도우에 모든 관련 리소스를 미리 로드해서 KV-cache 형태로 적재하고, 인퍼런스 과정에서 모델이 이 캐시를 사용해서 추가적인 검색 과정 없이 쿼리에 답변할 수 있다.

CAG는 검색 지연을 줄이고 검색 오류를 최소화하면서도 컨텍스트 관련성을 유지할 수 있다고 한다. 여러 벤치마크 데이터셋에서 CAG가 기존 RAG보다 더 나은 성능을 보인다고 주장한다.

제한된 지식 기반을 가진 경우, CAG가 RAG보다 더 괜찮은 선택지로 고려해볼 수 있다고 한다. 왜냐하면, CAG 는 필요한 모든 정보를 모델의 context 안에 캐시 형태로 밀어 넣기 때문에 대용량 데이터를 다루는 상황이라면 적절하지 않기 때문이다. 또한 LLM 이 long context prompt 를 처리하는데 어려움을 겪는 경우가 많기 때문에, 이러한 단점을 고려해서 CAG 를 사용하는 것이 좋다.

Implementation

저자의 코드를 확인해보니 아래와 같이 kv cache 를 생성한다.

from transformers.cache_utils import DynamicCache

# Initialize dynamic cache
past_key_values = DynamicCache()

# Generate KV cache with proper device placement
with torch.no_grad():
    outputs = model(
        input_ids=input_ids,
        past_key_values=past_key_values,
        use_cache=True,
        output_attentions=False,
        output_hidden_states=False
    )

# The model's device mapping will automatically place each layer's
# KV cache on the correct device
return outputs.past_key_values

생성한 kv cache 를 사용하는 것은 간단한데, vllm 같은 framework 를 사용하지 않고 naive 한 iteration 을 통해서 진행한다.

# Main generation loop
for _ in range(max_new_tokens):
    # Forward pass with proper device placement
    outputs = model(
        input_ids=next_token,  # Only process last token
        past_key_values=past_key_values,
        use_cache=True
    )

    # Get next token prediction (logits will be on the last device)
    next_token_logits = outputs.logits[:, -1, :]
    next_token = next_token_logits.argmax(dim=-1).unsqueeze(-1)
    
    ...

    # Update KV cache
    past_key_values = outputs.past_key_values

    # Append prediction
    output_ids = torch.cat([output_ids, next_token], dim=1)

Experiment

Baselines

Sparse Retrieval System (BM25): 일반적인 BM25. It ranks documents based on term frequency inverse document frequency (TF-IDF) and document length normalization.
Dense Retrieval System (OpenAI Indexes): LlamaIndex framework 에 openai Embedding 모델을 사용하여 구축한 vector search system.

Other Settings

Metrics: BERT-Score
Datasets: SQuAD, HotpotQA

Results

성능: CAG 와 기존 Sparse, Dense RAG 비슷한 성능을 보이는데, CAG 가 아주 약소하게 좋음
속도: CAG 가 압도적으로 빠른편이다. 데이터 사이즈를 small, medium, large 로 나눠서 모델 generation 속도를 비교하는데, 사이즈가 클수록 10배이상 빨라진다 (e.g. 2초 vs. 100초).

느낀점

평가 진행에 있어서 좀 아쉬운 점들이 있다.

BERT-Score 외에 Rouge, GPT-eval 같은 다양한 평가 지표로 진행하면 좋았을 것 같음.
CAG 는 KV-cache 를 생성하는 속도까지 고려해야 더 공평한 결과일 것이다.

이 접근에 착안해서 생각해봤는데 KV-cache 를 미리 생성하고, 해당 cache 를 검색하는 RAG 를 만드는것도 하나의 접근법일수도 있을것 같다.