Pristine Models

1 INTRODUCTION

• Excluding unsafe examples can discard useful visual information, while inference-time filtering leaves the model able to generate unsafe images and burdens guardrails.The introduction also notes that scaling post-hoc unlearning risks catastrophic forgetting as more concepts are removed.Source: §1 · p. 1
• The approach targets restricted concepts known before training and complements post-hoc unlearning for restrictions identified only afterward.The intended factorization retains benign information such as geometry, composition, lighting, and object interactions while isolating concept-specific information.Source: §1 · p. 2
• Pristine Attention separates safe, unsafe, and shared knowledge into independently parameterized branches; the unsafe branches can be detached at inference.Mode-specific prompts route safe or unsafe concepts to their respective branches, while shared prompts describe information intended to remain transferable.Source: §1 · p. 2
• Across ImageNet architectures and CC12M, Pristine suppresses restricted-concept generation while retaining more generation quality than safe-only training.The results indicate that pretraining-time safety can complement post-hoc methods when restricted concepts are known in advance.Source: §1 · p. 2

2 APPROACH

• The approach seeks to match the safe model’s safety while retaining the full model’s generation quality by separating restricted knowledge from useful safe knowledge.Pristine Attention implements this separation through independently parameterized safe, unsafe, and shared attention branches.Source: §2 · p. 3, §2.1 · p. 3

2.1.1 PROMPT DECOMPOSITION

• Prompt decomposition routes concept identity to a safe or unsafe prompt and transferable attributes and scene information to a shared prompt.The same decomposition is applied to safe examples so the shared branch plays the same role in both modes.Source: §2.1.1 · p. 3
• The shared branch learns properties such as material and scene context, while mode-specific branches specialize in safe or unsafe identities.This is intended to let properties learned from unsafe examples compose with safe identities after the unsafe branch is removed.Source: §2.1.1 · p. 3

2.1.2 ATTENTION MECHANISM

• Each branch attends to its own prompt; the shared update is combined with the active safe or unsafe update, while text representations are embedded independently.The design replaces existing attention layers and applies to cross-attention and MMDiT backbones; the cross-attention variant sums branch outputs into the shared spatial state.Source: §2.1.2 · pp. 3–4, Figure 3 · p. 4

2.2 TRAINING AND INFERENCE

• Prompt decomposition adapts to data structure: open-world CC12M uses a shared prompt for safe examples and adds an unsafe branch for unsafe examples, which is detached at inference.Unlike object-centric data, CC12M has no single dominant concept; its shared prompt can describe nearly the entire scene, and safe prompts require no decomposition.Source: §2.2 · pp. 4–5
• Table 1 reports retained-concept quality, restricted-concept generation, and balanced safety–utility metrics; Pristine has the strongest balanced accuracy and precision across both architectures.The caption specifies the ImageNet dog setting.Source: Table 1 · p. 5

2.3 TOY EXPERIMENT: MNIST DIGIT GENERATION

• In the toy digit setting, the shared branch learns thickness from both safe and unsafe digits, enabling thick safe digits while suppressing unsafe identities.Digits 0–4 are safe and 5–9 unsafe; thickness is routed to the shared branch and identity to the mode-specific branch.Source: §2.3 · p. 5

3.1 IMAGENET RESULTS

• Pristine leads Balanced Accuracy, Precision, and Recall on UNet (0.89, 0.56, 0.72), and Balanced Accuracy and Precision on MMJiT (0.95, 0.57), outperforming Safe and Full.On safe concepts, Pristine’s UNet FID/VQAScore are 8.35/0.63 and its MMJiT values are 7.32/0.75; implicit unsafe accuracy is 2.82% on UNet and 6.20% on MMJiT.Source: §3.1 · p. 6
• Activating the unsafe branch restores unsafe generation while degrading safe generation; Figure 4 illustrates the intended reversal across safe and unsafe prompts.Source: §3.1 · p. 6, Figure 4 · p. 6

3.2 EVALUATING ALTERNATE DESIGN CHOICES

