UniPinRec: Unifying Generative Retrieval and Ranking at Pinterest Scale

Resumo Técnico: UniPinRec

Definição do Problema
Sistemas de recomendação industriais modernos, como os do Pinterest, operam tipicamente como funis de múltiplos estágios, onde a recuperação (retrieval) e o ranqueamento (ranking) são treinados como modelos distintos. Apesar de ambos os estágios dependerem cada vez mais de grandes backbones de transformer para codificar os mesmos dados de comportamento do usuário, essa separação cria uma redundância significativa. O paradigma atual duplica parâmetros, computação de treinamento e custos de serviço porque a recuperação e o ranqueamento não conseguem compartilhar sinais aprendidos ou pesos do modelo. Embora trabalhos anteriores tenham tentado unificar arquiteturas de modelos (ex: HSTU, OnePiece), essas abordagens frequentemente falham em unificar o pipeline completo. Elas permanecem fragmentadas entre formatos de entrada, procedimentos de treinamento e stacks de serviço, ou dependem de métodos generativos de ponta a ponta que substituem todo o funil, dificultando a integração com fontes de candidatos existentes e o controle operacional.

Metodologia
Os autores propõem o UniPinRec, um sistema que alcança a unificação de stack completo da recuperação e do ranqueamento. Esta abordagem integra formatos de entrada, arquitetura de modelo, treinamento e infraestrutura de serviço em um único pipeline coeso implantado no ambiente de produção existente do Pinterest. A metodologia baseia-se em três inovações técnicas principais:

1. Modelagem de Ação Mascarada (MAM - Masked Action Modeling):
   Para permitir que um único modelo sirva tanto para recuperação (previsão do próximo item) quanto para ranqueamento (previsão de ação) sem dobrar o comprimento do contexto ou quebrar a compatibilidade de entrada, os autores introduzem o MAM.

   • Arquitetura: Em vez de intercalar tokens de ação entre itens (o que inflaria o comprimento da sequência), as ações são codificadas como embeddings densos concatenados com as representações dos itens ao longo da dimensão de características.
   • Mascaramento: Durante o treinamento, as ações no histórico do usuário são mascaradas aleatoriamente com uma probabilidade $p_{mask}$. Na inferência, ações futuras são sempre mascaradas. Isso permite que o modelo aprenda objetivos de previsão de ação (ranqueamento) na mesma sequência não intercalada usada para recuperação (previsão do próximo item).
   • Padrão de Atenção: O sistema utiliza um padrão de máscara causal modificado (padrão M-FALCON) onde as posições candidatas atendem a todo o histórico passado, mas não entre si. Isso reduz a complexidade de atenção de $O((n+k)^2)$ para $O(n^2 + nk)$ e permite o compartilhamento eficiente de KV-cache.

2. Treinamento Misturado e Perda Conjunta:
   O sistema constrói exemplos de treinamento unificados que pareiam o histórico de engajamento passado de um usuário com uma visualização de feed futura (incluindo resultados negativos).

   • Infraestrutura de Dados: Um join baseado em Ray dentro do treinador funde tabelas Iceberg separadas para o histórico do usuário e impressões de visualização de feed, evitando a duplicação de dados e permitindo proporções de amostragem flexíveis.
   • Objetivo: O modelo é treinado com uma função de perda conjunta: $L = L_{item} + L_{action}$.
     • $L_{item}$: Uma perda de softmax amostrada para previsão do próximo item (recuperação).
     • $L_{action}$: Uma perda de entropia cruzada binária para prever tipos de ação específicos (clique, salvar, ocultar) em posições mascaradas (ranqueamento).
   • Essa otimização conjunta permite que o codificador de itens aprenda com ambas as tarefas simultaneamente, melhorando a generalização em comparação com o pré-treinamento e o ajuste fino sequenciais.

3. Compartilhamento de KV-Cache entre Estágios:
   Para abordar a eficiência de serviço, o UniPinRec reutiliza a computação da codificação do histórico do usuário entre os estágios.

   • Mecanismo: O estágio de recuperação codifica o histórico do usuário e armazena o respectivo cache de Chave-Valor (KV) em um pool de memória GPU pré-alocado. O estágio de ranqueamento, rodando como um processo separado, mapeia essa mesma memória GPU e reutiliza o histórico cacheado sem re-codificá-lo.
   • Eficiência: Isso transforma o estágio de ranqueamento de um forward pass completo ($O(n^2)$) em um passo de decodificação incremental ($O(nk)$), onde $n$ é o comprimento do histórico e $k$ é a contagem de candidatos.
   • Implementação: O sistema utiliza NVIDIA Triton para o serviço, com Faiss para busca ANN. O índice Faiss retorna tanto os IDs dos candidatos quanto seus embeddings pré-computados, eliminando a necessidade de uma busca de embedding separada durante o ranqueamento.

Principais Contribuições

• Unificação de Stack Completo: O UniPinRec é o primeiro sistema a unificar recuperação e ranqueamento através de entradas, modelo, treinamento e infraestrutura de serviço em um sistema de recomendação de produção. Ele substitui a necessidade de dois modelos independentes por um único decodificador generativo.
• Arquitetura Eficiente: Ao introduzir o MAM e sequências não intercaladas, o sistema alcança o compartilhamento total de pesos e eficiência de parâmetros sem inflar o comprimento do contexto.
• Otimização de Serviço: O mecanismo de compartilhamento de KV-cache entre estágios reduz significativamente o total de FLOPs e a latência ao evitar a codificação duplicada do histórico do usuário.
• Compatibilidade Operacional: O design é um substituto "drop-in" que se compõe com geradores de candidatos existentes (ex: palavras-chave, tendências) e permite testes A/B e reversões (rollback) independentes dos estágios.

Resultados
O sistema foi implantado nas superfícies principais do Pinterest, gerando os seguintes resultados:

• Desempenho Offline: O UniPinRec igualou o Recall@10 do modelo de recuperação de produção, enquanto melhorou o Hit@3 do ranqueamento em +14,8% comparado ao ranqueador dedicado de produção (TransAct V2 + DCNv2). Também superou o HSTU e um baseline de PinRec ajustado.
• Eficiência de Serviço: A abordagem unificada reduziu a latência de serviço de ponta a ponta em 11,1% e aumentou as Consultas Por Segundo (QPS) em 63,6% em comparação com uma implantação separada ingênua. O reaproveitamento do KV-cache sozinho proporcionou um aumento de velocidade de ~2,4x para o passo de ranqueamento.
• Engajamento Online: Em testes A/B em "Board More Ideas" e "Notifications":
  • Board More Ideas: Alcançou um aumento de +0,95% em salvamentos na superfície e um aumento de +0,08% em salvamentos em todo o site.
  • Notifications: Alcançou um aumento de +0,91% em aberturas de notificações push, +1,72% para usuários inativos e +0,09% em Usuários Ativos Semanais (WAU).

Significância
O artigo afirma que o UniPinRec demonstra que o "problema mais difícil" de unificar formatos de entrada, treinamento e infraestrutura de serviço é solucionável em produção. Ao ir além da unificação arquitetural para a unificação de stack completo, o sistema elimina a redundância de codificar o histórico do usuário duas vezes, reduz custos e melhora a qualidade do modelo através do compartilhamento de sinais. Os autores posicionam isso como um caminho prático e incremental para sistemas industriais adotarem recuperação e ranqueamento generativos sem sacrificar a flexibilidade operacional ou exigir uma reformulação completa do funil de recomendação. Trabalhos futuros incluem estender essa unificação para o estágio de ranqueamento L2 e outras superfícies como Busca e Anúncios.