Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

이 논문은 기존 확산 모델 기반의 제거 기법과 호환되지 않는 한 단계 생성 모델의 안전성을 확보하기 위해, 불균형 최적 수송 (UOT) 을 활용하여 특정 클래스의 생성을 억제하면서도 전체 생성 품질을 유지하는 'UOT-Unlearn' 프레임워크를 제안합니다.

핵심

1. 배경: "빠른 화가"와 "나쁜 그림"의 문제

🎨 상황 설정:
예전에는 AI 가 그림을 그릴 때, "노이즈 (잡음) 에서 시작해서 천천히, 한 번에 한 번씩 선을 지우며 (반복 과정) 그림을 완성"하는 방식이었습니다. 이 방식은 그림이 잘 나왔지만, 시간이 너무 오래 걸렸습니다.

하지만 최근에는 **"한 번에 뿅! 하고 그림을 완성하는 화가 (One-step Generative Models)"**가 등장했습니다. 이 화가는 그림을 그리는 속도가 매우 빨라 인기가 많지만, 문제는 이 화가가 실수로 '금지된 그림 (예: 폭력적이거나 저작권이 있는 그림)'을 그릴 수도 있다는 점입니다.

🚫 문제:
만약 이 '빠른 화가'가 나쁜 그림을 그리기 시작하면, 기존 방식처럼 "그림을 다시 그리는 과정 (반복 denoising)"을 거치면서 나쁜 부분을 지우려 하면, 화가가 이미 한 번에 그리는 방식을 깨뜨리게 되어 오히려 그림이 망가질 수 있습니다. 즉, 기존에 개발된 '잊기 기술'들은 이 '빠른 화가'에게는 통하지 않습니다.

---

2. 해결책: "균형 잡힌 이동 (UOT)"을 이용한 잊기

이 논문은 이 문제를 해결하기 위해 **'불균형 최적 수송 (Unbalanced Optimal Transport, UOT)'**이라는 수학적 개념을 차용했습니다. 이를 쉽게 비유해 보겠습니다.

📦 비유: 물자 재분배 창고

• 상황: 가상의 거대한 창고에 '사과', '배', '오렌지'가 섞여 있습니다. 하지만 '사과'는 더 이상 팔고 싶지 않아서 (잊고 싶어서) 창고에서 완전히 없애야 합니다.
• 기존 방식 (기존 기술들): 사과를 없애려고 하면, 창고 전체를 뒤집어엎거나, 사과가 있던 자리에 빈 공간만 남기거나, 엉뚱한 곳에 사과를 던져버립니다. 결과적으로 배와 오렌지도 엉망이 되거나 (화질 저하), 사과가 완전히 사라지지 않거나 (잊기 실패) 합니다.
• 이 논문의 방식 (UOT-Unlearn):
  1. 사과 (잊을 대상) 에는 '엄청난 벌금'을 매깁니다. 사과를 창고에 두는 것 자체가 매우 비싸게 책정됩니다.
  2. 하지만, 사과를 버리는 대신 다른 과일 (배, 오렌지) 로 자연스럽게 대체할 수 있는 길을 열어줍니다.
  3. 핵심: 사과를 없애는 비용과, 다른 과일의 모양을 유지하는 비용 사이의 최적의 균형을 찾습니다.

이 방법을 쓰면, 사과 (나쁜 그림) 는 자연스럽게 사라지고, 그 자리는 배나 오렌지 (나쁜 그림이 아닌 다른 좋은 그림) 가 자연스럽게 채워집니다. 결과적으로 창고 전체의 질서는 유지되면서, 사과만 깔끔하게 사라지는 것입니다.

---

3. 이 방법의 핵심 장점

이 논문에서 제안한 **'UOT-Unlearn'**이라는 도구의 특징은 다음과 같습니다.

1. 실제 데이터가 필요 없습니다 (Data-Free):

   • 보통 AI 가 "이것은 잊어라"라고 가르치려면, 잊어야 할 그림 (사과) 과 남겨야 할 그림 (배) 을 모두 보여줘야 합니다.
   • 하지만 이 방법은 오직 '사과가 어떤 모양인지'에 대한 정보 (중심점) 만 있으면 됩니다. 실제 '배' 그림을 보여주지 않아도, AI 가 스스로 "사과를 없애고 남은 공간은 자연스럽게 채워야겠다"라고 판단하도록 유도합니다. 이는 개인정보나 저작권이 있는 데이터를 건드리지 않고도 잊기를 가능하게 합니다.

2. 화질은 그대로, 내용만 바뀝니다:

   • 기존 방법들은 나쁜 그림을 지우려다 보니, 나머지 좋은 그림들도 흐릿해지거나 뭉개지는 경우가 많았습니다.
   • 이 방법은 나쁜 그림을 지우는 대신, 그 자리에 다른 좋은 그림을 자연스럽게 채워 넣기 때문에 전체적인 그림의 선명도와 질감 (화질) 이 거의 떨어지지 않습니다.

