Balanced Latent Space of Diffusion Models for Counterfactual Generation

Abstract

Counterfactual generation has demonstrated impressive performance in tasks such as image editing and synthesis, largely due to the development of diffusion models. However, existing diffusion-based counterfactual generation models suffer from instability due to a lack of understanding of the latent space. These models either retain too much of the original information or make excessive modifications, sacrificing crucial details, leading to inefficiency and inauthenticity. In this paper, we propose a framework that balances the latent space by incorporating signals that facilitate the transition to new counterfactuals while preserving factual information. We first identify the cause of this imbalance as the uncontrolled signal from the counterfactuals. Based on this understanding, we introduce a balancing method within the diffusion process. Our approach is evaluated on the colored MNIST dataset, a modified version of the standard MNIST dataset, with experimental results showing significant improvements over previous latent space methods.