Pristine Models

AI summary

Overview

• Research area: Computer Vision — safety and controllability of text-to-image diffusion/flow models, specifically intervening at pretraining time rather than after training.
• Technical level: Intermediate (readers should know roughly how cross-attention and diffusion/flow training work, but the paper's central idea — splitting attention into removable branches — is intuitive).
• Scope: The paper proposes Pristine Models, a text-to-image architecture whose attention layers are split into separately parameterized safe, unsafe, and shared branches so that the unsafe branch can be deleted at inference, and evaluates it on ImageNet (1M images) and CC12M.

What This Paper Is About

Text-to-image models trained on broad, unfiltered data produce high-quality images, but they also learn to reproduce concepts that should not be generated (weapons, copyrighted characters, adult material). Simply deleting unsafe images from the training set makes the model safer but also throws away genuinely useful information those images contain — lighting, geometry, composition, object interactions — which degrades quality. The authors ask whether this knowledge can instead be compartmentalized during pretraining, so that the unsafe part can simply be switched off at generation time while everything else is retained.

Key Contributions

1. A pretraining-time framing of concept suppression. Rather than removing concepts from an already-trained model (post-hoc unlearning), the authors target concepts known in advance and try to prevent the unsafe knowledge from becoming entangled with useful knowledge in the first place.
2. Pristine Attention, an attention mechanism with three independently parameterized branches — safe, unsafe, and shared — where the shared branch's update is always added to the active safe or unsafe update. It replaces existing attention layers in both cross-attention (Stable Diffusion UNet) and MM-DiT-style self-attention (MMJiT) backbones without changing the rest of the network or the training objective.
3. Prompt decomposition, an LLM-based procedure that splits each training caption into a mode-specific prompt (routed to the safe or unsafe branch) and a shared prompt describing only what should remain transferable. This gives explicit control over what is inherited from unsafe data, which the authors show is critical compared with simpler parameter-separation schemes.
4. Validation at two scales: the object-centric 1M-image ImageNet benchmark (with a "Dog Setting" of 131 safe canine classes and 869 unsafe classes) and the larger open-world CC12M corpus, including comparisons against large-scale unlearning baselines (UCE, ScaPre, SPEED, SAFREE) and a test of compositional attribute transfer.

Main Findings

• Pristine beats both naive baselines on the safety–utility tradeoff (ImageNet). On the Stable Diffusion UNet, Pristine (with the safe branch) reaches Balanced Accuracy / Precision / Recall of 0.89 / 0.56 / 0.72, versus 0.88 / 0.53 / 0.70 for the Safe-only baseline and 0.47 / 0.37 / 0.31 for the Full-data baseline. On MMJiT it reaches 0.95 Balanced Accuracy and 0.57 Balanced Precision, the best of any configuration in Table 1.
• The unsafe branch really does hold the restricted concepts. Activating the unsafe branch instead restores unsafe generation while collapsing safe generation (e.g., MMJiT safe-concept accuracy drops to 0.00 and unsafe-concept accuracy rises to 0.62), evidence that the branch specialized in the concepts it was meant to hold.
• Quality is largely retained. On the UNet, Pristine reaches FID 8.35 and VQAScore 0.63, close to Full (7.92, 0.68) and well above Safe (10.72, 0.48); on MMJiT it reaches 7.32 and 0.75 versus Full (6.68, 0.79) and Safe (8.23, 0.63).
• Suppression holds under implicit prompts. Without naming the concept, Pristine keeps implicit unsafe accuracy at 2.82% on the UNet (Safe 2.80%, Full 32.30%) and 6.20% on MMJiT (Full 40.21%).
• Prompt decomposition matters. Removing it hurts: AttentionSwap (swapping the whole attention module between modes) reaches only 0.61 Balanced Accuracy, and MLPSwap (an extra MLP expert for unsafe examples) reaches 0.91 Balanced Accuracy and 0.50 Balanced Precision, both below Pristine's 0.95 / 0.57.
• Post-hoc use of the same mechanism is weaker. Pristine (Unlearn), which starts from the Full model and retrains with Pristine Attention, reaches 0.93 / 0.40 / 0.55 Balanced Accuracy / Precision / Recall, versus 0.95 / 0.57 / 0.68 when the factorization is learned during pretraining.
• Compositional transfer works. For 10 materials common in the unsafe split but rare in the safe split, TPIPS scores against Qwen-Image references were Full 0.7705, Pristine 0.7404, Pristine (Unlearn) 0.7484, AttentionSwap 0.7283, MLPSwap 0.7272, and Safe 0.6890 — so material knowledge from unsafe data remains usable on safe concepts.
• Pretraining-based separation outperforms large-scale unlearning. On the ImageNet Dog Setting, the strongest unlearning baseline (UCE) reaches 0.85 / 0.42 / 0.56 Balanced Accuracy / Precision / Recall versus Pristine's 0.89 / 0.56 / 0.72, while UCE's safe-generation quality drops to VQAScore 0.31 and FID 12.43, versus 0.63 and 8.35 for Pristine.
• The tradeoff persists but improves at open-world scale. On CC12M, with 11 restricted categories and roughly 38% of the data labeled restricted by an MLLM, Full generates the requested restricted concept on 56.42% of restricted prompts (DPG-Bench 84.21, PRISM-Bench 5.59), Safe on 5.87% (79.30, 4.79), and Pristine on 13.29% (82.31, 5.02).
• A toy MNIST check confirms the mechanism. Treating digits 0–4 as safe and 5–9 as unsafe, with thick digits appearing only among the unsafe classes, the model generates thick 0–4s (never seen in training) while not generating 5–9s.

Methodology in Plain English

The dataset is split into examples that contain a restricted concept (unsafe) and examples that do not (safe). For each training caption, an LLM rewrites it into two prompts: one that names only the central concept ("An assault rifle" or "A Chihuahua"), and one that describes everything else in neutral terms ("A dark object resting on sand"). The same decomposition is applied to safe examples, so the shared branch plays an identical role in both modes.

Each attention layer is then tripled. The shared branch always runs, attending between the shared prompt and the image tokens. Exactly one mode-specific branch runs alongside it — the safe branch for safe examples, the unsafe branch otherwise. Text tokens are embedded separately per branch to limit leakage; the spatial updates from the active branches are summed into a single shared spatial state, so the shared knowledge can compose with either mode. The rest of the network and the standard generative training objective are untouched.

Because the branches have their own parameters, the unsafe branch can simply be deleted at inference; prompts are decomposed the same way at test time, and generation proceeds through the shared (and, when relevant, safe) branches. On object-centric data like ImageNet the three-way split is used. On open-world data like CC12M there is no single dominant concept, so the shared prompt already describes nearly the whole scene; the authors find it sufficient to route safe examples through the shared branch alone, adding the unsafe branch only for unsafe examples. Evaluation uses a classifier for whether the intended concept appeared, precision/recall for quality and diversity, FID and VQAScore for overall quality and prompt alignment, and a "balanced" harmonic mean of safe-concept performance and unsafe-concept suppression.

Why This Matters

Impact on research. The paper reframes safety in generative models as an architectural and pretraining question rather than a purely post-hoc editing problem. It offers evidence that isolating restricted knowledge during training avoids the catastrophic forgetting that grows as more concepts are unlearned, and it shows the idea survives scaling from a controlled ImageNet setting to open-world CC12M — a step many concept-removal papers do not take.

Real-world applications:

• Deploying image generators that cannot reproduce copyrighted characters or branded IP, while still benefiting from the lighting, composition, and style cues in the broader training data.
• Building consumer or enterprise image tools with removable capability modules, so a single pretrained backbone can ship in permissiveness tiers for different markets or customers.
• Reducing reliance on inference-time guardrails and the costly reject-and-regenerate loops they cause, since the model is less likely to produce prohibited content in the first place.
• Supporting compliance workflows where the set of prohibited concepts is known in advance (weapons, adult material, protected figures) and needs to be auditable and switchable.

Industry relevance. The method drops into existing attention layers of both UNet-style and MM-DiT-style backbones, does not change the training objective, and produces a single artifact with a detachable module — all attractive properties for organizations training or fine-tuning large image models.

Future Directions

• Handling concepts that are identified after pretraining. The authors are explicit that Pristine complements rather than replaces post-hoc unlearning and inference-time guardrails, and their Pristine (Unlearn) result (0.93 / 0.40 / 0.55) suggests post-hoc reuse of the architecture is weaker than pretraining with it.
• Closing the remaining gap to the Safe model on suppression (13.29% restricted generation versus 5.87% on CC12M) without giving up the quality gains (82.31 versus 79.30 on DPG-Bench).
• Reducing dependence on a clean, reliable safe/unsafe partition and on LLM-based prompt decomposition, which the paper assumes; the CC12M setup relies on an MLLM to label roughly 38% of the corpus as restricted.
• Extending the factorization beyond attention — the experiments isolate unsafe knowledge in attention branches only — and testing whether the approach scales further in data, model size, and number of restricted concepts.

Target Audience

Researchers and engineers working on generative image models, AI safety, or machine unlearning, especially those who train or fine-tune diffusion/flow models and can modify architecture rather than only post-process outputs. It is also useful for policy- and product-oriented readers who want a concrete picture of what "safety built into pretraining" can and cannot buy you. Readers without familiarity with attention mechanisms and diffusion training will find the conceptual argument accessible but the experimental section dense.

Authors’ abstract

Training on broad, unfiltered data enables high-quality image generation, but can also teach models concepts that should not be generated; restricting the training data ensures safety, but reduces quality because examples containing unsafe concepts also provide useful knowledge that is not itself unsafe. Existing approaches largely address this trade-off post hoc through concept unlearning, which can lead to catastrophic forgetting as the number of concepts to remove grows. We instead ask whether unsafe concept knowledge can be compartmentalized during pretraining so that access to it can be disabled at inference. We introduce Pristine Models, a novel text-to-image generation architecture with an attention mechanism that encourages the separation of safe and unsafe knowledge into independently parameterized branches. The branch-specific parameters are learned during training, while the unsafe branch can be detached at inference. On ImageNet, Pristine Models enable safer generation while retaining much of the utility of full-data pretraining, achieving higher image quality and better text alignment than training only on safe data or applying post-hoc concept removal. We further scale our approach to CC12M, demonstrating that it remains effective in a substantially larger and more diverse open-world pretraining setting.