Bayesian Power Steering: An Effective Approach for Domain Adaptation of Diffusion Models

TL;DR

This work introduces Bayesian Power Steering (BPS), a Bayesian fine-tuning framework for domain adaptation of diffusion models that converts learning from a large probability space to a task-specific small probability space. By modeling the posterior denoise function $\bar{\boldsymbol{\epsilon}}^*(\mathbf{z}_t,t,\mathbf{c}_{text},\mathbf{c}_{add})$ as a combination of the pretrained denoise and a time-dependent steering term $M$, BPS implements learnable, time-aware interventions across a multi-scale feature hierarchy (CP, TP, EP) while keeping the base model frozen. The method demonstrates data-efficient, high-fidelity adaptation across layout-to-image, segmentation-to-image, artistic drawing, and sketch-to-image tasks, achieving state-of-the-art or competitive results (e.g., FID $=10.49$ on COCO17 sketch) and robust performance under data scarcity. This approach enables practical deployment of diffusion models for user-specific conditioning with limited labeled data, by leveraging structured, task-aware conditioning and a lightweight, discriminative integration strategy. Overall, BPS offers a compact, fast-converging, and versatile framework for domain adaptation in diffusion-based generative systems.

Abstract

We propose a Bayesian framework for fine-tuning large diffusion models with a novel network structure called Bayesian Power Steering (BPS). We clarify the meaning behind adaptation from a \textit{large probability space} to a \textit{small probability space} and explore the task of fine-tuning pre-trained models using learnable modules from a Bayesian perspective. BPS extracts task-specific knowledge from a pre-trained model's learned prior distribution. It efficiently leverages large diffusion models, differentially intervening different hidden features with a head-heavy and foot-light configuration. Experiments highlight the superiority of BPS over contemporary methods across a range of tasks even with limited amount of data. Notably, BPS attains an FID score of 10.49 under the sketch condition on the COCO17 dataset.

Figures (27)

• Figure 1: The relationship between the probability spaces learned by the pre-trained model (the prior probability space or large probability space) and the target task (the target probability space or small probability space). Different colors represent different densities at various locations on the manifold.
• Figure 2: The simulation of 2D measure. The pre-trained model is trained with 500K samples, and the generative results are shown in Fig. (b). In Figures (c) and (d), labeled subsets of 100 samples were utilized in the experiments. Fig. (c) illustrates the generation outcomes achieved through fine-tuning the pre-trained model with 2.5K epochs, whereas Fig. (d) displays the results obtained by training the same model from scratch with 20K epochs. The above legends represent the sampling of 5K samples on either the simulated or real measure.
• Figure 3: Classification accuracy of generated samples from distinct integration modes with $n=12$ training samples. For each mode, five experiments are conducted.
• Figure 4: Robustness analysis of the integrated Models of distinct modes with respect to sample size and number of iterations. The horizontal coordinate is the number of epochs and the vertical coordinate is the accuracy.
• Figure 5: Overview of the integration model. The gray dumbbell shapes represent the pre-trained model. The blue rectangles represent blocks with residual structures, and rectangles of the same color in the same scale represent interventions in the same functional unit sharing features from the BPS.

Theorems & Definitions (1)

• Lemma 3.1