3. 한 번에 그리는 AI 에 특화되었습니다:

   • 복잡한 반복 과정 없이, AI 가 그림을 그리는 마지막 한 번의 결정 단계에서 이 기술을 적용합니다. 그래서 어떤 '빠른 화가' 모델에도 쉽게 끼워 넣을 수 있습니다 (Plug-and-play).

---

4. 실험 결과: 얼마나 잘 될까요?

연구진은 **CIFAR-10(작은 그림 10 가지)**과 **ImageNet(고화질 그림)**으로 실험을 해보았습니다.

• 결과: 다른 방법들 (기존 기술들을 변형한 것들) 은 나쁜 그림을 지우려다 보니 나머지 그림들이 엉망이 되거나, 반대로 나쁜 그림이 여전히 나오거나 했습니다.
• 이 방법: 나쁜 그림은 90% 이상 완벽하게 지워졌고 (PUL 점수 높음), 동시에 나머지 그림들의 질은 **거의 그대로 유지 (FID 점수 낮음)**되었습니다.

---

5. 한 줄 요약

**"AI 화가가 실수로 나쁜 그림을 그리지 않도록, 나쁜 그림을 지우는 대신 다른 좋은 그림으로 자연스럽게 대체하게 만들어, 화질은 그대로 유지하면서 나쁜 내용만 깔끔하게 지우는 새로운 기술"**입니다.

이 기술은 AI 가 더 빠르고 안전하게 그림을 그릴 수 있도록 도와주는, 마치 **'AI 의 양심'**을 심어주는 것과 같은 역할을 합니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 최근 흐름 매핑 (Flow Map) 모델이나 일관성 모델 (Consistency Models) 과 같은 One-Step 생성 모델은 노이즈에서 데이터로의 직접적인 매핑을 통해 기존 확산 모델 (Diffusion Models) 의 수백 단계 반복적 디노이징 과정 없이도 고품질 이미지를 생성할 수 있어 추론 속도가 획기적으로 개선되었습니다.
• 문제점: 이러한 강력한 생성 모델이 보편화됨에 따라, NSFW(성적 콘텐츠) 나 저작권이 있는 이미지와 같은 유해 콘텐츠 생성 위험이 증가했습니다. 이를 해결하기 위한 머신 언러닝 (Machine Unlearning) 기술이 필요하지만, 기존 언러닝 방법론들은 대부분 다단계 확산 모델의 반복적 디노이징 과정에 의존하고 있습니다.
• 핵심 한계: One-Step 모델은 단일 순전파 (Single Forward Pass) 로 동작하므로, 중간 단계에서 노이즈 예측이나 그래디언트를 수정하는 기존 언러닝 기법들을 적용할 수 없습니다. 또한, 기존 방법들은 대부분 재학습을 위한 실제 데이터 (Retain Data) 가 필요하여 효율성이 떨어집니다.

2. 제안 방법론: UOT-Unlearn (Methodology)

저자들은 One-Step 생성 모델을 위한 최초의 플러그 앤 플레이 (Plug-and-Play) 클래스 언러닝 프레임워크인 UOT-Unlearn을 제안했습니다. 이 방법은 불균형 최적 수송 (Unbalanced Optimal Transport, UOT) 이론을 기반으로 합니다.

2.1 핵심 아이디어

• UOT 의 활용: 기존 최적 수송 (OT) 은 소스와 타겟 분포의 마진 (Marginal) 을 엄격하게 일치시키는 반면, UOT 는 마진 제약 조건을 완화하고 수송 비용과 분포 편차 (f-divergence) 사이의 균형을 최적화합니다.
• 언러닝의 재정의: 언러닝을 '기억할 클래스의 확률 질량을 잊을 클래스 (Forget Class) 에서 제거하고, 나머지 클래스 (Retain Classes) 로 부드럽게 재분배하는 과정'으로 정의합니다.

2.2 수학적 공식화

• 목표: 사전 학습된 One-Step 생성기 $G_{pre}$를 미세 조정하여, 잊을 클래스 $S_f$의 생성 확률을 0 으로 만들고, 나머지 클래스 $S_r$의 생성 품질을 유지하도록 합니다.
• 비용 함수 (Cost Function):
  • Forget Cost: 생성된 샘플이 잊을 클래스의 특징 공간 (Feature Space) 에 속할 경우, 해당 특징이 잊을 클래스의 중심 (Anchor) 에서 멀어지도록 강력한 페널티를 부과합니다.
  • Retain Cost: 나머지 샘플에 대해서는 사전 학습된 출력과 현재 출력 간의 L2 거리를 최소화하여 생성의 충실도 (Fidelity) 를 유지합니다.