• AttentionSwap reaches 0.61 Balanced Accuracy versus Pristine’s 0.95, while MLPSwap’s Balanced Accuracy/Precision are 0.91/0.50 versus Pristine’s 0.95/0.57.The authors conjecture that swapping the entire attention module makes AttentionSwap more brittle to prompting-context shifts; MLPSwap suppresses unsafe concepts less effectively.Source: §3.2 · p. 7
• Post-hoc Pristine (Unlearn) scores 0.93 Balanced Accuracy, 0.40 Precision, and 0.55 Recall, below pretraining-time Pristine’s 0.95, 0.57, and 0.68.The authors conjecture that replacing attention after pretraining encourages compatibility with surrounding representations, allowing some unsafe information to be relearned.Source: §3.2 · p. 7

3.3 COMPOSITIONAL GENERATION

• Pristine reaches TPIPS 0.7404, above Safe’s 0.6890 and below Full’s 0.7705, transferring materials learned from unsafe examples to safe dog breeds.Figure 5 qualitatively shows bronze and algae rendered more faithfully than Safe and comparably to Full.Source: §3.3 · p. 7, Figure 5 · p. 8

3.4 COMPARISON TO UNLEARNING BASELINES

• On the ImageNet Dog Setting, Pristine improves on UCE’s Balanced Accuracy/Precision/Recall (0.85/0.42/0.56) with 0.89/0.56/0.72, while reducing unsafe accuracy to 0%.Pristine also retains VQAScore 0.63 and FID 8.35, compared with UCE’s 0.31 and 12.43.Source: Table 2 · p. 8, §3.4 · p. 8

3.5 CC12M RESULTS

• Pristine reaches 82.31 on DPG-Bench and 5.02 on PRISM-Bench, while restricted generation falls to 13.29% versus Full’s 56.42% and Safe’s 5.87%.Figure 6 qualitatively shows suppression on unsafe prompts and Full-matching quality on safe concepts.Source: §3.5 · p. 9, Figure 6 · p. 9

• Inference-time filtering and safety guidance, alongside concept-erasure fine-tuning, suppress concepts the model has already learned; Pristine discourages learning restricted concepts in the first place.Source: §4 · p. 9
• As concept sets grow, post-hoc model editing risks catastrophic forgetting; Pristine instead separates concepts during pretraining.Source: §4 · p. 10
• Modular LLM pretraining (Roland et al., 2026) routes designated concepts through removable modules, while Pristine uses prompt decomposition to control shared versus detachable knowledge.The paper reports that analogous modular separation is less effective without explicit control over the factorization.Source: §4 · p. 10

• Pristine factorizes learned knowledge into safe, unsafe, and shared modules, allowing removal of the unsafe module while retaining useful information from the same data.Source: §5 · p. 10
• Across architectures, Pristine strongly suppresses restricted-concept generation while improving image quality, text alignment, and compositionality over safe-only training.Source: §5 · p. 10
• Post-hoc unlearning and inference-time guardrails remain necessary as new restrictions emerge and to mitigate harmful generations at deployment.Source: §5 · p. 10

• The appendix provides dataset and evaluation-set construction details, training configurations, and hyperparameters for the experiments.Source: A APPENDIX · p. 13
• It includes additional ablations and controls, along with extended qualitative results across a broader range of examples.Source: A APPENDIX · p. 13
• Across designated unsafe concepts, Pristine fails to reproduce the requested concept and instead generates semantically related safe alternatives, unlike Full.The Safe model is included for reference and was trained without the corresponding unsafe data.Source: Figure 7 · p. 14

• Pristine fails to generate designated unsafe concepts including copyrighted characters, weapons, and recognized artwork.Source: §A.1 · p. 14
• Pristine improves safe-prompt following over Safe, including counting, attribute binding, and imaginative generation.The figure contrasts diverse DPG-Bench and PRISM-Bench prompts; Safe often struggles with counting, compositionality, and imaginative details.Source: §A.1 · p. 14, Figure 8 · p. 15
• The qualitative results support the hypothesis that retaining access to unsafe training data can preserve useful generative capabilities.Source: §A.1 · p. 14

