Síntese de imagem em alta resolução

01 · CONTEXTO

Por que alta resolução é um problema de custo

Gerar imagens de alta resolução é tema do livro Generative AI System Design Interview: o capítulo 8 tem exatamente esse título. Este caso o trata como um problema de projeto de sistemas de machine learning.

Duas pesquisas abertas mostram onde o custo aparece. Transformers são expressivos, mas computacionalmente inviáveis para sequências longas, como imagens de alta resolução. Já modelos de difusão que operam direto no espaço de pixels costumam consumir centenas de dias de GPU na otimização, e a inferência é cara por exigir avaliações sequenciais.

02 · COMPRESSÃO LATENTE

Três maneiras de gerar sem tocar nos pixels

O VQ-VAE é a base da compressão discreta: o codificador emite códigos discretos, e não contínuos, e o prior é aprendido em vez de fixo. Com um prior autoregressivo sobre essas representações, o modelo gera imagens de alta qualidade.

O artigo Taming Transformers combina CNNs, que aprendem um vocabulário de constituintes da imagem, com transformers, que modelam a composição desses constituintes em imagens de alta resolução.

Os modelos de difusão latente aplicam a difusão no espaço latente de autoencoders pré-treinados, o que reduz de forma significativa os requisitos computacionais em comparação com a difusão em pixels. Esse é o argumento de custo em escala para trabalhar no espaço latente.

03 · SERVIÇO

Do pedido à imagem final

O diagrama abaixo mostra um serviço assíncrono que gera no espaço latente, decodifica e refina a imagem em cascata.

1. No diagrama, o cliente chama o gateway de API, que coloca o pedido na fila de jobs.

2. O worker de geração opera no espaço latente de um autoencoder pré-treinado, em vez de operar direto no espaço de pixels.

3. No diagrama, o decodificador do autoencoder recebe a saída do worker e a passa ao refinador.

4. Modelos generativos podem ser encadeados com modelos de super-resolução; é essa cascata que o refinador representa neste diagrama.

5. Um refinador baseado em difusão, como o SR3, parte de ruído gaussiano puro e refina a saída de forma iterativa, com uma U-Net treinada para remover ruído em vários níveis.

6. No diagrama, a imagem final vai para o armazenamento e é entregue ao cliente pela CDN.

04 · TREINO

Treinar em dois estágios

O diagrama abaixo separa o treino do autoencoder, que define o espaço latente, do treino do gerador, que opera nesse espaço.

1. No diagrama, o dataset de imagens alimenta o treino do autoencoder.

2. Em uma versão discreta, como no VQ-VAE, o codificador emite códigos discretos, e não contínuos, e o prior é aprendido.

3. No diagrama, as representações latentes ficam guardadas em armazenamento e servem de entrada para o treino do gerador.

4. No artigo Taming Transformers o gerador é um transformer que modela a composição dos constituintes da imagem; na difusão latente, o gerador é um modelo de difusão aplicado no espaço latente.

5. Operar no espaço latente reduz de forma significativa os requisitos computacionais, em comparação com a difusão em pixels.

6. No diagrama, os checkpoints alimentam a avaliação de amostras, cujas medidas seguem para as métricas de treino.

05 · DECISÕES

Pixels ou latente, e como condicionar

Para controlar o resultado, camadas de cross-attention tornam a difusão um gerador flexível para condicionamentos gerais, como texto ou caixas delimitadoras.