HybridStitch: Pixel and Timestep Level Model Stitching for Diffusion Acceleration

O artigo apresenta o HybridStitch, um novo paradigma de geração de imagens que acelera os modelos de difusão ao dividir a imagem em regiões simples e complexas, utilizando um modelo menor para esboçar as áreas fáceis e um modelo maior para refinar as áreas complexas, alcançando um aumento de velocidade de 1,83× no Stable Diffusion 3.

O essencial

Imagine que você pediu a um artista genial (o Modelo Grande) para desenhar uma paisagem complexa. Ele é incrível, mas muito lento e cansativo. Agora, imagine que você também tem um assistente rápido e ágil (o Modelo Pequeno), mas que às vezes comete erros em detalhes finos.

O problema das tecnologias atuais é que elas tratam a imagem inteira como um bloco único: ou o artista genial faz tudo, ou ele faz o começo e passa a bola para o assistente fazer o resto. Isso é ineficiente. Por que o artista genial precisa gastar tempo pintando o céu azul simples (que o assistente faria perfeitamente) se ele poderia focar apenas nas montanhas difíceis e nos detalhes da cara do personagem?

É aqui que entra o HybridStitch, a solução proposta neste artigo.

A Analogia da "Costura Híbrida"

Pense no processo de gerar uma imagem como se fosse costurar um vestido de noiva em vez de apenas desenhar.

1. O Esboço Rápido (O Assistente):
   No início, o "assistente" (Modelo Pequeno) faz um esboço rápido de todo o vestido. Ele define a forma geral, a cor do tecido e a silhueta. Isso é rápido e cobre tudo.

2. A Costura Inteligente (O Modelo Grande):
   Aqui está a mágica do HybridStitch. Em vez de o artista genial refazer todo o vestido, ele olha para o esboço e diz: "O assistente acertou o tecido do fundo, mas errou os detalhes do colarinho e das rendas."

   O HybridStitch cria uma máscara invisível. Ele "corta" apenas as partes difíceis (o colarinho, as rendas) e pede ao artista genial para refinar apenas essas áreas. O resto do vestido continua sendo trabalhado pelo assistente rápido.

3. A Costura Final (A Fusão):
   O HybridStitch pega o trabalho do assistente (o fundo) e o trabalho do artista genial (os detalhes) e os "costura" juntos perfeitamente. O resultado é uma imagem que parece ter sido feita 100% pelo artista genial, mas que foi produzida muito mais rápido.

Por que isso é diferente do que já existe?

• O Jeito Antigo (Naive Stitch): Era como se o artista genial fizesse os primeiros 10 minutos do desenho inteiro, e depois parasse e dissesse: "Ok, agora você (assistente) faz o resto do desenho inteiro". Se o céu estava pronto, o artista genial ainda gastou tempo nele. Se a montanha estava difícil, o assistente estragou.
• O HybridStitch: É como se o artista genial e o assistente trabalhassem ao mesmo tempo, mas em áreas diferentes da mesma tela. O assistente cuida do que é fácil, e o genial cuida do que é difícil, pixel por pixel.

Como eles sabem o que é "difícil"?

O sistema é como um inspetor de qualidade. A cada momento do desenho, ele compara o que o artista genial faria com o que o assistente fez.

• Se a diferença for pequena (ex: o céu azul), o sistema diz: "Pode deixar o assistente cuidar disso".
• Se a diferença for grande (ex: os olhos do personagem), o sistema diz: "Pare! O artista genial precisa intervir aqui".

Eles usam uma técnica chamada KV Cache (uma espécie de "memória de contexto"). Imagine que o artista genial precisa ver o resto do vestido para pintar o colarinho com coerência. Como ele só está pintando uma parte, o sistema "empresta" a memória do passo anterior para que ele não se perca e pinte algo que não combine com o resto.

Os Resultados (O "Pulo do Gato")

Os pesquisadores testaram isso no Stable Diffusion 3 (uma das melhores IAs de imagem do mundo) e descobriram:

