Pristine Models

Anirudh Sundara Rajan; Krishna Kumar Singh; Richard Zhang; Yong Jae Lee

cs.CVOct 8, 2026

Pristine Models

Authors: Anirudh Sundara Rajan, Krishna Kumar Singh, Richard Zhang, Yong Jae Lee

Organizations: University of Wisconsin–Madison · Adobe Research

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.

Figures & tables

Figure 1: An image is rarely unsafe in its entirety. A single image can entangle an unsafe concept that we wish to suppress at generation time, here a firearm, with safe , broadly useful properties like lighting, shadows, camera angle, and object interactions.

Figure 2: Goal. The Full Model ( Wang et al., 2026b ) , trained on both safe and unsafe data, generates high-quality images but can reproduce restricted concepts such as SpongeBob; the Safe Model, trained only on safe data, avoids them at the cost of useful visual knowledge, here reflected in incorrect gift-wrap placement. Pristine aims to retain the useful information learned from unsafe data while isolating restricted concept-specific knowledge in a detachable branch, achieving safe generation without sacrificing quality.

Figure 3: Pristine Cross Attention. Branch-specific cross-attention updates are summed into the shared spatial representation, while text tokens remain unchanged. MM-DiT-style self-attention follows the same high-level design; see Fig. 14 and Appendix A.4 for details.

Stable Diffusion UNet (Cross Attention)
	Quality (safe concepts)			Safety (unsafe concepts)			Balanced
Model	Acc $\uparrow$	P $\uparrow$	R $\uparrow$	Acc $\downarrow$	P $\downarrow$	R $\downarrow$	$\overline{\text{Acc}}\uparrow$	$\overline{\text{P}}\uparrow$	$\overline{\text{R}}\uparrow$
Safe baseline	0.80	0.54	0.57	0.02	0.49	0.09	0.88	0.53	0.70
Full baseline	0.82	0.63	0.68	0.67	0.74	0.80	0.47	0.37	0.31
Pristine (w/ safe branch)	0.81	0.59	0.70	0.00	0.46	0.25	0.89	0.56	0.72
Pristine (w/ unsafe branch)	0.04	0.31	0.29	0.42	0.75	0.67	0.07	0.27	0.31

Table 1: ImageNet Evaluation (Dog Setting). ‘Quality’ columns measure retained-concept generation quality (Accuracy/Precision/Recall

\uparrow

), while ‘Safety’ columns measure restricted-concept generation (Accuracy/Precision/Recall

\downarrow

). ‘Balanced’ columns summarize the safety–utility tradeoff. Pristine achieves the strongest Balanced Accuracy and Precision across both architectures while maintaining strong generative quality.

Figure 4: Given a safe prompt (“gray fox”), the safe branch succeeds while the unsafe branch fails; on the unsafe prompt (“garbage truck”), the pattern reverses, which is the intended behavior.

Figure 5: Compositional generation. For “bronze” and “algae”—materials rare on safe dogs but common on unsafe classes— Pristine (w/ safe branch) renders them far more faithfully than Safe and comparably to Full , showing that attributes learned from unsafe data transfer to safe concepts.

	Quality (safe concepts)			Safety (unsafe concepts)			Balanced
Model	Acc $\uparrow$	P $\uparrow$	R $\uparrow$	Acc $\downarrow$	P $\downarrow$	R $\downarrow$	$\overline{\text{Acc}}\uparrow$	$\overline{\text{P}}\uparrow$	$\overline{\text{R}}\uparrow$
ScaPre ( Deng et al., 2026 )	0.66	0.62	0.57	0.30	0.74	0.56	0.68	0.36	0.49
UCE ( Gandikota et al., 2024 )	0.79	0.62	0.47	0.07	0.68	0.31	0.85	0.42	0.56
SPEED ( Li et al., 2026 )	0.82	0.69	0.53	0.24	0.73	0.48	0.79	0.39	0.53
SAFREE ( Yoon et al., 2025 )	0.73	0.65	0.61	0.28	0.74	0.68	0.73	0.37	0.42
Pristine (w/ safe branch)	0.81	0.59	0.70	0.00	0.46	0.25	0.89	0.56	0.72

