Resumo Técnico: Deriva de Características do Professor (TFD)

Declaração do Problema

Modelos de difusão e correspondência de fluxo pré-treinados alcançam geração de imagens de alta fidelidade, mas sofrem com custos de inferência elevados, frequentemente exigindo dezenas ou centenas de passagens forward para gerar uma única imagem. Embora existam métodos de destilação para comprimir esses professores de múltiplos passos em geradores eficientes de um ou poucos passos, as abordagens atuais enfrentam complexidade significativa. Métodos como Destilação de Correspondência de Distribuição (DMD) e suas variantes dependem de redes auxiliares (por exemplo, redes de pontuação falsa, discriminadores adversariais) ou cronogramas de treinamento multiestágio. Outras abordagens, como Modelos de Consistência, podem sofrer com treinamento instável ou exigir produtos vetoriais de Jacobiano complexos. Além disso, objetivos recentes de "Modelo em Deriva", que oferecem uma perda direta de atração-repulsão para geração em um passo, tipicamente dependem de um extrator de características pré-treinado externo (por exemplo, modelos auto-supervisionados ou latent-MAE) para definir o espaço de representação. Isso introduz sobrecarga adicional de treinamento e complexidade arquitetônica.

O problema central abordado é como alcançar uma destilação de difusão de um passo eficiente e de alta qualidade usando um único objetivo, sem depender de redes auxiliares, pipelines de otimização complexos ou extratores de características externos.

Metodologia: Deriva de Características do Professor (TFD)

Os autores propõem a Deriva de Características do Professor (TFD), um framework que simplifica a destilação de um passo aproveitando as representações internas do próprio modelo professor pré-treinado.

1. Espaço de Representação Induzido pelo Professor

Diferentemente do Modelo em Deriva original, que requer um codificador de características externo, a TFD utiliza os estados ocultos intermediários do professor pré-treinado congelado como o espaço de características para calcular a perda de deriva.

• Mecanismo: Para uma amostra gerada $x$ e uma amostra real $x^+$, as características são extraídas de camadas específicas $\ell$ da rede do professor $T$: $\phi^\ell_T(x, c) = h^\ell_T(x, t_f, c)$.
• Perda de Deriva: O campo de deriva é computado diretamente neste espaço induzido pelo professor. O gerador $f_\theta$ é otimizado para minimizar a distância entre sua saída e um "alvo transportado congelado" definido pelo campo de atração-repulsão do espaço de características do professor. Isso alinha o objetivo de destilação com o modelo sendo destilado, eliminando a necessidade de uma rede de representação separada.

2. Extração de Características com Ruído

Uma descoberta empírica chave é que extrair características de entradas moderadamente ruidosas produz melhores resultados do que usar amostras limpas.

• Antes de alimentar uma amostra no professor para extração de características, ela é perturbada com o processo de adição de ruído forward do professor: $x_{tf} = x + \sigma_{tf} \xi$.
• Os autores hipotetizam que essa perturbação suaviza a geometria das características do professor, suprimindo discrepâncias de alta frequência e tornando o espaço de representação mais adequado para correspondência ao nível da distribuição. Experimentos indicam que um nível de ruído de $\sigma_{tf} = 0.1$ é ótimo.

3. Regularização de Margem de Âncora

Para abordar o problema de colapso de modo (onde o gerador falha em cobrir regiões diversas suportadas pelo professor), os autores introduzem uma leve perda de margem de âncora.

• Âncoras: Um conjunto de características de âncora $\{a_i\}$ é extraído de amostras reais ou geradas pelo professor.
• Mecanismo: A perda incentiva amostras geradas a ocuparem regiões ao redor dessas âncoras. Ela ajusta adaptativamente o "suporte" necessário para cada âncora com base na densidade local da própria distribuição de âncoras.
• Objetivo: Isso fornece um sinal de cobertura complementar ao objetivo local de deriva de atração-repulsão, garantindo que o gerador estudante cubra diversas regiões suportadas pelo professor sem introduzir novas redes.

A função objetivo final combina a perda de Deriva de Características do Professor ($L_{TFD}$) e a perda de margem de âncora ($L_{anchor}$):
$$L = L_{TFD} + \lambda_{anchor} L_{anchor}$$

Contribuições Principais

1. Framework Simplificado: A TFD propõe um framework de destilação de difusão de um passo que computa o objetivo de deriva diretamente no espaço de representação oculta do professor pré-treinado, removendo a necessidade de uma rede de representação adicional.
2. Extração de Características com Ruído: O trabalho demonstra que extrair características do professor de entradas moderadamente ruidosas melhora a geometria das características para destilação, levando a melhores pontuações FID.
3. Mitigação de Colapso de Modo: Uma leve perda de margem de âncora é introduzida para incentivar a diversidade de amostras e a cobertura de regiões de características suportadas pelo professor, abordando um modo de falha comum na destilação.
4. Desempenho: O método alcança qualidade e diversidade de imagens competitivas com uma única passagem forward, simplificando substancialmente o pipeline geral de destilação em comparação com baselines de última geração.

Resultados Experimentais

Os autores avaliaram a TFD na geração ImageNet-64×64 (condicional à classe) e SDXL (texto-para-imagem).

• ImageNet-64×64:

  • A TFD alcançou um FID de 1,58 com uma única passagem forward após 200k iterações de treinamento.
  • Este resultado é competitivo com o DMD2 (FID 1,51) e supera significativamente outras baselines de um passo como DMD (2,62), EMD (2,20) e CTM (1,92).
  • Convergência: A TFD atingiu FID $\le$ 10 em 1,0k atualizações, enquanto a baseline DMD2 exigiu 5,5k atualizações (5,5x mais rápido). Ela atingiu FID $\le$ 3 em 12,5k atualizações, comparado a 24,0k do DMD2.

• Texto-para-Imagem SDXL:

  • No conjunto de dados COCO 2014, a TFD alcançou um FID de 18,39 com uma passagem forward, superando outras variantes destiladas do SDXL (por exemplo, DMD2 em 19,01, SDXL-Turbo em 24,57).
  • Ela manteve uma pontuação CLIP competitiva de 0,332, indicando forte alinhamento texto-imagem.
  • Resultados qualitativos mostraram que a TFD preserva estrutura coerente, detalhes visuais e alinhamento de prompts comparáveis a baselines destiladas de 4 passos e ao professor de 50 passos, apesar de usar apenas um passo.

Significado e Alegações

O artigo afirma que a TFD oferece uma alternativa mais simples e eficiente aos frameworks de destilação existentes. Ao utilizar o próprio modelo professor como extrator de características, a TFD evita a complexidade de treinar redes auxiliares ou projetar pipelines multiestágio. A introdução da extração de características com ruído e da regularização de margem de âncora aborda desafios específicos na geometria de representação e na cobertura de modos, respectivamente.

Os autores posicionam a TFD como um método que alcança geração eficiente de um passo com qualidade e diversidade de imagens competitivas, demonstrando que um único objetivo no espaço de representação nativo do professor é suficiente para destilação de alta fidelidade. O trabalho sugere que o conceito de "deriva" pode ser efetivamente realizado sem codificadores de características externos, desde que os estados internos do professor sejam utilizados corretamente.

Limitações: Os autores observam que o trabalho atual foca em professores baseados em U-Net e ainda não explorou arquiteturas de Transformador de Difusão (DiT). Além disso, a seleção de camadas de características e níveis de ruído atualmente carece de orientação teórica, e o método não foi testado em conjuntos de dados de maior escala ou configurações de resolução mais alta além dos experimentos atuais.