단계 없는 생성 모델을 위한 잊기 기법 불균형 최적 수송 기반 접근
본 논문은 한 번에 이미지를 생성하는 플로우 맵·컨시스턴시 모델에 대해, 목표 클래스를 안전하게 삭제하면서 전체 생성 품질을 유지하는 새로운 머신 언러닝 프레임워크 UOT‑Unlearn을 제안한다. 불균형 최적 수송(Unbalanced Optimal Transport, UOT) 이론을 활용해 “잊기 비용”과 f‑다이버전스 패널티 사이의 트레이드오프를 정량화하고, 잊히는 클래스의 확률 질량을 다른 클래스에 부드럽게 재분배한다. CIFAR‑10·I…
저자: Hyundo Choi, Junhyeong An, Jinseong Park
1. 서론
최근 이미지 생성 분야에서 확산 모델이 높은 품질을 달성했지만, 수십~수백 단계의 반복 연산으로 인해 추론 속도가 느리다는 한계가 있다. 이를 극복하고자 플로우 맵, 컨시스턴시 트라젝터리 모델(CTM), MeanFlow 등 ‘한‑스텝’ 생성 모델이 등장했으며, 노이즈를 바로 데이터 공간으로 매핑해 초고속 샘플링을 가능하게 한다. 그러나 이러한 모델이 점점 강력해짐에 따라, 부적절하거나 저작권 침해가 가능한 콘텐츠를 빠르게 대량 생성할 위험도 커졌다. 기존 머신 언러닝(ML) 연구는 다중 단계 확산 모델을 대상으로 ‘잊기’를 설계했지만, 한‑스텝 모델은 중간 단계가 없으므로 기존 방법을 직접 적용할 수 없다. 따라서 한‑스텝 모델 전용 언러닝 프레임워크가 절실히 필요하다.
2. 관련 이론
2.1 한‑스텝 생성 모델
연속 시간 확률 흐름을 ODE 형태로 기술하고, 이를 직접 학습하거나 근사해 ψ(x_t,t,s) 와 같은 흐름 맵을 구한다. 최종적으로 x₁ = G_θ(x₀) 와 같이 한 번의 전방 연산만으로 샘플을 생성한다.
2.2 불균형 최적 수송(UOT)
전통 OT는 마진을 정확히 맞추는 제약이 강해, 질량을 제거하거나 재분배해야 하는 상황에 부적합하다. UOT는 마진 차이를 f‑다이버전스 형태의 패널티로 완화해, 전송 비용과 마진 위반 비용 사이의 트레이드오프를 최적화한다. 특히 KL‑역방향을 사용하면 데이터가 거의 없는 영역에 질량이 배치되는 것을 강하게 억제한다.
3. 제안 방법 – UOT‑Unlearn
3.1 문제 정의
사전 학습된 한‑스텝 생성기 G_pre 가 전체 데이터 p_data (잊기 클래스 S_f + 유지 클래스 S_r)를 학습했다. 목표는 G_θ 를 미세조정해 S_f 에 대한 생성 확률을 0에 가깝게 만들고, S_r 에 대한 품질은 유지하는 것이다. 실제 S_r 데이터에 접근할 수 없으므로, 전송 목표 분포 ν 를 p_pre 로 근사한다.
3.2 UOT 기반 목표 함수
UOT의 반이중(semi‑dual) 형태를 차용해 전송 맵 ΔT_θ 와 dual potential v_ϕ 를 동시에 학습한다. 최종 목표는 식 (15)와 같이 두 부분으로 구성된다: (1) c_ul (잊기 비용) – R_f 내 샘플에 대해 힌지형 λ·(m‑d_cos) 패널티를 부여, (2) c_ul (보존 비용) – R_f 외 샘플에 대해 τ·‖G_pre‑G_θ‖₂² 패널티를 부여. 이 두 비용은 각각 UOT의 전송 비용과 f‑다이버전스 패널티에 대응한다.
3.3 최적화 절차
알고리즘 1은 배치 B₁, B₂, B₃ 을 이용해 v_ϕ 와 G_θ 를 교대로 업데이트한다. 첫 번째 배치 B₁ 은 dual potential을 업데이트하는 데 사용되고, B₂ 는 ν≈p_pre 에 대한 KL‑역방향을 계산한다. 마지막 배치 B₃ 은 실제 생성기 파라미터 θ 를 업데이트한다. 전체 과정은 생성된 샘플만으로 진행되므로, 실제 데이터 없이도 학습이 가능하다.
4. 실험
4.1 설정
CIFAR‑10(10 클래스)과 ImageNet‑256(1000 클래스) 데이터셋에 대해 CTM·MeanFlow 모델을 사전 학습시킨 뒤, 특정 클래스를 ‘잊기’ 대상으로 지정한다. 평가 지표는 (i) PUL – 잊기 대상 클래스가 생성될 확률, (ii) u‑FID – 잊기 후 전체 이미지 품질을 기존 모델과 비교한 지표이다.
4.2 결과
UOT‑Unlearn은 모든 실험에서 PUL을 0.05 이하로 낮추며, u‑FID는 사전 모델 대비 평균 3~5 % 정도 상승에 그쳤다. 기존 diffusion‑based 언러닝 방법은 PUL이 0.2~0.3 수준에 머물렀고, u‑FID는 10 % 이상 악화되는 경우가 많았다. 시각적 샘플도 잊기 클래스가 거의 사라지고, 나머지 클래스는 원본과 거의 구분되지 않을 정도로 품질이 유지되었다.
4.3 분석
잊기 비용 λ와 보존 비용 τ의 비율을 조절하면, 잊기 강도와 보존 품질 사이의 트레이드오프를 세밀하게 제어할 수 있다. 또한, 마진 m 을 크게 잡으면 더 보수적인 잊기가 이루어지며, 작은 m 은 빠른 수렴을 돕지만 일부 잊기 클래스가 완전히 사라지지 않을 위험이 있다.
5. 결론 및 향후 연구
UOT‑Unlearn은 한‑스텝 생성 모델에 대한 최초의 데이터‑프리 클래스 언러닝 프레임워크로, 불균형 최적 수송을 통해 잊기 대상 질량을 다른 클래스에 자연스럽게 재분배한다. 이는 빠른 이미지 생성 모델의 안전성을 크게 향상시키며, 마진 제약을 완화한 OT가 머신 언러닝에서 새로운 활용 가능성을 보여준다. 향후 연구에서는 (i) 다중 클래스 동시 잊기, (ii) 다양한 f‑다이버전스 함수 적용, (iii) 실시간 스트리밍 환경에서의 온라인 언러닝 등으로 확장할 수 있다.
원본 논문
고화질 논문을 불러오는 중입니다...
댓글 및 학술 토론
Loading comments...
의견 남기기