CAPO: Causal-Adaptive Preference Optimization for Diffusion Models via Causal Inference Uncertainty

Post-training alignment of diffusion models based on human feedback uniformly applies equal weights to all preference pairs, ignoring the inherent aleatoric and epistemic uncertainty in feedback, leading to overfitting on noisy signals and insufficient attention to difficult samples. To address this problem, we propose the Causal-Adaptive Preference Optimization (CAPO) framework. This method is the first to leverage the causal perspective on uncertainty for preference learning to achieve intelligent decoupling. We design a causal uncertainty decomposer that reliably decomposes the total prediction uncertainty for aleatoric uncertainty and epistemic uncertainty in a single forward pass. Based on this decomposition, we propose a causal-adaptive preference optimization loss function. This loss function dynamically adjusts the optimization objective weights using dual uncertainty information: suppressing the loss contribution of high aleatoric uncertainty samples while enhancing the loss contribution of high epistemic uncertainty samples. The CAPO framework enables the model to intelligently filter noisy signals and focus on tackling difficult tasks, thereby significantly improving alignment efficiency and robustness. On multiple text-to-image generation benchmarks, our method achieves significant performance improvements compared to existing techniques.