ShortFT: Diffusion Model Alignment via Shortcut-based Fine-Tuning

TL;DR

ShortFT introduces a shortcut-based fine-tuning method for diffusion models, significantly improving alignment with reward functions by reducing computational costs and enhancing effectiveness over traditional backpropagation approaches.

Contribution

The paper proposes a novel shortcut-based fine-tuning strategy using a trajectory-preserving few-step diffusion model to improve alignment efficiency and performance.

Findings

Enhanced alignment with reward functions

Reduced computational costs during fine-tuning

Outperforms state-of-the-art methods

Abstract

Backpropagation-based approaches aim to align diffusion models with reward functions through end-to-end backpropagation of the reward gradient within the denoising chain, offering a promising perspective. However, due to the computational costs and the risk of gradient explosion associated with the lengthy denoising chain, existing approaches struggle to achieve complete gradient backpropagation, leading to suboptimal results. In this paper, we introduce Shortcut-based Fine-Tuning (ShortFT), an efficient fine-tuning strategy that utilizes the shorter denoising chain. More specifically, we employ the recently researched trajectory-preserving few-step diffusion model, which enables a shortcut over the original denoising chain, and construct a shortcut-based denoising chain of shorter length. The optimization on this chain notably enhances the efficiency and effectiveness of fine-tuning the foundational model. Our method has been rigorously tested and can be effectively applied to various reward functions, significantly improving alignment performance and surpassing state-of-the-art alternatives.