Boosting Alignment After Unlearning in Text-to-Image Models#

Why Unlearning Matters#

Machine unlearning has emerged as a critical capability for deployed AI systems for several reasons:

Regulatory Compliance: Regulations like the EU AI Act, GDPR’s “right to be forgotten,” and various copyright frameworks increasingly require the ability to remove specific learned behaviors or knowledge from AI models. When copyrighted training data must be removed, or when a model has learned to produce harmful content, unlearning provides a path to compliance without the enormous cost of full retraining.

Safety Alignment: Models trained on internet-scale data inevitably learn to produce harmful content. Unlearning provides a targeted mechanism for removing specific harmful capabilities — violence, explicit content, bias, misinformation — while preserving the model’s useful capabilities.

Rapid Response: When new safety concerns emerge — a novel misuse vector is discovered, a new type of harmful content is identified — unlearning provides a faster response mechanism than retraining from scratch or applying output filters.

Model Customization: Organizations may want to deploy models with specific capabilities removed — a model for educational use that cannot generate violent content, or a model for commercial use that cannot reproduce copyrighted styles.

The Problem with Current Unlearning#

Current machine unlearning methods for text-to-image models, while effective at removing targeted capabilities, suffer from several significant side effects:

Capability Degradation: Unlearning specific content often degrades the model’s ability to generate related but benign content. A model trained to unlearn nudity might become worse at generating portraits, medical imagery, or artistic nudes in the classical tradition. The “blast radius” of unlearning extends beyond the targeted capability.

Concept Bleeding: Unlearning one concept can affect the model’s representation of semantically related concepts. Removing the ability to generate images of a specific copyrighted character might degrade the model’s ability to generate images of similar but non-copyrighted characters.

Alignment Distortion: Perhaps most critically, unlearning can distort the model’s overall alignment. A model that was carefully aligned to refuse harmful requests might, after unlearning, become less reliable at refusing other types of harmful content. The unlearning process can disrupt the delicate balance of safety training.

Distribution Shift: Unlearning changes the model’s output distribution in ways that extend beyond the targeted concept. The overall quality, diversity, and coherence of generated images can be affected.

Catastrophic Forgetting: In extreme cases, aggressive unlearning can trigger catastrophic forgetting — the model loses broad categories of capability, not just the targeted one.

This constellation of side effects creates what Ko et al. term the post-unlearning alignment gap: the discrepancy between a model’s alignment after unlearning and its alignment before unlearning, across all dimensions of safety and quality.

Formalizing the Problem#

Ko et al. formalize the post-unlearning alignment problem with precision. Let us consider a text-to-image model that has been trained and aligned to satisfy a set of safety objectives — it refuses to generate violent content, explicit content, copyrighted content, and so on. When we apply unlearning to remove one specific capability (say, generating images in a copyrighted artist’s style), we want to achieve two goals simultaneously:

1. Effective Unlearning: The model should genuinely lose the ability to generate the targeted content.
2. Alignment Preservation: The model should maintain its performance and alignment on all other dimensions — other safety objectives, generation quality, and useful capabilities.

The challenge is that these two goals are in tension. Current unlearning methods tend to optimize for the first goal at the expense of the second. The question Ko et al. address is: can we boost alignment after unlearning, recovering the safety and quality properties that unlearning disrupts?

Why Alignment Degrades After Unlearning#

The paper provides a thoughtful analysis of why unlearning disrupts alignment:

Entangled Representations: In text-to-image models, concepts are not represented in isolation. The model’s internal representations entangle multiple concepts — style, content, composition, quality, safety boundaries. When unlearning modifies the representation of one concept, it inevitably affects the representations of entangled concepts.

Safety Training Erosion: Safety alignment in text-to-image models is typically achieved through fine-tuning on curated datasets, reinforcement learning from human feedback (RLHF), or other alignment techniques. These methods produce subtle adjustments to the model’s weights that encode safety boundaries. Unlearning, which also modifies weights, can overwrite or disrupt these carefully calibrated safety adjustments.