• Implicit prompts refer to a target class without naming it explicitly, and the examples compare how different models respond.Source: §A.2 · p. 15
• The section covers ImageNet dataset and evaluation details, along with training configurations.Source: §A.2 · p. 15
• It reports additional ablations and controls conducted for the ImageNet study.Source: §A.2 · p. 15

A.2.1 TRAINING DETAILS

• The Dog Setting designates 131 classes as safe and 869 as unsafe; a Random Setting instead samples 150 safe and 850 unsafe classes.The Dog Setting includes dog breeds, wild canids, foxes, hyena, and dogsled among its safe classes.Source: §A.2.1 · p. 15
• On implicit ImageNet prompts, Pristine avoids the restricted concept and falls back to a safe concept, unlike the Full baseline that generates it.The prompts evoke the concept without naming it, so this tests removal beyond exact class-name prompts.Source: Figure 9 · p. 16
• On unsafe prompts, Pristine suppresses the requested concepts, whereas MLPSwap reproduces them more faithfully.The examples use “coffeepot” and “daisy.”Source: Figure 10 · p. 16
• On safe prompts, Pristine preserves the requested dog breed, while AttentionSwap often drifts to another breed.Source: Figure 11 · p. 17
• The study trains from-scratch SD 1.5 U-Net and MMJiT denoisers, adding parallel safe/shared/unsafe projections to U-Net cross-attention.The pretrained SD 1.5 VAE and CLIP text encoder remain frozen; MMJiT uses a flow-matching objective.Source: §A.2.1 · p. 17

A.2.2 EVALUATION

• The evaluation measures generation quality, prompt alignment, restricted-concept leakage, and compositional generalization.Source: §A.2.2 · p. 17
• For quality and alignment, FID and VQAScore are measured on held-out ImageNet data, augmented with ImageNetV2 and matched in size across safe and unsafe distributions.The augmentation addresses the small number of safe-class validation images, especially for FID estimation.Source: §A.2.2 · p. 17
• Leakage evaluation uses an ImageNet classifier on plain-background prompts, including explicit class names and implicit descriptions that do not name the target.The explicit and implicit sets contain 10,000 and 15,000 prompts, respectively.Source: §A.2.2 · pp. 17, 19
• Qualitative examples show the safe branch suppressing designated unsafe concepts, which reappear when the unsafe branch is activated.Full and Safe models are included for reference.Source: Figure 12 · p. 18
• For compositional generalization, each safe concept is paired with a training-unseen material, and generated images are compared with synthesized references using TPIPS.The tested materials include moss, hay, gold, algae, resin, bronze, foam, ceramic, clay, and stone.Source: §A.2.2 · p. 19

A.2.3 CAPTION GENERATION

• Each ImageNet image is paired with an approximately 40-word dense caption generated using its class label and a fixed photographic-caption prompt.The prompt asks for visible subjects, appearance, lighting, framing, and scene relations.Source: §A.2.3 · p. 19
• A text-only Qwen3.5-27B pass decomposes each dense caption into a subject caption for identity and a residual caption for scene details.The model receives the ImageNet class name and dense caption and returns the two fields as a single-line JSON object.Source: §A.2.3 · p. 19
• The subject caption names only the category, excluding visual properties, background, pose, lighting, and other entities.The category must use the exact name from the category line.Source: §A.2.3 · p. 20
• The residual caption retains the scene and subject appearance while replacing subject mentions with “the object” and avoiding the category or its synonyms.It includes interactions, viewpoint, lighting, people, and background.Source: §A.2.3 · p. 20
• During training, the subject caption feeds the safe/unsafe branch and the residual caption feeds the shared branch.Source: §A.2.3 · p. 20

A.2.4 PROMPT SPLITTING FOR IMPLICIT DATA

