MIPS 성능 향상을 위한 학습 기반 검색 기법: SupportNet과 KeyNet

새로운 MIPS 접근 방식

Amortized MIPS는 일반적인 쿼리들이 어떻게 매치되는지 이해하기 위한 초기 상각 비용을 지불하는 대신, 이러한 소형 신경망을 통해 추론 시 더 빠르고 정확한 검색을 제공한다. 이는 추천 시스템, 검색, RAG(Retrieval-Augmented Generation) 등 MIPS가 핵심적인 역할을 하는 다양한 애플리케이션의 성능을 개선하는 데 기여한다.

기존 MIPS의 한계와 배경

MIPS 시스템은 쿼리 벡터가 주어졌을 때 수백만 개의 후보 키 벡터 데이터베이스에서 가장 잘 일치하는 키를 검색한다. 그러나 키의 수가 수백만 개에 달하고 벡터 차원이 수천 개에 이르면, 모든 키와 쿼리를 비교하는 계산 비용이 급격히 증가한다. 기존의 근사 방법들은 클러스터링 및 양자화를 사용하여 정확도를 속도와 교환하지만, 쿼리 세트가 키 세트와 유사할 것이라는 가정을 느슨하게 적용하거나 쿼리의 특성을 고려하지 않는 경향이 있다.

특히 RAG의 경우, 쿼리는 짧은 사용자 생성 질문에서 비롯되는 반면 키는 길고 정보가 풍부하며 선별된 텍스트 조각을 나타내어 키와 쿼리가 상당히 다를 수 있다. 또한, 대부분의 프로덕션 MIPS 사용 사례는 안정적이고 잘 특성화된 쿼리 분포를 가진 기존 사용자 기반을 대상으로 하지만, 일반적인 쿼리 불가지론적 방법은 이러한 신호를 활용하지 못한다. 이 연구는 이러한 한계를 극복하고, 알려진 쿼리 분포에 최적화된 방법을 통해 컴퓨팅 자원을 보다 효율적으로 재분배하는 현대 머신러닝의 핵심 원칙을 따른다.

SupportNet과 KeyNet의 작동 원리

이 연구는 MIPS가 볼록 분석에서 잘 연구된 객체인 집합의 서포트 함수(support function)라는 수학적 구조에 기반한다는 점에 주목한다. 서포트 함수는 볼록하며, 특정 유형의 신경망인 ICNN(Input Convex Neural Network)으로 근사할 수 있다. 또한, 서포트 함수의 기울기는 최적의 매칭 키 자체를 반환한다. 이 아이디어는 실제 서포트 함수(또는 그 기울기)를 훨씬 낮은 계산 비용으로 모방할 수 있는 신경망을 학습하는 것이다.

SupportNet 아키텍처는 느슨하게 제약된 ICNN을 사용하여 서포트 함수를 직접 학습한다. 이 접근 방식은 문제의 볼록 구조와 일치하며, 데이터 클러스터당 하나의 서포트 함수를 계산하여 라우팅 품질을 개선하는 데 효과적이다. KeyNet 모델은 쿼리가 주어졌을 때 최적의 매칭 키 벡터를 직접 출력하도록 학습된다. 동차 함수에 대한 오일러 정리에서 파생된 일관성 손실(consistency loss)은 명시적인 볼록성이 강제되지 않더라도 예측이 기본 서포트 함수에 기하학적으로 기반하도록 유지한다.

두 모델 모두 관련 쿼리 분포에서 추출된 예제를 사용하여 훈련된다. 각 샘플링된 쿼리에 대해, 최적의 매칭은 한 번의 완전 탐색을 통해 계산되고 저장되어, 실제로 중요한 쿼리 분포에 대한 MIPS 솔루션 환경을 포착하는 지도 학습 데이터셋을 생성한다.

성능 평가 및 구체적 수치

