Search papers, labs, and topics across Lattice.
The paper introduces SafeDIG, a framework for safety steering in Diffusion Transformers (DiTs) that addresses the challenge of controlling harmful content generation in text-to-image models. SafeDIG uses Sparse Autoencoders (SAEs) to identify robust intervention points within the DiT architecture and transfers safety-related features across different risk domains. By freezing the SAE encoder and adapting only the decoder, SafeDIG achieves improved generalization and reduces unsafe generation rates while maintaining image quality, as demonstrated on FLUX.1 Dev and Stable Diffusion 3.5 Large.
Freezing a Sparse Autoencoder's encoder creates a reusable "safety dictionary" that generalizes to new risks in text-to-image diffusion models, offering a more robust alternative to fixed-layer steering.
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.
DAMO