• 최적화 목적 함수:
  • UOT 의 반이중 (Semi-dual) 형식을 활용하여, 실제 Retain 데이터 없이도 생성된 샘플과 잊을 클래스의 중심 (Centroid) 만을 사용하여 목적 함수를 최적화합니다.
  • 식 (15) 에서 보듯, 타겟 분포 $\nu$를 사전 학습된 분포 $p_{pre}$로 근사화하여 실제 데이터 (Real Data) 가 전혀 필요 없는 (Data-free) 최적화 과정을 가능하게 합니다.

2.3 알고리즘 흐름

1. 잊을 클래스의 특징 공간 중심 $\mu_f$를 미리 계산합니다.
2. 생성된 샘플이 잊을 영역 ($R_f$) 에 속하는지 확인합니다.
3. UOT 목적 함수를 통해 생성기 $G_\theta$와 이중 잠재 (Dual Potential) $v_\phi$를 교대로 업데이트합니다.
4. 잊을 클래스의 확률 질량은 UOT 의 유연한 마진 제약 덕분에 저품질 노이즈로 붕괴되지 않고, 나머지 유효한 클래스 영역으로 자연스럽게 이동합니다.

3. 주요 기여 (Key Contributions)

1. One-Step 모델용 최초의 언러닝 프레임워크: 다단계 확산 모델에 국한되었던 기존 연구와 달리, 단일 순전파 아키텍처에 특화된 언러닝 방법을 최초로 제안했습니다.
2. UOT 기반의 새로운 목적 함수: f-divergence 페널티를 통해 잊을 클래스의 확률 질량을 나머지 클래스로 부드럽게 재분배하는 수학적 프레임워크를 정립했습니다. 이는 분포 붕괴를 방지하고 생성 품질을 유지합니다.
3. 실제 데이터 불필요 (Data-Free): 최적화 과정에서 실제 Retain 데이터에 접근할 필요가 없으며, 생성된 샘플과 잊을 클래스의 중심 벡터만으로 학습이 가능합니다. 이는 프라이버시 보호와 효율성 측면에서 큰 장점을 가집니다.

4. 실험 결과 (Results)

저자들은 CIFAR-10 과 ImageNet-256 데이터셋에서 CTM(Consistency Trajectory Model) 과 MeanFlow 와 같은 One-Step 아키텍처를 사용하여 실험을 수행했습니다.

• 성능 지표:
  • PUL (Percentage of Unlearning): 잊을 클래스 제거 성공률.
  • u-FID (Unlearned FID): 나머지 클래스의 생성 품질 (낮을수록 좋음).
• 결과 요약:
  • CIFAR-10: 모든 대상 클래스 (자동차, 개구리, 배) 에서 기존 Baseline(GA, SA, SalUn, VDU) 보다 더 높은 PUL과 더 낮은 u-FID를 기록했습니다. 특히 Gradient Ascent(GA) 는 PUL 은 높지만 u-FID 가 급격히 나빠지는 (생성 품질 저하) 반면, UOT-Unlearn 은 두 가지 목표를 동시에 잘 달성했습니다.
  • ImageNet-256: 'Goldfish' 클래스를 잊는 실험에서, Baseline 은 구조적 왜곡이 심해 u-FID 가 79.89 로 나빠진 반면, UOT-Unlearn 은 85.08% 의 높은 PUL을 달성하면서도 u-FID 를 20.16 수준으로 유지했습니다.
  • 2D 합성 데이터 시각화: 잊을 모드 (Forget Mode) 의 확률 질량이 유효하지 않은 영역으로 흩어지는 것이 아니라, 나머지 유효 모드 (Retain Modes) 로 부드럽게 재분배되는 것을 시각적으로 확인했습니다.

5. 의의 및 결론 (Significance)

• 효율성과 안전성의 균형: One-Step 생성 모델의 빠른 속도가 유해 콘텐츠 확산의 위험을 증가시킬 수 있다는 점을 고려할 때, 본 연구는 이러한 모델에 대한 안전장치 (Unlearning) 를 효율적으로 제공할 수 있는 길을 열었습니다.
• 구조적 재분배: 기존 방법들이 분포를 왜곡하거나 붕괴시키는 경향이 있었던 것과 달리, UOT 기반의 접근법은 확률 질량의 **구조적 재분배 (Structured Redistribution)**를 통해 생성 모델의 전반적인 품질을 유지하면서 특정 개념만 제거할 수 있음을 입증했습니다.
• 확장성: 실제 데이터 접근이 불필요하고 아키텍처 수정이 필요 없는 '플러그 앤 플레이' 방식이므로, 다양한 One-Step 생성 모델에 쉽게 적용 가능하여 미래의 생성형 AI 안전성 연구에 중요한 기여를 할 것으로 기대됩니다.

이 논문은 생성형 AI 의 윤리적 안전성을 확보하기 위해, 수송 이론 (Optimal Transport) 을 머신 언러닝에 적용한 독창적인 접근법을 제시하며, 특히 차세대 고속 생성 모델 (One-Step) 에 대한 필수적인 기술적 해법을 제공했습니다.