하드 네거티브를 넘어서: 밀집 검색을 위한 지식 증류에서 점수 분포의 중요성

기존의 밀집 검색 모델 학습은 주로 어려운 오답(Hard Negatives)을 찾는 데 집중했으나, 이는 모델이 교사 모델의 전체적인 선호 구조를 배우는 것을 방해할 수 있다. 이 논문은 교사 모델의 점수 분포를 골고루 반영하는 Stratified Sampling 기법을 통해 검색 모델의 정확도와 새로운 데이터에 대한 적응력을 동시에 높일 수 있음을 증명한다.

밀집 검색(Dense Retrieval) 모델은 문장의 의미를 하나의 벡터(Embedding)로 압축하여 유사도를 계산한다. 이때 더 강력한 성능을 가진 교사 모델(Cross-Encoder)의 지식을 학생 모델(Bi-Encoder)에게 전달하는 지식 증류(Knowledge Distillation)가 널리 사용된다. 기존 방식은 주로 학생 모델이 틀리기 쉬운 '어려운 오답'을 골라 학습시키는 데 집중해 왔으나, 이는 교사 모델이 가진 미세한 선호도 차이를 배우지 못하게 만드는 편향을 초래한다.

이 논문은 교사 모델의 점수 분포 자체가 중요한 학습 신호라는 점에 주목한다. 교사 모델이 특정 문서에 대해 내리는 상대적인 점수 차이의 전체적인 '골격'을 학생 모델이 그대로 모사해야만 다양한 난이도의 결정 경계를 정확히 학습할 수 있다. 이를 위해 점수 대역을 여러 구간으로 나누고 각 구간에서 대표 샘플을 뽑는 방식을 사용한다.

결과적으로 이러한 접근은 학생 모델이 특정 데이터에 과적합(Overfitting)되는 것을 방지한다. 교사 모델의 지식을 파편적으로 배우는 대신 전체적인 판단 기준을 학습함으로써, 학습 데이터와 성격이 다른 새로운 질문이나 문서가 들어와도 안정적인 검색 성능을 유지하게 된다.

전체 학습 과정은 2단계 프로토콜로 구성된다. 첫 번째 단계에서는 MLM(Masked Language Modeling) 기반의 백본 모델을 MS MARCO 데이터셋에서 대조 학습(Contrastive Learning)을 통해 기본 검색 능력을 갖추도록 적응시킨다. 이때는 배치 내 네거티브(In-batch negatives)만을 사용하여 기본적인 임베딩 공간을 형성한다.

두 번째 단계인 지식 증류에서는 교사 모델인 Qwen3-Reranker-8B의 점수를 활용한다. 먼저 각 쿼리당 200개의 후보 문서 풀을 구성하고, 교사 모델의 점수를 0에서 1 사이로 최소-최대 정규화(Min-max normalization)한다. 정규화된 점수 s̃ᵢ = (sᵢ - min(S)) / (max(S) - min(S))를 통해 점수 간의 상대적 거리를 보존한다.

핵심 기법인 Stratified Sampling은 원하는 샘플 수 K에 대해 pⱼ = (j-1)/(K-1)로 정의되는 균등 간격의 분위수 앵커를 설정한다. 각 앵커 점수 τⱼ에 대해 |s̃_d - τⱼ|를 최소화하는 문서를 선택하여, 무작위성에 의한 변동 없이 교사 모델의 점수 분포 골격을 구조적으로 모사하는 학습 세트를 구축한다. 이렇게 샘플링된 데이터는 KL Divergence 또는 MarginMSE 손실 함수를 통해 학생 모델 학습에 사용된다.

실험 결과, Stratified Sampling은 모든 백본 모델(BERT, DistilBERT, Co-condenser)과 손실 함수 조합에서 가장 우수한 성능을 기록했다. 특히 BEIR 벤치마크를 이용한 도메인 외(Out-of-domain) 평가에서 bert-base 모델 기준 nDCG@10 점수가 0.3140.318을 기록하며, Top-K 방식(0.1830.217) 대비 압도적인 일반화 성능을 보여주었다.

분포 통계 분석 결과, Stratified Sampling으로 구축된 데이터셋은 커버리지(0.990), 엔트로피(1.523), 표준편차(0.359) 모든 지표에서 가장 높은 수치를 기록했다. 이는 해당 전략이 특정 점수대에 편향되지 않고 교사 모델의 지식 전체를 고르게 반영하고 있음을 수치적으로 증명한다.

샘플 수 K의 변화에 따른 분석에서도 Stratified Sampling은 K가 커질수록 성능이 안정적으로 향상되는 견고함을 보였다. 반면 Top-K 방식은 K가 늘어날수록 오히려 성능이 급격히 하락하는 현상이 관찰되었는데, 이는 중복되거나 불필요하게 어려운 샘플이 추가되면서 학습 신호에 노이즈가 발생했기 때문으로 분석된다.

본 연구는 지식 증류에서 데이터 구성(Data Composition)이 모델 성능에 미치는 영향을 격리하여 분석하기 위해 고정된 후보 풀(Fixed Pool) 방식을 채택했다. 이는 복잡한 동적 마이닝(Dynamic Mining) 변수를 배제하고 순수하게 샘플링 전략의 효과만을 측정하기 위함이다.

수학적으로 MarginMSE 손실 함수는 교사가 판단한 두 문서 간의 점수 차이(Margin)를 학생 모델이 그대로 회귀(Regression)하도록 유도한다. 실험 결과 MarginMSE는 부적절한 샘플링(예: reranker-top) 시 학습이 완전히 붕괴(Collapse)되는 민감성을 보였으나, Stratified Sampling과 결합했을 때는 가장 높은 성능을 달성했다. 이는 회귀 기반 목적 함수일수록 데이터 분포의 다양성이 학습 안정성에 결정적인 역할을 함을 시사한다.

제안된 분위수 기반 샘플링은 결정론적(Deterministic) 알고리즘으로, 무작위 샘플링이 가질 수 있는 분산 문제를 해결하면서도 통계적 대표성을 확보한다. 이는 향후 대규모 검색 모델 학습 시 데이터 파이프라인의 표준적인 기준으로 활용될 가능성이 높다.

본 연구는 MS MARCO 데이터셋을 중심으로 실험이 진행되었으며, 더 거대한 규모의 데이터셋이나 다른 언어 환경에서의 효과에 대해서는 추가적인 검증이 필요할 수 있다. 또한 고정된 후보 풀에서의 샘플링 효과에 집중했으므로, 실시간으로 변하는 동적 마이닝 환경과의 결합 효과는 명시적으로 다루지 않았다.