ConceptAttention: Diffusion Transformers Learn Highly Interpretable Features

ConceptAttention is an interpretability method for multi-modal diffusion transformers. We implement it for the Flux DiT architecture in PyTorch. Check out the paper here.

Code setup

You will then need to install the code here locally by running

pip install -e .

Running the model

Here is an example of how to run Flux with Concept Attention

from concept_attention import ConceptAttentionFluxPipeline

pipeline = ConceptAttentionFluxPipeline(
    model_name="flux-schnell",
    device="cuda:0"
)

prompt = "A dragon standing on a rock. "
concepts = ["dragon", "rock", "sky", "cloud"]

pipeline_output = pipeline.generate_image(
    prompt=prompt,
    concepts=concepts,
    width=1024,
    height=1024,
)

image = pipeline_output.image
concept_heatmaps = pipeline_output.concept_heatmaps

image.save("image.png")
for concept, concept_heatmap in zip(concepts, concept_heatmaps):
    concept_heatmap.save(f"{concept}.png")

Examples

Example scripts are located in the examples/ directory:

To run an example:

cd examples
python generate_image_flux.py

Output images will be saved to examples/results/flux/ or examples/results/flux2/ depending on the model.

Concept Attention Also Works on Video!

https://github.com/user-attachments/assets/d4a0e100-1fc3-44c1-ae9e-09fb420bd003

ConceptAttention can also be applied to video generation models. Here's an example using CogVideoX:

from concept_attention.cogvideox import ModifiedCogVideoXTransformer3DModel, ModifiedCogVideoXPipeline
from diffusers.utils import export_to_video

# Load model
model_id = "THUDM/CogVideoX-5b"
transformer = ModifiedCogVideoXTransformer3DModel.from_pretrained(
    model_id, subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe = ModifiedCogVideoXPipeline.from_pretrained(
    model_id, transformer=transformer, torch_dtype=torch.bfloat16
).to("cuda")

# Generate video with concept attention
prompt = "A golden retriever with a ball by a tree in the grass."
concepts = ["dog", "grass", "sky", "tree", "ball"]

video, concept_attention_dict = pipe(
    prompt=prompt,
    concepts=concepts,
    num_frames=81,
    num_inference_steps=50,
    concept_attention_kwargs={
        "timesteps": list(range(0, 50)),
        "layers": list(range(0, 30)),
    }
)

# Save video
export_to_video(video.frames[0], "output.mov", fps=8)

# Access concept attention maps (shape: num_concepts, num_frames, height, width)
concept_attention_maps = concept_attention_dict["concept_attention_maps"]

See the full example at examples/generate_video_cogvideox.py.

Experiments

Each of our experiments are in separate directories in /experiments.

You can run one for example like this

cd experiments/qualitative_baseline_comparison
python generate_image.py # Generates test image using flux
python plot_flux_concept_attention.py # Generates concept attention maps and saves them in results. 

Data Setup

To use ImageNetSegmentation you will need to download gtsegs_ijcv.mat into /experiments/imagenet_segmentation/data.

cd experiments/imagenet_segmentation/data
wget http://calvin-vision.net/bigstuff/proj-imagenet/data/gtsegs_ijcv.mat

Bibtex

@misc{helbling2025conceptattentiondiffusiontransformerslearn,
    title={ConceptAttention: Diffusion Transformers Learn Highly Interpretable Features}, 
    author={Alec Helbling and Tuna Han Salih Meral and Ben Hoover and Pinar Yanardag and Duen Horng Chau},
    year={2025},
    eprint={2502.04320},
    archivePrefix={arXiv},
    primaryClass={cs.CV},
    url={https://arxiv.org/abs/2502.04320}, 
}