Robust and Generalizable Safety Steering for Text-t...

Robust and Generalizable Safety Steering for Text-t... | AI Research

Key Takeaways

Robust and Generalizable Safety Steering for Text-to-Image Diffusion Transformers Diffusion Transformers (DiTs) have become the standard for high-quality ima...
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.
During inference, SafeDIG combines Blend and Repel operations to steer unsafe activations toward transferred safety manifolds or away from harmful sparse directions.

Paper AbstractExpand

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.

Robust and Generalizable Safety Steering for Text-to-Image Diffusion Transformers
Diffusion Transformers (DiTs) have become the standard for high-quality image generation, but their complex internal structure makes traditional safety measures—like filtering prompts or checking final images—insufficient. Because harmful concepts evolve from text into visual patterns deep within the model, safety interventions must happen inside the model's "brain." This paper introduces SafeDIG, a framework designed to identify the most effective locations within a DiT to intervene and to transfer safety knowledge across different types of risks, ensuring models remain safe even when encountering new, unseen harmful content.

Identifying Where to Intervene

In a DiT, information flows through different stages: from text understanding to cross-modal binding (where text meets image) and finally to rendering. SafeDIG recognizes that a "one-size-fits-all" intervention layer is ineffective. Instead, it uses a "robustness-aware" routing system that analyzes these stages before training. By predicting which intervention sites will remain stable when the model faces new risks, SafeDIG prioritizes the most reliable locations for safety control, moving beyond simple trial-and-error layer selection.

Transferring Safety Knowledge

A major challenge in AI safety is ensuring that a model trained to avoid one type of harm can also avoid others. SafeDIG addresses this by using Sparse Autoencoders (SAEs), which act as a dictionary for safety-related features. The framework separates the "what" (the safety features) from the "where" (the specific way the model renders images). By freezing the encoder—the part that identifies safety features—and only adapting the decoder to the specific target domain, SafeDIG can learn to block new risks without needing to be fully retrained, effectively reusing its safety knowledge.

Steering During Generation

During the actual image generation process, SafeDIG uses two specific operations to keep the model on a safe path. The "Blend" operation gently pulls the model’s internal activations toward a safe, stable state. The "Repel" operation identifies specific "harmful" directions in the model's internal space and pushes the generation process away from them. These techniques allow the model to steer away from unsafe outputs while maintaining the high visual quality and creative flexibility that users expect from modern diffusion models.

Proven Performance

Experiments conducted on powerful models like FLUX.1 Dev and Stable Diffusion 3.5 Large demonstrate that SafeDIG significantly reduces the rate of unsafe image generation. By successfully transferring safety protocols to new, held-out risk categories, the framework proves that it can maintain protection across diverse scenarios. The results show that SafeDIG consistently outperforms existing safety methods, providing a more robust and adaptable way to keep generative AI models aligned with safety standards.