Exploring Diffusion Models' Corruption Stage in Few-Shot Fine-tuning and Mitigating with Bayesian Neural Networks

Este artigo investiga a fase de corrupção observada no ajuste fino de poucos exemplos de Modelos de Difusão, identifica sua causa na distribuição de aprendizado restrita e propõe o uso de Redes Neurais Bayesianas para mitigar esse problema, melhorando a fidelidade, qualidade e diversidade das imagens geradas sem custos adicionais de inferência.

O essencial

Imagine que você tem um chef de cozinha extremamente talentoso (o Modelo de Difusão) que aprendeu a cozinhar milhões de pratos diferentes com milhões de ingredientes. Ele sabe fazer de tudo: desde um bolo de chocolate até um sushi perfeito.

Agora, você quer que esse chef aprenda a fazer apenas um prato específico: a sua receita secreta de bolo de cenoura. Você só tem uma foto desse bolo para mostrar a ele (isso é o "Few-Shot", ou "poucas amostras").

O Problema: A Fase da "Corrupção" (O Pesadelo do Chef)

O artigo descobre algo estranho que acontece quando tentamos ensinar esse chef com apenas uma foto:

1. O Começo Promissor: No início, o chef olha a foto e começa a entender. Ele tenta copiar o bolo. Parece ótimo!
2. O Pesadelo (A Fase de Corrupção): De repente, algo dá errado. O chef começa a alucinar. Em vez de um bolo, ele começa a desenhar manchas de ruído, formas estranhas e padrões confusos na imagem. É como se ele tivesse perdido o foco e começado a misturar farinha com tinta. A imagem fica pior, cheia de "estática".
3. O Fim Trágico (Sobreajuste): Depois dessa fase de caos, o chef "desiste" de tentar entender a receita de verdade. Ele entra em um estado de memorização cega. Ele para de criar e apenas copia a foto original perfeitamente, pixel por pixel. Se você pedir um "bolo de cenoura em um prato azul", ele só te dá a foto original do bolo no prato original. Ele perdeu a criatividade e a capacidade de adaptar o prato a novas situações.

Os autores chamam essa fase do meio (onde aparecem as manchas e ruídos) de "Fase de Corrupção".

A Causa: O Chef com Visão de Túnel

Por que isso acontece? O artigo explica que, ao aprender com apenas uma foto, o chef fica com uma visão de túnel.

• Imagine que a mente do chef é um mapa gigante de todas as receitas possíveis.
• Quando ele aprende com milhões de fotos, o mapa é vasto e rico.
• Quando ele tenta aprender com uma única foto, ele tenta encolher todo esse mapa gigante para caber apenas naquele pequeno ponto.
• Como o mapa é muito pequeno e rígido, quando o chef tenta criar algo novo (mesmo que seja apenas um pequeno detalhe diferente), ele "quebra". Ele não tem espaço para variar, então a imagem fica distorcida (corrompida) ou ele simplesmente copia a foto original sem pensar.

A Solução: O "Chef com Mente Aberta" (Redes Neurais Bayesianas)

Para consertar isso, os autores propõem usar uma técnica chamada Redes Neurais Bayesianas (BNNs).

Pense nas Redes Neurais comuns como um chef que é obcecado por precisão absoluta. Ele quer acertar a receita exatamente como na foto, sem margem de erro.

As Redes Bayesianas, por outro lado, ensinam o chef a ser um pouco mais flexível e incerto. Em vez de dizer "Isso é 100% farinha", o chef Bayesian diz: "Isso é provavelmente farinha, mas talvez um pouco de amido, e tenho uma pequena chance de estar errado".

A Analogia da "Augmentação de Dados Invisível":

• Sem BNN: O chef tenta copiar a foto exata. Se a foto tiver um ponto de luz estranho, ele copia o ponto estranho.
• Com BNN: O chef é treinado para imaginar que a receita pode variar um pouquinho. Ele "alarga" o mapa mental dele. Em vez de focar apenas em uma versão exata do bolo, ele aprende a entender a essência do bolo, permitindo variações.

