LightRetriever: A LLM-based Text Retrieval Architecture with Extremely Faster Query Inference

기술 요약: LightRetriever

문제 정의

현재의 거대 언어 모델(LLM) 기반 텍text 검색 시스템은 일반적으로 문서와 쿼리를 모두 깊은 파라미터화된 LLM으로 처리하여 밀집(dense) 또는 희소(sparse) 표현을 생성하는 대칭적 이중 인코더(dual-encoder) 아키텍처를 채택합니다. 이러한 방식은 심층적인 문맥 모델링을 통해 검색 능력을 크게 향상시키지만, 온라인 배포 시 결정적인 병목 현상을 초래합니다. 문서는 오프라인에서 사전 인코딩 및 인덱싱될 수 있지만, 검색 쿼리는 실시간으로 도착하며 즉각적인 인코딩을 요구합니다. 온라인 서버(예: A800 GPU)에서 전체 크기의 LLM을 쿼리 추론용으로 서빙하는 것은 높은 지연 시간, 낮은 처리량, 과도한 자원 소비를 야기합니다. 본 논문은 근본적인 긴장 관계를 제기합니다. 즉, 문서는 헤비한 사전 계산의 혜택을 받는 반면, 쿼리 인코더는 엄격한 온라인 제약 조건 하에서 작동해야 함에도 불구하고, 기존 방식들은 쿼리가 문서와 동일하게 계산 비용이 많이 드는 심층 모델링을 거치도록 강제한다는 점입니다.

방법론: LightRetriever

저자들은 문서와 쿼리 인코더 사이의 대칭성을 명시적으로 깨뜨리는 새로운 하이브리드 검색 아키텍처인 LightRetriever를 제안합니다. 핵심 통찰은 문서에는 LLM의 전체 모델링 능력이 필요하지만, 쿼리는 강력한 검색 성능을 달성하기 위해 그만큼 무거운 처리를 할 필요가 없다는 것입니다. 이 아키텍처는 주요 모델링 비용을 문서 측으로 옮기는 대신 쿼리 측을 극도로 가벼운 연산으로 축소합니다.

1. 비대칭 아키텍처

• 문서 측: 밀집 및 희소 인코딩 모두를 위해 전체 크기의 LLM 인코더($Enc_d$)를 유지합니다. 문서는 오프라인에서 사전 계산됩니다.
• 쿼리 측: 온라인 서빙 중 심층 트랜스포머 추론을 제거합니다. 쿼리 인코더($Enc_q$)는 단순한 임베딩 룩업(lookup) 메커니즘으로 축소됩니다.

2. 밀집 검색 메커니즘 (Dense Retrieval Mechanism)

밀집 검색을 위해 LightRetriever는 3단계 프로세스를 채택합니다:

• 학습 (Training): 전체 크기의 LLM이 시스템 학습에 사용됩니다. 주어진 쿼리 토큰 $t_i$(태스크 지침이 앞에 붙음)에 대해, 모델은 토큰 레벨 표현 $v^{den}_{t_i} = Enc_q(Inst; t_i)$를 생성합니다. 최종 밀집 쿼리 벡터 $v^{den}_q$는 이 토큰 벡터들의 평균을 통해 계산됩니다:
  $$v^{den}_q = \frac{1}{n} \sum_{i=0}^{n-1} v^{den}_{t_i}$$
  시스템은 쿼리와 관련 문서 벡터를 정렬하기 위해 대조 손실(contrastive loss, 식 2)을 사용하여 최적화됩니다.
• 캐싱 (Caching): 추론 중에 쿼리 토큰들은 서로 깊게 상호작용하지 않으므로, 토큰 벡터는 캐싱이 가능합니다. 서빙 전, 시스템은 LLM의 전체 어휘 사전($V$)에 대한 토큰 임베딩을 사전 계산하여 단일 임베딩 룩업 행렬 $E$ (형태 $[V, H]$)에 저장합니다.
• 온라인 서빙 (Online Serving): 추론 시, 쿼리는 토큰화되며, 각 토큰의 캐시된 임베딩을 찾아 평균을 내는 방식으로 밀집 벡터가 생성됩니다. 이는 대규모 트랜스포머의 순전파(forward pass)를 단일 임베딩 룩업 및 평균 연산으로 대체하며, RAM에 상주하는 임베딩 레이어만을 필요로 합니다.

3. 희소 검색 메커니즘 (Sparse Retrieval Mechanism)

희소 검색의 경우 쿼리 측은 더욱 가볍습니다:

