Generative Retrieval via Diffusion Transformer with Metric-Ordered Sequence Training and Hybrid-Policy Preference Optimization

Definição do Problema: Recuperação de Atributos com Preservação de Padrão

O artigo aborda uma limitação específica em sistemas de recuperação de produção: a tensão entre satisfazer um atributo alvo (ex: qualidade, segurança, conformidade de política) e preservar um padrão granular (ex: um estilo visual específico, tópico ou modo de falha) definido por um conjunto semente (seed set).

Abordagens padrão falham em equilibrar esses objetivos:

• A média dos embeddings das sementes preserva o padrão, mas frequentemente cai em uma região de baixa densidade de atributos, recuperando itens comuns que carecem do atributo desejado.
• A recuperação de atributos globais encontra itens com o atributo alvo, mas frequentemente deriva para padrões não relacionados, inundando os resultados com variantes irrelevantes.

Os autores formalizam isso como Recuperação de Atributos com Preservação de Padrão. O objetivo é gerar um embedding de consulta que recupere vizinhos top-$K$ que sejam simultaneamente positivos para o atributo e do padrão da semente. Eles definem uma métrica de interseção primária, Joint@K, que penaliza tanto a falha de atributo quanto o desvio de padrão.

Metodologia: MO-DiT+HPPO

O framework proposto, MO-DiT+HPPO (Metric-Ordered Diffusion Transformer with Hybrid-Policy Preference Optimization), é um pipeline de treinamento em estágios projetado para navegar no espaço de embedding de regiões de baixa densidade para regiões de alta densidade enquanto permanece dentro de um cluster de padrão específico.

1. Arquitetura: Recuperador Generativo Contínuo

O modelo central é um Diffusion Transformer (DiT) operando em um espaço de embedding multimodal congelado.

• Entrada: Uma sequência de embeddings de itens (o conjunto semente) atua como tokens de condição.
• Saída: Um embedding de consulta contínuo gerado via flow matching.
• Mecanismo: O modelo aprende um campo de velocidade para transportar ruído Gaussiano para um embedding alvo ao longo de um caminho geodésico na hiperesfera unitária. Diferente da recuperação generativa discreta, ele sintetiza um vetor contínuo diretamente, evitando a alucinação de identificadores inválidos.

2. Construção de Sequência Ordenada por Métrica

Para ensinar o modelo a direção de "melhoria" sem rotulagem densa online, os autores constroem trajetórias de treinamento:

• Clustering: Itens positivos para o atributo são agrupados para definir proxies de padrão latentes.
• Rotulagem Esparsa: Um pequeno subconjunto de itens é rotulado online para "densidade de recall" (fração de vizinhos top-$K$ que satisfazem o atributo).
• Predição & Ordenação: Um preditor de métrica leve é treinado nesses rótulos esparsos para pontuar todo o pool. Itens dentro de um cluster de padrão são ordenados de baixa para alta densidade de atributo prevista.
• Resultado: As sequências de treinamento representam uma trajetória movendo-se de itens comuns dentro do padrão em direção a regiões de alta densidade de atributos.

3. Pipeline de Treinamento em Estágios

O framework procede através de quatro estágios, todos usando o mesmo backbone DiT e objetivo de flow-matching, mas com diferentes dados e alvos de supervisão:

