Robust and Generalizable Safety Steering for Text-to-Image Diffusion Transformers

2026-05-28

AI Summary

Existing text-to-image diffusion models struggle to achieve stable and generalizable safety control, as harmful semantics are dynamically coupled across textual and visual representations, rendering fixed interventions ineffective and safety mechanisms non-transferable. This work proposes SafeDIG, a novel framework that introduces position-aware sparse feature transfer into diffusion Transformers (DiTs) for safety control. SafeDIG constructs a sparse autoencoder with functionally distinct intervention sites and employs robustness-aware routing to select stable intervention points. The encoder is frozen as a reusable safety dictionary, while only the decoder is adapted to align with the target domain’s activation manifold. During inference, Blend and Repel operations jointly steer activations away from harmful directions. Evaluated on FLUX.1 Dev and Stable Diffusion 3.5 Large, SafeDIG significantly reduces unsafe generation rates while preserving source-domain safety and image fidelity, enabling effective safety generalization across diverse risk domains.

Abstract

Diffusion Transformers have become a powerful backbone for text-to-image generation, but their layered and cross-modal generation process makes safety control fundamentally different from prompt-level filtering or output-level detection. Harmful semantics may be weakly expressed in text representations, progressively bound to visual latents, and finally entangled with rendering dynamics. As a result, safety steering at a fixed layer can be unstable, and a steering mechanism learned from known risks may not transfer reliably to a shifted target risk domain. We propose SafeDIG, a safety steering framework that formulates DiT safety adaptation as position-aware sparse feature transfer. SafeDIG first constructs Sparse Autoencoders over functionally distinct DiT intervention positions and uses robustness-aware pre-training routing to prioritize intervention sites that are expected to remain stable under source-target risk shift. It then separates transferable safety features from domain-specific activation geometry by freezing the SAE encoder as a reusable sparse safety dictionary and adapting only the decoder to the target-domain activation manifold. During inference, SafeDIG combines Blend and Repel operations to steer unsafe activations toward transferred safety manifolds or away from harmful sparse directions. Experiments on FLUX.1 Dev and Stable Diffusion 3.5 Large show that SafeDIG consistently reduces target-domain and overall unsafe generation rates while preserving source-domain safety and image quality.

Problem

• Safety Steering
• Text-to-Image Generation
• Diffusion Transformers
• Domain Shift
• Harmful Semantics

Innovation

• Diffusion Transformers
• Safety Steering
• Sparse Autoencoders
• Feature Transfer
• Robustness