• Velocidade: O HybridStitch foi 1,83 vezes mais rápido do que usar apenas o modelo grande.
• Qualidade: A imagem final ficou tão boa quanto a feita pelo modelo gigante, sem perder detalhes importantes.
• Comparação: É mais rápido e melhor do que as técnicas anteriores que tentavam misturar modelos.

Resumo em uma frase

O HybridStitch é como ter uma equipe de construção onde o mestre pedreiro (Modelo Grande) só constrói as paredes difíceis e os detalhes artísticos, enquanto o ajudante (Modelo Pequeno) faz o reboco e a pintura das partes fáceis, tudo acontecendo ao mesmo tempo para terminar a casa mais rápido, sem perder a qualidade.

Isso permite que IAs gerem imagens incríveis em celulares ou computadores mais simples, sem precisar de supercomputadores caros para cada clique.

Resumo técnico

Resumo Técnico: HybridStitch

1. O Problema

Os modelos de difusão para geração de imagem a partir de texto (Text-to-Image ou T2I) têm demonstrado capacidades notáveis, mas sofrem de uma sobrecarga computacional pesada, especialmente em modelos grandes com dezenas de bilhões de parâmetros (como o Stable Diffusion 3 e 3.5). Isso resulta em alta latência, criando barreiras para aplicações sensíveis ao tempo.

Trabalhos anteriores tentaram acelerar a inferência utilizando uma abordagem de "Mistura de Modelos" (Mixture of Models), onde se troca de um modelo grande para um menor durante os passos de desruído (denoising). No entanto, essas abordagens existentes (como Naive Stitch, T-Stitch e SRDiffusion) operam em granularidade de imagem inteira. Elas assumem que todos os pixels da imagem têm a mesma demanda computacional e mudam o modelo simultaneamente para toda a imagem.

• Limitação: Isso ignora a heterogeneidade dentro de um único passo de tempo. Algumas regiões (ex.: fundo) são fáceis de renderizar e poderiam usar um modelo menor mais cedo, enquanto regiões complexas (ex.: objetos detalhados) precisam do modelo grande por mais tempo. Mudar o modelo para a imagem inteira cedo demais degrada a qualidade; mudar tarde demais aumenta a latência desnecessariamente.

2. Metodologia: HybridStitch

O HybridStitch propõe um novo paradigma de geração que trata a criação de imagem como um processo de edição, realizando a troca de modelos em nível de píxel e passo de tempo.

Principais Componentes da Arquitetura:

• Estágios Híbridos: O método divide o processo em três estágios dinâmicos:

  1. Primeiro Estágio: Apenas o modelo grande processa o ruído gaussiano inicial para estabelecer o layout global.
  2. Segundo Estágio (Estitching Híbrido): Ambos os modelos (grande e pequeno) operam simultaneamente.
     • O modelo pequeno processa o estado latente completo da imagem para construir um esboço global e manter a consistência.
     • O modelo grande processa apenas uma máscara de pixels (as regiões mais difíceis/complexas).
     • Combinação: A saída do modelo grande substitui a saída do modelo pequeno nas regiões mascaradas.
  3. Terceiro Estágio: Apenas o modelo pequeno processa a imagem restante até o final, uma vez que as diferenças entre os modelos tornam-se insignificantes.

• Seleção de Região (Máscara Dinâmica):

  • A máscara não é estática. Ela é atualizada a cada passo de desruído.
  • A região a ser processada pelo modelo grande é determinada calculando a diferença absoluta entre os estados latentes de passos consecutivos. Os pixels com as maiores diferenças (indicando instabilidade ou mudanças significativas) são selecionados (Top-K) para refinamento pelo modelo grande.

