💬 NLP

DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

DiffRetriever introduz um método de recuperação de tokens representativos paralelos para modelos de linguagem de difusão que anexa múltiplas posições mascaradas aos prompts e as lê em uma única passagem bidirecional, alcançando desempenho e eficiência superiores em relação às abordagens sequenciais autoregressivas de múltiplos tokens e aos recuperadores existentes ajustados por contraste.

Autores originais: Shuai Wang, Yin Yu, Shengyao Zhuang, Bevan Koopman, Guido Zuccon

Publicado 2026-05-11

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Shuai Wang, Yin Yu, Shengyao Zhuang, Bevan Koopman, Guido Zuccon

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando encontrar o livro perfeito em uma biblioteca massiva com base em uma única frase que você escreveu em um guardanapo. É isso que um "recuperador" de computador faz: ele pega sua consulta e encontra os documentos que melhor correspondem a ela.

Por muito tempo, a melhor maneira de fazer isso foi usar um "bibliotecário inteligente" (um modelo de IA) que lê sua frase e anota uma única palavra para representar toda a ideia. Isso é rápido, mas às vezes uma palavra não é suficiente para capturar um pensamento complexo.

Pesquisadores tentaram fazer o bibliotecário escrever múltiplas palavras para ser mais preciso. No entanto, com o tipo antigo de IA (chamado "autoregressivo"), o bibliotecário tinha que escrever essas palavras uma por uma, em sequência. Se você pedisse 20 palavras, eles teriam que escrever a primeira, esperar, escrever a segunda, esperar, e assim por diante. Isso era tão lento que as palavras extras não pareciam valer a espera, e os resultados não melhoravam muito.

Apresentando o DiffRetriever: O Bibliotecário do "Pensamento em Grupo"

Este artigo apresenta um novo tipo de bibliotecário chamado DiffRetriever, que usa um tipo diferente de IA chamado Modelo de Linguagem de Difusão.

Aqui está o truque mágico:
Em vez de escrever palavras uma por uma em sequência, o bibliotecário de difusão olha para uma página em branco com 20 espaços vazios (posições mascaradas) e preenche todos os 20 espaços exatamente ao mesmo tempo.

Pense nisso assim:

O Jeito Antigo (Autoregressivo): Você pede a um amigo para descrever um filme. Ele diz: "Foi..." (pausa) "um..." (pausa) "grande..." (pausa) "filme de..." (pausa) "ação." Eles têm que esperar cada palavra terminar antes de começar a próxima. Se você quiser 20 palavras, leva uma eternidade.
O Novo Jeito (Difusão): Você entrega a um amigo uma folha com 20 caixas em branco. Você diz: "Preencha essas caixas com a descrição." Eles olham para a folha inteira, pensam sobre o filme e simultaneamente preenchem todas as 20 caixas em um único olhar.

O Que o Artigo Encontrou

Velocidade vs. Qualidade: Os pesquisadores testaram isso em dois tipos de IA: os antigos escritores "um por um" e os novos escritores "todos de uma vez".
- Quando pediram aos escritores antigos para produzir múltiplas palavras, ficou cada vez mais lento, e a qualidade não melhorou realmente.
- Quando pediram aos novos escritores para produzir múltiplas palavras, foi tão rápido quanto escrever uma palavra (porque fizeram tudo de uma vez), mas a qualidade saltou significativamente. Foi como obter uma descrição muito melhor pelo mesmo tempo.
O Melhor Desempenho: Após treinar o novo bibliotecário "todos de uma vez" (especificamente um chamado "Dream"), ele se tornou o melhor em encontrar documentos relevantes em seus testes, superando todos os outros métodos, incluindo as antigas tentativas de múltiplas palavras e outros modelos de IA modernos.
O Orçamento "Perfeito": Os pesquisadores também notaram que perguntas diferentes precisam de números diferentes de palavras. Uma pergunta simples como "Qual é o tempo?" pode precisar apenas de 2 palavras, enquanto uma pergunta complexa pode precisar de 16.
- Atualmente, eles escolhem um número (como 4 ou 16) e o usam para cada pergunta.
- Eles descobriram que, se pudessem magicamente saber exatamente quantas palavras cada pergunta específica precisava antes de escrevê-las, o sistema ficaria ainda melhor — melhor até do que os modelos treinados.
- Eles também descobriram que pistas simples, como o tamanho da pergunta, podem prever quantas palavras são necessárias. Isso sugere que, no futuro, poderíamos construir um sistema que ajusta automaticamente quantas palavras escreve para cada pergunta, obtendo o melhor dos dois mundos.

