기술 요약: MixLM

문제 정의

대규모 언어 모델(LLM)은 전통적인 어휘 기반 방식이 놓치는 미묘한 의미적 신호를 포착함으로써 검색 및 추천 시스템에서 탁월한 관련성 순위 지정(relevance ranking) 능력을 입증해 왔습니다. 그러나 산업 현장에서 LLM을 크로스 인코더(cross-encoder) 랭커로 배포하는 데에는 결정적인 병목 현상인 컴퓨팅 비용이 존재합니다.

표준 크로스 인코더 접근 방식은 사용자 쿼리와 긴 후보 아이템 설명(종종 수천 개의 토큰)을 결합한 프롬프트를 구성합니다. 이는 다음과 같은 문제를 야기합니다:

1. 과도한 프리필(Prefill) 추론: 계산 비용이 어텐션 메커니즘으로 인해 컨텍스트 길이에 따라 제곱으로 증가하는 프리필 단계에 의해 지배됩니다.
2. 처리량(Throughput) 제약: 엄격한 레이턴시 예산(예: 500ms) 하에서, 시스템은 성능이 낮은 모델을 사용하거나 정보가 풍부한 특징(예: 텍스트 요약 또는 프루닝을 통해)을 버려야 하는 상황에 직면하며, 이는 랭킹 품질의 희생으로 이어집니다.

핵심 과제는 전체 문서의 풍부한 의미적 상호작용을 보존하면서 입력 컨텍스트 길이를 획기적으로 줄이는 것입니다.

방법론: MixLM 프레임워크

저자들은 긴 아이템 텍스트를 학습된 임베딩 토큰의 압축된 집합으로 대체하는 LLM 기반 랭킹 프레임워크인 MixLM을 제안합니다. 이 "텍ext-임베딩 혼합 상호작용(text-embedding mix-interaction)" 아키텍처는 아이템 설명의 처리 과정을 온라인 랭킹 단계와 분리합니다.

1. 아키텍처

MixLM은 두 가지 주요 구성 요소로 이루어집니다:

• 인코더 LLM (Encoder LLM): 오프라인(또는 니어라인)에서 전체 아이템 설명을 처리하여 압축된 표현을 생성합니다. 단일 벡터를 출력하는 대신, 인코더는 일련의 임베딩 토큰을 생성합니다. 프로덕션 환경에서 저자들은 아이템을 표현하기 위해 인코더의 은닉 상태(hidden states)로부터 소수의 토큰(예: $T_S=1$)을 선택하는 last-$N$ 샘플링을 사용합니다.
• 랭커 LLM (Ranker LLM): 온라인에서 작동합니다. 사용자 쿼리, 보조 텍스트, 그리고 사전 계산된 인코더 임베딩 토큰을 전달받습니다. 이들은 "혼합 입력(mixed input)" 프롬프트를 형성하기 위해 결합되며, 이는 전체 텍스트 프롬프트보다 훨씬 짧습니다.

랭커는 이 혼합 입력을 바탕으로 관련성($p_{yes}$)을 예측합니다. 아이템 임베딩은 미리 계산되어 캐싱되어 있기 때문에, 온라인 추론은 쿼리 프리픽스(prefix)와 소수의 임베딩 토큰만을 처리하게 되어 시퀀스 길이를 획기적으로 줄입니다.

2. 훈련 파이프라인

저자들은 인코더와 랭커를 정렬하기 위해 3단계 훈련 레시피를 채택했습니다:

• 단계 I (도메인 추론 미세 조정): 0.6B 사전 학습된 LLM을 7B 관련성 판별기(relevance judge)의 사고 사슬(chain-of-thought) 트레이스를 사용하여 도메인 특화 추론 데이터셋(180k 샘플)으로 미세 조정합니다. 이를 통해 관련성 근거를 이해할 수 있는 베이스 모델을 구축합니다.
• 단계 II (랭킹 교사 학습): 단계 I 모델을 전체 아이템 설명(10.9M 샘플)을 사용하는 순수 텍스트 랭커로 추가 미세 조정합니다. 이 모델은 고품질의 "교사(teacher)" 역할을 하지만, 온라인 서빙에는 너무 비쌉니다.
• 단계 III (공동 인코더-랭커 학습): 인코더와 랭커를 혼합 입력 아키텍처를 사용하여 공동 학습(co-training)합니다. 훈련 목적 함수는 다음을 결합합니다:
  • Soft-Label SFT 손실: 예측값을 실제 레이블과 정렬합니다.
  • 증류(Distillation) 손실: 혼합 입력 모델의 출력 분포를 순수 텍스트 교사 모델(단계 II)과 정렬합니다.
  • 자기 정렬(Self-Alignment) 손실: 인코더의 출력 임베딩이 랭커의 내부 은닉 상태와 호환되도록 보장합니다. 여기에는 은닉 상태 정렬(혼합 입력의 마지막 은닉 상태를 전체 텍스트 입력과 매칭) 및 예측 정렬이 포함됩니다.

단계 III에서는 커리큘럼 학습 전략을 사용하여, 초기에는 공간 정렬(인코더-랭커 호환성)을 우선시한 후 증류(관련성 예측 정확도)에 초점을 맞추도록 전환합니다.