• Geração com Máscara e KV Cache:

  • Para evitar que o modelo grande perca o contexto global ao processar apenas uma parte da imagem (o que causaria baixa qualidade), o HybridStitch utiliza o KV Cache (Chave-Valor) do passo anterior.
  • O modelo grande recebe a parte mascarada como entrada, mas concatena as chaves e valores do cache (do passo anterior) para preencher o contexto completo durante o cálculo de atenção. Isso garante coerência sem a necessidade de processar todos os tokens do modelo grande a cada passo.

• Estratégia de Troca (Switch Strategy):

  • A transição entre os estágios é controlada por uma métrica de diferença ($D_t$) baseada na distância L1 entre os estados latentes adjacentes. Quando a diferença cai abaixo de um limiar específico, o sistema avança para o próximo estágio (reduzindo a área de processamento do modelo grande).

3. Contribuições Chave

1. Granularidade Pixel a Pixel: É a primeira abordagem a realizar a troca de modelos em nível de pixel dentro de um único passo de tempo, reconhecendo que diferentes regiões da imagem têm diferentes necessidades de refinamento.
2. Técnica sem Treinamento (Train-free): O método não requer re-treinamento dos modelos; ele é uma técnica de aceleração de inferência que combina modelos existentes.
3. Mecanismo de Contexto via KV Cache: Resolve o problema de perda de contexto em geração parcial (masked generation) reutilizando o cache de atenção de passos anteriores, permitindo que o modelo grande refine apenas regiões específicas sem sacrificar a qualidade global.
4. Paradigma de "Edição": Reenquadra a geração de difusão como um processo iterativo de refinamento de áreas complexas, em vez de um processo linear uniforme.

4. Resultados Experimentais

Os testes foram realizados no conjunto de dados COCO utilizando o Stable Diffusion 3 (Modelo Grande: SD 3.5 Large; Modelo Pequeno: SD 3 Medium) com 50 passos de desruído.

• Aceleração:

  • O HybridStitch alcançou um speedup de 1.83× em relação ao modelo grande puro.
  • Superou os métodos existentes: T-Stitch (1.41×) e SRDiffusion (1.55×).
  • Representa uma redução de latência de 18.06% em relação ao SRDiffusion (o estado da arte anterior).

• Qualidade de Imagem:

  • FID (Frechet Inception Distance): O HybridStitch obteve o melhor FID entre as técnicas aceleradas (30.43 para a configuração de 30% de máscara), superando T-Stitch (31.87) e SRDiffusion (31.67).
  • CLIP Score e LPIPS: Mantiveram-se consistentes e superiores aos baselines, indicando boa semelhança semântica e perceptual com a imagem original.
  • A configuração de 30% de máscara mostrou o melhor equilíbrio entre eficiência e qualidade.

• Análise de Sensibilidade:

  • O método demonstrou robustez em diferentes GPUs (RTX 6000 Ada, H100, A100), mantendo ganhos significativos, especialmente em plataformas de desempenho moderado.
  • Estudos de ablação confirmaram que a atualização dinâmica da máscara é crucial; máscaras estáticas degradam a qualidade (aumento do FID).

5. Significado e Impacto

O HybridStitch representa um avanço significativo na eficiência de modelos de difusão. Ao demonstrar que a computação pode ser alocada de forma heterogênea dentro de uma única imagem e passo de tempo, ele quebra o paradigma de "tudo ou nada" das abordagens anteriores.

• Eficiência: Permite que modelos de última geração (com bilhões de parâmetros) sejam executados com latência muito menor, tornando-os viáveis para aplicações em tempo real ou em dispositivos com recursos limitados.
• Qualidade: Mantém a alta fidelidade visual exigida por aplicações comerciais, provando que a aceleração não precisa vir às custas da qualidade da imagem.
• Generalidade: A abordagem de usar KV cache para manter o contexto em gerações parciais pode ser aplicada a outros métodos de aceleração de difusão e até mesmo a geração de vídeo.

Em resumo, o HybridStitch oferece uma solução prática e eficaz para o gargalo de computação em modelos de difusão modernos, equilibrando a necessidade de alta qualidade com a demanda por baixa latência.