기술 요약: STORM (Stepwise Token Optimization with Reward-Guided Beam Search)

문제 정의

현대적인 정보 검색(Information Retrieval, IR)은 효율성과 효과성 사이의 트레이드오프에 직면해 있습니다. 밀집(Dense) 및 학습된 희소(Learned-sparse) 신경망 모델은 높은 검색 성능을 제공하지만, 모델이 변경될 때마다 재구축해야 하는 특수 인덱스를 필요로 하므로 상당한 인프라 비용을 발생시킵니다. 반면, BM25와 같은 어휘적(Lexical) 검색기는 표준 역색인(Inverted Index) 상에서 매우 효율적이고 투명하지만, 사용자의 질의에 포함되지 않은 용어를 관련 문서가 포함하고 있는 "어휘 불일치(Vocabulary Mismatch)" 문제로 인해 효과가 떨어집니다.

대규모 언어 모델(LLM)은 이 간극을 메우기 위해 질의를 재작성(Rewrite)할 수 있지만, 기존 방식들은 다음과 같은 한계를 가집니다:

1. **프롬프트 기반 재작성(Prompt-based rewriters)**은 형태는 잘 갖추었으나 검색에는 효과적이지 않은 용어나 긴 의사 문서(Pseudo-documents)를 생성하는 경우가 많으며, 생성 과정 중에 검색기를 참조하지 않기 때문에 높은 추론 지연 시간을 초과합니다.
2. 강화 학습(RL) 접근법은 검색 보상(Retrieval Reward)에 대해 학습할 때 일반적으로 전체 시퀀스가 생성된 후에만 보상을 관찰합니다. 이러한 시퀀스 수준의 감독은 탐색 병목 현상을 생성하여, 어떤 특정 어휘적 용어가 보상에 기여했는지 식별하기 어렵게 만들며, 이는 느린 수렴이나 일반적인 확장(Generic expansions)으로의 붕괴를 초래합니다.

방법론

저자들은 보상 유도 빔 서치(Reward-guided beam search)를 통해 검색 보상을 토큰 수준의 신호로 변환하는 자가 지도 학습 프레임워크인 STORM(Stepwise Token Optimization with Reward-guided beaM search)을 소개합니다.

핵심 메커니즘: 보상 유도 빔 서치 (Reward-Guided Beam Search)

확률에 따라 후보를 선택하는 표준 빔 서치와 달리, STORM은 생성 과정 중에 검색 시스템 자체를 가이드로 사용합니다.

• 단계별 스코어링 (Stepwise Scoring): 각 생성 단계에서 후보 키워드 시퀀스는 BM25 역색인에 대해 점수가 매겨집니다.
• 가지치기 (Pruning): 낮은 검색 보상을 보이는 브랜치는 즉시 가지치기 됩니다. 이는 모델이 LLM의 내부 파라미터 지식에만 의존하는 대신, 검색에 입증된 효과가 있는 어휘에 탐색을 집중하도록 합니다.
• 학습 루프 (Training Loop): 이 프레임워크는 GCN(Generative Cooperative Networks) 아키텍처를 적응시킵니다. 학습된 판별기(Discriminator) 대신, 정의된 검색 보상 함수를 사용합니다.
  • 제안 분포 (Proposal Distribution): 학습 샘플은 표준 뉴클리어스 샘플링(Nucleus sampling)과 보상 유도 빔 서치의 혼합물로부터 추출됩니다.
  • 목표 (Objective): 정책(Policy)은 검색 보상에 의해 가중치가 부여된 타겟 분포와 KL 발산(KL divergence)을 최소화하도록 업데이트됩니다.

보상 함수

보상 함수 $R(\hat{y})$는 효율성을 저해하면서 검색 품질을 최적화하기 위해 세 가지 구성 요소를 결합합니다:

1. 검색 품질 (Retrieval Quality): Cross-encoder 의사 레이블(OpenSearch Project MS-MARCO Hard Negatives에서 증류됨)에 대한 nDCG@100으로 측정됩니다.
2. 어휘 페널티 (Vocabulary Penalty): 생성된 용어의 문서 빈도(Document Frequency, DF)를 기반으로 한 페널티로, 모델이 고빈도이며 변별력이 낮은 용어를 방출하는 것을 억제합니다.
3. 길이 정규화 (Length Regularization): 토큰 선택을 위한 사전 확률(Prior) 역할을 하고 시퀀스 길이를 정규화하여, 지나치게 길거나 정보가 없는 문자열 생성을 방지하는 로그 가능도(Log-likelihood) 항입니다.

주요 기여