Table 2: Comparison against state-of-the-art unlearning baselines on the ImageNet Dog Setting. Pristine achieves the strongest overall safety–utility balance, substantially improving the balanced metrics while retaining strong safe-generation quality.

	Safety $\downarrow$	Generative Quality $\uparrow$
Model	Restricted Gen. (%)	DPG	PRISM
Safe	5.87	79.30	4.79
Full	56.42	84.21	5.59
Pristine	13.29	82.31	5.02

Table 3: Safety and generative quality on CC12M. Restricted Gen. is the fraction of prompts explicitly requesting a restricted concept for which that concept is generated (

\downarrow

); DPG-Bench and PRISM-Bench measure prompt following on safe prompts (

\uparrow

Panel of Figure 6 — Figure 6: CC12M generations. Left: on unsafe prompts, Full ( Wang et al., 2026b ) produces the requested concept while Safe and Pristine suppress it. Right: on safe concepts, Pristine matches Full ’s generation quality. Prompts are truncated for readability; all generations use the full prompts. More examples can be found in Appendix A.1 .

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Qualitative suppression of designated unsafe concepts. Across a diverse set of designated unsafe concepts, the Full model ( Wang et al., 2026b ) readily generates the requested concept, whereas our Pristine model successfully and consistently fails to reproduce it, instead generating semantically related but safe alternatives. The Safe model is shown for reference, trained without the corresponding unsafe data.

Figure 8: Improved quality on safe generations. Across diverse prompts from DPG-Bench and PRISM-Bench, the Safe model often struggles with counting, compositionality, and imaginative details, while our Pristine model generates them more faithfully. The Full model is shown as a reference for the highest achievable generation quality.

Figure 9: Robustness to implicit prompting on ImageNet (Dog Setting). Each model is prompted with descriptions that semantically evoke an unsafe class without naming it (e.g., “a bright orange, faceted vessel” for a coffee pot, “a bright green reptile” for a snake). The Full baseline generates the concept from these implicit cues, confirming the prompts are sufficient to elicit it; the Safe baseline, trained without the unsafe data, cannot; and our Pristine, though trained on that data, also does not—falling back to a safe concept (a dog). Removal therefore generalizes beyond prompts that contain the exact class name.

Figure 10: On unsafe prompts (“coffeepot”, “daisy”), Pristine suppresses the requested concepts, whereas MLPSwap reproduces them more faithfully, consistent with its higher unsafe-concept leakage.

Figure 11: On safe prompts, Pristine preserves the requested dog breed, while AttentionSwap often drifts to a different breed, illustrating its weaker and more brittle safe-concept generation.

Figure 12: Qualitative compartmentalization of designated unsafe concepts. Across diverse unsafe prompts, the Pristine safe branch suppresses the requested concept, while activating the unsafe branch restores it, suggesting that the unsafe branch has specialized to generating the restricted content as intended. The Full and Safe models are shown for reference.

Model	FID $\downarrow$	VQAScore $\uparrow$	P $\uparrow$	R $\uparrow$
Safe baseline (original, $d=17$ )	8.23	0.63	0.76	0.56
Safe baseline (compute-equalized, $d=28$ )	8.46	0.62	0.79	0.52
Pristine (w/ safe branch)	7.32	0.75	0.77	0.60

Appendix

Table 4: Generative quality on the Safe (retained-concept) split. Giving the Safe baseline enough extra depth to match Pristine’s parameter count does not close the quality gap. FID/VQAScore/Precision/Recall stay essentially flat (within noise) between

d=17

and

d=28

, while Pristine still leads on VQAScore and Recall.

Figure 13: Given a safe prompt (“electric guitar”), our Pristine safe branch succeeds while the unsafe branch fails; on the unsafe prompt (“cauliflower”), the pattern reverses, which is the intended behavior. The Safe baseline struggles to generate the safe “electric guitar” well.

MMJiT (Self Attention)
	Quality (safe concepts)			Safety (unsafe concepts)			Balanced
Model	Acc $\uparrow$	P $\uparrow$	R $\uparrow$	Acc $\downarrow$	P $\downarrow$	R $\downarrow$	$\overline{\text{Acc}}\uparrow$	$\overline{\text{P}}\uparrow$	$\overline{\text{R}}\uparrow$
Safe baseline	0.25	0.68	0.39	0.01	0.69	0.15	0.40	0.42	0.54
Full baseline	0.57	0.69	0.68	0.53	0.80	0.79	0.51	0.31	0.32
Pristine (w/ safe branch)	0.39	0.66	0.61	0.01	0.68	0.41	0.56	0.43	0.60
Stable Diffusion UNet (Cross Attention)

Appendix

Table 5: ImageNet Evaluation (Random Setting). ‘Quality’ columns measure retained-concept generation quality (Accuracy/Precision/Recall

\uparrow

), while ‘Safety’ columns measure restricted-concept generation (Accuracy/Precision/Recall

\downarrow

Prompt Decomposition	Overall Similarity $\uparrow$
Attribute-Shared	0.7112
Attribute-Restricted	0.6814

Appendix

Table 6: Effect of prompt decomposition on compositionality. Overall similarity is measured on held-out object–material compositions; higher is better.

Figure 14: Pristine Self-Attention. For MMDiT-style architectures, each active branch jointly updates its text and spatial tokens. Spatial updates are combined into the shared representation, while text states remain branch-specific.

Category	Representative instance
Dogs	golden retriever
Dinosaurs	Tyrannosaurus rex
Mythical / fantasy creatures	dragon
Children / infants	toddler
Firearms & gun-like objects	handgun
Bladed weapons	katana

Appendix

Table 7: The 11 restricted (unsafe) categories used in the CC12M study, each with a representative instance. A text-image pair is excluded from the safe training pool if a strong MLLM tags it with any of these categories.

Safe T2I Generation

Safe text-to-image generation develops methods to prevent harmful images while preserving useful, faithful outputs.

See all 53 papers

No Safe Dose: How Training Data Drives Unsafe Image Generation
May 27, 2026Felix Friedrich, Lukas Helff, Niharika Hegde +2Safe T2I Generation Data Contamination
1Black Forest Labs · 2TU Darmstadt & hessian.AI · 3DFKI +1
Haste Makes Waste? Towards Safe Few-Step Text-to-Image Models
Oct 27, 2026Lior Ziv, Mahmood SharifT2I Generation Safe T2I Generation
Tel Aviv University
Introspective Attention Modulation for Safe Text-to-Image Generation
Jul 16, 2026Basim Azam, Hossein Rahmani, Naveed AkhtarDiffusion Transformer T2I Generation
The University of Melbourne, Australia · Lancaster University, United Kingdom

May 27, 2026J/K move · Enter open · S save

No Safe Dose: How Training Data Drives Unsafe Image Generation

Felix Friedrich, Lukas Helff, Niharika Hegde, Patrick Schramowski+1

1Black Forest Labs · 2TU Darmstadt & hessian.AI · 3DFKI+1

Safe T2I Generation Data Contamination T2I Generation Evaluation T2I Diffusion Models

Open paper AskPDF arXiv

Text-to-image models trained on large-scale data often inevitably ingest unsafe content. While some people observe input-output amplifications, it remains unclear whether and how training data composition directly drives model output safety or by other factors. We shed light on this question by isolating this variable: we train the same text-to-image model on datasets that differ only in their fraction of unsafe images (0% to 9.6%), across several dataset scales (100K to 8M). Then we generate images with the resulting models, and evaluate them with four independent safety classifiers. Output unsafety rises monotonically from 16.6% at 0% contamination to 25.5% at 5%. A factorial design reveals that the proportion, not the absolute count, of unsafe training images is the operative variable. The 16.6% irreducible baseline at zero contamination implicates the other components, e.g. frozen text encoder, as a residual safety risk -- confirmed by a text encoder ablation showing that SafeCLIP reduces this floor to 9.6%, while the dose-response effect persists across all three encoders tested. Critically, no quality degradation in terms of FID, CLIPscore and ImageReward accompanies safety filtering. These results establish that data curation and text encoder safety are complementary and independently effective interventions. At the same time, the remaining level of unsafety poses questions for future research about emerging capabilities and compositionality.

Abstract

Safe T2I Generation

Explore similar work

No Safe Dose: How Training Data Drives Unsafe Image Generation

Haste Makes Waste? Towards Safe Few-Step Text-to-Image Models

Introspective Attention Modulation for Safe Text-to-Image Generation