Test-Time Scaling for Safe Text-Guided Image Generation via Intermediate Clean Estimates

2026-08-04

AI Summary

This work addresses the vulnerability of current text-to-image diffusion models to prompt-induced generation of harmful content—such as nudity or copyrighted intellectual property—and the limitations of existing defenses, which are either computationally expensive or rely solely on textual modifications while neglecting visual signals. The authors propose a weight-invariant, test-time safety mechanism that leverages intermediate clean image estimates during generation to detect prohibited concepts via sparse boundary-aware objectives. Upon detection, the method performs real-time intervention by optimizing a structured low-rank residual in the text conditioning space through truncated backpropagation. This approach uniquely integrates visual signals from the generative process for efficient detection and correction, enables flexible safety tuning, and incurs negligible overhead on standard inference latency. Experiments on Stable Diffusion v1.4 and v3.5 demonstrate substantial improvements over existing weight-invariant baselines in nudity removal, IP protection, and style erasure, effectively suppressing violations while better preserving output quality and semantic fidelity.

Abstract

Ensuring safety and policy compliance in text-to-image diffusion models remains a critical challenge, as benign or adversarial prompts can often elicit prohibited content, e.g. nudity and protected intellectual property. While training-based unlearning methods are effective, they are computationally expensive and prone to catastrophic interference with general capabilities. Conversely, existing test-time defenses are primarily prompt-centric, relying on modifying textual descriptions only, and overlook the visual signals for detection. In this paper, we propose to leverage the intermediate clean image estimated during the generation process and employ a sparse margin objective to detect prohibited concepts. When a violation is detected, we immediately intervene by optimizing a structured low-rank residual in the text-conditioning space via truncated backpropagation. This design allows weight-preserving detection, keeps non-violating inference latency nearly unchanged as the maximum budget increases, and offers flexibility in safety performance via test-time scaling. Extensive experiments on Stable Diffusion v1.4 and v3.5 across nudity removal, IP protection, and style erasure demonstrate superior performance across suppression, fidelity and preservation compared to prior weight-preserving baselines, providing a scalable and flexible solution for safe generative deployment.

Problem

• safe text-to-image generation
• prohibited content
• policy compliance
• diffusion models
• content safety

Innovation

• test-time scaling
• intermediate clean estimate
• low-rank residual
• weight-preserving safety
• diffusion model safety