Pristine Models

More data, better models. But at what cost?

Train an image generation model on more data and it gets better at generating images. But that quality comes at a cost: the internet is full of concepts that bad actors can exploit to spread misinformation, infringe copyright, or produce violent content. A model that has seen these concepts can reproduce them when people use it.

So why not just delete the unsafe images?

Take a look at this photo.

The only problem here is the gun. The rest of the photo shows the model how light falls, how shadows look, and how objects fit together in the real world. Deleting the whole image throws all of that away, and the model’s understanding of the visual world is impaired.

Goal: leverage the whole dataset and learn everything except the unsafe concepts.

Given an unsafe model, what can we do?

Discarding the unsafe images is off the table, since we don't want to lose the quality they carry. So suppose we have a model that has already learned from everything. What can we do after the fact?

Safety at pretraining time

Ideally, we would find the parts of the network responsible for an unsafe concept and switch them off. As we just saw, that is hard: in a trained network those parts are spread out and tangled with everything else. So instead of searching for them after the fact, why not train the network to put them somewhere we can find?

This does not replace post-hoc methods; it complements them. Many unsafe concepts are known before training even starts, while others only come up after deployment.

The figure below illustrates exactly this goal.

Here, “knowledge” means the ability to reproduce the restricted concept as a whole. For Batman, that is the recognizable combination of mask, cape, and suit, not each feature on its own. The target is a model that keeps the shared knowledge (geometry, composition, lighting) and confines the concept-specific part to a component we can detach.

Pristine Attention

It comes down to controlling where information goes. From every unsafe example, we want two things:

Pristine Attention gets these separate modules from the attention layers, which are where words in the prompt get tied to parts of the image. It replaces each attention layer with three parallel copies, which we call branches: a safe branch, an unsafe branch, and a shared branch. Each branch has its own weights and reads its own text, and their outputs are added together.

The branches are trained to separate restricted concept information from knowledge we want to retain. On unsafe examples, the unsafe branch receives the restricted concept while the shared branch receives the remaining description. Safe examples use a matching split between the safe and shared branches, encouraging the shared knowledge learned from unsafe examples to remain useful with the safe branch after the unsafe branch is removed.

To illustrate, we start with object-centric, ImageNet-like data. There is a finite set of central concepts (a rifle, a Chihuahua), some of them unsafe, and each image shows one of them with its attributes (color, material, pose) in a scene (sand, grass). Open-world data like CC12M is covered at the end of this section.

Split the prompt

We control what each module learns by controlling what it reads. Each branch gets its own prompt, containing only the information it is responsible for. An LLM splits every training prompt in two: the concept goes to a detachable unsafe branch, and everything else goes to a shared branch.

Together the two prompts still describe the whole image, but each branch can only account for its own part. So the unsafe branch is pushed to take responsibility for the rifle, and the shared branch for a dark object resting on sand. That covers separate.

Make it useful: add a safe branch

In the rifle example, the shared branch learns attributes and setting: the color of an object and what it rests on. For that knowledge to be useful, the shared branch should play the same role when generating safe content. So we split safe prompts the same way, sending the identity to its own safe branch and the attributes and setting to the shared branch:

Now the shared branch has one job on every image, safe or unsafe: attributes and setting, never identity. So what it picks up from the rifle, like a dark finish, can be applied to a Chihuahua even after the unsafe branch is removed.

How it flows

Inside an attention layer, the image being generated is a grid of spatial tokens x, and the prompt is a sequence of text tokens t. With weights W, attention updates both:

In MM-DiT-style self-attention, the text tokens are updated too. In cross-attention, the text is only read, so (0, Δx) = Attn(t, x; W). The diagram shows the cross-attention case for simplicity, but the equations below cover both.

Pristine runs this once per branch, each with its own text and its own weights. The shared branch is always on, and exactly one mode-specific branch is active: safe for safe examples, unsafe otherwise.

Each branch’s text is updated independently:

Their spatial updates are added together:

Safe examples train the shared and safe branches; unsafe examples train the shared and unsafe branches. Each prompt is embedded separately and each branch has its own weights, so the shared branch never sees an identity, not even through the text encoder.

That is what makes the unsafe branch detachable. It is the only place the model ever linked “assault rifle” to the image, so deleting W_unsafe removes that link. Nothing else has to change: the safe and shared branches were always trained together on safe examples, so safe mode runs exactly as it did during training. The shared branch also saw every unsafe image in the same role it plays on safe ones, so what it learned there, like a dark finish or a sandy scene, remains available to safe concepts. The rest of the network and the training loss are unchanged.

