Beyond Generative Priors: Minority Sampling with JEPA-Guided Diffusion

Declaração do Problema

A amostragem de minoria visa gerar instâncias de baixa densidade em uma variedade de dados, o que é crítico para aplicações como diagnóstico médico, detecção de anomalias e IA criativa. No entanto, as abordagens existentes definem "minoria" em relação ao prior generativo aprendido pelo próprio modelo (ou seja, a distribuição dos dados de treinamento). Isso cria uma limitação fundamental: as amostras identificadas como "raras" são frequentemente apenas artefatos do conjunto de treinamento específico ou dos vieses indutivos do modelo, e não instâncias genuinamente atípicas no mundo real. Por exemplo, um modelo treinado em conjuntos de dados padrão pode considerar um cão em um fundo branco "raro" simplesmente porque está sub-representado nos dados de treinamento, enquanto uma aeronave furtiva pode ser genuinamente rara no mundo real, mas não ser capturada como tal se o prior do modelo for estreito. O artigo argumenta que os métodos atuais falham em alinhar-se com priors do mundo real, que são necessários para capturar a verdadeira raridade semântica.

Metodologia: Orientação JEPA

Os autores propõem a Orientação JEPA, um framework que desloca a amostragem de minoria de uma perspectiva centrada no gerador para uma perspectiva centrada no mundo. Isso é alcançado guiando modelos de difusão usando a densidade implícita de uma Arquitetura Preditiva de Incorporação Conjunta (JEPA), uma classe de modelos de mundo (por exemplo, DINOv2) conhecida por codificar representações amplas e semanticamente ricas.

Mecanismo Central:

1. Definição Centrada no Mundo: Em vez de minimizar a probabilidade sob a densidade do modelo de difusão $p_\theta(x)$, o método busca minimizar a probabilidade sob a densidade induzida pela JEPA $p_\phi(x)$. Amostras de minoria são definidas como aquelas que residem em regiões de baixa densidade do espaço de representação da JEPA.
2. JEPA-SCORE: O método utiliza o JEPA-SCORE, uma medida de densidade derivada do Jacobiano de um codificador JEPA pré-treinado $f_\phi$. Especificamente, é a soma dos logaritmos dos valores singulares do Jacobiano $J_f(x)$. Valores mais baixos de JEPA-SCORE indicam menor densidade (maior atipicidade) no prior do mundo.
3. Aproximação via SVD Randomizada: Calcular o Jacobiano completo e sua Decomposição em Valores Singulares (SVD) em cada etapa da difusão é computacionalmente proibitivo. Os autores empregam SVD Randomizada para aproximar eficientemente os $k$ maiores valores singulares do Jacobiano comprimido. Eles fornecem um limite superior teórico para o erro de aproximação, decompondo-o em erro de estimação da SVD randomizada e erro de truncamento.
4. Teorema do Envelope para Eficiência: Para evitar a retropropagação através do loop de otimização interna da SVD randomizada (que depende da entrada $x_t$), os autores invocam o Teorema do Envelope. Isso permite tratar a matriz de projeção ótima $Q^*$ como uma constante durante o cálculo do gradiente, reduzindo significativamente o uso de memória e a sobrecarga computacional, preservando ao mesmo tempo gradientes de primeira ordem corretos.
5. Orientação Adiada: Para abordar a lacuna de domínio entre os latentes intermediários ruidosos do modelo de difusão e as entradas limpas esperadas pelo codificador JEPA, os autores introduzem a orientação adiada. A orientação JEPA é aplicada apenas após um passo de tempo intermediário $\tau T$ (por exemplo, $t=0.8T$), garantindo que as estimativas de remoção de ruído $\hat{x}_{0|t}$ estejam suficientemente claras para que o codificador JEPA avalie efetivamente. Isso também permite que o método se estenda à geração condicional (por exemplo, texto para imagem), pois a condição é estabelecida cedo no processo antes do início da orientação centrada no mundo.

Principais Contribuições

• Amostragem de Minoria Centrada no Mundo: Uma mudança de paradigma redefinindo amostras de minoria com base em priors do mundo real (via JEPAs) em vez dos priors estreitos do modelo generativo.
• Framework de Orientação JEPA: Um método de amostragem de difusão principista que direciona trajetórias para regiões de baixa densidade sob priors induzidos por JEPA.
• Cálculo Eficiente: O desenvolvimento de estratégias de aproximação principistas (SVD Randomizada) e a aplicação do Teorema do Envelope para tornar a orientação JEPA computacionalmente prática, acompanhada de limites teóricos de erro.
• Integração sem Treinamento: A abordagem é inteiramente sem treinamento, exigindo apenas modelos de difusão e JEPA pré-treinados, e pode ser integrada no momento da inferência sem modificar os modelos subjacentes.

Resultados Experimentais

Os autores avaliaram a orientação JEPA em tarefas de geração incondicional, condicional por classe e texto para imagem (usando CelebA, ImageNet, Stable Diffusion 1.5 e SDXL-Lightning).

• Atipicidade: A orientação JEPA alcança consistentemente o menor JEPA-SCORE (indicando maior atipicidade sob priors do mundo) em comparação com bases de referência como Sehwag et al. (2022), MG, SGMS e MinorityPrompt.
• Fidelidade e Validade: Apesar de mirar em regiões raras, o método mantém alta fidelidade. Em tarefas de texto para imagem, alcança pontuações competitivas de alinhamento com texto (CLIPScore, PickScore, ImageReward) enquanto gera amostras semanticamente distintas (por exemplo, pessoal militar idoso, jogadores de beisebol afro-americanos) que outros métodos falham em produzir.
• Comparação Qualitativa: Visualizações mostram que métodos centrados no gerador frequentemente produzem amostras com contextos atípicos (por exemplo, uma águia perto de um guindaste), enquanto a orientação JEPA produz amostras com semânticas ou comportamentos atípicos (por exemplo, águias acasalando, formas animais incomuns), alinhando-se melhor com a raridade do mundo real.
• Utilidade a jusante: Em uma tarefa de aumento de dados para treinamento de classificadores, o aumento com amostras de minoria guiadas por JEPA produziu desempenho superior ou comparável a outros métodos usando menos amostras, sugerindo que essas amostras fornecem sinais de treinamento mais informativos.

Significado e Alegações

O artigo afirma que a orientação JEPA oferece uma solução robusta para a limitação dos amostradores de minoria existentes de ficarem presos na própria distribuição do modelo generativo. Ao alavancar modelos de mundo (JEPAs), o método gera amostras que são genuinamente raras no mundo real, e não apenas raras em relação a um conjunto de dados específico.

Os autores enfatizam que esta abordagem:

1. Desacopla a noção de minoria do prior generativo.
2. Ponteia a lacuna entre modelagem generativa e modelagem de mundo, um componente chave na busca pela inteligência artificial geral.
3. Fornece uma ferramenta prática e sem treinamento para aplicações que exigem dados diversos e atípicos, como testes de robustez, melhoria de equidade e geração de conteúdo criativo.

Os autores também reconhecem uma preocupação potencial de uso duplo: reverter a direção da orientação poderia reforçar vieses existentes em priors do mundo, e eles incentivam o uso responsável. Eles observam que, embora o método introduza sobrecarga computacional em comparação com amostradores padrão, ele permanece mais rápido do que algumas bases de referência complexas (como MinorityPrompt) e oferece uma compensação razoável pelo ganho significativo em atipicidade centrada no mundo.