Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

Resumo Técnico: Generalização Composicional em Tempo de Teste em Modelos de Difusão via Descoberta de Conceitos

Declaração do Problema

A generalização composicional — a capacidade de recombinar partes familiares em configurações novas — é uma marca distintiva da inteligência humana, mas permanece um desafio significativo para modelos de aprendizado profundo. Embora redes neurais frequentemente interpolem bem dentro das distribuições de treinamento, elas têm dificuldade quando primitivas familiares aparecem em combinações desconhecidas. Abordagens existentes para geração composicional em modelos de difusão geralmente dependem de uma biblioteca pré-definida de conceitos, sinais de condicionamento ou rótulos externos para guiar a geração de composições novas. Isso cria uma lacuna entre a descoberta de conceitos e a generalização composicional: métodos anteriores assumem que as primitivas relevantes já são conhecidas e disponíveis, falhando em abordar cenários onde um modelo deve descobrir conceitos específicos da consulta do zero, em tempo de teste, sem um vocabulário fixo ou supervisão externa.

Metodologia

O artigo propõe um framework que reaproveita um modelo de difusão pré-treinado como uma hierarquia implícita de modos de densidade indexados pelo tempo para descobrir e compor conceitos em tempo de teste. O método opera em três etapas principais:

1. Descoberta de Conceitos via Ascensão de Modo:
   Em vez de usar o modelo de difusão apenas para amostragem reversa, os autores utilizam a função de pontuação aprendida indexada pelo tempo $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$ para localizar modos locais de densidade das marginais ruidosas $p_t(x_t)$. Começando de uma única consulta fora da distribuição (OOD) $x_q$, o método realiza ascensão de gradiente na pontuação incondicional em múltiplos passos de tempo de ruído $t$ para recuperar modos locais.

   • Base Teórica: Aproveitando a fórmula de Tweedie, os autores demonstram que esses modos correspondem a "protótipos de conceito" livres de ruído no espaço de dados limpos ($x_0$), atuando efetivamente como centróides discretos em vários níveis de abstração, dependendo do nível de ruído.
   • Representação: Cada modo descoberto é mapeado de volta ao espaço $x_0$ e aproximado como um especialista Gaussiano local $q_j(x_0) = \mathcal{N}(x_0; m_j, \Sigma_j)$, onde a média é derivada do modo e a covariância é estimada por meio de sondas de Hutchinson.

2. Composição por Produto de Especialistas (PoE):
   Os especialistas Gaussianos descobertos são combinados em um único modelo professor analítico usando uma formulação de Produto de Especialistas (PoE).

   • Seleção: Um algoritmo de otimização gananciosa submodular seleciona o subconjunto de protótipos mais relevante para a consulta, maximizando um objetivo de cobertura de verossimilhança logarítmica por dimensão.
   • Composição: Os especialistas selecionados são combinados em uma distribuição PoE ponderada $q_T$. Como os especialistas são Gaussianos, a composição resultante também é Gaussiana, com uma função de pontuação de forma fechada. Isso permite a derivação de uma previsão de ruído analítica $\epsilon_{q_T}(x_t, t)$ sem a necessidade de um token condicional treinado para a composição nova.

3. Amostragem Composicional e Distilação:

   • Amostragem Analítica: A pontuação analítica do professor PoE é integrada a um framework de orientação sem classificador (CFG). O modelo contrasta a previsão do PoE com a previsão incondicional para guiar a amostragem em direção à composição descoberta.
   • Distilação LoRA: Para melhorar a geração no manifold e criar uma condição reutilizável, o método opcionalmente distila o professor PoE no modelo base. Isso envolve treinar um novo embedding de classe e um adaptador de Adaptação de Baixo RANK (LoRA) em um conjunto de amostras geradas pelo professor PoE. Esta etapa absorve a composição descoberta de volta no manifold de pontuação do modelo.

Contribuições Principais

1. Ponte entre Descoberta e Composição: O trabalho conecta a descoberta de conceitos e a generalização composicional ao tratar um modelo de difusão pré-treinado como uma hierarquia implícita de modos de densidade. Isso permite a recuperação de protótipos de conceitos discretos em tempo de teste sem um pool pré-definido de primitivas.
2. Mecanismo de Composição em Tempo de Teste: Os autores introduzem um método que seleciona gananciosamente modos relevantes usando um objetivo de verossimilhança submodular e os combina em um modelo professor PoE com pontuação de forma fechada, permitindo amostragem direta sem re-treinar o modelo base na composição específica.
3. Validação Empírica: O framework é avaliado em dois benchmarks, ColorMNIST e CelebA, demonstrando que os protótipos descobertos e sua composição PoE codificam informações mais ricas sobre composições não vistas do que a amostragem apenas da consulta ou a recuperação da classe treinada mais próxima.

Resultados

Experimentos em ColorMNIST (identidade do dígito, cor, fundo) e CelebA (atributos faciais) mostram que o método proposto supera várias linhas de base:

• Linhas de Base: O método é comparado com "Apenas Consulta" (amostragem ao redor da consulta com variância fixa), "Top-1/Top-3" (recuperando as classes treinadas mais próximas) e CFG padrão.
• Desempenho: Tanto o amostrador PoE analítico quanto a variante distilada por LoRA alcançam métricas superiores de FID, similaridade CLIP e Precisão/Recall em comparação com as linhas de base.
  • Em ColorMNIST, a variante LoRA alcança uma pontuação F1 de 87,6% (Fidelidade) e 76,3% (Generalização), superando significativamente a linha de base Top-1 (~6,5% F1).
  • Em CelebA, o método PoE alcança o melhor FID de Generalização (140,3) e pontuações CLIP, enquanto o LoRA melhora ainda mais o equilíbrio entre cobertura e fidelidade, alcançando o maior Recall de Generalização e F1.
• Análise Qualitativa: Resultados visuais indicam que, enquanto as linhas de base frequentemente falham em preservar a combinação específica de atributos (por exemplo, alterando a identidade do dígito ou desviando para combinações de atributos diferentes), o método proposto reconstrói com sucesso a composição OOD e generaliza para outras instâncias da mesma classe.

Significado e Alegações

O artigo afirma que a geometria da pontuação indexada pelo tempo de modelos de difusão contém conceitos de modo de densidade reutilizáveis que suportam a geração composicional em tempo de teste sem uma biblioteca de conceitos pré-definida.

• Insight Teórico: Os resultados sugerem que modelos de difusão codificam implicitamente uma hierarquia de conceitos, onde modos de marginais ruidosas servem como protótipos em diferentes níveis de abstração.
• Implicação Prática: O método permite a geração de configurações novas a partir de uma única consulta OOD, oferecendo um caminho para generalização composicional que não depende de supervisão externa ou vocabulários pré-especificados.
• Limitações: Os autores observam modestamente que o protocolo assume que os fatores primitivos são conhecidos (configuração composicional padrão) e que o método depende dos vieses indutivos do backbone pré-treinado. Embora possa interpolar entre primitivas conhecidas, não pode fabricar primitivas genuinamente novas que extrapolem além do suporte de treinamento (por exemplo, um fundo branco em um modelo treinado apenas em fundos escuros).

O trabalho fornece evidências teóricas e empíricas de que modelos de difusão podem ser reaproveitados para descoberta e composição de conceitos, potencialmente estendendo-se a aplicações mais amplas em modelagem de linguagem e controle robótico.