Isso funciona como se o chef estivesse praticando a receita em diferentes condições (luz diferente, ângulo diferente) sem você precisar tirar mais fotos. Isso preenche os "buracos" no mapa mental dele, evitando que ele entre na fase de corrupção (as manchas) e evitando que ele caia na memorização cega.

O Resultado: O Melhor dos Dois Mundos

Com essa técnica, o chef consegue:

1. Não alucinar: As manchas estranhas e ruídos desaparecem.
2. Não copiar cegamente: Ele consegue criar novos bolos baseados na sua receita, mas em cenários diferentes (em um prato azul, com velas, etc.), mantendo a qualidade e a fidelidade ao original.

Resumo Simples

O artigo diz: "Quando ensinamos uma IA a aprender com poucas fotos, ela passa por um momento de confusão onde a imagem fica cheia de ruído antes de virar uma cópia perfeita e sem vida. Para resolver isso, usamos uma técnica matemática (Bayesiana) que ensina a IA a ser um pouco mais 'flexível' e menos obcecada pela cópia exata. Isso faz com que ela aprenda melhor, crie imagens mais bonitas e não fique 'travada' em erros estranhos."

É como dar ao chef uma mente mais aberta para que ele entenda o conceito da receita, em vez de apenas decorar a foto.

Resumo técnico

Resumo Técnico: Explorando a Fase de Corrupção em Ajuste Fino de Poucos Exemplos de Modelos de Difusão e Mitigação com Redes Neurais Bayesianas

1. O Problema: A "Fase de Corrupção" no Ajuste Fino

Os autores identificam um fenômeno inesperado e prejudicial durante o processo de ajuste fino (fine-tuning) de poucos exemplos (few-shot) em Modelos de Difusão (DMs), como o Stable Diffusion. Ao tentar personalizar um modelo pré-treinado com um conjunto pequeno de imagens (ex: 4 a 6 imagens), o desempenho não evolui de forma monótona. Em vez disso, observa-se uma trajetória de três fases:

1. Melhoria Inicial: A fidelidade da imagem aumenta rapidamente nas primeiras iterações.
2. Fase de Corrupção (O Problema): A fidelidade da imagem cai abruptamente. As imagens geradas começam a exibir padrões ruidosos e distorcidos (artefatos visuais), embora o modelo ainda tente aprender o conceito.
3. Recuperação com Overfitting: A fidelidade volta a subir, mas o modelo entra em um estado de overfitting severo, onde passa a gerar apenas cópias idênticas das imagens de treinamento, perdendo a capacidade de criar variações diversificadas.

O artigo demonstra que a gravidade e o momento de início dessa fase de corrupção dependem do número de imagens de treinamento: quanto menos dados, mais cedo e mais severa é a corrupção.

2. Análise e Modelagem Heurística

Para entender a causa raiz, os autores desenvolveram uma modelagem heurística do processo de ajuste fino, começando com um cenário de "um único exemplo" (one-shot) e generalizando.

• Hipótese Central: A corrupção ocorre devido a uma distribuição aprendida excessivamente restrita. Em poucos exemplos, o modelo tenta modelar a distribuição de dados com um conjunto muito pequeno de amostras ($I_\theta$).
• Mecanismo de Erro: Durante a fase intermediária, o modelo possui uma "confiança" ($\sigma_1$) moderada na regeneração da amostra de treinamento. Quando o processo de difusão encontra ruído que não corresponde exatamente à amostra de treinamento, o erro de predição ($\delta_t$) é amplificado em vez de ser corrigido. Isso resulta na geração de padrões de ruído (corrupção).
• Conclusão da Análise: O modelo falha porque não consegue generalizar para além das amostras exatas de treinamento sem introduzir ruído, devido à falta de variabilidade na distribuição aprendida.

3. Metodologia: Aplicação de Redes Neurais Bayesianas (BNNs)

Para mitigar esse problema, os autores propõem incorporar Redes Neurais Bayesianas (BNNs) no processo de ajuste fino.