Beyond object-centric data. In open-world data such as CC12M there is no single central concept: once the unsafe concept is removed, the shared prompt still describes nearly the whole scene. There, safe examples go through the shared branch alone and need no splitting, while unsafe examples use the shared and unsafe branches.

ImageNet

We test Pristine on ImageNet using the MMJiT self-attention architecture. In this controlled experiment, the 131 canine classes (dogs, foxes, and wolves) are designated safe and the remaining 869 classes unsafe. Full trains on all images; Safe trains only on safe images; Pristine trains on both and then detaches its unsafe branch.

Disentangled concepts

Pristine learns concepts in a disentangled way: the mode branch learns which object it is, and the shared branch learns everything else. So the model generates according to which branch is attached. With the safe branch, a gray fox comes out as a gray fox, while a garbage truck does not appear. With the unsafe branch, the pattern flips.

Bronze and algae dogs

The shared branch is there so that attributes learned from unsafe images stay usable on safe concepts. Bronze and algae appear mostly on unsafe classes and rarely on dogs. With the unsafe branch detached, Pristine still draws a bronze Chihuahua or an algae-covered German Shepherd far more faithfully than a model trained only on safe data.

Safety and quality, measured together

For safe prompts, we want the requested concept to appear; for unsafe prompts, we want it suppressed. A classifier measures how often the requested class is generated. Safe-class accuracy should be high, while unsafe-class accuracy should be low.

Our balanced scores combine those two goals with the harmonic mean. For a metric M, we first turn unsafe performance into a desirable quantity: M_safe should be high, and 1 − M_unsafe should also be high. The harmonic mean then rewards a model only when both are high. If either safe generation or unsafe suppression is weak, the balanced score remains low instead of allowing a strong result on one side to compensate for failure on the other.

We apply this to accuracy and to generative precision and recall. The table below is the MMJiT portion of the paper’s ImageNet evaluation, presented in the same three groups: quality on safe concepts, safety on unsafe concepts, and their balance.

The comparison is direct. Full has strong safe-concept quality, but it also retains unsafe concepts: its unsafe accuracy is 0.56. Safe brings unsafe accuracy to 0.00, but its safe accuracy falls to 0.79. Pristine also reaches 0.00 unsafe accuracy while raising safe accuracy to 0.91. It therefore achieves the highest balanced accuracy (0.95) and balanced precision (0.57).

Conclusion

The Pristine Model shows promise when scaling to CC12M. There is still room for further research in achieving the best of both worlds between the Safe and Full models, while developing safe systems with high generation quality.

Citation

Anirudh Sundara Rajan, Krishna Kumar Singh, Richard Zhang, and Yong Jae Lee. Pristine Models. arXiv:2610.14091, 2026.

@misc{sundararajan2026pristine,
  title         = {Pristine Models},
  author        = {Sundara Rajan, Anirudh and Singh, Krishna Kumar and Zhang, Richard and Lee, Yong Jae},
  year          = {2026},
  eprint        = {2610.14091},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV},
  doi           = {10.48550/arXiv.2610.14091},
  url           = {https://arxiv.org/abs/2610.14091}
}

References

1. Y. Cao and J. Yang. Towards making systems forget with machine unlearning. IEEE Symposium on Security and Privacy, 2015.
2. N. Kumari, B. Zhang, S.-Y. Wang, E. Shechtman, R. Zhang, and J.-Y. Zhu. Ablating concepts in text-to-image diffusion models. ICCV, 2023.
3. R. Gandikota, J. Materzynska, J. Fiotto-Kaufman, and D. Bau. Erasing concepts from diffusion models. ICCV, 2023.
4. R. Gandikota, H. Orgad, Y. Belinkov, J. Materzyńska, and D. Bau. Unified concept editing in diffusion models. WACV, 2024.
5. K. Deng, G. Li, Y. Xiao, B. Hui, and X. Ma. Forget many, forget right: Scalable and precise concept unlearning in diffusion models. arXiv:2601.06162, 2026.
6. O. Li, Y. Wang, X. Hu, H. Jiang, Y. Hao, J. Ma, F. Feng, et al. SPEED: Scalable, precise, and efficient concept erasure for diffusion models. ICLR, 2026.
7. N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, et al. Toy models of superposition. arXiv:2209.10652, 2022.
8. M. McCloskey and N. J. Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. Psychology of Learning and Motivation, 24, 1989.