3. 서빙 최적화

처리량을 극대화하기 위해 저자들은 여러 시스템 수준의 최적화를 구현했습니다:

• 공유-프리픽스 프리필 (Shared-Prefix Prefill): 랭킹 리스트 내의 모든 아이템에 대해 쿼리 프리픽스가 동일하므로, 프리픽스의 키-값(KV) 캐시는 한 번만 계산되어 모든 아이템에 걸쳐 재사용됩니다.
• 인-배치 프리픽스 캐싱 (In-Batch Prefix Caching): 추론 엔진(SGLang)을 수정하여 어텐션 계산을 병합함으로써, 접미사 토큰(임베딩)이 단일 순방향 패스(forward pass) 내에서 공유 프리픽스에 어텐션을 가할 수 있도록 합니다.
• 멀티 프로세스 아키텍처: 서빙 스택은 Python의 글로벌 인터프리터 락(GIL) 병목 현상을 우회하기 위해 멀티프로세싱을 사용하여 gRPC 프런트엔드와 추론 엔진을 분리합니다.
• KV 캐시 관리: 프리필 전용 워크로드에 대해 중복된 캐시 작업을 제거하여 CPU 오버헤드를 줄이고 유효 GPU 메모리 용량을 높입니다.

주요 기여

1. MixLM 프레임워크: 긴 텍스트 문서를 압축된 임베딩 토큰으로 대체하여, 의미적 표현력을 유지하면서도 프롬프트 길이를 몇 자릿수나 줄이는 새로운 혼합 입력 랭킹 아키텍처를 제시했습니다.
2. 엔드 투 엔드 훈련 파이프라인: 고품질의 전체 텍스트 교사로부터 혼합 입력 학생 모델로 지식을 성공적으로 전달하는 커스텀 증류 및 자기 정렬 손실을 포함한 공동 학습 전략을 제공합니다.
3. 프로덕션 서빙 스택: (공유-프리픽스 캐싱 및 멀리 프로세스 핸들링을 포함한) 추론 엔진의 포괄적인 최적화를 통해 고처리량 LLM 서빙을 가능하게 했습니다.
4. 산업적 배포: LinkedIn의 구직(Job Search) 서비스에 대규모 A/B 테스트를 통해 검증된, LLM 기반 의미적 랭킹 시스템의 첫 번째 풀 트래픽 배포 사례를 보여주었습니다.

실험 결과

오프라인 평가

• 관련성 품질: MixLM은 NDCG@10 0.9239를 달성하였으며, 이는 순수 텍스트 교사 모델(0.9432)과 유사하고, 임베딩 전용 검색(0.8380) 및 요약 텍스트 베이스라인(0.9218)을 크게 상회하는 수치입니다.
• 처리량: 고정된 500ms 레이턴시 예산 하에서:
  • MixLM은 22,000 items/s/GPU를 기록했습니다.
  • 이는 요약 텍스트 LLM 랭킹(2,200 items/s/GPU) 대비 10.0배 개선된 수치입니다.
  • 이는 전체 텍스트 LLM 랭킹(290 items/s/GPU) 대비 75.9배 개선된 수치입니다.

온라인 A/B 테스트

• MixLM은 LinkedIn의 AI 기반 구직 서비스 트래픽의 100%에 배포되었습니다.
• 시스템은 기존의 제한된 규모 프로덕션 베이스라인과 대등한 관련성 지표를 유지했습니다.
• 향상된 처리량 덕분에 LLM 기반 랭킹의 전체 배포가 가능해졌으며, 이는 일간 활성 사용자 수(DAU)의 0.47% 증가를 이끌어냈습니다.

의의 및 주장

본 논문은 MixLM이 실제 환경에서 LLM 기반 랭킹 시스템을 배포하기 위한 실용적이고 확장 가능한 경로를 제공한다고 주장합니다. 전체 텍스트 LLM의 의미적 풍부함과 산업적 검색의 효율성 요구 사이의 간극을 메움으로써, MixLM은 다음을 입증합니다:

1. 임베딩을 통한 컨텍스트 압축: 랭킹 품질의 큰 손실 없이 수천 개의 토큰을 몇 개의 임베딩 토큰으로 압축하는 것이 가능하다.
2. 시스템-모델 공동 설계 (System-Model Co-Design): 높은 처리량은 단순히 모델 아키텍처의 변화뿐만 아니라, 혼합 입력(mixed-input) 설계와 인프라 최적화(KV 캐싱, 배칭)의 긴밀한 통합을 통해 달성된다.
3. 비즈니스 임팩트: 효율성 개선은 제한된 규모 또는 요약 텍스트 실험에서 전체 트래픽 LLM 배포로의 전환을 가능하게 하여, 대규모 전문 네트워크에서 사용자 참여 지표(DAU)를 직접적으로 견인한다.

저자들은 이 연구를 이전에 광범위한 채택을 가로막았던 "레이턴시-처리량" 트레이드오프에 대한 해결책으로 포지셔닝하며, 크로스 인코더 LLM을 프로덕션 검색 시스템에 적용하는 실질적인 방안을 제시합니다.