Resumo Técnico: Expansão de Domínio de Difusão (DDE)

Declaração do Problema

Modelos de difusão pré-treinados alcançaram resultados state-of-the-art na geração de imagens, áudio e outras modalidades de dados. No entanto, re-treinar esses modelos do zero para lidar com domínios de saída maiores (por exemplo, imagens panorâmicas, faixas de áudio longas) ou entradas de condicionamento mais complexas (por exemplo, múltiplas localizações de objetos) é computacionalmente caro. Abordagens existentes para expandir domínios generativos frequentemente dependem de mecanismos de coordenação algorítmicos fixos (por exemplo, MultiDiffusion) que não exigem treinamento adicional, mas são limitados a tarefas específicas e carecem de flexibilidade. Outros métodos, como o treinamento de módulos difusores dinâmicos, frequentemente exigem modalidades fixas no momento do treinamento ou falham em generalizar para domínios maiores do que aqueles observados durante o treinamento.

O desafio central abordado por este artigo é como estender eficientemente o domínio generativo de modelos de difusão pré-treinados para gerar objetos maiores e lidar com condicionamentos mais complexos sem re-treinar os modelos base, garantindo ao mesmo tempo que a solução generalize para instâncias maiores do que aquelas vistas durante o treinamento do coordenador.

Metodologia: Expansão de Domínio de Difusão (DDE)

Os autores propõem a Expansão de Domínio de Difusão (DDE), um framework que emprega uma rede "coordenadora" compacta e treinável para orquestrar as saídas de modelos de difusão pré-treinados.

Formulação Geral

Dado um modelo de difusão condicional pré-treinado $p(x|y)$ com uma rede de remoção de ruído $D(x, y, t)$, o objetivo é gerar um objeto expandido $X^{[L]}$ condicionado a uma entrada expandida $Y^{[L]}$, onde $L$ representa a escala (por exemplo, número de patches ou rótulos de condicionamento).

1. Decomposição: O par expandido $(X^{[L]}, Y^{[L]})$ é decomposto em $L$ partes menores $([x_1, \dots, x_L], [y_1, \dots, y_L])$ usando uma função de mapeamento $F$. Essas partes caem dentro do domínio original do modelo base.
2. Geração Base: Cada parte $x_i$ é processada pelo removedor de ruído pré-treinado $D(x_i, y_i, t)$.
3. Coordenação: Uma rede coordenadora treinável $C[\cdot]$ recebe como entrada a sequência de saídas desruidadas $[D(x_i, y_i, t)]_{i=1}^L$ e as informações de condicionamento $[y_i]_{i=1}^L$. Ela aprende a coordenar essas saídas para produzir uma estimativa coerente do objeto expandido completo $X^{[L]}(0)$.
4. Objetivo de Treinamento: O coordenador é treinado minimizando o erro de remoção de ruído do desruidor composto $D^{[L]}$ em um conjunto de dados de tamanho $L_{train}$. Crucialmente, o método é projetado para generalizar para tamanhos de teste $L_{test} \ge L_{train}$.

Arquitetura do Coordenador

Os autores utilizam uma arquitetura Vision Transformer (ViT) para o coordenador, adaptada de Peebles & Xie (2023).

• Processamento de Entrada: Os patches desruidados dos modelos pré-treinados são "patchificados" e aumentados com codificações posicionais globais relativas ao objeto maior. Os autores empregam Embeddings Posicionais Rotacionais (RoPE) para lidar efetivamente com dimensões espaciais maiores.
• Reconciliação: Durante o processo forward, o transformer processa todos os patches. Para lidar com patches sobrepostos (usados para garantir continuidade), a saída final é construída calculando a média dos valores nas regiões sobrepostas, inspirado pelo método de amostragem MultiDiffusion.
• Eficiência: O coordenador é significativamente menor que o modelo base (por exemplo, 66M vs. 405M parâmetros na geração de música) e requer menos épocas de treinamento.