Concept Manifold Distortion: The model’s concept space — the manifold of possible generations — is a complex, high-dimensional structure. Unlearning creates a “hole” in this manifold for the targeted concept, but this hole distorts the surrounding region, affecting nearby concepts and potentially creating unexpected pathways to harmful content.

Loss Landscape Disruption: The optimization process of unlearning moves the model’s weights to a new region of the loss landscape. This new region may not be well-aligned with the safety constraints that were satisfied in the original weight configuration.

Overview#

Ko et al. propose a principled approach to boosting alignment after unlearning. Their method, which I will refer to as alignment boosting, operates in two phases:

1. Unlearning Phase: Apply an existing unlearning method to remove the targeted capability from the model.
2. Alignment Recovery Phase: Apply targeted techniques to recover and even improve the model’s alignment on other dimensions.

The key insight is that alignment recovery is not simply re-doing the original alignment training — it requires techniques specifically designed to repair the distortions introduced by unlearning while preserving the unlearning itself.

Alignment Boosting Techniques#

The paper explores several techniques for boosting alignment after unlearning:

Integration and Optimization#

The paper presents the complete alignment boosting framework as an integrated optimization problem that jointly optimizes:

• Unlearning effectiveness: Ensuring the targeted capability is genuinely removed
• Safety preservation: Maintaining the model’s refusal behavior for other harmful content
• Quality preservation: Maintaining generation quality for benign content
• Distinguishing capability: Enabling the model to correctly distinguish between unlearned and related concepts

The optimization uses a multi-objective formulation with adaptive weighting that dynamically adjusts the relative importance of each objective based on the model’s current state.

Setup#

The paper evaluates the alignment boosting framework on several popular text-to-image models, including Stable Diffusion variants and other open-source models. The evaluation covers multiple unlearning scenarios:

Style Unlearning: Removing the model’s ability to generate images in the style of specific artists (relevant for copyright compliance).

Object Unlearning: Removing the model’s ability to generate specific objects or entities (relevant for privacy and safety).

Concept Unlearning: Removing the model’s ability to generate images related to specific concepts, such as violence or explicit content (relevant for safety alignment).

Metrics#

The evaluation uses a comprehensive set of metrics:

Unlearning Effectiveness: How completely has the targeted capability been removed? Measured through targeted generation attempts and classifier-based evaluation.

Safety Preservation: How well does the model maintain its refusal behavior for non-targeted harmful content? Measured through safety benchmarks and red-teaming evaluations.

Generation Quality: How does the quality of generated images compare before unlearning, after unlearning without alignment boosting, and after alignment boosting? Measured through FID scores, CLIP scores, and human evaluation.

Concept Distinguishment: Can the model correctly generate images of concepts related to but distinct from the unlearned concept? Measured through targeted generation tests on semantically proximate concepts.

Alignment Gap: The overall measure of how much alignment is preserved or recovered through the alignment boosting process.

Key Results#

The experimental results demonstrate that alignment boosting significantly improves post-unlearning alignment across all evaluated scenarios:

The Broader Unlearning-Safety Connection#

The CodeHalu paper’s findings have implications beyond their specific technical contributions. They highlight a broader principle: safety interventions in AI systems are not independent. Modifying a model to address one safety concern can affect its safety profile on other dimensions. This interconnectedness means that AI safety must be approached holistically, not as a collection of independent patches.

Regulatory Implications#

As regulations increasingly mandate the removal of specific capabilities from deployed models (copyright compliance, privacy protection, harmful content removal), the post-unlearning alignment problem becomes a regulatory concern. A model that has been modified for regulatory compliance should not become less safe on other dimensions. Alignment boosting provides a framework for ensuring that regulatory compliance does not compromise safety.

For Model Developers#

The work provides practical guidance for model developers:

Plan for Alignment Recovery: Unlearning should not be implemented as a standalone operation. It should be followed by alignment verification and recovery to ensure that the model remains safe and useful.

