A arquitetura por trás dos sistemas capazes de gerar imagens a partir de prompts de texto é chamada de modelo de difusão. No treinamento desses modelos, as imagens utilizadas têm ruído adicionado, e junto com essa distorção, o modelo recebe também um embedding do texto que descreve a imagem, e também um embedding temporal. Com base nessas informações, o sistema aprende a remover o ruído e produzir a imagem original, guiada pelo prompt, em passos sucessivos. Assim, em modo de inferência, o modelo pode produzir imagens a partir de ruído baseado em um prompt de texto.
Assim como nas aplicações tradicionais de visão computacional, os modelos de difusão populares atualmente utilizam uma rede neural convolucional chamada U-Net para processar as imagens. Para explorar uma nova possibilidade, pesquisadores americanos desenvolveram um sistema que eles batizaram de Diffusion Transformers (DiTs). Como o nome deixa evidente, os DiTs substituem o módulo U-Net por um transformer. Esta arquitetura, que aplica mecanismos de atenção, foi a responsável pela recente melhoria impressionante na inteligência artificial voltada para o processamento de linguagem natural.
Para realizar a nova tarefa, o sistema precisou receber algumas adaptações. Os embeddings da imagem com ruído, por exemplo, foram transformados em tokens. Blocos modificados do transformer aprendem a processar os tokens para produzir estimativas do ruído a ser removido. Além disso, antes de cada camada de atenção e camada densa, os tokens são multiplicados por um vetor, produzido por uma rede densa independente, que representa a descrição da imagem e dos embeddings temporais.
Os resultados foram avaliados através de uma métrica chamada FID (Fréchet Interception Distance), que mede a distância entre a distribuição da imagem gerada com aquela da imagem original. Em comparação com um modelo de difusão que utiliza a U-Net, a versão mais robusta do DiT teve uma melhoria de cerca de 9%. Alguns exemplos e comparações podem ser visualizados na página do projeto.
O artigo que descreve o trabalho está disponível para consulta e já existe até uma implementação oficial baseada em PyTorch para quem quiser testar o novo sistema.
