Pristine Models가 유해 개념을 사전학습에서 분리해 추론 시 떼어낸다

무엇인가

텍스트-이미지 생성 모델은 넓고 필터링되지 않은 데이터로 학습할수록 품질이 좋아지지만, 생성되어서는 안 되는 개념까지 함께 배운다. 반대로 안전한 데이터만 골라 학습하면 품질이 떨어진다. 유해 개념이 담긴 이미지에도 기하학, 구도, 조명, 재질, 객체 간 상호작용 같은 그 자체로는 무해한 지식이 들어 있기 때문이다. 기존 접근은 대부분 학습이 끝난 뒤 개념 언러닝으로 이 문제를 사후 처리하는데, 제거할 개념 수가 늘어날수록 파국적 망각(catastrophic forgetting)이 발생한다. 저자들은 여기서 관점을 바꿔, 유해 개념 지식을 사전학습 단계에서 구획화해 두고 추론 시점에 접근만 차단할 수 있는지 묻는다.

어떻게 동작하나

제안 방법의 핵심은 Pristine Attention이다. 어텐션 레이어를 안전(safe), 유해(unsafe), 공유(shared) 세 개의 독립 파라미터 브랜치로 나눈다. 각 학습 프롬프트는 두 갈래로 분해된다. 모드 특정 프롬프트는 안전 또는 유해 개념의 정체성을 담당하고, 공유 프롬프트는 전이 가능한 속성과 장면 정보를 담당한다. 예를 들어 "A dark assault rifle resting on sand."는 유해 브랜치용 "An assault rifle."과 공유 브랜치용 "A dark object resting on sand."로 쪼개진다. 안전 예시도 같은 방식으로 분해해("A brown Chihuahua standing on grass." → "A Chihuahua." + "A brown object standing on grass.") 공유 브랜치가 양쪽 모드에서 같은 역할을 하도록 만든다. 정보 누출을 막기 위해 각 프롬프트는 독립적으로 토큰화·임베딩되고(t_b = Embed(c_b)), 각 브랜치가 공통 공간 상태 x에 대해 어텐션을 수행한다. 텍스트 상태는 브랜치별로 따로 갱신되고, 공간 기여는 x ← x + Δx_shared + Δx_mode로 합산된다. cross-attention 기반 SD UNet과 MM-DiT 스타일 self-attention 양쪽에 적용 가능하며, 어텐션 레이어만 교체하고 나머지 네트워크와 학습 목적함수는 그대로 둔다.

무엇과 다른가

학습은 안전 데이터와 유해 데이터에 대해 각각 손실을 계산하되 공유 파라미터를 함께 최적화하는 형태다. θ_pristine = argmin [L_gen({φ_shared, φ_safe}; D_safe) + L_gen({φ_shared, φ_unsafe}; D_unsafe)]. 추론 시에는 유해 브랜치를 떼어내고 나머지만 사용한다. 객체 중심 데이터에서는 3분할을 쓰지만, CC12M 같은 오픈월드 데이터에서는 지배적인 단일 개념이 없으므로 안전 예시는 W_shared만 통과시키고 유해 예시만 W_unsafe를 추가로 사용한다. 프롬프트 분해는 LLM으로 수행한다. 개념 검증용 토이 실험도 있다. 숫자 0~4를 안전, 5~9를 유해로 가정하고 굵은 숫자가 유해 클래스에만 등장하게 만들면, 공유 브랜치가 굵기라는 속성을 학습해 학습 중 한 번도 본 적 없는 굵은 0~4를 생성하면서 5~9는 생성하지 않는다.

어떻게 쓰나

