Concept Removal for Frontier Image Generative Models
BLOCK replaces model bottleneck layers with trained transcoders for persistent concept removal in SD3.5, Flux, Infinity, outperforming state-of-the-art methods.
Key Findings
Methodology
The proposed BLOCK method replaces the bottleneck layer in modern diffusion and autoregressive models (e.g., SD3.5, Flux, Infinity) with a trained transcoder. This module learns to replicate the original layer while structuring activation features, creating an internal filter that selectively disables concept-specific signals. Training involves sparse regularization and TopK activation to ensure feature disentanglement. The intervention is embedded permanently, enabling persistent concept removal under white-box access without retraining the entire model. Empirical evaluations on UnlearnCanvas benchmark demonstrate superior performance in style and object removal, with robustness against adversarial prompts and support for sequential concept removal across multiple models and scales.
Key Results
- In Flux (12B parameters), the method surpasses baseline concept removal accuracy by up to 21 percentage points, maintaining high visual fidelity with minimal FID increase.
- On Infinity-8B, multiple concepts are removed sequentially with negligible quality loss, CLIPScore remains high, and the method outperforms weight-editing approaches.
- Across models, BLOCK supports multi-concept, multi-step removal, demonstrating robustness against adversarial prompts and outperforming external modules like SAE-based approaches.
Significance
This approach addresses core limitations of existing concept removal techniques by internalizing the filter within the model architecture, ensuring persistence and efficiency. It enables large-scale models to be safely deployed with content controls, reducing risks of undesired outputs and bias propagation. The method's scalability and robustness make it a promising standard for content moderation, model safety, and bias mitigation in AI-generated imagery, with broad industry and research implications.
Technical Contribution
The main innovation is the integration of a trained transcoder into the model’s bottleneck layer, replacing weight-based edits or external modules. This creates a permanent, internal filter that can be precisely targeted at concept-specific features. The training leverages a novel TopK sparsity mechanism with a straight-through estimator, enabling effective feature disentanglement without full model backpropagation. The approach is architecture-agnostic, scalable to models with billions of parameters, and supports sequential, multi-concept removal while maintaining high image quality.
Novelty
This work is the first to embed a trained transcoder directly into the core architecture of large diffusion and autoregressive models for concept removal. Unlike prior weight editing or external modules, it offers a persistent, internal filtering mechanism that is robust against sequential and adversarial challenges. The use of structured activation features for targeted concept suppression represents a significant departure from existing methods, providing both theoretical guarantees and practical advantages.
Limitations
- The method relies on white-box access, limiting applicability to open-source or accessible models. Black-box scenarios remain challenging.
- Training the transcoder requires curated prompt datasets and hyperparameter tuning, which may limit scalability in some contexts.
- Effectiveness in removing highly entangled or complex concepts, especially when multiple concepts overlap, needs further validation. Future work should explore multi-objective optimization and unsupervised training strategies.
Future Work
Future directions include developing unsupervised or few-shot training methods for transcoder generalization, extending the approach to multi-modal content (video, 3D), and integrating with content moderation pipelines. Additionally, exploring adaptive, context-aware concept removal and theoretical analysis of feature disentanglement will further enhance robustness and applicability.
AI Executive Summary
The rapid advancement of deep generative models, notably diffusion and autoregressive architectures like SD3.5, Flux, and Infinity, has revolutionized visual content creation, enabling the production of highly realistic and diverse images. However, these models are trained on vast, uncurated internet datasets, which inevitably contain undesirable concepts such as inappropriate objects, copyrighted styles, or biased representations. Existing methods for concept removal—either external modules like SAE or weight editing—face significant challenges: external modules can be bypassed or detached, weight edits often degrade over multiple iterations, and retraining entire models is computationally prohibitive.
This paper introduces BLOCK, a novel internal concept removal framework that directly modifies the model’s core architecture. By replacing the bottleneck layer—an architecturally mandated transformation point— with a trained transcoder, BLOCK creates a persistent internal filter. This filter can selectively disable signals associated with specific concepts, effectively removing them from generated images without compromising overall quality. The training process employs a sparse activation mechanism, TopK, combined with a carefully designed loss function to ensure feature disentanglement and precise control.
Empirical evaluations on multiple large-scale models demonstrate that BLOCK outperforms existing techniques in both object and style removal tasks. In the Flux model with 12 billion parameters, it improves concept removal accuracy by up to 21 percentage points, while maintaining high visual fidelity as measured by FID and CLIPScore. The method also supports sequential concept removal, showing robustness against adversarial prompts and multiple concept combinations. These results highlight the method’s scalability, effectiveness, and potential for real-world deployment.
Overall, BLOCK offers a significant step forward in the internal control of large generative models. Its architecture-embedded, persistent filtering mechanism ensures long-term robustness and efficiency, making it a practical tool for content moderation, bias mitigation, and safe AI deployment. Future work will explore unsupervised training, multi-modal extensions, and further theoretical insights into feature disentanglement, aiming to broaden its impact across AI content generation and safety domains.
Deep Dive
Abstract
Image generative models are trained on massive, largely uncurated internet-scale datasets that contain undesirable visual concepts. Efficiently removing such concepts from the model generations without degrading the quality of output images remains challenging. We introduce a novel concept removal method for frontier diffusion and image autoregressive models, such as SD3.5, Flux, and Infinity. Our intervention replaces the internal bottleneck layer present in all these modern models with a transcoder that is trained to replicate the original layer while structuring it into distinct activation features. This in-place substitution creates an integrated filter through which concept-specific signals can be selectively disabled while preserving the rest of the model's behavior. Since the intervention modifies the model backbone rather than attaching an external component, it remains persistent under white-box access. Empirically, the approach achieves state-of-the-art concept removal performance across modern diffusion and autoregressive models, maintains visual generation quality, provides robustness against adversarial prompts, and supports sequential removal of diverse concepts. This positions our method as a practical approach for concept removal in frontier image generative models.