UniPinRec: Unifying Generative Retrieval and Ranking at Pinterest Scale

기술 요약: UniPinRec

문제 정의
Pinterest와 같은 현대적인 산업용 추천 시스템은 일반적으로 검색(Retrieval)과 랭킹(Ranking)이 별개의 모델로 학습되는 다단계 퍼널(multi-stage funnel) 구조로 작동합니다. 두 단계 모두 동일한 사용자 행동 데이터를 인코딩하기 위해 대규모 트랜스포머 백본을 점점 더 많이 사용하고 있음에도 불구하고, 이러한 분리는 상당한 중복을 초양합니다. 현재의 패러다임은 검색과 랭킹이 학습된 신호나 모델 가중치를 공유할 수 없기 때문에 파라미터, 학습 연산량, 서빙 비용을 중복해서 소모합니다. 이전 연구들이 모델 아키텍처를 통합하려는 시도(예: HSTU, OnePiece)를 해왔으나, 이러한 방식들은 종종 전체 파이프라인을 통합하는 데 실패했습니다. 이들은 입력 형식, 학습 절차, 서빙 스택 측면에서 여전히 파편화되어 있거나, 기존 후보 소스와의 통합 및 운영 제어를 어렵게 만드는 엔드 투 엔드 생성형 방법론에 의존합니다.

방법론
저자들은 검색과 랭킹의 **풀스택 통합(full-stack unification)**을 달죄하는 시스템인 UniPinRec를 제안합니다. 이 접근 방식은 Pinterest의 기존 프로덕션 환경 내에 배포되는 단일하고 응집된 파이프라인으로서 입력 형식, 모델 아키텍처, 학습 및 서빙 인프라를 통합합니다. 이 방법론은 세 가지 핵심 기술 혁신에 기반합니다:

1. 마스크 액션 모델링 (Masked Action Modeling, MAM):
   단일 모델이 컨텍스트 길이를 두 배로 늘리거나 입력 호환성을 깨뜨리지 않으면서 검색(다음 아이템 예측)과 랭킹(액션 예측) 모두를 수행할 수 있도록 하기 위해, 저자들은 MAM을 도입했습니다.

   • 아키텍처: 액션 토큰을 아이템 사이에 끼워 넣는 대신(이는 시퀀스 길이를 팽창시킴), 액션을 아이템 표현과 함께 피처 차원을 따라 결합되는 밀집 임베딩(dense embeddings)으로 인코딩합니다.
   • 마스킹: 학습 중에 사용자의 히스토리 내 액션은 확률 $p_{mask}$로 무작위로 마스킹됩니다. 추론 시에는 미래의 액션이 항상 마스킹됩니다. 이를 통해 모델은 검색(다음 아이템 예측)에 사용되는 것과 동일한 비인터리빙(non-interleaved) 시퀀스 상에서 액션 예측(랭킹) 목적 함수를 학습할 수 있습니다.
   • 어텐션 패턴: 시스템은 수정된 인과적 마스크(M-FALCON 패턴)를 사용하며, 여기서 후보 위치들은 과거의 모든 히스토리에는 어텐션을 수행하지만 서로에게는 수행하지 않습니다. 이는 어텐션 복잡도를 $O((n+k)^2)$에서 $O(n^2 + nk)$로 줄여주며 효율적인 KV-캐시 공유를 가능하게 합니다.

2. 혼합 학습 및 공동 손실 (Blended Training and Joint Loss):
   시스템은 사용자의 과거 참여 히스토리를 미래의 피드뷰 임프레션 슬레이트(부정적 결과 포함)와 쌍으로 구성하는 통합 학습 예시를 구축합니다.

   • 데이터 인프라: Ray 기반의 인-트레이너 조인(in-trainer join)을 통해 별도의 Iceberg 테이블을 병합하여 데이터 중복을 피하고 유연한 샘플링 비율을 허용함으로써 사용자 히스토리와 피드뷰 임프레션을 결합합니다.
   • 목적 함수: 모델은 다음과 같은 공동 손실 함수로 학습됩니다: $L = L_{item} + L_{action}$.
     • $L_{item}$: 다음 아이템 예측(검색)을 위한 샘플드 소프트맥스(sampled softmax) 손실입니다.
     • $L_{action}$: 마스킹된 위치에서의 특정 액션 유형(클릭, 저장, 숨기기)을 예측하기 위한 바이너리 크로스 엔트로피 손실입니다(랭킹).
   • 이러한 공동 최적화는 아이템 임베더가 순차적인 사전 학습 및 미세 조정보다 더 나은 일반화 성능을 갖도록 두 태스크로부터 동시에 학습할 수 있게 합니다.