Em Resumo

O artigo prova que o problema de usar múltiplas palavras para busca não era a ideia de usar múltiplas palavras; era a forma antiga de escrevê-las (uma por uma). Ao mudar para uma nova IA que escreve todas as palavras de uma vez, eles tornaram a busca com múltiplas palavras rápida, barata e muito mais precisa. É como perceber que você não precisa esperar um trem lento entregar seu pacote; você só precisa de um drone que deixa tudo de uma vez.

Resumo Técnico: DiffRetriever

Declaração do Problema

Trabalhos recentes, especificamente PromptReps (Zhuang et al., 2024), demonstraram que modelos de linguagem autoregressivos (AR LMs) poderiam funcionar como recuperadores zero-shot ao solicitá-los a gerar tokens representativos (vetores densos e logits esparsos) para consultas e passagens. No entanto, estender essa abordagem para recuperação de múltiplos tokens (gerar $K$ representantes em vez de um) provou ser ineficiente e frequentemente ineficaz com modelos AR. Como os modelos AR geram tokens sequencialmente, produzir $K$ representantes requer $K$ passagens para frente, fazendo com que a latência de codificação escale linearmente com $K$ . Tentativas anteriores de PromptReps de múltiplos tokens não mostraram ganhos consistentes de desempenho sobre a decodificação de token único, apesar desse custo significativo de latência, deixando incerto se a limitação era o próprio conceito de múltiplos tokens ou o gargalo da geração sequencial.

Metodologia

Os autores propõem o DiffRetriever, um framework de recuperação projetado especificamente para Modelos de Linguagem de Difusão (DiffLMs) para superar o gargalo da geração sequencial.

Mecanismo Central

Diferentemente dos modelos AR que geram tokens um por um, os DiffLMs preenchem posições mascaradas conjuntamente sob atenção bidirecional. O DiffRetriever aproveita isso por meio de:

Construção do Prompt: Adição de $K$ posições mascaradas ([MASK]) a um prompt de token representativo (por exemplo, "As palavras são [MASK]...").
Decodificação Paralela: Alimentação de toda a sequência (prompt + $K$ máscaras) no DiffLM em uma única passagem para frente bidirecional.
Extração de Representação: Extração simultânea de $K$ estados ocultos densos e $K$ vetores de logits esparsos das $K$ posições mascaradas.

Pontuação e Treinamento

Pontuação: O framework adota a receita híbrida de pontuação do PromptReps:
- Pontuação Densa: Calculada via interação tardia estilo ColBERT (MaxSim) entre estados ocultos de consulta e passagem.
- Pontuação Esparsa: Calculada via logits ativados por ReLU e agrupados por máximo (max-pooled).
- Pontuação Híbrida: Uma interpolação linear das pontuações densas e esparsas normalizadas.
Ajuste Fino: Os autores empregam ajuste fino contrastivo (perda InfoNCE) em ambas as pontuações densa e esparsa, aplicável tanto a backbones AR quanto DiffLM.

Configuração Experimental

O estudo compara quatro backbones (7B–8B parâmetros):

Autoregressivo (AR): LLaMA3-8B, Qwen2.5-7B.
Difusão (Diff): Dream-7B (inicializado a partir do Qwen2.5), LLaDA-8B (treinado do zero).
Baselines: BM25, PromptReps (AR), DiffEmbed (DiffLM como encoder agrupado por média) e RepLLaMA (AR ajustado finamente de forma contrastiva).
Conjuntos de Dados: MS MARCO, TREC DL 2019/2020 e BEIR-7 (sete conjuntos de dados diversos fora do domínio).