• Abordagem: Em vez de tratar os parâmetros do modelo como valores fixos, os parâmetros são modelados como variáveis aleatórias com distribuições posteriores.
• Mecanismo de Ação:
  • O uso de BNNs impede que o modelo aprenda a distribuição exata e rígida do conjunto de treinamento.
  • A inferência variacional introduz ruído intrínseco durante o treinamento, atuando como uma aumentação de dados implícita.
  • Isso força o modelo a aprender uma distribuição mais ampla e robusta ($I_\theta$ expandido), cobrindo casos de baixa confiança e evitando a amplificação de erros que causa a corrupção.
• Função de Perda: O objetivo de aprendizado é decomposto em duas partes:
  1. A expectativa da perda de difusão (fidelidade aos dados).
  2. Um termo de regularização (KL Divergence) que restringe a distribuição dos parâmetros a permanecer próxima da distribuição do modelo pré-treinado (prior).
• Eficiência: O método é altamente compatível com técnicas existentes (DreamBooth, LoRA, OFT) e não introduz custos adicionais de inferência, pois durante a inferência, os parâmetros são substituídos apenas pelos seus valores médios ($\mu_\theta$).

4. Resultados Experimentais

Os experimentos foram realizados em tarefas de geração orientada a objetos (Dataset DreamBooth) e orientada a sujeitos (CelebA-HQ), utilizando diferentes arquiteturas de DMs (SD v1.4, v1.5, v2.0) e métodos de ajuste fino.

• Mitigação da Corrupção: A aplicação de BNNs eliminou quase completamente a fase de corrupção, permitindo que o modelo transicionasse suavemente da melhoria inicial para um estado de alta qualidade sem passar pela degradação de ruído.
• Métricas Quantitativas:
  • Fidelidade de Texto (Clip-T) e Imagem (Dino, Clip-I): Melhoria consistente em todos os métodos (DreamBooth, LoRA, OFT).
  • Qualidade de Imagem (Clip-IQA): Aumento significativo na qualidade visual.
  • Diversidade (Lpips): Aumento na diversidade das imagens geradas, indicando que o modelo não sofreu overfitting severo.
• Estudo com Usuários: Em uma avaliação humana com 101 participantes, os modelos ajustados com BNNs foram preferidos em quase todos os cenários (melhor caso e caso médio), especialmente em alinhamento com o texto e qualidade geral da imagem.
• Eficiência Computacional: A aplicação de BNNs apenas em camadas específicas (ex: camadas de normalização ou blocos superiores) reduziu drasticamente o custo de memória e tempo de treinamento, mantendo a eficácia.

5. Contribuições Principais

1. Descoberta do Fenômeno: Identificação e nomeação da "Fase de Corrupção" no ajuste fino de poucos exemplos, um problema anteriormente não documentado detalhadamente.
2. Modelagem Teórica: Desenvolvimento de um modelo heurístico que explica a corrupção como resultado de uma distribuição aprendida restrita e alta variância de erro em estágios intermediários.
3. Solução Inovadora: Proposta de usar BNNs para expandir implicitamente a distribuição aprendida, mitigando a corrupção sem custos de inferência.
4. Validação Abrangente: Demonstração de que a solução é genérica, funcionando bem em diferentes modelos de difusão, métodos de ajuste fino (LoRA, OFT, DreamBooth) e tamanhos de conjuntos de dados.

6. Significado e Impacto

Este trabalho é significativo porque resolve um obstáculo fundamental na personalização de IA generativa: a instabilidade do treinamento com poucos dados. Ao eliminar a fase de corrupção, a técnica permite que aplicações de "poucos exemplos" sejam mais confiáveis e de maior qualidade, facilitando o uso de modelos de difusão em cenários reais onde grandes conjuntos de dados não estão disponíveis. Além disso, a abordagem de BNNs oferece um novo paradigma para regularização em modelos generativos, equilibrando fidelidade e diversidade de forma mais eficaz do que as técnicas atuais.