Pristine Models

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.

한 줄 요약

Pristine Models는 텍스트-이미지 생성 모델에서 안전하지 않은 개념을 사전 학습 단계에서 분리하여 추론 시 제거할 수 있도록 설계된 새로운 아키텍처를 제안한다.

핵심 기여도

Pristine Attention이라는 새로운 주의 메커니즘을 도입하여 안전한/위험한 지식을 독립적인 브랜치로 분리.
ImageNet에서 안전성과 생성 품질을 동시에 개선, 전체 데이터 학습 모델보다 높은 텍스트-이미지 정렬 성능.
CC12M 데이터셋에서도 효과적으로 확장 가능, 대규모 오픈월드 사전 학습 환경에서 안정적 성능 유지.
기존의 post-hoc unlearning 접근 대비 카타스트로픽 포리겟 문제를 완화.

핵심 아이디어

기존 텍스트-이미지 생성 모델은 안전하지 않은 개념을 학습하면서 유용한 시각 정보(예: 조명, 구성)도 함께 배우게 되는데, 이를 제거하면 생성 품질이 저하된다. Pristine Models는 사전 학습 단계에서 안전한/위험한 개념을 독립적인 브랜치로 분리하여 추론 시 위험한 브랜치를 제거할 수 있도록 설계한다. 이는 post-hoc unlearning과는 달리, 학습 초기부터 지식 분리가 이루어지므로 카타스트로픽 포리겟을 방지할 수 있다. 주의 메커니즘(Pristine Attention)을 통해 텍스트 프롬프트를 안전/위험/공유 세 가지 브랜치로 분리하고, 공유 브랜치는 유용한 정보만 학습하도록 유도한다.

기술적 접근법

**Pristine Attention**: 텍스트-이미지 매핑을 위한 세 가지 브랜치(안전, 위험, 공유)로 분리.
**모드별 프롬프트 분해**: 각 훈련 샘플의 프롬프트를 안전/위험 브랜치 전용 또는 공유 브랜치 전용으로 분리.
**공유 브랜치**: 위험 데이터에서도 유용한 정보(예: 구성, 조명)를 학습.
**추론 시 분리**: 위험 브랜치를 제거하여 안전한 생성 가능.
**학습 데이터셋**: ImageNet, CC12M.
**모델 아키텍처**: Cross-attention 및 self-attention 기반 모델에서 검증.

주요 결과

**ImageNet**: Pristine Models는 안전한 데이터만 학습한 모델 대비 +12.3% 높은 텍스트-이미지 정렬 성능, 전체 데이터 학습 모델 대비 안전성 73.2% 개선.
**CC12M**: 대규모 오픈월드 데이터에서도 안전성 68.4% 유지, 생성 품질은 전체 데이터 학습 모델과 유사.
**성능 비교**: post-hoc unlearning 대비 1.6× 빠른 추론 속도, 카타스트로픽 포리겟 발생률 40% 감소.

의의 및 한계

Pristine Models는 사전 학습 단계에서 안전성과 품질을 동시에 보장하는 새로운 패러다임을 제시하며, 기존 post-hoc unlearning 접근의 한계를 극복한다. 특히, 대규모 데이터셋에서의 확장성을 보여주며, 텍스트-이미지 생성 모델의 실용적 안전성 강화에 기여할 수 있다. 그러나 사전에 제한 개념이 알려져야 하며, 추후 식별된 개념에는 여전히 post-hoc 방법이 필요하다는 한계가 있다.