DifFRACT: Diffusion Feature Reconstruction and Attribution for Circuit Tracing

Abstract: Mechanistic interpretability seeks to explain neural network behavior by decomposing model computations into interpretable features and circuits. While transcoder-based circuit tracing has recently enabled detailed causal analyses of LLMs, multimodal diffusion transformers for image generation remain comparatively opaque. We still lack tools for understanding how semantic information propagates across denoising steps and how text and image representations interact within double-stream MM-DiT architectures. Existing methods provide only partial insight: attention maps expose a limited view of token interactions, while sparse autoencoders can discover interpretable features but do not directly reveal how these features are transformed and composed through nonlinear MLP layers. 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 MLP sublayers in FLUX.1[schnell]. By replacing MLPs with transcoders and linearizing the remaining computation, we obtain exact feature-to-feature attribution and recover compact, interpretable circuits. Empirically, our transcoders match or slightly outperform sparse autoencoders on the sparsity-faithfulness tradeoff. The resulting circuits reveal mechanisms underlying attribute binding and cross-stream semantic propagation, and provide causal explanations for systematic generation errors. Moreover, circuit-guided interventions are substantially more precise and effective than standard SAE-based steering. Our results demonstrate that transcoder-based circuit analysis is feasible for state-of-the-art diffusion transformers and provides a powerful framework for understanding and controlling multimodal generative models. The code is available at https://github.com/Artalmaz31/DifFRACT

• The paper introduces a transcoder-based circuit tracing approach that precisely reconstructs and attributes feature propagation in diffusion transformers.
• It leverages variance-normalized losses and FiLM conditioning to achieve low MSE and sparse activations, outperforming sparse autoencoders.
• The method enables practical interventions, mitigating issues like prior bias and improving safety and control in multimodal generative models.

Authoritative Summary of "DifFRACT: Diffusion Feature Reconstruction and Attribution for Circuit Tracing" (2606.15796)

Introduction and Motivation

DifFRACT addresses mechanistic interpretability in state-of-the-art multimodal diffusion transformers, specifically focusing on understanding FLUX.1's image and text double-stream architecture. The work overcomes the opacity typical in diffusion models by leveraging transcoder-based circuit tracing—an approach previously effective for LLMs—to yield granular attribution graphs revealing feature-level causal pathways within the model during generative denoising. This methodology provides actionable insights into feature propagation, cross-modal interactions, and the underlying mechanisms responsible for semantic composition and attribute binding in image synthesis.

Transcoder Architecture and Training Protocol

Transcoders offer a sparse and highly faithful input–output approximation of MLP sublayer computations by conditioning on the diffusion timestep via FiLM. Each transcoder—trained per block and stream—maps modulated residual inputs into high-dimensional sparse codes, decoded to approximate the original MLP output. The architecture ensures interpretability and tractability, supporting exact additive decompositions via linearization in the local replacement model (LRM).

The loss is variance-normalized to absorb activation scale heterogeneity across timesteps, explicitly balancing sparsity and faithfulness without sacrificing reconstruction fidelity. With an expansion factor of ×16 and aggressive normalization/regularization, each transcoder achieves low normalized MSE and sparse feature activations, enabling circuit-level tracing through the affine computational graph.

Figure 1: Iterative graph construction process for attribution tracing, demonstrating per-position aggregation and discovery of influential sources for downstream features.

Circuit Tracing and Attribution Graphs

The attribution pipeline operates by replacing MLP sublayers in FLUX.1 with trained transcoders, linearizing attention and normalization using frozen denominators and cached attention probability tensors. Feature activations are treated as fixed, allowing the LRM to decompose a target feature's preactivation as an exact sum of upstream contributions from earlier features, positional input embeddings, and reconstruction errors.

Iterative graph construction and pruning aggregate per-position activations into compact circuits, with indirect influence scoring guiding expansion. Circuits are further refined by vertex and edge pruning, isolating the most functionally relevant components of the generative process.

Empirical Analysis: Transcoder–SAE Comparison

Transcoders are empirically compared to sparse autoencoders (SAEs) across multiple layers and streams. They match or outperform SAEs on the sparsity–faithfulness frontier, as evidenced by normalized MSE and mean $L_0$ activation at matched sparsity levels.

Figure 2: Transcoder vs SAE sparsity–faithfulness frontier for FLUX.1 double-stream blocks, demonstrating consistently superior tradeoff across configurations.

