CausalEmbed: Auto-Regressive Multi-Vector Generation in Latent Space for Visual Document Embedding

Resumo Técnico: CausalEmbed

1. O Problema

A Recuperação de Documentos Visuais (VDR - Visual Document Retrieval) é fundamental para sistemas de busca empresarial e RAG (Geração Aumentada por Recuperação) em documentos ricos em conteúdo visual.

• Limitação das Abordagens Atuais: Métodos de ponta baseados em Multimodal Large Language Models (MLLMs), como o ColPali e o ColQwen, utilizam representações de multi-vetores (geralmente centenas ou milhares de tokens visuais por página) para capturar detalhes finos e layout. Embora eficazes, essa abordagem gera uma sobrecarga de armazenamento proibitiva e custos de inferência elevados, dificultando a escalabilidade em ambientes de produção.
• Falhas das Técnicas de Compressão Existentes: Métodos atuais de compressão (como pruning baseado em agrupamento/clustering ou pooling) tentam reduzir o número de vetores selecionando ou fundindo os tokens existentes. No entanto, essas técnicas frequentemente resultam em perda significativa de precisão na recuperação, pois dependem da qualidade da representação densa original e não conseguem gerar novos semânticos compactos eficientemente.

2. Metodologia: CausalEmbed

O CausalEmbed propõe uma mudança de paradigma: em vez de codificar paralelamente patches de imagem em vetores e depois comprimir, ele utiliza uma geração auto-regressiva para sintetizar vetores de embedding latentes sequencialmente.

• Geração Auto-Regressiva: O modelo ajusta um MLLM pré-treinado para gerar uma sequência de vetores latentes ($z_1, ..., z_L$) token a token. Cada novo vetor depende do contexto original (imagem ou texto) e de todos os vetores gerados anteriormente.
  • Isso permite que o modelo "construa" uma representação compacta (apenas dezenas de tokens) que captura a semântica global e detalhada da página, em vez de apenas selecionar subconjuntos de tokens existentes.
• Objetivo de Treinamento (Função de Perda): Para evitar o colapso da representação e garantir a diversidade, o treinamento combina três componentes:
  1. Perda de Alinhamento Contrastivo ($L_m$): Garante que pares consulta-documento relevantes tenham pontuações mais altas que pares não relevantes (usando o mecanismo Late Interaction com MaxSim).
  2. Perda de Refinamento Progressivo ($L_d$): Incentiva que cada passo adicional na geração auto-regressiva traga uma melhoria marginal na pontuação de recuperação, garantindo que a sequência inteira seja útil.
  3. Regularização de Diversidade ($L_q$): Penaliza a similaridade entre vetores gerados na mesma sequência (especialmente no lado da consulta) para evitar homogeneização e garantir que cada token capture aspectos distintos da semântica.
• Análise Teórica de Gradiente: O artigo demonstra teoricamente que a abordagem auto-regressiva permite uma propagação de gradiente mais eficiente através de toda a cadeia de geração, comparada aos métodos de forward-pass paralelo, onde apenas tokens selecionados recebem sinais de gradiente diretos.

3. Principais Contribuições

1. Novo Paradigma de Codificação: Introdução de uma abordagem auto-regressiva para gerar embeddings de documentos visuais, abandonando a codificação paralela baseada em patches.
2. Alta Eficiência com Alta Precisão: O método alcança uma compressão de 30x a 155x no número de tokens (reduzindo de milhares para apenas ~32 vetores) mantendo ou superando o desempenho de métodos de base completos.
3. Escalabilidade no Tempo de Inferência (Test-Time Scaling): O CausalEmbed exibe uma propriedade única onde a precisão da recuperação pode ser ajustada dinamicamente alterando o número de tokens gerados durante a inferência (sem necessidade de retreinamento), permitindo um trade-off flexível entre latência e precisão.
4. Generalização Robusta: O método funciona consistentemente bem em diferentes backbones (desde modelos menores como PaliGemma até modelos maiores como Qwen2.5-VL), superando métodos de pruning mesmo em modelos com desempenho base inferior.

4. Resultados Experimentais

Os experimentos foram conduzidos nos benchmarks ViDoRe V1, V2 e V3, que cobrem diversos domínios (científico, jurídico, industrial, multilíngue).

• Desempenho vs. Compressão:
  • O CausalQwen (baseado em Qwen2.5-VL) com apenas 32 tokens superou significativamente os melhores métodos de pruning (como SemCluster e K-Means) e modelos de vetor único (BiQwen2) nos benchmarks ViDoRe V2 e V3.
  • No CausalPali (baseado em PaliGemma), o método superou o modelo completo ColPali (com 1000+ tokens) em quase 5% de desempenho em todos os benchmarks, demonstrando que a geração de vetores é superior à seleção de vetores densos.
• Eficiência:
  • Redução drástica no custo de armazenamento e latência de interação tardia (Late Interaction). Embora a geração auto-regressiva adicione um pequeno custo de codificação de consulta, a redução no número de vetores acelera a fase de comparação (MaxSim) em aproximadamente 11.8x.
• Escalabilidade: A curva de desempenho mostra que aumentar o número de tokens de 1 para 32 melhora consistentemente a pontuação nDCG@5, aproximando-se do limite superior do modelo completo, validando a propriedade de Matryoshka emergente.

5. Significado e Impacto

O CausalEmbed representa um avanço crucial para a viabilidade prática de sistemas de recuperação de documentos visuais em larga escala.

• Viabilidade Econômica: Ao reduzir o armazenamento de vetores em mais de 30 vezes sem sacrificar a precisão, torna-se economicamente viável implantar sistemas de VDR de alta qualidade em ambientes com recursos limitados.
• Flexibilidade Operacional: A capacidade de ajustar a precisão em tempo de execução (escalabilidade no tempo de teste) oferece uma flexibilidade inédita para aplicações que exigem diferentes níveis de latência e precisão.
• Mudança de Paradigma: O trabalho sugere que a geração auto-regressiva não é apenas para texto ou imagens, mas é uma ferramenta poderosa para a criação de representações latentes compactas e ricas em semântica, abrindo caminho para futuras pesquisas em embeddings generativos multimodais.

Em suma, o CausalEmbed resolve o dilema entre alta precisão e baixo custo de armazenamento na recuperação de documentos visuais, substituindo a abordagem de "selecionar o melhor" por uma abordagem de "gerar o ideal".