JEDI: The Force of Jensen-Shannon Divergence in Disentangling Diffusion Models

TL;DR

JEDI is a test-time method that improves subject separation and compositional alignment in diffusion models by minimizing semantic entanglement in attention maps using Jensen-Shannon divergence, without retraining.

Contribution

JEDI introduces a novel Jensen-Shannon divergence based objective for test-time adaptation in diffusion models, enhancing disentanglement and prompt alignment without retraining or external supervision.

Findings

Improves subject separation and compositional alignment in diffusion models.

Reduces number of update steps needed through adversarial optimization.

Provides a lightweight, CLIP-free disentanglement score for benchmarking.

Abstract

We introduce JEDI, a test-time adaptation method that enhances subject separation and compositional alignment in diffusion models without requiring retraining or external supervision. JEDI operates by minimizing semantic entanglement in attention maps using a novel Jensen-Shannon divergence based objective. To improve efficiency, we leverage adversarial optimization, reducing the number of updating steps required. JEDI is model-agnostic and applicable to architectures such as Stable Diffusion 1.5 and 3.5, consistently improving prompt alignment and disentanglement in complex scenes. Additionally, JEDI provides a lightweight, CLIP-free disentanglement score derived from internal attention distributions, offering a principled benchmark for compositional alignment under test-time conditions. Code and results are available at https://ericbill21.github.io/JEDI/.