3. 교차 단계 KV-캐시 공유 (Cross-Stage KV-Cache Sharing):
   서빙 효율성을 해결하기 위해, UniPinRec는 단계 전반에 걸쳐 사용자 히스토리 인코딩의 계산을 재사용합니다.

   • 메커니즘: 검색 단계는 사용자 히스토리를 인코딩하고 결과물인 키-값(KV) 캐시를 사전 할당된 GPU 메모리 풀에 저장합니다. 별도의 프로세스로 실행되는 랭킹 단계는 이 동일한 GPU 메모리를 매핑하여 인코딩을 다시 수행하지 않고 캐시된 히스토리를 재사용합니다.
   • 효율성: 이를 통해 랭킹 단계를 전체 포워드 패스($O(n^2)$)에서 증분 디코드 단계($O(nk)$)로 변환합니다 (여기서$n$은 히스토리 길이,$k$는 후보 수).
   • 구현: 시스템은 Faiss를 사용하여 ANN 룩업을 수행하는 NVIDIA Triton을 사용하여 서빙합니다. Faiss 인덱스는 후보 ID와 함께 사전 계산된 임베딩을 반환하여, 랭킹 중에 별도의 임베딩 페치를 수행할 필요를 없앱니다.

주요 기여

• 풀스택 통합: UniPinRec는 프로덕션 추천 시스템에서 입력, 모델, 학습 및 서빙 인프라 전반에 걸쳐 검색과 랭킹을 통합한 최초의 시스템입니다. 이는 두 개의 독립적인 모델을 하나의 생성형 디코더로 대체합니다.
• 효율적인 아키텍처: MAM과 비인터리빙 시퀀스를 도입함으로써, 시스템은 컨텍스트 길이를 팽창시키지 않고도 완전한 가중치 공유와 파라미터 효율성을 달성합니다.
• 서빙 최적화: 교차 단계 KV-캐시 공유 메커니즘은 사용자 히스토리를 중복 인코딩하는 것을 방지함으로써 총 FLOPs와 지연 시간을 크게 줄입니다.
• 운영 호환성: 이 설계는 기존의 후보 생성기(예: 키워드, 트렌딩)와 결합 가능한 "드롭인(drop-in)" 교체 모델이며, 각 단계의 독립적인 A/B 테스트 및 롤백을 허용합니다.

결과
시스템은 Pinterest의 핵심 서피스에 배포되어 다음과 같은 결과를 얻었습니다:

• 오프라인 성능: UniPinRec는 프로덕션 검색 모델의 Recall@10을 유지하면서, 전용 프로덕션 랭커(TransAct V2 + DCNv2) 대비 랭킹 Hit@3를 +14.8% 개선했습니다. 또한 HSTU 및 미세 조정된 PinRec 베이스라인보다 우수한 성능을 보였습니다.
• 서빙 효율성: 통합된 접근 방식은 단순한 별도 배포 방식에 비해 엔드 투 엔드 서빙 지연 시간을 11.1% 줄였고, QPS(초당 쿼리 수)를 63.6% 증가시켰습니다. KV-캐시 재사용만으로 랭킹 패스에서 약 2.4배의 속도 향상을 제공했습니다.
• 온라인 참여도: "Board More Ideas" 및 "Notifications"에 대한 A/B 테스트 결과:
  • Board More Ideas: 서피스 저장(saves)에서 +0.95% 리프트, 사이트 전체 저장에서 +0.08% 리프트를 달성했습니다.
  • Notifications: 푸시 오픈에서 +0.91% 리프트, 휴면 사용자(dormant users)에서 +1.72% 리프트, 주간 활성 사용자(WAU)에서 +0.09% 리프트를 달성했습니다.

의의
본 논문은 UniPinRec가 입력 형식, 학습 및 서빙 인프라를 통합하는 "더 어려운 문제"가 프로덕션 환경에서도 해결 가능하다는 것을 보여준다고 주장합니다. 아키텍처 통합을 넘어 풀스택 통합으로 나아감으로써, 시스템은 사용자 히스토리를 두 번 인코딩하는 중복을 제거하고, 비용을 절감하며, 신호 공유를 통해 모델 품질을 향ale합니다. 저자들은 이를 산업용 시스템이 운영의 유연성을 희생하거나 추천 퍼널을 완전히 개편할 필요 없이 생성형 검색 및 랭킹을 채택할 수 있는 실질적이고 점진적인 경로로 제시합니다. 향후 연구에는 이 통합을 L2 랭킹 단계 및 Search, Ads와 같은 다른 서피스로 확장하는 것이 포함됩니다.