• 쿼리 인코딩: 희소 쿼리 벡터 $v^{spr}_q$는 토큰화로부터 직접 유도되는 용어 기반의 학습 불가능한(unlearnable) 벡터이며, 토큰 $t$의 값은 단순히 쿼리 내 해당 토큰의 빈도수입니다. 별도의 인코더는 필요하지 않습니다.
• 문서 인코딩: 문서는 전체 LLM에 의해 처리됩니다. 최종 레이어의 은닉 상태(hidden states)는 모델의 언어 모델링 헤드를 통해 어휘 공간으로 투영됩니다. 희소성을 유도하기 위해 ReLU 활성화 함수와 로그 포화(log-saturation) 함수가 적용된 후, 시퀀스 길이에 대해 맥스 풀링(max-pooling)이 수행됩니다.
• 하이브리드 스코어링: 밀집 점수와 희소 점수를 선형 보간하여 최종 검색 순위를 산출합니다.

주요 기여

1. 극한의 쿼리 추론 속도: 심층 쿼리 모델링을 임베딩 룩업으로 대체함으로써, LightRetriever는 A800 GPU에서 전체 Llama-8b 모델을 서빙하는 것과 비교하여 쿼리 인코딩에서 1000배의 속도 향상을 달성했습니다.
2. 처리량 개선: 이 아키텍처는 엔드 투 엔드(end-to-end) 검색 처리량(초당 쿼리 수, QPS)을 10배 증가시킵니다.
3. 성능 유지: 쿼리 측의 급격한 계산 감소에도 불구하고, 이 방법은 다양한 벤치마크에서 전체 대칭적 LLM 기반 검색기와 거의 유사한 95%의 검색 성능을 유지합니다.
4. 일반화 능력: 이 접근 방식은 다양한 파운데이션 모델(Llama-1B/3B/8B, Qwen-1.5B/3B/7B)과 다양한 검색 태스크(BeIR, CMTEB)에서 검증되었으며, 태스크별 아키텍처 변경 없이도 견고함을 입증했습니다.

실험 결과

본 논문은 15개의 영어 태스크(BeIR)와 8개의 중국어 태스크(CMTEB Retrieval)에서 LightRetriever를 평가합니다.

• 속도: 1M 개의 패시지에 대해 65,536개의 쿼리를 대상으로 한 벤치마크에서, LightRetriever는 전체 엔드 투 엔드 검색 시간을 약 119초(Full-Llama8b 기준)에서 약 9.36초로 단축하여, 전체 모델의 549 QPS 대비 6,999 QPS의 처리량을 달애했습니다.
• 효과성:
  • LightRetriever-Llama3.1-8b는 BeIR에서 54.4, CMTEB-R에서 63.0의 nDCG@10을 달려했습니다.
  • 이는 전체 대칭적 Llama3.1-8b 베이스라인(56.8 및 67.6)과 비교했을 때 각각 -2.4 및 -4.6 포인트의 미미한 저하를 나타내며, 동시에 무거운 쿼리 추론 비용을 피할 수 있게 해줍니다.
  • 이 방법은 BGE-m3(BeIR에서 49.6)와 같은 다른 효율적인 베이스라인보다 우수한 성능을 보이며, E5-Mistral 및 LLM2Vec의 성능에 근접합니다.
• 절제 연구 (Ablation Studies):
  • 전체 문서 인코더를 제거했을 때(대칭적 경량화 모델) 성능이 심각하게 저하되었습니다(예: Beir에서 -13.8). 이는 심층 문서 모델링의 필요성을 확인시켜 줍니다.
  • 쿼리 학습 시에만 임베딩 백(embedding bag)을 사용했을 때(트랜스포머 없이)도 상당한 성능 손실(-11.2 on Beir)이 발생했습니다. 이는 효과적인 표현을 학습하기 위해서는 비록 추론 시에는 버리더라도 학습 단계에서는 심층 모델링이 필수적임을 나타냅니다.

의의 및 주장

본 논문은 LightRetriever가 지연 시간에 민감하고 높은 처리량이 요구되는 검색 시나리오를 위한 확장 가능한 패러다임을 제공한다고 주장합니다. 이는 LLM 기반 검색의 의미론적 이해 이점을 희생하지 않으면서, 실시간 쿼리 인코딩을 위한 거대 LLM 배포의 비효식성을 해결합니다. 저자들은 이 작업이 계산 부담을 온라인 쿼리 측에서 오프라인 문서 측으로 전이함으로써, 생산 환경의 검색 엔진에 필요한 응답성을 유지하면서도 거대 모델의 힘을 활용할 수 있게 한다고 강조합니다. 이 방법은 LLM 기반 검색기의 쿼리 측 온라인 효율성 병목 현상을 극단적인 디커플링(decoupling)을 통해 명시적으로 다룬 첫 번째 사례로 제시됩니다.