Official PyTorch implementation of "Where and How to Perturb: On the Design of Perturbation Guidance in Diffusion and Flow Models".

HeadHunter is a framework designed to steer diffusion models more effectively by selecting and combining individual attention heads. It addresses the observation that when perturbation guidance (like PAG) is applied at the head level, effects vary significantly, with some heads exhibiting interpretable characteristics related to specific visual concepts.

This repository provides a multi-GPU distributed inference pipeline for HeadHunter. The pipeline performs iterative attention head perturbation and ranking using image scoring models, with per-iteration validation and memory-efficient model handling.

• 🧠 Head-level Specialization: Exploits the distinct visual effects (structure, lighting, texture, style) captured by individual attention heads.
• 🔀 Compositional Power: Combines multiple heads to compose complex visual effects.
• 🎯 Objective-Aware Guidance: Iteratively finds the best set of heads to maximize arbitrary objectives (e.g., PickScore, CLIPScore).
• 🌫️ SoftPAG: Implements interpolation-based perturbation to control strength and prevent oversaturation.

• ⚡ Distributed Generation/Inference: Efficiently distributes attention head perturbation tasks across layers and GPUs (Supports SD3 & FLUX.1-Dev).
• 📊 Iterative Top-K Selection: Automatically selects the best heads using:
  • PickScore (Used in paper)
  • CLIPScore
  • Harmonic mean of Pick + CLIP
• 🧪 Validation: Validation performed before and after every iteration.

conda env create -f headhunter.yaml
conda activate headhunter
cd diffusers
pip install -e ".[torch]"

We provide scripts for both Stable Diffusion 3 (SD3) and FLUX.1-Dev.

bash run.sh      # For SD3
bash run_flux.sh # For FLUX.1-Dev (512x512)

# Please modify 'prompts_style.txt' as needed.
# Please add  <your_wandb_api_key> in shell script.

output_iterative/
  _style_prompt/
    iter0/
      content_prompt/
        layerL_headH.png
    iter1/
      ...
    validation/
      iter0/
        content0.png
      ...
      final_val.png   # Grid image of all validation results
    final/
      content_prompt/
        final_perturb_heads.pkl
        final_perturb_heads.txt

The perturbation logic is implemented in diffusers/src/diffusers/models/attention_processor.py. You can specify the behavior using the perturb_type string.

These define the target attention map ($A'$) that acts as the negative guidance or perturbation target.

Modifies the attention map after the softmax operation.

Modifies the attention scores (logits) before the softmax operation.

You can control the strength of the perturbation by interpolating between the original attention map ($A$) and the perturbed target ($A'$).

Syntax: Append @scale={u} to the perturb_type, where u is the interpolation factor ($0.0 \le u \le 1.0$).

• $u=0.0$: Original Attention ($A$).
• $u=1.0$: Full Perturbation ($A'$).

By default, Linear Interpolation is used. You can specify other geodesic paths by appending keywords:

Example: [PROB_PERTURB]attention_identity@scale=0.5 implements SoftPAG with 50% strength.

You can target specific attention heads for perturbation, as introduced in HeadHunter.

• Via Pipeline Argument: Pass perturb_heads=[0, 1, 5] to the pipeline call.

• PAG: Self-Rectifying Diffusion Sampling with Perturbed-Attention Guidance
• SEG: Smoothed Energy Guidance
• HeadHunter / SoftPAG: Where and How to Perturb: On the Design of Perturbation Guidance in Diffusion and Flow Models

If you find this work useful, please cite our paper:

@inproceedings{ahn2025headhunter,
  title={Where and How to Perturb: On the Design of Perturbation Guidance in Diffusion and Flow Models},
  author={Ahn, Donghoon and Kang, Jiwon and Lee, Sanghyun and Kim, Minjae and Jang, Wooseok and Min, Jaewon and Lee, Sangwu and Paul, Sayak and Kim, Seungryong},
  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
  year={2025}
}