Test-Time Alignment of Text-to-Image Diffusion Models via Null-Text Embedding Optimisation

Resumo Técnico: Alinhamento de Modelos de Difusão de Texto-para-Imagem via Null-Text em Tempo de Teste

Declaração do Problema
Modelos de difusão de texto-para-imagem pré-treinados são frequentemente treinados em conjuntos de dados da web de larga escala e não curados, que contêm conteúdo desalinhado ou indesejável. Alinhar esses modelos com valores humanos específicos ou objetivos alvo (ex: qualidade estética, preferências específicas) durante a inferência é crítico para uma implantação confiável. As abordagens de alinhamento existentes geralmente se dividem em duas categorias, ambas sofrendo de limitações significativas:

1. Métodos de Fine-tuning: Estes atualizam os parâmetros do modelo para otimizar uma função de recompensa. Embora eficazes, são computacionalmente caros e propensos à sobre-otimização de recompensa (reward hacking), onde o modelo sofre overfitting a recompensas de aproximação (proxies), perdendo a generalização para outras recompensas e a diversidade de saída.
2. Métodos de Alinhamento em Tempo de Teste (TTA): Estes otimizam variáveis latentes ou de ruído, amostram de posteriors condicionadas por recompensa ou utilizam busca discreta sem atualizar os pesos do modelo. No entanto, esses métodos frequentemente sofrem de sub ou sobre-otimização de recompensa devido à natureza altamente não estruturada dos espaços latentes e de ruído. A otimização nesses espaços pode levar ao "reward hacking", onde o modelo explora padrões de ruído não semânticos para maximizar a recompensa alvo enquanto degrada a qualidade visual, a coerência semântica e o desempenho em outras métricas.

Metodologia: Null-Text Test-Time Alignment (Null-TTA)
Os autores propõem o Null-TTA, um framework livre de treinamento que alinha modelos de difusão otimizando o embedding incondicional (null-text embedding) dentro do Classifier-Free Guidance (CFG), em vez de manipular variáveis latentes ou de ruído.

• Insight Central: No CFG, o embedding incondicional ($\phi$) serve como uma âncora geométrica para a distribuição generativa do modelo. Diferente dos espaços de ruído não estruturados, o espaço de embedding de texto possui uma natureza semântica estruturada. Otimizar dentro deste espaço garante que o alinhamento ocorra em um manifold semanticamente coerente, evitando que o modelo explore artefatos de ruído não semânticos.
• Função Objetivo: O método formula uma objetivo regularizado que direciona a distribuição generativa do modelo em direção a uma recompensa alvo enquanto penaliza o desvio de distribuição em relação ao modelo pré-treinado. O objetivo maximiza a recompensa esperada $R(x_0)$ enquanto minimiza a divergência KL entre a distribuição induzida pelo embedding otimizado $\phi'$ e a distribuição original induzida por $\phi$:
  $$\max_{\phi'} \left( \lambda_1 \mathbb{E}_{p(x_0|\phi')}[R(x_0)] - \lambda_2 \text{KL}(p(x_{0:T}, \phi') || p(x_{0:T}, \phi)) \right)$$
  O termo KL é decomposto em divergências locais entre passos de denoising consecutivos e uma divergência entre as distribuições de embedding. Isso é simplificado para uma expressão de forma fechada envolvendo a diferença no ruído previsto ($\tilde{\epsilon}$) e a diferença nos embeddings ($\phi' - \phi$).
• Estratégia de Otimização:
  • Regularização Adaptativa: O peso de regularização $\lambda_2$ é anelado de forma adaptativa conforme o processo de denoising progride ($t \to 0$). Uma regularização forte nos estágios iniciais estabiliza a otimização, enquanto uma regularização mais fraca nos estágios posteriores permite um alinhamento mais refinado.
  • Busca Gananciosa (Greedy Search): Em cada passo de denoising, o método emprega uma busca gananciosa leve. Ele amostra $K$ candidatos latentes do kernel de transição, estima a amostra limpa via fórmula de Tweedie, pontua-os com o modelo de recompensa e seleciona deterministicamente o candidato de maior pontuação para o próximo passo.
  • Recompensas Não-Diferenciáveis: Para cenários onde a função de recompensa é uma caixa-preta (não-diferenciável), o método emprega estimativa de gradiente de ordem zero usando perturbações aleatórias do null-text embedding.

Principais Contribuições

1. Framework Null-TTA: Um novo framework de TTA livre de treinamento que realiza o alinhamento no espaço semântico estruturado do condicionamento de texto, evitando as armadilhas da otimização de ruído/latente não estruturada.
2. Objetivo Principiado: Um objetivo regularizado por KL que reorienta a própria distribuição generativa do modelo em direção à recompensa alvo, garantindo um alinhamento estável e eficiente sem atualizações de parâmetros. Isso evita o reward hacking e mantém a coerência semântica.
3. Desempenho de Estado da Arte: Experimentos extensos demonstrando que o Null-TTA atinge uma performance superior de recompensa alvo enquanto mantém forte generalização de recompensa cruzada e qualidade perceptual.

Resultados Experimentais
Os autores avaliaram o Null-TTA no Stable Diffusion v1.5 e SDXL usando várias recompensas alvo (HPSv2, PickScore, Aesthetic Score, ImageReward) e objetivos não-diferenciáveis (compressibilidade JPEG).

• Otimização de Recompensa vs. Generalização: O Null-TTA consistentemente desloca a fronteira de Pareto para fora em comparação com os baselines (incluindo DNO, DAS, DSearch e TDPO). Ele alcança recompensas alvo mais altas enquanto degrada as recompensas mantidas (held-out) significativamente menos do que os competidores, indicando um alinhamento bem-sucedido sem sobre-otimização.
• Desempenho Qualitativo: Em categorias desafiadoras (contagem, composicionalidade, raciocínio espacial, cores incomuns, transferência de estilo, cenas impossíveis), o Null-TTA produz imagens que satisfazem mais fielmente as restrições do prompt enquanto preservam a coerência global, evitando modos de falha como contagem incorreta de objetos ou relações emaranhadas.
• Estudo com Usuários: Um estudo cego com 16 participantes mostrou uma preferência estatisticamente significativa pelo Null-TTA sobre os baselines, confirmando que os ganhos de recompensa se traduzem em qualidade perceptual.
• Eficiência Computacional: O Null-TTA alcança um desempenho de recompensa alvo mais forte com menor uso de memória GPU e tempos de execução (wall-clock time) comparáveis ou mais rápidos que os baselines. Essa eficiência decorre de realizar o backpropagation apenas através das camadas de cross-attention (atualizando o embedding) em vez de todo o U-Net.
• Recompensas Não-Diferenciáveis: O método otimizou com sucesso o tamanho de arquivo JPEG (um objetivo de caixa-preta) usando gradientes de ordem zero, alcançando taxas de compressão competitivas enquanto preserva melhor a qualidade visual do que os baselines baseados em gradiente.

Significância e Alegações
O artigo estabelece a otimização de espaço semântico como um paradigma novo, eficaz e principiado para o Alinhamento em Tempo de Teste. Ao alavancar a natureza estruturada do espaço de embedding de texto, o Null-TTA fornece uma regularização inerente que evita o reward hacking e mantém o conhecimento pré-treinado do modelo. Os autores alegam que esta abordagem oferece uma alternativa simples, porém poderosa, aos métodos de TTA existentes, permitindo um alinhamento em tempo de inferência confiável, eficiente e generalizável para modelos de difusão sem o custo computacional do fine-tuning.