Trained timestep-conditioned transcoders and SAE baselines for the FLUX.1[schnell] text-to-image diffusion transformer (MM-DiT), accompanying the paper DifFRACT: Diffusion Feature Reconstruction and Attribution for Circuit Tracing (PDF).

A transcoder decomposes an MLP sublayer into a sparse linear combination of interpretable features; conditioning it on the denoising timestep lets a single transcoder track how a feature behaves across the diffusion trajectory. Substituting these transcoders into a frozen Local Replacement Model yields the attribution graphs and circuit-guided interventions studied in the paper. The code to load and use these weights is at the companion repository (GitHub).

40 PyTorch checkpoints in two folders. Every checkpoint is a state_dict for a TemporalAwareTranscoder module (the SAE baseline shares the identical architecture).

The 32 transcoders for layers 0–15 (both streams) are the set analysed by the Local Replacement Model; layer 18 (and the SAEs at 6/12/18) support the sparsity–faithfulness comparison.

Each module maps an MLP input x to its output ŷ, conditioned on the diffusion timestep t:

The SAE baseline is architecturally identical but autoencodes the MLP output (input = target), so its reconstruction error is directly comparable to a transcoder's.

Install the companion code (GitHub), then:

Loading an individual SAE baseline:

The end-to-end pipeline (Local Replacement Model, attribution graph, intervention) is demonstrated in walkthrough.ipynb in the companion repository.