Lightweight Chunk Selection for Mobile Retrieval-Augmented Generation

기술 요약: 모바일 RAG를 위한 경량 청크 선택(Lightweight Chunk Selection)

문제 정의
검색 증강 생성(RAG)은 외부 지식을 통합하여 대규모 언어 모델(LLM)을 강화하지만, 모바일 및 에지 기기에서 검색된 컨텍스트를 처리하는 데 따르는 계산 및 메모리 비용이 배포를 저해합니다. 일반적인 완화 전략은 검색된 단일 청크로 생성을 제한하는 것입니다. 그러나 표준 밀집 검색(dense retrieval)에 의해 반환된 최상위 순위의 청크가 반드시 증거로서 충분하다는 것을 의미하지는 않습니다. 검색은 사실적 지원이 아닌 의미론적 유사성에 최적화되어 있기 때문입니다. 프롬프트 압축이나 크로스 인코더 재순위화(cross-encoder reranking)와 같은 기존의 컨텍스트 축소 방법은 종종 엄격한 모바일 리소스 예산을 초과하는 추가적인 대형 모델이나 압축기를 필요로 합니다. 또한, 정답 문자열 포함 여부에 의존하는 표준 평가 지표는 노이즈가 많습니다. 즉, 어떤 청크는 정답 문자열을 포함하고 있지만 필요한 관계적 증거를 제공하지 못할 수 있고, 그 반대의 경우도 발생할 수 있습니다.

방법론
저자들은 이 문제를 증거 정렬(evidence-alignment) 작업으로 취급하는 경량화된 LLM 결합형 청크 선택 프레임을 제안합니다. 이 시스템은 LLM의 질문 전용 웜업(warmup) 단계와 병렬로 작동하며, 추가적인 추론 비용을 피하기 위해 계산을 재사용합니다.

1. 특징 통합(Feature Integration): 선택기는 세 가지 상호 보완적인 특징 소스를 집계합니다:

   • 질문 은닉 상태($h$): 생성기의 질의에 대한 의미론적 해석을 나타내는 웜업 패스의 최종 토큰 은닉 상태입니다.
   • MoE 라우팅 신호($r$): 혼합 전문가(Mixture-of-Experts, MoE) 아키텍처에서 유도된 라우팅 가중치 기반 전문가 응답 점수입니다. 이는 모델이 전문가 간에 계산을 어떻게 분산하는지에 대한 구조화된 정보를 포착하며, 은닉 상태를 보완하는 신호 역할을 합니다.
   • 청크 임베딩($E$): 상위 $K$개의 검색 후보들에 대한 밀집 임베딩으로, 후보 측의 증거 기하학적 구조를 보존합니다.

2. 선택기 아키텍처(Selector Architecture): 컴팩트한 다층 퍼셉트론(MLP)은 이러한 결합된 특징들을 검색된 청크들과 동일한 임베딩 공간 내의 "증거 프로토타입(evidence prototype)"으로 매핑합니다. 최종 청크는 이 예측된 프로토타입과 후보 임베딩 사이의 코사인 유사도를 기반으로 선택됩니다. 이 방식은 질의와 모든 후보 사이의 명시적인 토큰 수준 상호작용을 피함으로써 계산 비용을 절감합니다.

3. 예산 인식 차원 축소(Budget-Aware Dimension Reduction): 더 엄격한 배포 제약 조건을 위해 선택적 기능 선택 모듈이 도입됩니다. 비지도 학습 방식의 압축 대신, 본 방법은 훈련된 전체 차원 선택기에 대한 그래디언트 기반 민감도 분석을 사용하여 가장 정보량이 많은 입력 차원을 식별하고 유지합니다. 이를 통해 선택기의 크기와 정확도 사이의 제어 가능한 트레이드오프를 가능하게 합니다.

4. 의미론적 충분성을 통한 지도 학습(Supervision via Semantic Sufficiency): 표준 데이터셋의 노이즈를 해결하기 위해, 저자들은 정답 문자열 포함 여부가 아닌 "증거적 충분성(evidential sufficiency)"에 기반하여 그라운드 트루스 레이블을 구축합니다. 자동화된 어노테이터를 사용하여, 청크가 질문에 답하기 위한 충분한 정보를 포함하고 있는 경우에만 양성으로 레이블링함으로써, (정답 문자열은 있으나 컨텍스트가 없는) 거짓 양성과 (의미는 통하나 문자열은 없는) 거짓 음성을 필터링합니다.

주요 기여

• 정식화: 본 논문은 모바일 RAG 청크 선택을 생성기 측의 질의 표현(은닉 상태 및 MoE 라우팅)을 재사용하는 경량화된 증거 정렬 문제로 정식화합니다.
• 모델 설계: 질문 은닉 상태, MoE 라우팅 유도 전문가 신호, 그리고 청크 임베딩을 융합하여 증거 프로토타입을 예측하는 컴팩트한 선택 모델을 제안합니다.
• 데이터셋 구축: 새로운 지도 학습 전략을 도입하여, 생성 전 청크 선택을 위한 지도 학습이 가능하도록 의미론적 청크 정답성을 기반으로 한 증거적 충분성 레이블을 생성했습니다.
• 효율성: 입력 차원을 줄이면서도 경쟁력 있는 성능을 유지하는 태스크 인식형 특징 선택 전략을 포함하여, 에지 기기 배포에 적합하도록 설계되었습니다.

결과
실험은 TriviaQA, PopQA, MS MARCO Passage Ranking의 세 가지 데이터셋을 대상으로 수행되었습니다.

• 성능: 제안된 선택기는 평가된 경량 베이스라인들 중 세 데이터셋 모두에서 가장 높은 Rank-1 청크 선택 정확도를 달고했습니다. 이는 가장 강력한 베이스라인인 KNRM보다 평균 2.49%포인트 향상된 수치이며, 구체적으로는 TriviaQA에서 4.60포인트, PopQA에서 9.13포인트, MS MARCO PR에서 4.61포인트의 이득을 보였습니다.
• 절제 연구(Ablation): 검색된 청크 임베딩이 주요 성능 신호를 제공했습니다. 그러나 이를 은닉 상태 및 라우팅 신호와 결합했을 때 가장 견고한 결과를 얻었으며, 특히 여러 청크가 표면적인 엔티티를 공유하지만 단 하나만이 사실적 지원을 제공하는 PopQA에서 효과적이었습니다.
• 예산 분석: 축소된 특징 변형(예: 입력 차원의 20%만 사용)은 경쟁력 있는 성능을 유지하면서도, 가장 강력한 컴팩트 신경망 베이스라인(KNRM)보다 우수한 성능을 보이면서 훨씬 적은 파라미터를 사용했습니다 (7M vs. 약 48M).

의의
본 논문은 LLM 측의 질의 표현(은닉 상태 및 MoE 라우팅)을 활용하고 이를 검색 측의 후보 임베딩과 정렬하는 것이 모바일 적용 가능한 RAG를 위한 효과적이고 파라미터 효율적인 전략임을 입증합니다. 검색 순위에서 증거 정렬으로 초점을 전환하고, 추가적인 대형 모델 없이 경량 선택기를 활용함으로써, 본 방법은 자원이 제한된 환경에서의 컨텍스트 관리라는 핵심 병목 현상을 해결합니다. 이 연구는 크로스 인코더나 생성형 압축기의 무거운 추론 오버헤드 없이도 고품질의 청크 선택이 가능하다는 것을 보여주며, 이를 통해 에지 배포를 위한 견고한 RAG 구현을 가능하게 합니다.