Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale

기술 요약: 언어 모델이 실제로 문맥 내(In-Context) 검색을 수행할 수 있는가?

문제 제식 (Problem Statement)
언어 모델(LM)은 관련 문서 식별자를 직접 생성하기 위해 문맥 내 코퍼스(in-context corpus)에 조건화됨으로써 벡터 기반 검색의 잠재적인 대안을 제공하지만, 이 능력은 현실적인 규모에서는 여전히 미개척 영역으로 남아 있다. 기존 연구들은 독점적 시스템이나 소규모 재순위화(reranking) 작업에 집중해 왔으며, 두 가지 핵심 과제를 해결하지 못했다: (C1) 수백만 토큰을 포함하는 코퍼스로의 확장성, 그리고 (C2) 훈련 시 보았던 크기를 훨씬 초과하는 코퍼스 크기에 대한 일반화 능력이다. 현재의 LM은 백만 토큰 규모의 어텐션 비용과 긴 문맥으로의 외삽(extrapolation)에 대한 취약성 문제로 어려움을 겪고 있다. 핵심 질문은 LM이 붕괴 없이 이러한 조건 하에서 코퍼스 규모의 검색기 역할을 신뢰성 있게 수행할 수 있는가 하는 점이다.

방법론: BLOCKSEARCH 및 분석
저자들은 문맥 내 코로스 규모 검색(ICR)을 해결하기 위해 설계된 체계적인 연구이자 모델 아키텍처인 BLOCKSEARCH를 소개한다.

• 아키텍처 및 훈련: 0.6B 파라미터 백본(Qwen3-0.6B)을 기반으로 구축된 BLOCKSEARCH는 분리된 문서들을 효율적으로 처리하기 위해 블록 희소 어텐션(block-sparse attention)을 채택한다. 주요 아키텍처 및 훈련 수정 사항은 다음과 같다:

  • 무작위 문서 식별자 (Randomized Document Identifiers): 절대적 위치에 대한 과적합을 방기하기 위해, 순차적 ID 대신 각 훈련 단계마다 문서에 무작위 4자리 코드를 할당한다.
  • 인-배치 네거티브(In-Batch Negatives) 및 온-폴리시 손실(On-Policy Loss): 모델은 쿼리당 1개의 포지티브와 15개의 하드 네거티브로 훈련된다. 노출 편향(inference 시 모델이 자신의 잠재적으로 틀린 접두사에 조건화되는 현상)을 완화하기 위해 온-폴리시 보조 손실이 추가된다. 이는 모델 자신의 분포로부터 코드를 생성(rollout)하고, 인-배치 점수를 바탕으로 경로를 수정하도록 모델을 훈련하는 과정을 포함한다.
  • 평가 설정: BLOCKSEARCH는 BEIR 데이터셋(MS MARCO, NQ, HotpotQA)을 통해 평가되며, 코퍼스 크기는 모델의 고유 컨텍스트 윈도우인 32k 토큰을 훨씬 초과하는 500개에서 10,000개의 문서(최대 약 1백만 토큰)까지 다양하다.

• 실패의 메커니즘 분석: 연구는 BLOCKSEARCH가 외삽 성능을 개선함에도 불구하고, 대규모 척도(1M 토큰 근처)에서 여전히 붕괴함을 확인했다. 저자들은 이것이 관련 문서를 식별하는 능력(ranking)의 실패가 아니라, 어텐션 희석 효과(attention dilution effect) 때문임을 밝혀냈다:

  • 검색 vs 생성 간극 (Retrieval vs. Generation Gap): 10,000개의 문서가 있는 상황에서도, 개별 어텐션 헤드는 종종 골드(gold) 문서에 가장 높은 pre-softmax 점수를 할당한다 ($R_{any} \approx 1.0$).
  • 소프트맥스 붕괴 (Softmax Collapse): 그러나 코퍼스($N$)가 증가함에 따라, 무관한 문서들의 합산 기여도가 골드 문서의 점수보다 더 빠르게 소프트맥스 분모(denominator)에 영향을 미친다. 결과적으로, 정규화된 어텐션 질량(post-softmax weight)이 골드 문서에서 붕괴하며, 이로 인해 레이어의 출력이 관련 콘텐츠가 아닌 방해 요소(distractors)들의 평균값이 되어버린다.

