🤖 AI

Safe Autoregressive Image Generation with Iterative Self-Improving Codebooks

This paper proposes an iterative self-improving framework that leverages a unified multimodal model's own judgment to identify unsafe generations and adaptively refine its visual codebook, thereby enhancing safety in autoregressive image generation without requiring external human feedback.

Original authors: Yunqi Xue, Zhijiang Li, Philip Torr, Jindong Gu

Published 2026-06-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yunqi Xue, Zhijiang Li, Philip Torr, Jindong Gu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very talented, all-knowing artist who can paint pictures based on your descriptions. This artist doesn't paint with a brush on a canvas; instead, they build images one tiny "pixel-block" at a time, like a master builder snapping together Lego bricks from a giant box. This box of bricks is called a codebook.

The problem is that sometimes, if you ask this artist to paint something specific (even if you don't mean to be harmful), they might accidentally pick a "bad brick" from the box that results in a scary, violent, or inappropriate picture.

This paper proposes a clever way to fix the artist's box of bricks so they can only build safe pictures, without needing a human teacher to constantly check their work. Here is how they do it, using a simple three-step story:

1. The Artist Becomes Its Own Teacher (Self-Reflection)

Usually, to teach an artist not to make bad pictures, you need a human to look at every painting and say, "No, that's too violent," or "Yes, that's safe." That takes forever.

Instead, this method lets the artist look at its own work.

The artist tries to paint a picture based on a prompt.
The artist then looks at the picture and asks, "Does this look dangerous or inappropriate?"
If the artist thinks, "Yes, that's bad," it creates a "bad example."
Then, it tries to paint a similar but safe version of that picture.
Now, the artist has a pair: a "Bad Picture" and a "Good Picture."

2. Finding the "Bad Bricks" (The Harmful Space)

The researchers realized that the difference between the "Bad Picture" and the "Good Picture" isn't random; it's hidden in the specific "bricks" (tokens) the artist chose.

They take the "Bad" bricks and the "Good" bricks and compare them.
They use math to find the specific direction in the box of bricks where the "badness" lives. Think of this as finding a specific shadowy corner in the storage room where all the dangerous Lego pieces are hiding.
They then lock that corner. They make sure the artist can never reach into that shadowy corner to grab a brick again. This is called projecting the codebook into a "harmless space."

3. Polishing the Art (Adaptive Fine-Tuning)

Here is the tricky part: If you just lock away the bad bricks, the artist might start making ugly or blurry pictures because they are missing some tools they need for normal art.

To fix this, the researchers let the artist practice painting safe pictures again, but with a strict rule: You can only use the tools that are NOT in the shadowy corner.
The artist learns to rebuild its skills using only the "safe" tools, smoothing out the rough edges and making the pictures look beautiful again.
This process is iterative, meaning they repeat the cycle: Try to paint -> Check for badness -> Lock the bad corner -> Practice safe painting -> Repeat. Each time, the artist gets better at avoiding the bad stuff while keeping the art high-quality.

The Result

By the end of this process, the artist has a new, upgraded box of bricks.

It can still paint beautiful, complex, and emotional masterpieces (like the "19th-century painting" or "dark emotional portrait" examples in the paper).
But if you ask it to paint something harmful (like violence or nudity), the "bad bricks" simply aren't there to be picked. The artist will instead paint a safe version of the scene.

In short: The paper teaches an AI artist to police itself. It finds its own mistakes, identifies exactly which "mental tools" cause those mistakes, removes those tools, and then relearns how to be a great artist using only the safe tools. The best part? It does all this without needing a human to label thousands of bad images for it.

Technical Summary: Safe Autoregressive Image Generation with Iterative Self-Improving Codebooks

Problem Statement

Autoregressive unified multimodal models generate images by sequentially predicting discretized visual tokens derived from a codebook, rather than operating in continuous latent spaces like diffusion models. While this architecture offers advantages in speed and text-conditioning consistency, it raises critical safety concerns: how safe are the images generated by these autoregressive systems? Existing safety mitigation methods for diffusion models (e.g., latent space guidance, post-generation filtering) rely on continuous representations and do not generalize well to the discrete token-based generation of autoregressive models. Furthermore, current approaches often require extensive human annotation or external feedback loops, which are resource-intensive and difficult to scale.

Methodology

The paper proposes Iterative Self-Improving Codebooks, a method that leverages the unified multimodal model's own capabilities to identify and eliminate harmful concepts without external human annotation. The approach consists of two core steps repeated iteratively until convergence:

1. Constructing the Harmful Space and Projection

Self-Generated Data Pairing: The model generates images from a set of prompts containing specific harmful attributes. Using the model's own understanding capabilities, it identifies which generated images are harmful. For each harmful prompt ( $x_u$ ), the model generates a semantically similar but safe alternative prompt ( $x_s$ ) by replacing unsafe terms with benign ones. This creates paired image-text datasets: $(x_u, y_u)$ for harmful and $(x_s, y_s)$ for safe.
Feature Difference Extraction: The model extracts embedding features ( $F_u$ and $F_s$ ) for the tokens corresponding to these prompts from its codebook. The difference between these feature vectors ( $F_u - F_s$ ) is computed to isolate the specific features associated with the harmful concept.
Subspace Construction: These difference vectors are aggregated and processed via Singular Value Decomposition (SVD) to extract the principal components. The top- $k$ singular vectors define a Harmful Subspace ( $P$ ) that captures the core features of the harmful concept (e.g., nudity, violence).
Hard Projection: The original codebook ( $W_{ori}$ ) is projected to remove components lying within this harmful subspace. The projected codebook ( $W_{proj}$ ) is calculated as:
$W_{proj} = W_{ori} (I - \lambda \cdot \text{Proj}_P(W_{ori}))$
This ensures the new embeddings are orthogonal to the harmful information.

2. Adaptive Fine-Tuning in the Harmless Space

Quality Restoration: The hard projection in Step 1 may degrade image quality by removing benign details. To counteract this, the method performs adaptive fine-tuning on the projected codebook.
Null Space Constraint: A perturbation matrix ( $\Delta$ ) is optimized to enhance image fidelity. Crucially, this perturbation is constrained to the null space of the difference matrix $E$ (the space orthogonal to the harmful subspace). This ensures that the fine-tuning process improves visual quality without reintroducing harmful information.
Final Codebook: The final safe codebook is constructed as $W_{safe} = W_{proj} + \Delta \cdot Q$ , where $Q$ is the projection matrix onto the null space.

The entire process is iterative: the updated codebook is used to generate new data, refine the harmful space, and further improve safety and quality.

Key Contributions

First Systematic Exploration: The authors are the first to systematically address the safety of image generation within autoregressive unified multimodal frameworks, distinguishing it from the well-studied diffusion-based safety landscape.
Iterative Self-Improving Mechanism: The paper introduces a method that uses the model's internal understanding and judgment capabilities to construct harmful/safe pairs and update the codebook, eliminating the need for annotated datasets or human feedback.
Codebook-Centric Safety: Unlike methods that modify generation trajectories or latent spaces, this approach directly modifies the discrete visual token representations (the codebook), offering a fundamental solution for autoregressive models.

Experimental Results

The method was evaluated on eight harmful-prompt datasets (including I2P, CoPro, ViSU, P4D, and MPUP) across five unified multimodal models (Janus series, VILA-U, Emu3, LlamaGen, OmniMamba).

Safety Improvement: The method significantly reduced the generation of inappropriate content across seven categories (Sexual, Violence, Hate, Self-harm, Illegal activity, Harassment, Shocking). For example, on the I2P dataset, the proportion of sexual content dropped from 0.12 (Baseline) to 0.04 (Safe-CB).
Iterative Efficacy: Experiments demonstrated that iterative removal outperforms single-turn removal. Safety improvements generally reached saturation after three iterations, with the harmful space becoming more refined over time.
Preservation of Capabilities: The method preserved the model's inherent image generation and text-understanding capabilities. Benchmarks like Geneval (generation) and MMMU (understanding) showed no significant degradation.
Semantic Fidelity: While the initial hard projection caused a drop in prompt-image alignment (measured by CLIP-Score and TIFA), the adaptive fine-tuning stage successfully recovered semantic consistency, bringing Safe-CB performance close to the baseline on benign prompts.
Generalization: The approach generalized well to out-of-distribution (OOD) data and various model sizes (from 1B to 7B parameters).

Significance and Claims

The paper claims that this work establishes a new paradigm for safety in autoregressive image generation. By leveraging the model's own self-judgment to iteratively refine its internal codebook, the method achieves safety enhancements without relying on external human annotation or retraining the entire model from scratch. The authors emphasize that this approach effectively eliminates harmful mappings while maintaining high-quality visual details, offering a scalable and self-sustaining solution for safe multimodal generation. The work highlights that safety in autoregressive models requires distinct strategies from diffusion models, specifically targeting the discrete nature of visual tokens.