ImageNet 실험은 131개 견종을 안전, 나머지 869개 클래스를 유해로 지정한 Dog Setting과 무작위 150개 안전 클래스의 Random Setting으로 구성했다. 캡션은 Gemini 3.1 Pro Preview로 생성했고, 베이스라인은 전체 데이터로 학습한 Full과 안전 데이터만 학습한 Safe다. 안전 개념 성능과 유해 개념 억제를 동시에 보기 위해 조화평균 기반 Balanced-M = H(M_safe, 1 − M_unsafe)를 정의했다. Pristine(w/ safe branch)은 UNet에서 Balanced Accuracy/Precision/Recall 0.89/0.56/0.72, MMJiT에서 Balanced Accuracy/Precision 0.95/0.57로 두 베이스를 모두 앞섰다. 안전 개념 품질도 UNet에서 FID 8.35, VQAScore 0.63으로 Full(7.92, 0.68)에 근접하고 Safe(10.72, 0.48)를 크게 웃돌았다. MMJiT에서는 7.32/0.75로 Full(6.68, 0.79)에 근접, Safe(8.23, 0.63)를 앞선다. 클래스 이름을 숨긴 암시적 프롬프트에서도 유해 정확도가 UNet 2.82%(Safe 2.80%, Full 32.30%), MMJiT 6.20%(Full 40.21%)로 억제가 유지됐다.

전제와 한계

설계 선택을 검증하는 절제 실험도 있다. 프롬프트 분해 없이 어텐션 모듈 전체를 모드별로 스왑하는 AttentionSwap은 유해 억제는 되지만 Balanced Accuracy가 0.61로 떨어졌고, 유해 전용 MLP 전문가를 두는 MLPSwap은 0.91/0.50에 그쳤다. 이미 학습된 Full에서 시작해 Pristine Attention으로 교체·재학습한 Pristine (Unlearn)은 0.93/0.40/0.55로, 처음부터 분해 구조로 학습한 0.95/0.57/0.68보다 약했다. 유해 데이터에서만 흔한 재질 10종을 안전한 견종에 합성하는 조성 생성 평가(TPIPS)에서는 Full 0.7705, Safe 0.6890, Pristine 0.7404로 전이 효과를 확인했다. 대규모 언러닝 베이스라인 UCE, ScaPre, SPEED, SAFREE와 비교하면 UCE가 가장 좋은 0.85/0.42/0.56이었고 Pristine은 0.89/0.56/0.72에 유해 정확도 0%를 기록했다. UCE는 안전 VQAScore가 0.31로 떨어지고 FID가 12.43으로 오르는 대가를 치렀다.

오픈월드 확장은 CC12M에서 이뤄졌다. 무기, 지식재산, 아동/유아를 아우르는 유해 6개 범주와 무해 5개 범주를 합쳐 11개 제한 카테고리를 지정했고, MLLM이 CC12M의 약 38%를 제한 데이터로 라벨링했다. Full은 DPG-Bench 84.21, PRISM-Bench 5.59를 기록했지만 제한 프롬프트의 56.42%에서 요청된 제한 개념을 생성했다. Safe는 제한 생성률을 5.87%로 낮추는 대신 품질이 79.30과 4.79로 떨어졌다. Pristine은 82.31과 5.02를 유지하면서 제한 생성률을 13.29%로 낮췄다. 안전 데이터만 쓸 때 잃는 품질의 상당 부분을 되찾으면서 억제 효과도 대부분 보존한 셈이다.

실무 관점에서 이 논문이 유용한 지점은 안전성 확보 시점을 추론이 아니라 사전학습으로 옮긴다는 발상이다. 배포 전에 금지 개념 목록이 확정되어 있는 서비스라면, 데이터 필터링으로 품질을 깎거나 사후 언러닝으로 기존 표현을 훼손하는 대신 학습 구조를 바꾸는 선택지가 생긴다. 다만 도입 전에 확인할 것이 있다. 프롬프트 분해를 LLM에 의존하므로 분해 품질이 결과를 좌우하고, 브랜치가 늘어나는 만큼 파라미터가 증가한다(저자들은 부록 A.3.1에서 Pristine이 Safe보다 파라미터가 많지만 결론에 영향을 주지 않는다고 밝혔다). 또 cross-attention과 MM-DiT 각각에 대해 어텐션 레이어를 직접 교체해야 하므로 사용 중인 백본에 맞는 구현이 필요하다.

저자들이 명시한 전제는 분명하다. 이 방법은 사전에 알려진 제한 개념에만 적용되며, 사후 언러닝과 추론 시점 가드레일을 대체하지 않고 보완한다. 학습 이후에 새로 식별된 개념은 여전히 사후 방법으로 처리해야 한다. 윤리 성명에서는 악의적 행위자가 시스템 설계에 대한 통찰을 얻어 우회 수단을 개발하는 데 악용될 가능성도 인정한다.