UniPinRec: Unifying Generative Retrieval and Ranking at Pinterest Scale

Resumen Técnico: UniPinRec

Planteamiento del Problema
Los sistemas de recomendación industriales modernos, como los de Pinterest, suelen operar como embudos de múltiples etapas donde la recuperación (retrieval) y el ranking se entrenan como modelos separados. A pesar de que ambas etapas dependen cada vez más de grandes arquitecturas transformer para codificar los mismos datos de comportamiento del usuario, esta separación crea una redundancia significativa. El paradigma actual duplica parámetros, cómputo de entrenamiento y costos de servicio porque la recuperación y el ranking no pueden compartir señales aprendidas o pesos del modelo. Si bien trabajos previos han intentado unificar las arquitecturas de los modelos (por ejemplo, HSTU, OnePiece), estos enfoques a menudo no logran unificar el flujo completo. Permanecen fragmentados en cuanto a formatos de entrada, procedimientos de entrenamiento y pilas de servicio, o dependen de métodos generativos de extremo a extremo que reemplazan todo el embudo, lo que dificulta la integración con las fuentes de candidatos existentes y el control operativo.

Metodología
Los autores proponen UniPinRec, un sistema que logra la unificación de pila completa (full-stack) de la recuperación y el ranking. Este enfoque integra los formatos de entrada, la arquitectura del modelo, el entrenamiento y la infraestructura de servicio en un único flujo cohesivo desplegado dentro del entorno de producción existente de Pinterest. La metodología se basa en tres innovaciones técnicas principales:

1. Modelado de Acción Enmascarada (MAM - Masked Action Modeling):
   Para permitir que un solo modelo sirva tanto para la recuperación (predicción del siguiente ítem) como para el ranking (predicción de acciones) sin duplicar la longitud del contexto ni romper la compatibilidad de la entrada, los autores introducen MAM.

   • Arquitectura: En lugar de intercalar tokens de acción entre los ítems (lo que inflaría la longitud de la secuencia), las acciones se codifican como embeddings densos concatenados con las representaciones de los ítems a lo largo de la dimensión de características.
   • Enmascaramiento: Durante el entrenamiento, las acciones en el historial del usuario se enmascaran aleatoriamente con una probabilidad $p_{mask}$. En la inferencia, las acciones futuras siempre están enmascaradas. Esto permite que el modelo aprenda objetivos de predicción de acciones (ranking) en la misma secuencia no intercalada utilizada para la recuperación (predicción del siguiente ítem).
   • Patrón de Atención: El sistema utiliza un patrón de máscara causal modificado (patrón M-FALCON) donde las posiciones candidatas atienden a todo el historial pasado pero no entre sí. Esto reduce la complejidad de la atención de $O((n+k)^2)$ a $O(n^2 + nk)$ y permite un intercambio eficiente de la caché KV.

2. Entrenamiento Mezclado y Pérdida Conjunta:
   El sistema construye ejemplos de entrenamiento unificados que emparejan el historial de interacción pasado de un usuario con una serie de impresiones de feedview futura (incluyendo resultados negativos).

   • Infraestructura de Datos: Un join basado en Ray dentro del entrenador fusiona tablas separadas de Iceberg para el historial del usuario y las impresiones de feedview, evitando la duplicación de datos y permitiendo ratios de muestreo flexibles.
   • Objetivo: El modelo se entrena con una función de pérdida conjunta: $L = L_{item} + L_{action}$.
     • $L_{item}$: Una pérdida de softmax muestreado para la predicción del siguiente ítem (recuperación).
     • $L_{action}$: Una pérdida de entropía cruzada binaria para predecir tipos de acciones específicos (clic, guardar, ocultar) en las posiciones enmascaradas (ranking).
   • Esta optimización conjunta permite que el codificador de ítems aprenda de ambas tareas simultáneamente, mejorando la generalización en comparación con el pre-entrenamiento secuencial y el ajuste fino (fine-tuning).