• For implicit prompts, the model splits an already-anonymized caption without receiving the true class name; this decomposition is not used for open-world data.The splitter reuses a generic hypernym already present in the caption rather than re-identifying the class.Source: §A.2.4 · p. 20
• The subject caption contains only the caption’s existing generic noun, while the residual replaces its mentions with “the object” and preserves other words.The instruction forbids guessing the specific subject or removing attached descriptive details.Source: §A.2.4 · p. 21
• Examples preserve modifiers attached to the subject, leave incidental brand names untouched, and keep plural counts in the residual caption.Source: §A.2.4 · p. 21
• With matched parameter count, the Safe baseline’s FID, VQAScore, Precision, and Recall remain essentially flat from d = 17 to d = 28, while Pristine still leads on VQAScore and Recall.Source: Table 4 · p. 22

A.3.1 COMPARISON WITH COMPUTE EQUALIZED SAFE BASELINE

• At matched capacity, Safe and compute-equalized Safe remain comparable across FID (8.23 vs. 8.46), VQAScore (0.63 vs. 0.62), and precision/recall (0.76/0.56 vs. 0.79/0.52).The compute-equalized Safe baseline has 416.0M parameters versus Pristine’s 420.8M; TwoBranch matches depth-17 Safe inference compute and exhibits the same quality gains over Safe.Source: §A.3.1 · p. 22

A.3.2 RANDOM SETTING

• MMJiT Pristine’s Balanced Accuracy, Precision, and Recall are 0.56, 0.43, and 0.60, while explicit unsafe accuracy falls from 53% to near zero.On U-Net, Balanced Recall improves to 0.64, and Balanced Accuracy and Precision match Safe (0.70, 0.43); safe-concept quality remains close to Full and better than Safe on both architectures.Source: §A.3.2 · p. 22, Table 5 · p. 23

A.3.3 PROMPT SPLITTING CONTROLS INFORMATION SHARING

• On held-out object–material compositions, Attribute-Shared achieves higher compositional similarity than Attribute-Restricted (0.7112 vs. 0.6814), supporting greater attribute transferability when attributes are explicitly routed through the shared prompt.Source: §A.3.3 · p. 23, Table 6 · p. 24

A.3.4 PRISTINE (TWOBRANCH)

• TwoBranch stays close to three-branch: safe FID is 7.38 vs. 7.32, VQAScore 0.7206 vs. 0.7479, TPIPS 0.7373 vs. 0.7404, and explicit unsafe accuracy 1.06% vs. 0.00%.Against Safe, TwoBranch improves safe FID from 8.23 to 7.38, VQAScore from 0.6286 to 0.7206, and TPIPS from 0.6890 to 0.7373; it also exceeds AttentionSwap and MLPSwap TPIPS scores of 0.7283 and 0.7272.Source: §A.3.4 · p. 24

A.5 CC12M STUDY

• The CC12M study uses dense LLaVA-NeXT recaptions at 512×512 and divides the data into safe and unsafe pools using 11 MLLM-tagged restricted categories.Roughly 38% of image–caption pairs are tagged with at least one restricted category.Source: §A.5 · pp. 24–25, Table 7 · p. 25

A.4 PRISTINE SELF-ATTENTION

• In Pristine self-attention, each active branch jointly updates text and spatial tokens, keeps text states branch-specific, and combines spatial updates into a shared representation.Figure 14 depicts the spatial updates being summed into that common representation.Source: §A.4 · p. 24, Figure 14 · p. 25

A.5.1 PROMPT DECOMPOSITION

• For cluttered open-world images, CC12M decomposition uses handles to caption restricted concepts separately from the shared scene description, which retains transferable details.Safe examples use the whole prompt in the shared branch; unsafe examples route the unsafe and shared captions to their respective branches, while inference needs no prompt decomposition.Source: §A.5.1 · pp. 25–26

A.5.2 EVALUATION DETAILS; A.5.3 JUDGE DETAILS

• Evaluation tests 5,500 held-out prompts across 11 categories, counting a requested category as generated only when a multilabel judge clearly identifies it.The judge considers all categories together, requires positive justification for selected labels, and performs a second self-consistency pass.Source: §A.5.2 · p. 26, §A.5.3 · p. 26