Os modelos de difusão surgiram recentemente como o padrão de fato para gerar saídas complexas e de alta dimensão. Você pode conhecê-los por sua capacidade de produzir arte de IA impressionante e imagens sintéticas hiper-realistasmas também obtiveram sucesso em outras aplicações, como design de medicamentos e controle contínuo. A ideia-chave por trás dos modelos de difusão é transformar iterativamente o ruído aleatório em uma amostra, como uma imagem ou estrutura de proteína. Isso é tipicamente motivado como um estimativa de máxima verossimilhança problema, onde o modelo é treinado para gerar amostras que correspondam o mais próximo possível aos dados de treinamento.
No entanto, a maioria dos casos de uso de modelos de difusão não está diretamente preocupada com a correspondência dos dados de treinamento, mas sim com um objetivo downstream. Não queremos apenas uma imagem que se pareça com imagens existentes, mas uma que tenha um tipo específico de aparência; não queremos apenas uma molécula de medicamento que seja fisicamente plausível, mas uma que seja o mais eficaz possível. Nesta publicação, mostramos como os modelos de difusão podem ser treinados nesses objetivos downstream diretamente usando aprendizado por reforço (RL). Para fazer isso, ajustamos Difusão estável em uma variedade de objetivos, incluindo compressibilidade de imagem, qualidade estética percebida pelo ser humano e alinhamento de imagem rápida. O último desses objetivos usa suggestions de um grande modelo de visão-linguagem para melhorar o desempenho do modelo em prompts incomuns, demonstrando como modelos poderosos de IA podem ser usados para melhorar uns aos outros sem nenhum humano envolvido.
Um diagrama ilustrando o objetivo de alinhamento de imagem de immediate. Ele usa LLaVAum grande modelo de visão-linguagem, para avaliar imagens geradas.
Otimização da política de difusão de redução de ruído
Ao transformar a difusão em um problema de RL, fazemos apenas a suposição mais básica: dada uma amostra (por exemplo, uma imagem), temos acesso a uma função de recompensa que podemos avaliar para nos dizer o quão “boa” essa amostra é. Nosso objetivo é que o modelo de difusão gere amostras que maximizem essa função de recompensa.
Os modelos de difusão são normalmente treinados usando uma função de perda derivada da estimativa de máxima verossimilhança (MLE), o que significa que eles são encorajados a gerar amostras que fazem os dados de treinamento parecerem mais prováveis. Na configuração RL, não temos mais dados de treinamento, apenas amostras do modelo de difusão e suas recompensas associadas. Uma maneira de ainda usarmos a mesma função de perda motivada por MLE é tratar as amostras como dados de treinamento e incorporar as recompensas ponderando a perda para cada amostra por sua recompensa. Isso nos dá um algoritmo que chamamos de regressão ponderada por recompensa (RWR), após algoritmos existentes da literatura RL.
No entanto, há alguns problemas com essa abordagem. Um deles é que o RWR não é um algoritmo particularmente exato — ele maximiza a recompensa apenas aproximadamente (veja Nair e outros.Apêndice A). A perda inspirada em MLE para difusão também não é exata e, em vez disso, é derivada usando um limite variacional na verdadeira probabilidade de cada amostra. Isso significa que o RWR maximiza a recompensa por meio de dois níveis de aproximação, o que achamos que prejudica significativamente seu desempenho.
Avaliamos duas variantes de DDPO e duas variantes de RWR em três funções de recompensa e descobrimos que o DDPO consistentemente atinge o melhor desempenho.
O perception principal do nosso algoritmo, que chamamos de otimização de política de difusão de denoising (DDPO), é que podemos maximizar melhor a recompensa da amostra ultimate se prestarmos atenção a toda a sequência de etapas de denoising que nos levaram até lá. Para fazer isso, reformulamos o processo de difusão como um processo de várias etapas Processo de decisão de Markov (MDP). Na terminologia MDP: cada etapa de redução de ruído é uma ação, e o agente só recebe uma recompensa na etapa ultimate de cada trajetória de redução de ruído quando a amostra ultimate é produzida. Essa estrutura nos permite aplicar muitos algoritmos poderosos da literatura RL que são projetados especificamente para MDPs multietapas. Em vez de usar a probabilidade aproximada da amostra ultimate, esses algoritmos usam a probabilidade exata de cada etapa de redução de ruído, o que é extremamente fácil de calcular.
Optamos por aplicar algoritmos de gradiente de política devido à sua facilidade de implementação e sucesso passado no ajuste fino do modelo de linguagem. Isso levou a duas variantes do DDPO: DDPOSFque usa o estimador de função de pontuação simples do gradiente de política, também conhecido como REFORÇAR; e DDPOÉque usa um estimador amostrado de importância mais poderoso. DDPOÉ é o nosso algoritmo de melhor desempenho e sua implementação segue de perto a de otimização de política proximal (PPO).
Ajuste fino da difusão estável usando DDPO
Para nossos principais resultados, nós ajustamos Difusão Estável v1-4 usando DDPOÉ. Temos quatro tarefas, cada uma definida por uma função de recompensa diferente:
- Compressibilidade: Quão fácil é comprimir a imagem usando o algoritmo JPEG? A recompensa é o tamanho negativo do arquivo da imagem (em kB) quando salva como JPEG.
- Incompressibilidade: Quão difícil é comprimir a imagem usando o algoritmo JPEG? A recompensa é o tamanho positivo do arquivo da imagem (em kB) quando salva como JPEG.
- Qualidade estética: Quão esteticamente atraente é a imagem para o olho humano? A recompensa é a saída do Preditor estético LAIONque é uma rede neural treinada nas preferências humanas.
- Alinhamento de Immediate-Imagem: Quão bem a imagem representa o que foi pedido no immediate? Este é um pouco mais complicado: nós alimentamos a imagem em LLaVApeça para ele descrever a imagem e, em seguida, calcule a similaridade entre essa descrição e o immediate unique usando Pontuação BERTS.
Como Steady Diffusion é um modelo de texto para imagem, também precisamos escolher um conjunto de prompts para dar a ele durante o ajuste fino. Para as três primeiras tarefas, usamos prompts simples do formulário “um(a) (animal)”. Para alinhamento de prompt-imagem, usamos prompts do formulário “uma(um) (animal) (atividade)”onde as atividades são “lavando louça”, “jogando xadrez”e “andar de bicicleta”. Descobrimos que a Difusão Estável muitas vezes tinha dificuldade para produzir imagens que correspondessem ao immediate para esses cenários incomuns, deixando bastante espaço para melhorias com o ajuste fino do RL.
Primeiro, ilustramos o desempenho do DDPO nas recompensas simples (compressibilidade, incompressibilidade e qualidade estética). Todas as imagens são geradas com a mesma semente aleatória. No quadrante superior esquerdo, ilustramos o que a Steady Diffusion “vanilla” gera para nove animais diferentes; todos os modelos RL-finetuned mostram uma clara diferença qualitativa. Curiosamente, o modelo de qualidade estética (canto superior direito) tende a desenhos minimalistas em preto e branco, revelando os tipos de imagens que o preditor estético LAION considera “mais estéticos”.
Em seguida, demonstramos o DDPO na tarefa mais complexa de alinhamento de prompt-imagem. Aqui, mostramos vários instantâneos do processo de treinamento: cada série de três imagens mostra amostras para o mesmo immediate e semente aleatória ao longo do tempo, com a primeira amostra vindo do Steady Diffusion vanilla. Curiosamente, o modelo muda para um estilo mais parecido com desenho animado, o que não foi intencional. Nossa hipótese é que isso ocorre porque os animais que fazem atividades semelhantes às humanas têm mais probabilidade de aparecer em um estilo parecido com desenho animado nos dados de pré-treinamento, então o modelo muda para esse estilo para se alinhar mais facilmente com o immediate, aproveitando o que ele já sabe.
Generalização inesperada
Foi descoberto que uma generalização surpreendente surge ao ajustar grandes modelos de linguagem com RL: por exemplo, modelos ajustados com base em instruções apenas em inglês muitas vezes melhoram em outras línguas. Descobrimos que o mesmo fenômeno ocorre com modelos de difusão de texto para imagem. Por exemplo, nosso modelo de qualidade estética foi ajustado usando prompts que foram selecionados de uma lista de 45 animais comuns. Descobrimos que ele se generaliza não apenas para animais não vistos, mas também para objetos cotidianos.
Nosso modelo de alinhamento prompt-imagem usou a mesma lista de 45 animais comuns durante o treinamento, e apenas três atividades. Descobrimos que ele generaliza não apenas para animais não vistos, mas também para atividades não vistas, e até mesmo novas combinações dos dois.
Superotimização
É bem sabido que o ajuste fino de uma função de recompensa, especialmente uma aprendida, pode levar a recompensar a superotimização onde o modelo explora a função de recompensa para atingir uma alta recompensa de uma forma não útil. Nosso cenário não é exceção: em todas as tarefas, o modelo eventualmente destrói qualquer conteúdo de imagem significativo para maximizar a recompensa.
Também descobrimos que o LLaVA é suscetível a ataques tipográficos: ao otimizar o alinhamento em relação aos prompts do formulário “(n) animais”o DDPO conseguiu enganar o LLaVA com sucesso, gerando um texto vagamente semelhante ao número correto.
Atualmente, não há um método de uso geral para evitar a superotimização, e destacamos esse problema como uma área importante para trabalhos futuros.
Conclusão
Modelos de difusão são difíceis de superar quando se trata de produzir saídas complexas e de alta dimensão. No entanto, até agora eles têm sido bem-sucedidos principalmente em aplicações onde o objetivo é aprender padrões de muitos e muitos dados (por exemplo, pares de imagem-legenda). O que descobrimos é uma maneira de treinar efetivamente modelos de difusão de uma forma que vai além da correspondência de padrões — e sem necessariamente exigir nenhum dado de treinamento. As possibilidades são limitadas apenas pela qualidade e criatividade da sua função de recompensa.
A maneira como usamos DDPO neste trabalho é inspirada pelos sucessos recentes de finetuning de modelos de linguagem. Os modelos GPT da OpenAI, como Steady Diffusion, são primeiro treinados em grandes quantidades de dados da Web; eles são então finetuned com RL para produzir ferramentas úteis como ChatGPT. Normalmente, sua função de recompensa é aprendida a partir de preferências humanas, mas outros tem mais recentemente descobrimos como produzir chatbots poderosos usando funções de recompensa baseadas em suggestions de IA. Comparados ao regime de chatbot, nossos experimentos são de pequena escala e limitados em escopo. Mas considerando o enorme sucesso desse paradigma de “pré-treinamento + ajuste fino” na modelagem de linguagem, certamente parece que vale a pena prosseguir no mundo dos modelos de difusão. Esperamos que outros possam desenvolver nosso trabalho para melhorar grandes modelos de difusão, não apenas para geração de texto para imagem, mas para muitas aplicações interessantes, como geração de vídeo, geração musical, edição de imagem, síntese de proteínas, robóticae muito mais.
Além disso, o paradigma “pré-treinamento + ajuste fino” não é a única maneira de usar DDPO. Contanto que você tenha uma boa função de recompensa, não há nada que o impeça de treinar com RL desde o início. Embora essa configuração ainda seja inexplorada, este é um lugar onde os pontos fortes do DDPO podem realmente brilhar. O RL puro tem sido aplicado há muito tempo a uma ampla variedade de domínios que vão desde jogando jogos para manipulação robótica para fusão nuclear para design de chip. Adicionar a expressividade poderosa dos modelos de difusão à mistura tem o potencial de levar as aplicações existentes de RL para o próximo nível — ou até mesmo descobrir novas.
Este put up é baseado no seguinte artigo:
Se você quiser saber mais sobre DDPO, você pode conferir o papel, website, código uniqueou pegue o pesos modelo em Hugging Face. Se você quiser usar DDPO em seu próprio projeto, confira meu Implementação PyTorch + LoRA onde você pode ajustar a difusão estável com menos de 10 GB de memória da GPU!
Se o DDPO inspira seu trabalho, cite-o com:
@misc{black2023ddpo,
title={Coaching Diffusion Fashions with Reinforcement Studying},
writer={Kevin Black and Michael Janner and Yilun Du and Ilya Kostrikov and Sergey Levine},
12 months={2023},
eprint={2305.13301},
archivePrefix={arXiv},
primaryClass={cs.LG}
}