1. 자가 지도 학습 (Self-Supervised Training): 인간이 작성한 재작성문 없이도 작동하는 어휘적 질의 확장 방법으로, 감독은 전적으로 생성된 확장에 대한 검색기의 점수로부터 파생됩니다.
2. 토큰 수준 탐색 (Token-Level Exploration): 보상 유도 확률적 빔 서치를 도입하여, 지연된 시퀀스 수준의 검색 피드백을 즉각적인 토큰 수준의 신호로 변환함으로써 모델이 생성 중에 검색에 효과적인 어휘를 발견하도록 유도합니다.
3. 인프라 경량 효율성 (Infrastructure-Light Efficiency): STORM은 짧은 키워드 문자열을 생성하므로, 훨씬 더 많은 용어(최대 8배)를 포함하는 질의를 발행함에도 불구하고 일반 BM25와 유사한 추론 지연 시간을 유지하며 경쟁력 있는 효과성을 달ach합니다.

실험 결과

저자들은 Qwen3 백본(0.6B에서 8B 파라미터 범위)을 사용하여 인도메인(In-domain), 아웃도메인(Out-of-domain), 다국어 벤치마크에서 STORM을 평가했습니다.

• 인도메인 (TREC DL, MS-MARCO): STORM은 모든 모델 크기에 걸쳐 프롬프트 기반 재작성기(HyDE, MuGI, W2P) 및 RL 기반 베이스라인(QUESTER)보다 일관되게 우수한 성능을 보입니다. 8B 모델의 경우, 경쟁력 있는 LLM 재작성기들과 대등하거나 이를 능가합니다.
• 아웃도메인 (BEIR): STORM은 12개의 BEIR 컬렉션에서 가장 높은 평균 nDCG@10을 달성했습니다. 특히, 0.6B STORM 모델조차 모든 0.6B 베이스라인을 평균적으로 능가하여 강력한 일반화 능력을 입증했습니다.
• 제로샷 다국어 (MIRACL): 영어 MS-MARCO 데이터로만 학습되었음에도 불구하고, STORM-8B는 18개 언어에 제로샷으로 전이되어 평균 48.4의 nDCG@10을 달성했습니다. 이는 순수하게 어휘적 백엔드에 의존함에도 불구하고, 전용 다국어 밀집 검색기(mDPR, mContriever, mColBERT)를 평균적으로 능가하는 수치입니다.
• 독점 모델과의 비교: STORM-8B(Qwen3)는 훨씬 더 큰 GPT-4.1 기반의 질의 재작성기들(예: MuGI-GPT4, Q2D-GPT4)보다 인도메인 및 BEIR 데이터셋 전반에서 평균적으로 더 우수한 성능을 보였으며, 이는 보상 신호가 검색 작업을 위한 Chain-of-Thought 프롬프팅보다 더 효과적임을 시사합니다.
• 효율성: STORM은 훨씬 더 많은 용어(최대 8배)를 포함하는 질의를 발행함에도 불구하고 일반 BM25와 동등한 검색 지연 시간(~0.01s)을 유지합니다. 이는 보상 함수가 비용이 많이 드는 인덱스 탐색을 유발할 수 있는 고빈도 용어를 억제하기 때문입니다. 생성 지연 시간 또한 경쟁하는 LLM 재작성기(예: W2P의 34.2s 대비 STORM-8B의 4.39s)보다 낮습니다.

의의 및 주장

본 논문은 STORM을 밀집 신경망 검색의 경쟁력 있고 인프라가 가벼운 대안으로 포지셔닝합니다. 주요 의의는 다음과 같습니다:

• **어휘적 확장(Lexical expansion)**이 밀집 인덱스나 거대한 독점 모델 없이도 효과적으로 최적화될 수 있음을 보여주었습니다.
• **보상 유도 디코딩(Reward-guided decoding)**이 비미분 가능한 검색 목적 함수와 생성을 정렬하기 위한 유효한 범용 도구임을 입증하여, 시퀀스 수준 RL의 탐색 병목 현상을 해결했습니다.
• **제로샷 전이(Zero-shot transfer)**가 순수 어휘적 접근 방식으로도 가능하다는 것을 보여줌으로써, 교차 언어 검색을 위해 밀집 의미 매칭이 반드시 필요하다는 가설에 도전했습니다.

저자들은 결론적으로 STORM이 효과성과 효율성 사이의 견고한 균형을 제공하며, 인덱스 재구축 비용이 높거나 지연 시간이 중요한 검색 엔진 및 RAG(Retrieval-Augmented Generation) 시스템에 실용적인 솔루션이라고 주장합니다.

한계점

저자들은 STORM이 인도메인 작업에서 학습된 희소 검색기(SPLADE-v2 등)와의 격차를 좁히기는 하지만, 용어 가중치 메커니즘을 통해 풍부한 문맥적 관계를 인코딩하는 밀집 모델의 능력을 완전히 대체하지는 못한다는 점을 인정하며, 이는 아마도 밀집 모델의 특성 때문일 것이라고 언급했습니다. 또한, 최적의 생성 길이는 데이터셋에 따라 달라지므로 지연 시간과 효과성 사이의 트레이드오프가 필요합니다.