Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

SHIFT: Steering Hidden Intermediates in Flow Transformers

About

Diffusion models have become leading approaches for high-fidelity image generation. Recent DiT-based diffusion models, in particular, achieve strong prompt adherence while producing high-quality samples. We propose SHIFT, a simple but effective and lightweight framework for concept removal in DiT diffusion models via targeted manipulation of intermediate activations at inference time, inspired by activation steering in large language models. SHIFT learns steering vectors that are dynamically applied to selected layers and timesteps to suppress unwanted visual concepts while preserving the prompt's remaining content and overall image quality. Beyond suppression, the same mechanism can shift generations into a desired \emph{style domain} or bias samples toward adding or changing target objects. We demonstrate that SHIFT provides effective and flexible control over DiT generation across diverse prompts and targets without time-consuming retraining.

Nina Konovalova, Andrey Kuznetsov, Aibek Alanov• 2026

Related benchmarks

TaskDatasetResultRank
Nudity ErasureI2P (dev)
Common Count155
10
Style ErasureOld artists
AccE30
7
Style ErasureModern artists
AccE60
7
Image Generation EvaluationMS-COCO 5K (sampled images)
FID33.9
7
Concept PreservationMickey prompts
Concept Score (CS)26.06
3
Concept PreservationDog prompts
CS24.18
3
Concept PreservationLegislator prompts
CS21.77
3
Concrete Object ErasurePikachu Concrete Object Erasure
CS25.14
3
Concept PreservationSpongebob prompts
Concept Score (CS)27.35
3
Concept PreservationPikachu prompts
Concept Score (CS)26.25
3
Showing 10 of 16 rows

Other info

Follow for update