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

TL;DR

This paper investigates the corruption stage in few-shot fine-tuning of diffusion models, models this phenomenon theoretically, and proposes a Bayesian neural network approach to mitigate corruption and enhance image quality.

Contribution

The paper introduces a theoretical model of the corruption stage in diffusion model fine-tuning and applies Bayesian neural networks to mitigate this issue without extra inference costs.

Findings

Bayesian neural networks significantly reduce corruption in generated images.

The method improves fidelity, quality, and diversity in object- and subject-driven tasks.

The approach is compatible with existing fine-tuning methods and requires no additional inference cost.

Abstract

Few-shot fine-tuning of Diffusion Models (DMs) is a key advancement, significantly reducing training costs and enabling personalized AI applications. However, we explore the training dynamics of DMs and observe an unanticipated phenomenon: during the training process, image fidelity initially improves, then unexpectedly deteriorates with the emergence of noisy patterns, only to recover later with severe overfitting. We term the stage with generated noisy patterns as corruption stage. To understand this corruption stage, we begin by theoretically modeling the one-shot fine-tuning scenario, and then extend this modeling to more general cases. Through this modeling, we identify the primary cause of this corruption stage: a narrowed learning distribution inherent in the nature of few-shot fine-tuning. To tackle this, we apply Bayesian Neural Networks (BNNs) on DMs with variational inference to implicitly broaden the learned distribution, and present that the learning target of the BNNs can be naturally regarded as an expectation of the diffusion loss and a further regularization with the pretrained DMs. This approach is highly compatible with current few-shot fine-tuning methods in DMs and does not introduce any extra inference costs. Experimental results demonstrate that our method significantly mitigates corruption, and improves the fidelity, quality and diversity of the generated images in both object-driven and subject-driven generation tasks. Code is available at https://github.com/Nicholas0228/BNN-Finetuning-DMs.