• Estágio 1: Pré-treinamento de Sequência Bruta em Larga Escala: Estabelece um prior de recuperação contínua geral usando sequências naturais (ex: fluxos de comportamento do usuário) sem ordenação métrica.
• Estágio 2: Pré-treinamento de Continuação (CPT) Multidomínio Compartilhado Ordenado por Métrica: Treina nas sequências ordenadas por métrica construídas através de todos os domínios. O modelo aprende a prever o próximo item em uma trajetória de baixa para alta densidade, internalizando a "direção de melhoria métrica" como uma capacidade geral.
• Estágio 3: Ajuste Fino Supervisionado (SFT) de Centroide de Cauda Específico do Domínio: Mapeia o prefixo de entrada (primeiros $m$ itens) para o centroide da cauda de alta densidade (últimos $L-m$ itens). O uso de um centroide, em vez de um único item do topo, reduz a sensibilidade a ruídos e outliers.
• Estágio 4: Otimização de Preferência de Política Híbrida (HPPO): Alinha o modelo com o objetivo real online (Joint@K).
  • Pool de Candidatos Híbrido: Gera candidatos de duas fontes: (1) construções de centroide de cauda determinísticas (estáticas, preservadoras de padrão) e (2) consultas geradas pela política amostradas com um "leque de escala de orientação" (guidance-scale fan) (variando o equilíbrio entre padrão e atributo).
  • Rotulagem: Os candidatos são pontuados usando recuperação online real (Joint@K e Same@K).
  • Filtro de Par Pareto: Um componente crítico que filtra pares de preferência. Um "vencedor" (maior Joint@K) só é aceito se não reduzir a Participação de Mesmo Padrão (Same-Pattern Share/pureza) em relação ao "perdedor". Isso evita que o modelo ganhe pontuações de atributo ao derivar para padrões não relacionados.
  • Otimização: Utiliza uma perda de preferência ancorada em referência (semelhante ao DPO), onde a pontuação é o negativo da perda de flow-matching, regularizada por um modelo de referência SFT congelado.

Principais Contribuições

1. Formalização: Definiu a Recuperação de Atributos com Preservação de Padrão e a métrica de interseção Joint@K.
2. Sequências Ordenadas por Métrica: Introduziu um método para converter rótulos esparsos online em trajetórias de treinamento em larga escala dentro do padrão que codificam a direção de melhoria do atributo.
3. Framework em Estágios: Demonstrou que o CPT multidomínio compartilhado seguido pelo SFT de centroide de cauda específico do domínio melhora a recuperação em diversos domínios.
4. HPPO: Adaptou a otimização de preferência para recuperação de difusão contínua, utilizando um pool de candidatos híbrido e um filtro de par Pareto consciente de trade-off para garantir que ganhos de atributo não venham ao custo da pureza do padrão.
5. Protocolo de Avaliação: Estabeleceu um protocolo de avaliação rigoroso com divisões de holdout de item e de padrão, além de testes de significância de bootstrap pareado.

Resultados

O artigo avalia a abordagem em quatro grandes domínios de atributos (D1–D4) sob protocolos rigorosos de holdout.

• Desempenho: O pipeline MO-DiT+HPPO supera significativamente um recuperador generativo pré-treinado forte e baselines de média simples.
• Significância: Os ganhos em Joint@K são grandes e estatisticamente significativos (via paired-bootstrap) em sete de oito células de divisão de domínio. No split de padrão do domínio mais difícil, o método atinge um empate marginal, indicando robustez em cenários complicados.
• Ablações:
  • Filtro Pareto: Remover o filtro faz com que o modelo deslize ao longo da fronteira atributo-padrão (ganhando atributos, mas perdendo pureza). O filtro empurra a fronteira para fora, melhorando o Joint@K enquanto mantém ou aumenta a Participação de Mesmo Padrão.
  • Política de Candidatos: Mostra-se que os candidatos estáticos (centroide de cauda) e gerados pela política (leque de orientação) são complementares; o pool híbrido produz melhores resultados do que cada um isoladamente.
  • Ordenação: O treinamento ordenado por métrica supera os baselines não ordenados ou de domínio único.

Significância e Alegações

O artigo afirma que o MO-DiT+HPPO resolve com sucesso a tensão entre a satisfação do atributo e a preservação do padrão na recuperação generativa contínua. Ao tratar a recuperação como um problema de geração em um espaço contínuo e usar uma abordagem em estágios que vai de priors gerais para trajetórias específicas ordenadas por métrica, o sistema aprende a sintetizar consultas que visam regiões de alta densidade de padrões específicos.

Os autores enfatizam que o filtro de par Pareto é o ingrediente fundamental que evita o "reward hacking" (sacrificar o padrão pelo atributo). O framework permite atualizações práticas por domínio, pois o pré-treinamento caro é compartilhado, enquanto o alinhamento final (HPPO) é computacionalmente leve e depende de métricas reais online em vez de recompensas aproximadas aprendidas. O trabalho preenche a lacuna entre a modelagem generativa offline e os objetivos de recuperação online, fornecendo um método para expandir conjuntos semente para tarefas downstream como descoberta de conteúdo, detecção de variantes e curadoria de datasets sem perder a intenção específica do padrão da semente.