Diffusion Models for Attribution

Abstract

In high-stakes domains such as healthcare, finance, and law, the need for explainable AI is critical. Traditional methods for generating attribution maps, including white-box approaches relying on gradients and black-box techniques that perturb inputs, face challenges like gradient vanishing, blurred attributions, and computational inefficiencies. To overcome these limitations, we introduce a novel approach that leverages diffusion models within the framework of Information Bottleneck (IB) theory. By utilizing the Gaussian noise from diffusion models, we connect the information bottleneck with the Minimum Mean Squared Error (MMSE) from classical information theory, enabling precise calculation of mutual information. This connection leads to a new loss function that minimizes the Signal-to-Noise Ratio (SNR), facilitating efficient optimization and producing high-resolution, pixel-level attribution maps. Our method achieves greater clarity and accuracy in attributions than existing techniques, requiring significantly fewer pixel values to reach the necessary predictive confidence. This work demonstrates the power of diffusion models in advancing explainable AI, particularly in identifying critical input features with high precision.