DifFRACT: Diffusion Feature Reconstruction and Attribution for Circuit Tracing

Abstract

While transcoder-based circuit tracing has recently enabled detailed causal analyses of large language models, it remains unclear how semantic information propagates across denoising steps and how text and image representations interact within multimodal diffusion transformers (MM-DiT) for image generation. Attention maps expose only a limited view of token interactions, and sparse autoencoders find interpretable features but do not reveal how nonlinear MLP layers transform and compose them. In this work, we extend transcoder-based circuit tracing to multimodal diffusion transformers. We train timestep-conditioned transcoders that faithfully approximate the input-output behavior of the MLP sublayers in the double-stream blocks of FLUX.1[schnell]. By replacing MLPs with transcoders and linearizing the remaining computation, we obtain exact feature-to-feature attribution and build compact attribution graphs over interpretable features. Empirically, our transcoders match or slightly outperform sparse autoencoders on the sparsity-faithfulness tradeoff. The graphs recover the known shift from prompt-driven to image-driven computation across denoising steps, point to candidate mechanisms of attribute binding and cross-stream semantic propagation, and yield intervention-testable explanations for systematic generation errors. On prompt grids labeled blind, graph-guided interventions reach effects that single-feature steering does not; at a matched budget, features chosen from the graph fix a color prior more often than the target feature alone and destroy the object less often than features chosen by co-activation. FLUX.1[dev] reproduces the temporal and color-prior results, and SD3.5-medium reproduces them in direction. Our results demonstrate that transcoder-based circuit analysis is feasible for state-of-the-art multimodal diffusion transformers.

Abstract: arXiv descriptive metadata (CC0). Thank you to arXiv for use of its open access interoperability. This service was not reviewed or approved by, nor does it necessarily express or reflect the policies or opinions of, arXiv.