Evaluate Comprehensively: After unlearning, models should be evaluated not just on the targeted capability but on the full range of safety and quality metrics. The post-unlearning alignment gap should be explicitly measured.

Use Integrated Approaches: Treating unlearning and alignment as separate concerns leads to suboptimal outcomes. Integrated approaches that jointly optimize for unlearning effectiveness and alignment preservation produce better results.

For the Unlearning Research Community#

The paper opens several important research directions:

Theoretical Understanding: Why does unlearning disrupt alignment? What are the theoretical limits of post-unlearning alignment recovery? Can we develop theoretical frameworks that predict the alignment impact of specific unlearning operations?

Efficient Recovery: The current alignment boosting process requires significant computational resources. Can we develop more efficient recovery techniques that achieve comparable results with less computation?

Proactive Alignment Preservation: Rather than recovering alignment after it is lost, can we develop unlearning methods that intrinsically preserve alignment? This proactive approach would be more efficient than the reactive recovery approach.

Scalability: The current work focuses on text-to-image models. Do the same principles apply to other modalities — language models, video generation models, multimodal systems? Extending the framework to other model types is an important direction.

Machine Unlearning Landscape#

Ko et al.’s work sits within a rapidly growing body of research on machine unlearning:

Exact Unlearning: Approaches that guarantee the complete removal of specific training data influence, typically requiring significant architectural support or data organization. While theoretically clean, exact unlearning is often impractical for large-scale generative models.

Approximate Unlearning: Approaches that approximate the effect of removing training data without guaranteeing complete removal. Ko et al.’s work focuses on this category, recognizing that practical unlearning is inherently approximate.

Task-Level Unlearning: Removing specific capabilities rather than specific training data influence. This is the most relevant category for safety applications — we want to remove the capability to generate harmful content, not necessarily the influence of specific training examples.

Certified Unlearning: Providing formal guarantees about the effectiveness of unlearning. This is an active research area that could provide important assurances for regulatory compliance.

Alignment in Generative Models#

The paper also contributes to the broader discourse on alignment in generative models:

Multi-Objective Alignment: The challenge of simultaneously satisfying multiple safety objectives (no violence, no explicit content, no copyright infringement, etc.) while maintaining model utility. The post-unlearning alignment problem is a specific instance of this broader challenge.

Alignment Stability: The question of whether alignment is a stable property of trained models or a fragile equilibrium that can be disrupted by perturbations (including unlearning). The paper’s findings suggest that alignment is more fragile than commonly assumed.

Alignment Verification: The need for comprehensive, ongoing verification of model alignment, not just at deployment time but after any modification (including unlearning). The alignment gap metric introduced in the paper provides a framework for such verification.

Computational Cost#

The alignment boosting process adds computational overhead to the unlearning process. For organizations that need to frequently modify deployed models (responding to new copyright claims, emerging safety concerns, etc.), this overhead may be significant. Developing more efficient recovery techniques is an important practical priority.

Evaluation Coverage#

The evaluation focuses on a limited set of safety dimensions and quality metrics. A truly comprehensive evaluation would cover a broader range of safety concerns (bias, fairness, misinformation, etc.) and quality dimensions (composition, text rendering, spatial reasoning, etc.).

Long-Term Stability#

The paper evaluates alignment boosting immediately after application. Whether the recovered alignment is stable over time and under distribution shift (changes in user prompts, new adversarial techniques) is an important open question.

Interaction with Other Modifications#

In practice, models may undergo multiple modifications — sequential unlearning operations, fine-tuning for new capabilities, safety updates, etc. How alignment boosting interacts with these subsequent modifications is not yet well understood.

The Completeness of Recovery#

While alignment boosting significantly reduces the post-unlearning alignment gap, it does not eliminate it entirely. There remains a residual gap between the original model’s alignment and the post-boosting model’s alignment. Whether this gap can be closed completely — or whether there is an inherent trade-off between unlearning and alignment — is a fundamental open question.