3. Intercambio de Caché KV entre Etapas (Cross-Stage KV-Cache Sharing):
   Para abordar la eficiencia del servicio, UniPinRec reutiliza la computación de la codificación del historial del usuario a través de las etapas.

   • Mecanismo: La etapa de recuperación codifica el historial del usuario y almacena la caché de Clave-Valor (KV) resultante en un pool de memoria GPU preasignado. La etapa de ranking, que se ejecuta como un proceso separado, mapea esta misma memoria GPU y reutiliza la caché del historial sin volver a codificarlo.
   • Eficiencia: Esto transforma la etapa de ranking de un pase hacia adelante completo ($O(n^2)$) en un paso de decodificación incremental ($O(nk)$), donde $n$ es la longitud del historial y $k$ es el conteo de candidatos.
   • Implementación: El sistema utiliza NVIDIA Triton para el servicio, con Faiss para la búsqueda ANN. El índice de Faiss devuelve tanto los IDs de los candidatos como sus embeddings precomputados, eliminando la necesidad de una búsqueda de embeddings separada durante el ranking.

Contribuciones Clave

• Unificación de Pila Completa: UniPinRec es el primer sistema que unifica la recuperación y el ranking a través de los formatos de entrada, el modelo, el entrenamiento y el servicio en un sistema de recomendación de producción. Reemplaza la necesidad de dos modelos independientes con un único decodificador generativo.
• Arquitectura Eficiente: Al introducir MAM y secuencias no intercaladas, el sistema logra un intercambio total de pesos y eficiencia de parámetros sin inflar la longitud del contexto.
• Optimización del Servicio: El mecanismo de intercambio de caché KV entre etapas reduce significativamente los FLOPs totales y la latencia al evitar la codificación duplicada del historial del usuario.
• Compatibilidad Operativa: El diseño es un reemplazo directo ("drop-in") que se compone con los generadores de candidatos existentes (por ejemplo, palabras clave, tendencias) y permite realizar pruebas A/B y reversiones (rollbacks) independientes de las etapas.

Resultados
El sistema se desplegó en las superficies principales de Pinterest, obteniendo los siguientes resultados:

• Rendimiento Offline: UniPinRec igualó el Recall@10 del modelo de recuperación de producción mientras mejoró el Hit@3 del ranking en un +14.8% en comparación con el ranker dedicado de producción (TransAct V2 + DCNv2). También superó a HSTU y a la línea base de PinRec ajustada.
• Eficiencia de Servicio: El enfoque unificado redujo la latencia de servicio de extremo a extremo en un 11.1% y aumentó las Consultas Por Segundo (QPS) en un 63.6% en comparación con un despliegue separado ingenuo. El uso de la caché KV por sí solo proporcionó una aceleración de ~2.4x para el pase de ranking.
• Engagement Online: En pruebas A/B en "Board More Ideas" y "Notifications":
  • Board More Ideas: Logró un aumento de +0.95% en los guardados de la superficie y un aumento de +0.08% en los guardados en todo el sitio.
  • Notifications: Logró un aumento de +0.91% en las aperturas de notificaciones push, un aumento de +1.72% para usuarios inactivos y un aumento de +0.09% en los Usuarios Activos Semanales (WAU).

Significancia
El artículo afirma que UniPinRec demuestra que el "problema más difícil" de unificar los formatos de entrada, el entrenamiento y la infraestructura de servicio es soluble en producción. Al ir más allá de la unificación arquitectónica hacia la unificación de la pila completa, el sistema elimina la redundancia de codificar el historial del usuario dos veces, reduce los costos y mejora la calidad del modelo mediante el intercambio de señales. Los autores posicionan esto como un camino práctico e incremental para que los sistemas industriales adopten la recuperación y el ranking generativos sin sacrificar la flexibilidad operativa ni requerir una revisión completa del embudo de recomendación. El trabajo futuro incluye extender esta unificación a la etapa de ranking L2 y otras superficies como Búsqueda y Anuncios.