• 제안된 완화 방법: 어텐션 희석을 방지하기 위해 저자들은 두 가지 길이 인식 조정을 제안한다:

  1. 길이 인식 소프트맥스 조정 (Length-Aware Softmax Adjustments):
     • 가산 싱크 (Additive Sinks): 확산된 어텐션 질량을 흡수하기 위해 학습된 상수를 소프트맥스 분모에 추가한다 (Streaming-LLM에서 영감을 받음).
     • 곱셈 재스케일링 (SSMax): pre-softmax 점수에 $s \cdot \log N$을 곱하여 골드 점수와 방해 요소 사이의 간격을 넓힘으로써, $N$이 증가함에 따라 골드 가중치가 사라지는 것을 방지한다.
  2. 문서 단위 희소 어텐션 (Document-Level Sparse Attention): 중간 레이어(L16)에서 모든 문서를 스코어링하고, 이후 레이어들을 위해 상위 $B$(예: 256)개의 문서에 대해서만 밀집 어텐션을 제한하는 라우팅 메커니즘을 사용하여 최종 생성 전 컨텍스트를 효과적으로 가지치기(pruning)한다.

주요 결과

• 베이스라인 성능: 초기 BLOCKSEARCH 모델은 작은 코퍼스 크기(~45k 토큰)에서 밀집 검색(dense retrieval) 성능과 일치하며, ~500k 토큰(훈련 컨텍스트의 10배)까지 일반화되지만, 1M 토근 근처에서 붕괴한다.
• 수정 사항의 영향:
  • SSMax: 전체 $N$ 스윕(sweep)에 걸쳐 성능을 유의미하게 향상시킨다. MS MARCO 1M 토큰에서 16.5%의 Recall@1을 달s achieve하며(수정되지 않은 모델의 0.2% 대비), 밀집 베이스라인(20.2%)과의 격차를 대부분 좁힌다.
  • 라우팅 (Routing): 문서 단위 라우팅은 성능을 더욱 개선하여, MS MARCO 1M 토큰에서 18.8% Recall@1을 달성하며 밀집 베이스라인에 거의 근접한다.
  • 결합 접근법: SSMax와 라우팅을 결합한 모델(BLOCKSEARCH-SSMax-routing)은 MS MARCO 1M 토큰에서 20.5% Recall@1을 달성하여, 밀집 베이스라인(20.2%)을 약간 앞지르고 동시기에 발표된 7배 더 큰 MSA-4B 모델(16.0%)을 능가한다.
• 분포 외 데이터 (LIMIT): 의미적 임베딩 유사성이 아닌 어휘적 유사성을 요구하는 LIMIT 벤치마크에서, 수정된 BLOCKSEARCH는 밀집 검색 모델이 고전하는 지점에서 밀집 베이스라인을 크게 압도한다. 850k 토큰에서 BLOCKSEARCH-SSMax-routing은 14.9% Recall@1을 달성하여 밀집 검색기(3.5%) 대비 약 3배의 성능 향상을 보여준다.
• 어텐션 천장 (Attention Ceiling): 모든 실험에서 "어텐션 천장"(적어도 하나의 헤드가 골드 문서를 1위로 랭크하는 쿼리의 비율)은 100%에 가깝게 유지되었다. 이는 검색 신호가 내부적으로는 존재하지만, 어텐션 희석으로 인해 읽기/생성 단계에서 손실됨을 확인시켜 준다.

의의 및 주장
본 논문은 백만 토큰 규모에서의 문맥 내 검색에 대한 최초의 체계적인 연구를 제공하며, 특정 조건 하에서 이를 고전적 밀집 검색의 실행 가능한 대안으로 확립한다.

• 실행 가능성: 아키텍처적 수정(SSMax 및 라우팅)을 통해, LM은 표준 벤치마크에서 밀집 검색과 대등하거나 이를 능가할 수 있으며, 비의미적 유사성을 요구하는 작업에서 이를 크게 앞지를 수 있다.
• 효율성: 제안된 0.6B 모델은 동시기의 대규모 컨텍스트 모델(MSA-4B)보다 7배 적은 파라미터로 이러한 결과를 달성하며, 생성적 검색(generative retrieval)과 달리 코퍼스가 변경될 때 재학습할 필요가 없다.
• 근본적 병목 현상: 본 연구는 어텐션 희석(softmax dispersion)을 확장 가능한 문맥 내 검색의 주요 병목 현상으로 식별했다. 단순한 개입으로 이를 완화할 수는 있지만, 본 논문은 극단적인 컨텍스트 성장 하에서 어텐션을 제어하는 것이 미래의 검색 아키텍처를 위한 근본적인 과제로 남아 있음을 강조한다.

저자들은 자신들의 방법이 더 큰 코퍼스로 기능적 검색을 확장시키기는 하지만, $N$이 증가함에 따라 성능이 여전히 감소하고 일부 데이터셋에서 밀집 검색과의 잔여 격차가 존재한다는 점에서, 극한 규모에서의 검색 문제를 완전히 "해결"한 것은 아니라고 겸허하게 밝히고 있다.