Update README.md

README.md

@@ -1,3 +1,107 @@
 ---
 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
+
+<a href="https://arxiv.org/abs/2606.15796"><img src="https://img.shields.io/badge/Paper-PDF-b31b1b.svg" height=22.5></a>
+<a href="https://huggingface.co/Artalmaz31/DifFRACT"><img src="https://img.shields.io/badge/🤗-Weights-yellow" height=22.5></a>
+<a href="https://colab.research.google.com/github/Artalmaz31/DifFRACT/blob/main/walkthrough.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>
+<a href="./LICENSE"><img src="https://img.shields.io/badge/license-Apache--2.0-green" height=22.5></a>
+
+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