Contribuições Principais

1. Framework Inovador: Introdução da DDE, um método que usa um coordenador pequeno e treinável para estender as capacidades generativas de modelos de difusão pré-treinados para objetos maiores e condicionamentos mais complexos.
2. Capacidade de Generalização: Demonstração de que o coordenador pode generalizar para gerar objetos e lidar com entradas de condicionamento maiores em escala do que aquelas presentes em seus dados de treinamento ($L_{test} \ge L_{train}$).
3. Design de Arquitetura: Validação de uma arquitetura baseada em ViT como um coordenador adequado, que aprende dependências de longo alcance enquanto impõe consistência local através da média de sobreposição.
4. Eficiência: A abordagem alcança geração de alta qualidade com significativamente menos parâmetros e tempo de treinamento em comparação com re-treinar modelos base ou usar arquiteturas de coordenador maiores.

Resultados Experimentais

Os autores avaliaram a DDE em três domínios:

1. Geração de Faixas de Áudio Longas (Slakh2100)

• Tarefa: Gerar faixas de música de comprimento $10l$ (120 segundos) usando um modelo base treinado no comprimento $l$ (12 segundos).
• Baselines: Concatenação e MultiDiffusion.
• Resultados: A DDE alcançou uma Distância de Áudio Fréchet (FAD) de 2.112 para faixas de 4l e 2.142 para faixas de 10l. Isso superou significativamente o MultiDiffusion (4.796) e a Concatenação (4.596) para a tarefa de 10l. O coordenador (66M parâmetros) foi muito menor que o modelo base (405M).

2. Geração de Imagens Condicionais (CLEVR)

• Tarefa: Gerar imagens condicionadas a múltiplas localizações de objetos ($L$ rótulos), onde o modelo base lida apenas com uma localização.
• Baselines: RRR (Reduce, Reuse, Recycle) e MultiDiffusion.
• Resultados: A DDE demonstrou precisão superior em satisfazer múltiplas restrições de condicionamento. Para 5 condicionamentos, a DDE alcançou 44.5% de precisão (amostrador Heun), comparado a 33.6% para RRR e 36.3% para MultiDiffusion. O coordenador generalizou do treinamento em $L_{train}=2$ para o teste em $L_{test}=5$.

3. Geração de Imagens de Satélite

• Tarefa: Gerar grandes imagens de satélite condicionadas a mapas esquemáticos.
• Baselines: MultiDiffusion.
• Resultados: A DDE alcançou uma pontuação FID mais baixa que o MultiDiffusion tanto para tamanhos in-distribution ($N=96$, FID 31.75 vs 37.81) quanto out-of-distribution ($N=128$, FID 27.37 vs 35.01). A inspeção visual mostrou que a DDE produziu imagens mais coerentes, sem as costuras visíveis observadas nas amostras do MultiDiffusion.

Significado e Alegações

O artigo alega que a DDE oferece uma solução prática e eficiente para escalar modelos de difusão sem o custo proibitivo de re-treinar. Ao aprender a coordenar processos pré-treinados, o método permite:

• Expansão de Domínio: Geração de objetos maiores (áudio mais longo, imagens maiores) e tratamento de condições mais complexas (múltiplos rótulos) além das capacidades do modelo original.
• Generalização: A capacidade de extrapolar para escalas maiores do que aquelas vistas durante o treinamento do coordenador.
• Qualidade e Consistência: O coordenador baseado em ViT aprende efetivamente dependências de longo alcance, resolvendo conflitos entre etapas de desruido locais para produzir saídas globalmente consistentes.

Os autores observam que a limitação primária é a exigência de dados adicionais e tempo para treinar o modelo coordenador, embora esse custo seja substancialmente menor do que treinar um modelo de difusão completo do zero. Trabalhos futuros são propostos para aplicar a DDE à geração de sequências de proteínas e criação de vídeos a partir de clipes curtos.