SHIFT: Steering Hidden Intermediates in Flow Transformers

1. Problem Statement

Diffusion models, particularly modern Diffusion Transformers (DiTs) like FLUX and SD3.5, have achieved state-of-the-art image generation quality. However, their scale and training on large, imperfect datasets introduce significant risks, including the generation of harmful, copyrighted, or prohibited content.

Existing solutions for Concept Erasure (CE) face critical limitations when applied to DiTs:

• Computational Cost: Most methods (e.g., ESD, Concept Ablation) rely on gradient-based fine-tuning. While feasible for smaller UNet-based models (<1B parameters), they are prohibitively expensive for massive DiTs (e.g., FLUX with 12B parameters).
• Guidance Dependency: Many optimization-based methods require negative guidance during training. This fails for distilled models (like FLUX.1[schnell]) which are optimized to operate at a fixed guidance scale (often 0.0) and cannot utilize the guidance mechanisms required for these erasure techniques.
• Architectural Mismatch: Previous activation steering methods (inspired by LLMs) often rely on cross-attention mechanisms specific to UNet architectures. DiTs use Multimodal Diffusion Transformers (MM-DiT) with unified attention mechanisms where text and image tokens share a latent space, making traditional intervention points less effective.

2. Methodology: SHIFT Framework

The authors propose SHIFT, a lightweight, inference-time framework that steers DiT generation by manipulating intermediate activations without modifying model weights.

Core Architecture & Intervention Points

SHIFT targets the MM-DiT architecture (specifically Flux), focusing on two key intervention points where text semantics are propagated:

1. Text Encoder Pooled Embedding: Specifically the pooled vector (e.g., CLIP token).
2. Diffusion Transformer Text Tokens: The text token representations ($Y_t$) after the shared attention layer in the DiT backbone.

The steering operation is defined as:
$$\tilde{a} = a + \alpha \times v_{\text{steering}}$$
Where $a$ is the activation, $v_{\text{steering}}$ is the learned steering vector, and $\alpha$ is the steering strength.

Vector Construction

To generate a steering vector for a target concept (e.g., "nudity" or "Van Gogh style"):

1. Dataset Creation: A small set of contrastive prompt pairs is collected (e.g., $p^-$ without the concept vs. $p^+$ with the concept).
2. Activation Collection: The model runs on these prompts, and activations are recorded at the chosen intervention points.
3. Vector Estimation: Two methods are proposed:
   • Mean Difference: The simple difference between the mean activations of positive and negative prompts.
   • SVM-Based: A linear SVM is trained on the pooled activations to find the separating hyperplane. The steering vector is the normalized normal vector of this hyperplane, providing a more robust direction for separation.

Inference-Time Control & Regularization

SHIFT employs a dynamic mechanism to prevent "over-steering" (where the image quality degrades or non-target concepts are lost):

• Text Encoder Strength: Scaled by the cosine similarity between the initial prompt embedding and the target concept embedding.
• DiT Strength: Regularized by a lightweight classifier (SVM) signal. If the classifier detects the target concept is still present in the current activation ($p_{cls}$ is high), the steering strength is amplified. If the concept is already suppressed, steering is reduced.
• Temporal Invariance: A key finding is that a single, time-independent steering vector can be applied across all diffusion timesteps, simplifying the process significantly compared to timestep-specific interventions.

3. Key Contributions

1. Unified Steering Framework: SHIFT is the first comprehensive framework for steering both original and distilled DiT models (Flux.1[dev] and Flux.1[schnell]) via latent activation shifts, eliminating the need for retraining.
2. Temporal Consistency Discovery: The authors demonstrate that semantic concepts in DiTs are encoded as stable, global directions in the latent manifold. A single vector works effectively across all timesteps, streamlining the steering process.
3. Cross-Distillation Transfer: The method validates that steering vectors computed on a distilled model (Flux.1[schnell]) can be successfully transferred to the non-distilled version (Flux.1[dev]), enabling efficient control in resource-constrained settings.
4. Versatility: Beyond erasure, the framework supports domain stylization (e.g., changing to cyberpunk style), object addition (e.g., adding a hat), and concept switching (e.g., cat $\to$ dog).

4. Experimental Results

The authors evaluated SHIFT on Flux.1[dev] and Flux.1[schnell] across multiple tasks:

• Abstract Concept Erasure (Nudity):

  • On the I2P benchmark, SHIFT achieved 3x to 4x stronger suppression of nudity compared to baselines like ESD, CA, and UCE.
  • It maintained high prompt adherence (CLIP scores remained stable) and only incurred a minor degradation in FID compared to the best baselines.
  • Crucially, it worked effectively on Flux.1[schnell] (distilled, guidance=0.0), where optimization-based methods fail.

• Concrete Concept Erasure:

  • Tested on removing specific objects (e.g., "Snoopy") while preserving related concepts (e.g., "Mickey", "Pikachu").
  • SHIFT outperformed baselines in suppressing the target concept (lower CLIP for the target) while maintaining better image quality (lower FID) for non-target concepts.

• Style and Object Manipulation:

  • Successfully removed artist styles (Van Gogh) while preserving others (Picasso, Rembrandt).
  • Demonstrated the ability to add small objects (glasses, hats) and switch concepts (young $\to$ old, dog $\to$ cat) with high fidelity.

• Ablation Studies:

  • Block Selection: Steering is most effective when applied to the early stages of the diffusion trajectory (blocks 0–9). Applying it only to later steps yielded poor results.
  • Regularization: The use of the SVM-based classifier for dynamic strength adjustment was critical for preserving image quality and preventing the loss of non-target attributes.

5. Significance

SHIFT represents a paradigm shift in controlling large-scale generative models:

• Efficiency: It replaces computationally heavy fine-tuning with a lightweight, inference-time intervention, making it scalable for 10B+ parameter models.
• Distilled Model Compatibility: It solves the critical challenge of controlling distilled models that cannot use standard guidance-based erasure techniques.
• Generalizability: By treating DiTs similarly to LLMs (via activation steering), it opens the door for applying representation engineering techniques to the rapidly evolving field of diffusion transformers.
• Safety: It provides a practical, deployable tool for mitigating safety risks in high-fidelity image generation without compromising the model's core generative capabilities.

The code is publicly available, facilitating further research into activation-based control for next-generation diffusion models.