REVIEW 4 major objections 5 minor 20 references
JEDI: The Force of Jensen-Shannon Divergence in Disentangling Diffusion Models
T0 review · 4 major / 5 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read JEDI claims that minimizing a Jensen-Shannon divergence over internal attention maps at inference time separates subjects in diffusion model images better than the cosine-based CONFORM baseline, with fewer updates and no retraining.
desk verdict A plausible JSD-based test-time attention disentanglement method whose central quantitative claim currently rests on a circular metric and post hoc block selection. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is the normalized Jensen-Shannon divergence over sets of attention maps, treated as discrete probability distributions: $D_{JS}(P) = \frac{1}{|P|} \sum_{p \in P} D_{KL}(p \| m)$ with $m$ the mean distribution, divided by $\log n$ so it lies in $[0,1]$. The paper combines three terms built from this object: intra-group coherence, which minimizes JSD within a subject's attention maps; inter-group separation, which maximizes JSD between subjects' mixture distributions; and Shannon entropy regularization, which keeps the maps spatially spread. Updates use the Fast Gradient Sign Method, $x_t \leftarrow x_t - \alpha \cdot \operatorname{sign}(\nabla_{x_t} \text{score}(A_t))$, keeping the latent close to the base model. The mechanism's work is to turn the qualitative failure of attribute mixing into a quantitative overlap penalty on the model's own spatial probability fields.
What would settle it
Run JEDI with the same objective on attention maps taken from a deliberately wrong source, such as blocks outside 5 to 15 or unsymmetrized attention rows, and compare the generated images: if the inter-group JSD still rises without visual disentanglement, the attention-to-subject link is not carrying the result. More directly, collect generated images judged entangled versus disentangled by human raters across many prompts and check that the paper's inter-group JSD score separates the two sets; if it does not, the proposed CLIP-free score is not measuring what the paper claims.
Extended reading notes
Core claim
On the paper's own terms, the central discovery is that semantic entanglement can be formulated as overlap between attention distributions and minimized directly at inference time. For each subject, JEDI collects its token attention maps, minimizes their within-group Jensen-Shannon divergence, maximizes the divergence between subjects' mixture distributions, and adds a small entropy regularizer to prevent degeneracy. The loss is backpropagated into the latent with sign gradients, so only the first 18 denoising steps are adjusted. In Stable Diffusion 3.5, where self-attention mixes image and text tokens in DiT blocks, the paper extracts a symmetrized attention map per prompt token, bypassing softmax during attention and applying it only at extraction time, and optimizes blocks 5 to 15; in SD1.5 and LoRACLR it uses cross-attention maps. The reported outcome is comparable or better subject separation than the CONFORM baseline with 18 updates versus 69 and a smaller learning rate, plus a disentanglement score from inter-group JSD that rises when attribute mixing drops.
Load-bearing premise
The paper's load-bearing premise is that the attention maps it optimizes, in particular the symmetrized self-attention expression over DiT blocks 5 to 15 in Stable Diffusion 3.5, faithfully show which image regions each prompt subject controls; if those maps do not track subject location, minimizing divergence on them will not transfer to real image disentanglement.
Editorial extensions
If this is right
- Any diffusion model that exposes attention maps can be retrofitted with JEDI at test time, with no retraining, no external models, and no hand-made layouts.
- JEDI uses only 18 latent updates in the first 18 timesteps, roughly 67 percent fewer than CONFORM's 69, and its smaller learning rate keeps generated images closer to the base model's distribution.
- The inter-group JSD value doubles as a lightweight, CLIP-free disentanglement score, with higher values tracking visually disentangled subjects during sampling.
- The same method works on U-Net style SD1.5, DiT-style SD3.5, and the multi-concept LoRACLR model, indicating the objective is architecture-agnostic.
- The objective can in principle be moved into training or fine-tuning, which the paper identifies as a route to removing the inference-time overhead.
Reading between the lines
- Beyond the paper: shifting the JEDI objective into training or fine-tuning, rather than test-time optimization, would directly penalize entangled attention maps during learning and remove the doubling of inference time, though the paper only sketches this direction.
- Beyond the paper: because the method needs only internal attention maps, it should transfer to other transformer-based generators and even joint-attention video models; this is a prediction, not a paper claim.
- Beyond the paper: the inter-group JSD score could serve as a cheap diagnostic for attribute mixing across any text-to-image pipeline, testable by comparing it against human perceptual judgments on models the paper never ran.
- Beyond the paper: the blocks 5 to 15 choice for Stable Diffusion 3.5 suggests a principled block-selection criterion could be learned or derived, since the paper selects these blocks by visual inspection and computational efficiency.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. JEDI is a test-time adaptation method for diffusion models. At each early denoising step it extracts attention maps, groups them by prompt subject, and updates the latent by signed gradient descent on a loss of three terms: normalized Jensen-Shannon divergence within subject groups, separation between subject-mixture distributions, and a Shannon-entropy diversity regularizer. The paper proves upper bounds on JSD and entropy, demonstrates on toy Gaussians that JEDI preserves support structure better than NT-Xent, and reports qualitative comparisons on Stable Diffusion 1.5, Stable Diffusion 3.5, and LoRACLR. It also proposes an internal attention-based disentanglement score and claims roughly 67% faster updates than CONFORM.
Significance. If the empirical claims held, JEDI would be a simple, model-agnostic test-time intervention with a principled probabilistic objective and meaningful efficiency gains. The mathematical core is sound: Lemmas B.1 and B.2 are correct, and the synthetic comparison in Figure 2 is an instructive illustration of the difference between JSD and cosine/NT-Xent objectives. The paper also ships code and a project page. However, the empirical support for the central claim of consistent disentanglement is not yet established: the main score is the optimized objective itself, and the SD3.5 attention extraction protocol is selected post hoc on the evaluated examples. External, independent metrics and a predefined extraction protocol are needed before the central claims can be accepted.
major comments (4)
- [Section 5, Figure 10] The claimed disentanglement score is not independent of the optimized objective. The inter-group loss term in Section 3 is exactly 1 − DJS(M), where M contains the per-subject mixture attention maps, and JEDI minimizes this term during optimization. Figure 10 then reports that JEDI produces higher inter-group JSD than the base model. This improvement is guaranteed by construction, because JEDI directly maximizes the reported quantity. To support the claim of image-level disentanglement, the paper needs an external metric (e.g., instance segmentation, object detection, CLIP-based alignment, or a blinded human study) that is not a function of the optimized attention maps.
- [Appendix E] The SD3.5 attention extraction protocol is selected post hoc on the very images used for evaluation. The paper states that DiT blocks 5 to 15 are chosen 'based on visual analysis and computational efficiency' and that bypassing softmax during attention 'empirically yields more stable and consistent results.' Because the block range and softmax handling are tuned on the same prompts and images shown in Figures 4, 8, and 9, the reported improvements may reflect selection effects. A predefined protocol, or validation on a separate set of prompts and seeds, is necessary to rule out this confound.
- [Section 4] The section header promises 'quantitative experiments' on Stable Diffusion 1.5, including a comparison to CONFORM, but the results are exclusively qualitative: no metric table, no error bars over seeds, and no user study are reported anywhere in the paper or appendices. The only quantitative claim in the CONFORM comparison is the number of update steps (18 versus 69), which is an efficiency measure, not a quality measure. The abstract's statement that JEDI 'consistently improv[es] prompt alignment and disentanglement' is therefore not supported by the experiments as presented. Please add quantitative evaluation or soften the claim.
- [Section 5, 'Unbiased Disentanglement Score'] The label 'unbiased' is not justified for a score computed from the same internal attention distributions that JEDI directly optimizes. Even if the score is not identical to the full JEDI loss, it shares the same mechanism and is therefore circular for evaluating disentanglement. This is not merely a naming issue: Section 5 uses this score as the primary evidence that the method works. The paper should either supply an external evaluation or explicitly characterize the score as a diagnostic of the optimized objective, not as an independent benchmark.
minor comments (5)
- [Figure 10] The caption says 'Intra-Group JSD' while the surrounding text and the discussion in Section 5 refer to inter-group JSD; these labels should be reconciled.
- [Appendix B] In the proof of Lemma B.2, 'Langrage multiplier' should be 'Lagrange multiplier'.
- [Appendix C, Algorithm 1] The attention maps A_t appear in the update step but are never returned by Model in the pseudocode; the extraction step should be made explicit.
- [Section 4 and Appendix E] The number of prompts in the 'custom dataset' for SD3.5 is not stated; without a sample size it is difficult to interpret the claim that improvements are consistent.
- [Section 3 and Appendix A] The method relies on several free hyperparameters (learning rate alpha, diversity weight lambda, optimization window K, and the DiT block range), but only alpha is analyzed; please report how the other fixed choices were made or provide a robustness check.
Circularity Check
The SD3.5 disentanglement score is the JEDI objective itself, so the reported quantitative improvement is guaranteed by construction rather than independently measured.
-
fitted input called prediction
[Section 5 (Discussion and Future Work) and Appendix F, Figure 10]
"We find that the inter-group loss term in the JEDI objective naturally serves as an effective metric for measuring subject disentanglement during generation. For the images in Figure 1, the disentangled image achieves a mean JSD of 0.40 ± 0.15, compared to 0.17 ± 0.10 for the entangled counterpart, with full progression over time shown in Figure 10."
The JEDI objective's Inter-group Separation term minimizes 1 - D_JS(M) over subject-mixture attention distributions (Section 3). Appendix F defines Figure 10 as 'inter-group JSD between the two subjects ... computed across DiT blocks 7 to 15', which is exactly D_JS(M) over the same attention extraction used for optimization (blocks 5-15 optimized; 7-15 evaluated). Therefore JEDI directly maximizes the reported metric, so the higher Figure 10 value for JEDI relative to the base model is expected by construction. This is not an independent test of image-level disentanglement; the quantity reported as an evaluation metric is the same quantity being minimized as a loss.
full rationale
The strongest circularity is in the SD3.5 quantitative evaluation: the paper proposes the inter-group JSD as an unbiased disentanglement score, but this score is precisely the inter-group separation term of the JEDI loss, computed over the same attention blocks used for optimization. Thus the headline quantitative claim that JEDI 'consistently improves' disentanglement on SD3.5 is, in its only quantitative form, a restatement of the optimization objective rather than an independent outcome. The synthetic experiment and the SD1.5/CONFORM visual comparisons provide some independent qualitative content, and the mathematical lemmas are self-contained, which prevents a maximal circularity score. However, the central SD3.5 claim lacks an independent metric, so the paper's central quantitative evidence reduces by construction to its own loss. The Appendix E block selection ('Based on visual analysis and computational efficiency, we select blocks 5 to 15') is a post hoc choice that compounds the concern, but it is a selection risk rather than a separate definitional circularity.
Assumptions & free parameters
free parameters (4)
- alpha (learning rate) =
0.003
- lambda (diversity regularizer weight) =
0.01
- optimization window K =
18 (SD1.5 and SD3.5), 30 (LoRACLR)
- DiT block range =
blocks 5 to 15 of 24
assumptions (4)
- domain assumption Attention maps can be treated as discrete probability distributions over image positions.
- domain assumption The symmetrized self-attention expression in SD3.5 approximates the spatial influence of a prompt token.
- domain assumption Cross-attention maps for a subject token group localize that subject in the image.
- standard math JSD and Shannon entropy upper bounds hold as stated.
Cite this review
Pith. "Pith review of JEDI: The Force of Jensen-Shannon Divergence in Disentangling Diffusion Models." pith.science (2026). https://pith.science/paper/FG4UPCGT
@misc{pith2026250519166,
author = {Pith},
title = {Pith review of: JEDI: The Force of Jensen-Shannon Divergence in Disentangling Diffusion Models},
year = {2026},
howpublished = {\url{https://pith.science/paper/FG4UPCGT}},
note = {Machine review of arXiv:2505.19166}
}
read the original 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/.
Figures
Figures from the paper (11 more)
Reference graph
Works this paper leans on
-
[1]
write newline
" write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 global.max substring 't := if while FUNCTION format.date year duplicate empty "emp...
-
[2]
Make it count: Text-to-image generation with an accurate number of objects
Binyamin, L., Tewel, Y., Segev, H., Hirsch, E., Rassin, R., and Chechik, G. Make it count: Text-to-image generation with an accurate number of objects. arXiv preprint arXiv:2406.10210, 2024
arXiv 2024
-
[3]
Attend-and-excite: Attention-based semantic guidance for text-to-image diffusion models
Chefer, H., Alaluf, Y., Vinker, Y., Wolf, L., and Cohen-Or, D. Attend-and-excite: Attention-based semantic guidance for text-to-image diffusion models. ACM transactions on Graphics (TOG), 42 0 (4): 0 1--10, 2023
work page 2023
-
[4]
A simple framework for contrastive learning of visual representations
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp.\ 1597--1607. PmLR, 2020
2020
-
[5]
Scaling rectified flow transformers for high-resolution image synthesis
Esser, P., Kulal, S., Blattmann, A., Entezari, R., M \"u ller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024
work page 2024
-
[6]
J., Shlens, J., and Szegedy, C
Goodfellow, I. J., Shlens, J., and Szegedy, C. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572, 2014
arXiv 2014
-
[7]
Z., Shi, Y., Chen, Y., Fan, Z., Xiao, W., Zhao, R., Chang, S., Wu, W., et al
Gu, Y., Wang, X., Wu, J. Z., Shi, Y., Chen, Y., Fan, Z., Xiao, W., Zhao, R., Chang, S., Wu, W., et al. Mix-of-show: Decentralized low-rank adaptation for multi-concept customization of diffusion models. Advances in Neural Information Processing Systems, 36: 0 15890--15902, 2023
work page 2023
-
[8]
Momentum contrast for unsupervised visual representation learning
He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp.\ 9729--9738, 2020
2020
Show all 20 references
-
[9]
Denoising diffusion probabilistic models
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33: 0 6840--6851, 2020
2020
-
[10]
C., and Heilbron, F
Kwon, G., Jenni, S., Li, D., Lee, J.-Y., Ye, J. C., and Heilbron, F. C. Concept weaver: Enabling multi-concept fusion in text-to-image models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.\ 8880--8889, 2024
2024
-
[11]
T., Ben-Hamu, H., Nickel, M., and Le, M
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022
2022 arXiv
-
[12]
Liu, N., Li, S., Du, Y., Torralba, A., and Tenenbaum, J. B. Compositional visual generation with composable diffusion models. In European Conference on Computer Vision, pp.\ 423--439. Springer, 2022
2022
-
[13]
Meral, T. H. S., Simsar, E., Tombari, F., and Yanardag, P. Conform: Contrast is all you need for high-fidelity text-to-image diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.\ 9005--9014, 2024
2024
-
[14]
Oord, A. v. d., Li, Y., and Vinyals, O. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748, 2018
2018 arXiv
-
[15]
and Xie, S
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp.\ 4195--4205, 2023
2023
-
[16]
Sdxl: Improving latent diffusion models for high-resolution image synthesis
Podell, D., English, Z., Lacey, K., Blattmann, A., Dockhorn, T., M \"u ller, J., Penna, J., and Rombach, R. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952, 2023
2023 arXiv
-
[17]
Hierarchical text-conditional image generation with clip latents
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 1 0 (2): 0 3, 2022
2022 arXiv
-
[18]
High-resolution image synthesis with latent diffusion models
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp.\ 10684--10695, 2022
2022
-
[19]
Loraclr: Contrastive adaptation for customization of diffusion models
Simsar, E., Hofmann, T., Tombari, F., and Yanardag, P. Loraclr: Contrastive adaptation for customization of diffusion models. arXiv preprint arXiv:2412.09622, 2024
2024
-
[20]
Enhancing mmdit-based text-to-image models for similar subject generation
Wei, T., Chen, D., Zhou, Y., and Pan, X. Enhancing mmdit-based text-to-image models for similar subject generation. arXiv preprint arXiv:2411.18301, 2024
2024 arXiv
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.