연구진은 Natural Questions 데이터셋을 사용하여 SupportNet과 KeyNet의 성능을 평가했다. 이 데이터셋은 128개 클러스터로 분할된 약 260만 개의 고유한 문서 구절을 포함한다. SupportNet은 다양한 크기(XS, S, M, L)로 훈련되었으며, 라우팅 성능은 리콜(recall)과 지연 시간(latency) 또는 스캔 허용 데이터베이스 비율로 측정되었다.

SupportNet의 경우, 소형 SupportNet 모델은 95% 라우팅 리콜을 기존 센트로이드 라우팅보다 약 1.9배 빠르게 달성했다(1.74ms 대 3.36ms). 대형 모델은 1.74ms 내에 97% 리콜을 달성했다. 라우팅 품질 측면에서는, 데이터베이스의 5%만 스캔했을 때 학습된 라우터가 실제 파티션을 90~93%의 확률로 복구했으며, 중간 크기(M) 모델은 92.6%를 기록하여 센트로이드 기준선인 87.5%를 능가했다.

KeyNet은 기존 IVF(inverted file index) 시스템에 쿼리 재작성기(query rewriter)로 통합되어 평가되었다. KeyNet은 다양한 크기(XS, S, M, L)로 훈련되었으며, 99% 리콜을 달성하는 데 가장 큰 모델의 경우 원본 쿼리보다 최대 약 8배 적은 시간(0.23ms 대 1.8ms)이 소요되었다. 데이터베이스의 0.001%만 스캔했을 때 KeyNet으로 매핑된 쿼리는 실제 키를 89~97%의 확률로 검색했으며, 가장 큰 모델(L)은 96.7%를 기록하여 원본 쿼리의 86%를 상회했다. 99% 리콜을 달성하는 데 필요한 키 수는 최대 약 11배 적었다(0.009% 대 0.10% 원본).

개발자 활용 및 적용 조건

SupportNet과 KeyNet 중 어떤 신경망을 사용할지는 실제적인 고려 사항에 따라 결정된다. 추론 속도가 병목 현상이고 최적의 키를 근사하는 것만 필요하다면 KeyNet이 적합하다. 반면, 재순위 지정, 임계값 설정 또는 후속 사용을 위한 보정된 점수(calibrated score)가 필요하거나 소형 모델 크기에서 컴퓨팅 자원이 제한적일 때는 SupportNet이 경쟁력을 유지한다.

두 모델 모두 유사하게 훈련되며 기존의 근사 검색 인프라와 통합될 수 있다. 또한, 쿼리 패턴의 가벼운 분포 변화에 대해 합리적으로 견고하다. 매우 대규모의 MIPS 시스템에서는 두 가지 접근 방식을 모두 활용할 수 있다. 즉, 모든 클러스터에 대해 SupportNet 라우터를 훈련하고, 각 클러스터에 KeyNet을 사용하여 일반 인덱스가 최적의 키를 더 빠르게 찾아내도록 돕는 방식이다.

향후 방향과 고려사항

Amortized MIPS는 쿼리 분포가 미리 알려져 있는 경우(이는 프로덕션 ML 시스템에서 흔한 조건이다) 검색에 접근하는 방식의 변화를 반영한다. MIPS 솔루션을 신경망 가중치에 직접 인코딩함으로써, 이 접근 방식은 한 번의 훈련(상각) 비용을 빠르고 가벼운 추론과 교환한다. 이는 서비스할 쿼리에 명시적으로 맞춤화된 학습 기반 검색 시스템으로 나아가는 실용적인 길을 열어준다.

이러한 학습 기반 검색 시스템은 기존 기하학적 접근 방식이 가질 수 있는 한계, 즉 센트로이드 기반 공간 분할로 인해 가장 가까운 이웃을 놓칠 수 있는 문제를 해결한다. 이 연구는 MIPS의 효율성과 정확성을 크게 향상시킬 잠재력을 보여주며, 미래의 지식 기반 및 검색 시스템 개발에 중요한 방향을 제시한다.