MMD Guidance: Training-Free Distribution Adaptation for Diffusion Models via Maximum Mean Discrepancy Guidance

Resumo Técnico: Orientação por MMD para Adaptação de Distribuição sem Treinamento

Problema
Modelos de difusão pré-treinados servem como poderosos priors generativos para geração de amostras incondicionais e condicionais. No entanto, seus resultados frequentemente se desviam das distribuições alvo específicas do usuário, um problema crítico em tarefas de adaptação de domínio onde apenas um número limitado de exemplos de referência (ex: 50–100) está disponível. Retreinar ou realizar o ajuste fino (fine-tuning) desses modelos é frequentemente computacionalmente inviável ou impraticável devido a restrições de memória e à falta de grandes conjuntos de dados alvo. Os métodos existentes de orientação em tempo de inferência, como a orientação por classificador (classifier guidance) ou orientação livre de classificador (classifier-free guidance), ajustam as trajetórias de amostragem otimizando objetivos substitutos (ex: probabilidades de classe ou força de condicionamento) em vez de alinhar diretamente a distribuição gerada com uma distribuição de referência alvo. Consequentemente, esses métodos não podem alinhar explicitamente uma distribuição alvo arbitrária definida por um pequeno conjunto de amostras de referência.

Metodologia
O artigo propõe a Orientação por MMD (MMD Guidance), um mecanismo sem treinamento que aumenta o processo de difusão reversa com os gradientes da Discrepância de Média Máxima (MMD) entre as amostras geradas e um conjunto de dados de referência.

1. Mecanismo Central: O método modifica a atualização padrão de amostragem de difusão reversa em cada passo temporal $t$ subtraindo o gradiente da MMD empírica ao quadrado em relação ao estado atual da amostra. A regra de atualização é formulada como:
   $$z_{t-1} = \text{sampler}(z_t, t, \epsilon_\theta) - \lambda_t \nabla_{z_t} \widehat{\text{MMD}}^2(\hat{P}_t, \hat{Q})$$
   onde $\hat{P}_t$ é a distribuição empírica do lote (batch) gerado no passo temporal $t$, $\hat{Q}$ é a distribuição empírica das amostras de referência, e $\lambda_t$ controla a força da orientação.

2. Escolha da Divergência: Os autores selecionam a MMD em vez de $f$-divergências (como KL) e distâncias de Wasserstein porque a MMD fornece estimativas não enviesadas e de baixa variância a partir de dados limitados sem sofrer com a maldição da dimensionalidade. Além disso, sua formulação baseada em kernel permite o cálculo eficiente de gradientes em forma fechada.

3. Implementação no Espaço Latente: Para garantir eficiência computacional, a orientação é aplicada diretamente no espaço latente dos Modelos de Difusão Latente (LDMs). As imagens de referência são codificadas uma única vez em representações latentes usando o codificador VAE. Isso reduz a dimensionalidade do cálculo da MMD e melpar a relação sinal-ruído ao filtrar o ruído ao nível de pixel.

4. Adaptação Sensível ao Prompt: Para geração condicional, o framework se estende para a adaptação sensível ao prompt usando um kernel de produto. Este kernel, $k_{\otimes}([p, z], [p', z']) = k_{\text{text}}(p, p') \cdot k_{\text{image}}(z, z')$, mede conjuntamente a similaridade tanto no espaço semântico (prompt) quanto no espaço visual (latente). Isso permite que a orientação alinhe as amostras geradas tanto com o prompt textual quanto com o estilo visual de pares de referência sem retreinar o denoiser.

5. Dinâmica de Gradiente: O gradiente da MMD ao quadrado consiste em dois termos: um termo intra-lote que cria repulsão entre as amostras geradas (promovendo diversidade) e um termo transversal que cria atração em direção às amostras de referência (incentivando o alinhamento da distribuição). A análise teórica fornece garantias de concentração para o gradiente do termo transversal, garantindo que ele estime confiavelmente o gradiente da população mesmo com amostras de referência limitadas.

Principais Contribuições

• Orientação por MMD: Introdução de um método de tempo de inferência, sem treinamento, para adaptar modelos de difusão a distribuições alvo especificadas por conjuntos de referência limitados.
• Extensão Sensível ao Prompt: Uma extensão prática usando kernels de produto para permitir o alinhamento conjunto com prompts de texto e estilos visuais na geração condicional.
• Implementação Eficiente: Uma implementação escalável operando diretamente no espaço latente de LDMs, reduzindo significativamente o custo computacional em comparação com a orientação no espaço de pixels.
• Validação Empírica: Experimentos abrangentes demonstrando que a Orientação por MMD alcança um alinhamento de distribuição superior enquanto preserva a fidelidade das amostras.

Resultados Experimentais
Os autores avaliaram a Orientação por MMD em Modelos de Mistura Gaussiana (GMMs) sintéticos e em benchmarks de imagens do mundo real (FFHQ, CelebA-HQ e subconjuntos de estilo específico de saídas do Stable Diffusion).

• Dados Sintéticos: Em GMMs de 8 e 25 modalidades, a Orientação por MMD recuperou com sucesso as proporções da mistura alvo e corrigiu distribuições desbalanceadas, superando as linhas de base em Distância de Fréchet (FD) e Distância de Kernel (KD).
• Imagens do Mundo Real: Em experimentos com FFHQ (ex: gerando rostos com estilos específicos de óculos) e CelebA-HQ, a Orientação por MMD alcançou as maiores pontuações de fidelidade e cobertura de distribuição em comparação com a amostragem não orientada, orientação por classificador e até mesmo linhas de base treinadas pelo usuário que exigiam ajuste fino.
• Geração Condicional: Em configurações condicionadas por texto (usando Stable Diffusion XL e PixArt), o método adaptou com sucesso o estilo visual das imagens geradas (ex: "gatos em estilo anime" ou "carros em estilo Pixar") para corresponder aos conjuntos de referência enquanto mantinha a adesão aos prompts de texto.
• Eficiência: A análise de tempo de execução indica que o custo computacional adicional da orientação por MMD no espaço latente é negligenciável em comparação com a amostragem não orientada, mesmo para modelos de grande escala como o SDXL.
• Robustez: O método demonstrou robustez através de variações no número de amostras de referência (convergindo efetivamente com apenas 50–100 amostras) e diferentes lacunas de domínio.

Significância e Alegações
O artigo afirma que a Orientação por MMD oferece uma solução prática para implantar modelos generativos em cenários onde os recursos computacionais para retreinamento são limitados ou onde os dados específicos do usuário são escassos. Ao operar inteiramente em tempo de inferência sem modificar os parâmetros do modelo, o método permite implantação imediata e preserva a privacidade (já que nenhum dado do usuário é usado para treinamento). Os autores enfatizam que a abordagem efetivamente preenche a lacuna entre as capacidades dos modelos de fundação pré-treinados e os requisitos específicos do usuário, alcançando o alinhamento de distribuição enquanto mantém a diversidade e a qualidade das amostras geradas. O trabalho posiciona a Orientação por MMD como uma ferramenta complementar às técnicas existentes de ajuste fino e engenharia de prompt, sendo particularmente valiosa em ambientes com restrições de privacidade ou de recursos.