---
license: apache-2.0
base_model: black-forest-labs/FLUX.1-schnell
datasets:
- yvdao/midjourney-v6
library_name: pytorch
tags:
- mechanistic-interpretability
- sparse-autoencoder
- transcoder
- circuit-tracing
- diffusion
- flux
- text-to-image
---
# DifFRACT: Diffusion Feature Reconstruction and Attribution for Circuit Tracing
Trained **timestep-conditioned transcoders** and **SAE baselines** for the [FLUX.1[schnell]](https://huggingface.co/black-forest-labs/FLUX.1-schnell) text-to-image diffusion transformer (MM-DiT), accompanying the paper *DifFRACT: Diffusion Feature Reconstruction and Attribution for Circuit Tracing* ([PDF](https://arxiv.org/abs/2606.15796)).
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](https://github.com/Artalmaz31/DifFRACT)).
## Contents
40 PyTorch checkpoints in two folders. Every checkpoint is a `state_dict` for a `TemporalAwareTranscoder` module (the SAE baseline shares the identical architecture).
| Folder | Files | Naming | Streams | Layers |
|---|---|---|---|---|
| `temporal-aware-transcoders/` | 34 | `transcoder_{stream}_{layer}.pt` | `img`, `txt` | 0-15 and 18 |
| `temporal-aware-saes/` | 6 | `sae_{stream}_{layer}.pt` | `img`, `txt` | 6, 12, 18 |
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.
## Model architecture
Each module maps an MLP **input** `x` to its **output** `Å·`, conditioned on the diffusion timestep `t`:
- a sinusoidal timestep embedding → 2-layer SiLU MLP → linear head producing FiLM `(scale, shift)`;
- modulation `x_mod = x ⊙ (1 + scale) + shift`;
- a ReLU encoder `z = ReLU(W_enc x_mod + b_enc)` (sparse code);
- a unit-norm linear decoder `Å· = W_dec z + b_dec`.
| Hyperparameter | Value |
|---|---|
| Base model | FLUX.1[schnell], MM-DiT, `d_model = 3072` |
| Expansion factor | 16 (`d_feat = 49152`) |
| Timestep embedding dim | 256 |
| Sparsity | L1, `λ_img = 3e-4`, `λ_txt = 5e-5` |
| Reconstruction loss | variance-normalized MSE |
| Optimizer | AdamW, lr `2e-4`, weight decay `0`, CosineAnnealingLR |
| Activation buffer / batch | `1e6` / `4096` |
| Inference steps / guidance / resolution | 4 / 0 / 512×512 |
| Training prompts | [`yvdao/midjourney-v6`](https://huggingface.co/datasets/yvdao/midjourney-v6) |
The **SAE baseline** is architecturally identical but autoencodes the MLP **output** (input = target), so its reconstruction error is directly comparable to a transcoder's.
## Usage
Install the companion code ([GitHub](https://github.com/Artalmaz31/DifFRACT)), then:
```python
from huggingface_hub import snapshot_download
from transcoder_training.transcoder import load_transcoders
transcoders = load_transcoders(
f"{path}/temporal-aware-transcoders",
layers=range(16),
d_model=3072,
expansion_factor=16,
time_embed_dim=256,
)
```
Loading an individual SAE baseline:
```python
import torch
from transcoder_training.transcoder import TemporalAwareSAE
sae = TemporalAwareSAE(d_model=3072, expansion_factor=16, time_embed_dim=256)
sae.load_state_dict(torch.load(f"{path}/temporal-aware-saes/sae_img_12.pt", map_location="cpu"))
sae.eval()
```
The end-to-end pipeline (Local Replacement Model, attribution graph, intervention) is demonstrated in `walkthrough.ipynb` in the companion repository.
## Citation
```bibtex
@misc{mazur2026diffractdiffusionfeaturereconstruction,
title={DifFRACT: Diffusion Feature Reconstruction and Attribution for Circuit Tracing},
author={Artyom Mazur and Nina Konovalova and Aibek Alanov},
year={2026},
eprint={2606.15796},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2606.15796},
}
```
- Base model: [`black-forest-labs/FLUX.1-schnell`](https://huggingface.co/black-forest-labs/FLUX.1-schnell)
- License: Apache-2.0