Q-RAG: 가치 기반 임베더 학습을 통한 긴 컨텍스트 다단계 검색

기존의 다단계 RAG는 LLM 자체를 파인튜닝해야 하므로 비용이 매우 높고 대형 모델 적용이 어려웠다. 이 논문은 LLM 대신 가벼운 Embedder 모델만 강화학습으로 최적화하여, 훨씬 적은 비용으로도 초장문 컨텍스트에서 복잡한 추론 검색을 수행할 수 있음을 입증했다.

기존 RAG의 핵심인 Attention Mechanism은 컨텍스트가 길어질수록 관련 없는 정보에 주의력이 분산되는 Attention Dilution 문제를 겪는다. 특히 여러 단계의 정보를 조합해야 하는 다단계 검색에서는 첫 번째 검색 결과가 부정확할 경우 이후 단계가 모두 실패하는 연쇄 오류가 발생하기 쉽다.

Q-RAG는 이 문제를 해결하기 위해 검색 과정을 강화학습의 Markov Decision Process(MDP)로 모델링한다. 에이전트는 현재까지 찾은 정보(State)를 바탕으로 다음에 어떤 청크를 가져오는 것이 최종 정답에 도달할 가능성(Q-value)이 높은지 판단한다. 이는 단순히 유사도만 측정하는 기존 방식과 달리, 정답을 맞히기 위해 '필요한 조각'을 전략적으로 찾아 나서는 탐색 과정과 같다.

결과적으로 LLM은 전체 문서를 읽을 필요 없이 에이전트가 골라준 핵심 요약 정보만 처리하면 되므로, 1,000만 토큰에 달하는 방대한 데이터에서도 추론의 정확도를 유지하면서 연산 비용을 획기적으로 줄일 수 있다.

다단계 검색을 유한 시간 지평 Markov Decision Process(MDP)로 정의한다. 상태 $s_t$는 초기 쿼리와 이전에 선택된 청크들의 순서 리스트로 구성되며, 액션 $a_t$는 문서 내의 특정 텍스트 청크를 선택하는 행위이다.

Q-함수는 State Embedder $E_s$와 Action Embedder $E_a$의 내적으로 근사한다. $Q_\theta(s, a^i) = \langle E_s(s; \theta_1), E_a(a^i, i; \theta_2) \rangle$ 계산을 통해 각 청크의 유용성을 수치화한다. [상태 벡터와 청크 벡터 입력 → 내적 연산 수행 → 해당 청크 선택 시 기대 보상 출력 → 가장 높은 값의 청크 선택].

학습에는 최근 제안된 PQN(Parallel Q-Network) 알고리즘을 기반으로 한 Soft Q-learning을 사용한다. Replay Buffer 없이 온폴리시(On-policy) 방식으로 학습하여 메모리 효율을 극대화했으며, $\lambda$-return을 사용하여 학습의 안정성을 높였다. [보상 및 다음 상태 가치 입력 → 가중치 합산 연산 → 타겟 가치 생성 → 예측값과의 오차를 줄이는 방향으로 가중치 갱신].

BabiLong 벤치마크의 QA3(3개의 사실 확인 필요) 작업에서 1,000만 토큰 컨텍스트까지 성능 저하 없이 높은 정확도를 유지했다. 이는 GPT-4나 Gemini 2.0 같은 상용 모델들이 컨텍스트가 길어질수록 성능이 급격히 하락하는 것과 대조적이다.

RULER 벤치마크의 NIAH(Needle-in-a-Haystack) 테스트에서 100만 토큰 기준 평균 99.7%의 정확도를 기록하며 사실상 완벽한 검색 성능을 보였다. 특히 다중 쿼리(MQ) 및 다중 값(MV) 추출 상황에서도 기존 SOTA 모델들을 압도했다.

효율성 측면에서 Q-RAG는 단일 A100 GPU에서 12시간 이내에 학습이 완료된다. 이는 수십 대의 GPU 클러스터가 필요한 기존 LLM 파인튜닝 방식 대비 수십 배 이상의 비용 절감 효과를 의미한다.

Q-RAG의 핵심 아키텍처는 State Embedder와 Action Embedder의 이중 구조다. Action Embedder에는 Rotary Position Embedding(RoPE)을 적용하여 청크의 위치 정보를 보존하며, 본문에서는 RoPE가 임의의 연속 함수를 근사할 수 있음을 수학적으로 증명(Universal Approximation Theorem 변형)하여 이론적 토대를 마련했다.

시간적 추론을 위해 도입된 '상대적 위치 매핑' $\rho_t(i)$는 선택된 청크들을 기준으로 문서 전체를 구간화한다. 이는 절대적 위치 값이 컨텍스트 길이에 따라 커지는 문제를 방지하고, 에이전트가 '이미 찾은 정보의 바로 앞이나 뒤'를 탐색해야 한다는 논리적 관계를 학습하게 한다.

추론 시에는 학습된 Q-함수를 기반으로 Beam Search를 수행할 수 있다. 이는 단순 탐욕적 선택(Greedy Selection)보다 더 넓은 탐색 공간을 고려하게 하여 복잡한 다단계 추론의 성공률을 높인다. 또한 Q-value 임계값을 활용한 조기 종료(Early Stopping) 메커니즘을 통해 불필요한 검색 단계를 줄여 추론 속도를 최적화했다.

현재 연구는 정답이 명시된 Support-fact 신호를 기반으로 보상을 설계했기 때문에, 정답만 주어지는 일반적인 QA 데이터셋에서 LLM의 피드백만으로 학습하는 방식에 대해서는 추가적인 연구가 필요하다.