Resumo Técnico: Sobre o Design de Difusão de Passo Único via Atalhos de Caminhos de Fluxo

Declaração do Problema

Modelos generativos baseados em difusão tornaram-se o paradigma dominante para modelagem generativa profunda, mas sofrem com velocidades de inferência lentas, tipicamente exigindo dezenas ou centenas de avaliações de função neural (NFEs). Embora modelos de consistência e técnicas de destilação tenham permitido a geração de passo único, eles frequentemente dependem de processos de treinamento de dois estágios custosos (pré-treinamento de um modelo de difusão seguido de destilação). Recentemente, "modelos de atalho" treinados do zero (ex: Treinamento de Consistência, Difusão de Atalho, Correspondência de Momento Indutivo) surgiram para aprender mapeamentos diretos entre estados intermediários em um único passo. No entanto, o espaço de design para esses modelos permanece obscuro. Derivações teóricas, implementações práticas e escolhas de componentes específicos (como amostradores de tempo, caminhos de fluxo e métricas de perda) são frequentemente fortemente acoplados, tornando difícil identificar melhorias sistematicamente ou compreender a necessidade de módulos específicos.

Metodologia

Os autores propõem um framework de design comum para desmembrar os componentes de modelos de atalho representativos (variantes de tempo discreto e de tempo contínuo). Este framework unifica esses métodos sob o princípio de aproximar um alvo de mapa de fluxo de dois passos com uma predição parametrizada de um passo.

1. Framework de Design Unificado

O artigo formaliza modelos de atalho usando resolvedores de mapa de fluxo. O objetivo central é treinar uma rede neural $F_\theta$ para prever um mapa de fluxo de um passo $X^\theta_{t,r}(x_t)$ que aproxime um alvo de dois passos construído via propriedades de consistência:
$$X^\theta_{t,r}(x_t) \approx \text{sg}(\hat{X}_{s,r} \circ \hat{X}_{t,s}(x_t))$$
onde $r \le s \le t$, $\text{sg}$ denota o operador de parada de gradiente (stop-gradient), e o alvo é construído usando velocidades condicionais ou predições de rede intermediárias. O framework categoriza os métodos existentes (CT, SCD, IMM, sCT, MeanFlow) por suas escolhas específicas em:

• Caminhos de Fluxo: Caminhos Lineares vs. Cosseno.
• Amostradores de Tempo: Discretos (intervalos fixos) vs. Contínuos (intervalos infinitesimais).
• Parametrização da Rede: Velocidade instantânea ($v_\theta$) vs. Velocidade média ($u_\theta$).
• Métricas de Perda: LPIPS, $L_2$ ao quadrado ou métricas baseadas em kernel.

2. Justificativa Teórica

Os autores fornecem limites teóricos sobre a distância Wasserstein-2 entre a distribuição de dados e a distribuição gerada. Eles provam que os objetivos de treinamento de modelos de atalho de tempo discreto (DTSC) e de tempo contínuo (CTSC) servem como limites superiores válidos para o erro de geração. Crucialmente, eles demonstram que caminhos lineares são ótimos para modelos de atalho treinados do zero sob configurações condicionadas a dados, oferecendo custos de transporte convexos mais baixos e redução da curvatura da trajetória em comparação com caminhos de cosseno.

3. Melhorias Propostas

Com base no framework desmembrado, o artigo introduz três refinamentos técnicos principais para aprimorar o treinamento de modelos de atalho de tempo contínuo (especificamente variantes MeanFlow):

• Velocidade de Plug-in (Plug-in Velocity): Em vez de depender exclusivamente da velocidade condicional $v_{t|0}$ (que possui alta variância devido à amostragem de um mini-batch), os autores propõem o uso de uma velocidade de plug-in $v^*_t$. Esta é computada como uma média ponderada das velocidades condicionais dentro do mini-batch, aproximando efetivamente a velocidade marginal. Isso reduz a variância da supervisão da velocidade por $O(1/B)$ (onde $B$ é o tamanho do batch) com um custo computacional negligenciável, estabilizando o treinamento.
• Mini-batching Consistente por Classe: Para manter a eficácia do Classifier-Free Guidance (CFG) ao usar a velocidade de plug-in, os autores propõem amostrar mini-batches de tal forma que todos os exemplos dentro de um batch pertençam à mesma classe. Isso evita que os sinais específicos de classe sejam diluídos durante o processo de média.
• Amostrador de Tempo Gradual: Uma estratégia de aprendizado por currículo é introduzida onde o tempo terminal $r$ é fixado em 0 durante os estágios iniciais do treinamento (imitando uma tarefa de denoising) e transiciona gradualmente para a estratégia padrão de amostragem de tempo contínuo. Isso auxilia em uma convergência inicial mais rápida.

Resultados Principais

As melhorias propostas foram avaliadas no ImageNet-256$\times$256 e CIFAR-10 usando arquiteturas SiT e U-Net.

• Desempenho de Estado da Arte: O modelo resultante, denominado ESC (Explicit & Easier Shortcut), alcança um FID50k de 2.85 no ImageNet-256$\times$256 com geração de passo único sob classifier-free guidance. Com treinamento estendido (480 épocas), atinge 2.53.
• Comparação: Estes resultados estabelecem um novo estado da arte entre modelos de atalho treinados do zero, superando significativamente métodos anteriores de passo único como MeanFlow (3.43) e SCD (10.60). O desempenho supera inclusive a fidelidade de geração de dois passos do MeanFlow (2.93).
• Benefícios de Escala: Os ganhos de desempenho das técnicas propostas são mais pronunciados em modelos maiores (SiT-XL/2) em comparação com modelos menores (SiT-B/2), sugerindo que a redução de variância é crítica para escalar modelos de atalho.
• Eficiência: O cálculo da velocidade de plug-in adiciona um overhead negligenciável (aumento de aproximadamente 0,7% no tempo de iteração) enquanto fornece estabilidade significativa.

Significância e Alegações

O artigo alega que sua principal contribuição é baixar a barreira para a inovação em nível de componente em modelos de atalho. Ao fornecer um framework de design comum e justificativa teórica, os autores:

1. Desmembraram o espaço de design: Eles mostram que muitos módulos anteriormente considerados "indispensáveis" são, na verdade, escolhas independentes, permitindo exploração e combinação sistemáticas.
2. Validaram caminhos lineares: Eles justificam teoricamente o uso de caminhos lineares para o treinamento de modelos de atalho a partir do zero, desafiando a preferência prévia por caminhos de cosseno em certos contextos.
3. Possibilitaram a exploração principiada: As melhorias propostas (velocidade de plug-in, batching consistente por classe) oferecem um caminho para geração de passo único de alta fidelidade sem a necessidade de pré-treinamento, destilação ou aprendizado por currículo complexo.

Os autores enfatizam que seu trabalho facilita uma abordagem mais sistemática para o design de modelos de difusão de passo único, indo além da "caixa preta" de derivações teóricas acopladas e truques práticos. Eles reconhecem limitações, como a atual inflexibilidade no ajuste de hiperparâmetros de CFG e o custo computacional de produtos vetor-Jacobiano (JVP) em arquiteturas baseadas em transformer, identificando estas como direções para trabalhos futuros.