Principais Resultados

1. Eficiência e Eficácia de Múltiplos Tokens

Difusão vs. Autoregressão: O DiffRetriever de múltiplos tokens melhora substancialmente sobre o DiffRetriever de token único em todos os backbones de difusão testados. Em contraste, modelos AR de múltiplos tokens mostram ganhos planos ou negativos sobre modelos AR de token único.
Latência: O custo de codificação para o DiffRetriever de múltiplos tokens permanece aproximadamente constante, independentemente de $K$ (passagem única para frente), enquanto o custo de múltiplos tokens AR escala linearmente. No cenário zero-shot ( $K \le 20$ ), o AR de múltiplos tokens incorre em uma penalidade de latência de $\approx 15\times$ em comparação com a difusão, sem ganho consistente de desempenho.
Reversão de Desempenho: Em cenários zero-shot, os modelos AR inicialmente superam os DiffLMs em $K=1$ . No entanto, para $K>1$ , os DiffLMs (especificamente Dream e LLaDA) superam seus equivalentes AR. Por exemplo, o Dream de múltiplos tokens quase duplica o desempenho no MS MARCO em comparação com sua versão de token único, superando o Qwen2.5.

2. Desempenho de Ponta (State-of-the-Art)

Zero-Shot: O DiffRetriever no LLaDA alcança a maior pontuação média BEIR-7 entre sistemas zero-shot na comparação, superando o PromptReps em backbones AR e a baseline DiffEmbed.
Ajuste Fino: Após ajuste fino supervisionado, o DiffRetriever no Dream torna-se o recuperador BEIR-7 mais forte no geral, superando o PromptReps ajustado finamente, o DiffEmbed e o RepLLaMA ajustado finamente de forma contrastiva.
Sinergia de Backbone: O backbone Dream (inicializado a partir do Qwen2.5) mostra uma oscilação de desempenho maior do zero-shot para o ajuste fino em comparação com o LLaDA, sugerindo que priores de pré-treinamento AR combinados com decodificação de difusão bidirecional são altamente eficazes para recuperação.

3. Análise de Orçamento e Margem do Oráculo

Fixo vs. Adaptativo: O estudo identifica que o uso de um orçamento fixo $(K_q, K_p)$ em todas as consultas deixa desempenho significativo na mesa.
Análise do Oráculo: Um "oráculo por consulta" (que seleciona o $K$ ótimo para cada consulta usando a verdade fundamental) no modelo base congelado supera o ajuste fino contrastivo em um orçamento fixo em todas as combinações.
Previsibilidade: O orçamento ótimo correlaciona-se positivamente com características de consulta baratas (comprimento e entropia de Shannon), sugerindo que um roteador aprendido para seleção adaptativa de orçamento é uma direção futura viável.

Significado e Alegações

O artigo afirma resolver a ambiguidade em torno da recuperação de múltiplos tokens:

O Gargalo é a Geração Sequencial: A falha de recuperadores AR de múltiplos tokens anteriores não se deveu ao conceito de múltiplos tokens em si, mas ao custo de decodificação sequencial. O DiffRetriever prova que a recuperação de múltiplos tokens é eficaz e barata quando desacoplada da geração sequencial via modelos de difusão.
Objetivos Alinhados ao Treinamento: Consultar o modelo de difusão em sua forma pré-treinada (prever posições mascaradas em paralelo) transfere-se melhor para tarefas de recuperação do que reaproveitar o modelo como um encoder estilo BERT padrão (como feito no DiffEmbed).
Potencial Adaptativo: A existência de uma grande lacuna entre o desempenho de orçamento fixo e o desempenho do oráculo por consulta destaca a seleção adaptativa de orçamento como uma via crítica e inexplorada para melhorar a eficiência e eficácia da recuperação.

Os autores concluem que o DiffRetriever estabelece um novo estado da arte para recuperação baseada em difusão, demonstrando que tokens representativos paralelos podem superar tanto a difusão de token único quanto abordagens autoregressivas sequenciais, mantendo baixa latência.