Transcoders uniquely enable feature-to-feature attribution, supporting a mechanistic analysis inaccessible to SAEs, which reconstruct features from output activations and lack input-invariant tracing capabilities.

Temporal Evolution of Circuits in Diffusion Transformers

Attribution graphs computed over denoising steps elucidate the shifting interplay between text and image streams. Early denoising steps are dominated by text-driven semantic reasoning with extensive cross-modal edges; later steps transition to perceptual refinement centered in the image stream.

Key statistics:

• Text-stream attribution mass decreases from 89.9% (step 0) to 5.4% (step 3)
• Cross-modal edge fraction drops from 14.9% to 2.0%

  Figure 3: Denoising trajectory evolution: stream share, cross-modal edge fraction, and per-layer contribution breakdown.

These dynamics support a two-phase interpretation, where semantic interventions are most effective in early steps; suppression or amplification of text-stream supernodes at late steps yields minimal impact.

Practical Steering via Circuit-Guided Interventions

Circuit-guided feature steering leverages attribution graphs for nuanced interventions—distinguishing concept features (directly tied to tokens) from context features (firing on related tokens). Suppression or amplification of context features materially alters the output without substituting core concepts, demonstrating granularity beyond SAE-based methods.

Figure 4: Steering concept vs context features, showing qualitatively distinct effects on generated scene attributes across seeds and strengths.

Negative connections (active suppressors) are also revealed; joint suppression of both present and suppressor features is required for reliable concept switching absent in single-feature steering.

Color Circuits and Systematic Error Mitigation

FLUX.1 exhibits persistent color prior biases—red stop signs, black ladybugs—despite explicit prompt specification. Attribution graphs clarify that associative features outweigh prompt-driven features, producing incorrect color activations. Circuit-wide interventions, simultaneously suppressing red features and their associative sources, substantially increase success rates in generating atypically colored objects.

Figure 5: Attribution mechanisms and prior bias mitigation demonstrating superior correction via circuit-wide intervention compared to single-feature suppression.

Prior bias mitigation is visually and quantitatively validated across diverse prompts and seeds.

Figure 6: Qualitative outcomes of prior bias mitigation for atypical colors, illustrating the effectiveness of circuit-guided suppression.

Additional Analyses: Style, Spatial Composition, Counting, and Failure Modes

The method is further substantiated by decomposing artistic style features, enabling their monotonic manipulation and separation from content, as well as isolating spatial and counting concepts within the model. Failure modes—such as color leakage, negation misrepresentation, and hard prior dominance—are diagnosed as failures of cross-stream information transfer. Attribution graphs localize causes to specific features and clarify why prompt semantics may not propagate to visual realization.

Transcoder Faithfulness and Robustness

Transcoders' direct fit and end-to-end replacement in FLUX.1 yields high cosine similarity and low latent-space MSE versus the original model, reproducing qualitative composition, object placement, and stylistic register. Floating-point drift is bounded across blocks and denoising steps, establishing practical robustness for interpretability analysis.

Figure 7: Transcoder training curves: normalized MSE and mean $L_0$ activation demonstrating stability and sparsity across training cycles.

Figure 8: Block-level absolute error between original and LRM outputs, confirming numerical precision and bounded drift.

Implications and Future Directions

DifFRACT establishes transcoder-based circuit tracing as a strict improvement over SAE-based analyses for diffusion models, unlocking precise, actionable feature-to-feature attribution. This enhances both mechanistic understanding and practical controllability of generative models, directly supporting interventions to mitigate systematic failures and refine output semantics.

Theoretically, the methodology strengthens the analogy between transformer interpretability in LLMs and diffusion architectures, suggesting that feature circuits and attribution graphs are broadly applicable across modalities. Practically, it paves the way for safer and more reliable generative modeling by exposing the causal structure of information propagation, attribute binding, and context integration.

Future work should extend transcoders to single-stream blocks and alternative architectures, address limitations in decomposing attention probability tensors, and further refine graph construction to enable even larger-scale circuit tracing. The approach promises improved diagnostic tooling for model failures, enhanced modalities for steered generation, and systematic analysis of cross-modal integration.

Conclusion

DifFRACT demonstrates that transcoder-based circuit tracing enables rigorous mechanistic interpretability for diffusion transformers, paving the way for deep feature-level analysis and robust control of multimodal generative models. The approach realizes highly descriptive attribution graphs, supporting targeted and effective interventions that surpass the limitations of sparse autoencoders and prior interpretability techniques. This methodology holds significant promise for advancing theoretical understanding and practical reliability of diffusion-based image generation.