불균형 최적 수송을 이용한 단일 단계 생성 모델의 역학습

Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

최근 플로우 맵 모델과 같은 단일 단계 생성 프레임워크의 발전은 노이즈-데이터 매핑을 단일 순방향 패스에서 학습함으로써 이미지 생성 효율성을 크게 향상시켰습니다. 그러나 이러한 강력한 생성 모델의 안전성을 보장하기 위한 머신 언러닝은 아직 연구되지 않았습니다. 기존의 디퓨전 언러닝 방법은 본질적으로 이러한 단일 단계 모델과 호환되지 않는데, 이는 다단계 반복적인 노이즈 제거 프로세스에 의존하기 때문입니다. 본 연구에서는 불균형 최적 수송(UOT)을 기반으로 단일 단계 생성 모델을 위한 새로운 플러그 앤 플레이 클래스 언러닝 프레임워크인 UOT-Unlearn을 제안합니다. 저희의 방법은 언러닝을, 대상 클래스를 억제하는 '망각 비용'과 전체 생성 품질을 유지하는 'f-다이버전스 페널티' 사이의 균형 잡힌 절충으로 공식화합니다. 저희의 방법은 UOT를 활용하여, 잊혀진 클래스의 확률 질량을 저품질 또는 노이즈와 유사한 샘플로 붕괴시키는 대신, 나머지 클래스로 원활하게 재분배할 수 있도록 합니다. CIFAR-10 및 ImageNet-256 데이터셋에 대한 실험 결과는 저희 프레임워크가 뛰어난 언러닝 성공률(PUL)과 유지 품질(u-FID)을 달성하며, 기존 방법들을 크게 능가한다는 것을 보여줍니다.

Recent advances in one-step generative frameworks, such as flow map models, have significantly improved the efficiency of image generation by learning direct noise-to-data mappings in a single forward pass. However, machine unlearning for ensuring the safety of these powerful generators remains entirely unexplored. Existing diffusion unlearning methods are inherently incompatible with these one-step models, as they rely on a multi-step iterative denoising process. In this work, we propose UOT-Unlearn, a novel plug-and-play class unlearning framework for one-step generative models based on the Unbalanced Optimal Transport (UOT). Our method formulates unlearning as a principled trade-off between a forget cost, which suppresses the target class, and an $f$-divergence penalty, which preserves overall generation fidelity via relaxed marginal constraints. By leveraging UOT, our method enables the probability mass of the forgotten class to be smoothly redistributed to the remaining classes, rather than collapsing into low-quality or noise-like samples. Experimental results on CIFAR-10 and ImageNet-256 demonstrate that our framework achieves superior unlearning success (PUL) and retention quality (u-FID), significantly outperforming baselines.