역증류 프레임워크가 한 단계 생성 모델의 망각과 증류를 동시에 처리한다

무엇인가

다단계 확산·플로우 매칭 모델은 품질 좋은 샘플을 만들지만 추론 비용이 크고, 학습 데이터의 원치 않는 부분을 그대로 재현할 위험이 있다. 기존 기계적 망각(machine unlearning) 연구는 조건부 모델의 클래스 단위 삭제나 다단계 모델의 데이터 삭제에 집중했고, 한 단계 생성기(one-step generator)에서 특정 데이터 샘플을 지우는 연구는 UOT-Unlearn 정도뿐이었다. 이 논문은 그 공백을 겨냥해 망각과 증류를 한 번에 처리하는 IDU(Inverse Distillation Unlearning)를 제안한다.

어떻게 동작하나

방법의 뼈대는 역증류(inverse distillation)다. 교사 모델 f*가 데이터 분포 p0*로부터 UM 손실을 최소화해 얻어졌다면, 그 손실의 min-max를 뒤집으면 교사가 담고 있는 데이터 분포를 복원할 수 있다. IDU는 이때 최적화 대상 분포 p0를 망각 데이터 p0^F와 생성 분포 p0^θ의 혼합 p0^mix = ρ·p0^F + (1-ρ)·p0^θ로 매개화한다. UM 손실이 데이터 분포에 대해 선형이기 때문에 목적함수는 ρ·[L(f*,pF) - L(f,pF)] + (1-ρ)·[L(f*,pθ) - L(f,pθ)]로 분해된다. 논문의 정리 1은 ρ를 실제 망각 데이터 비율 π로 두면 최적 생성기가 정확히 남은 데이터 분포 p0^R만 학습한다고 주장한다. 망각 성분은 이미 주어진 망각 샘플이 담당하므로, 생성기는 남은 데이터만 흉내 내면 교사 분포를 복원할 수 있다는 논리다.

무엇과 다른가

학습은 두 단계를 교대한다. 먼저 가짜 모델 f를 망각 샘플과 생성 샘플을 섞은 혼합 분포에 맞추는 손실 ρ·L_UM(f,pF) + (1-ρ)·L_UM(f,pθ)로 갱신한다. 다음으로 f를 고정하고 생성기 θ를 갱신하는데, 기본 이론 손실 대신 SiD(Score identity Distillation) 프레임워크의 변형 손실을 쓰고 두 번째 항을 α_SiD(0.5~1.2)로 스케일한다. FM에는 α_SiD=0.5, SiD에는 1.2를 사용했다. 새로 도입되는 하이퍼파라미터는 ρ 하나뿐이며 이론값은 망각 데이터 비율 π다. 학습에는 교사 모델과 망각 샘플만 필요하고, 남은 데이터·특징 추출기·분류기는 손실이나 그래디언트 갱신에 들어가지 않는다.

어떻게 쓰나

실험은 MNIST와 CIFAR-10에서 flow matching(FM)과 SiD가 쓰는 EDM-VP 스코어 기반 백본으로 수행했다. 주 실험은 MNIST 숫자 {3,7}, CIFAR-10 클래스 {1,9}(자동차·트럭)을 잊게 하는 것이다. 지표는 남은 데이터에 대한 FID(Retain FID)와, 생성 이미지 5만 장을 기성 분류기로 판정했을 때 망각 클래스로 분류된 비율(FGR)이다. 네 가지 데이터셋-백본 조합 모두에서 망각 모드를 억제했고, FM 실험의 망각 클래스 FGR은 최대 0.56%, SiD 실험은 최대 1.25%였다. Retain FID는 세 설정에서 남은 데이터만으로 재학습한 뒤 증류한 오라클과 비슷했고, FM-MNIST에서는 그 기준보다 좋았다.

전제와 한계

클래스별 강건성도 확인했다. MNIST 전 클래스(FM), CIFAR-10 전 클래스(FM, SiD)에 대해 단일 클래스 망각을 수행한 결과 평균 Retain FID/FGR은 FM-MNIST 4.54/0.16%, FM-CIFAR-10 7.73/0.84%, SiD-CIFAR-10 3.61/0.75%였다. ρ 민감도 실험에서 FM CIFAR-10은 ρ=0.9에서 즉시 발산했고, 낮은 ρ는 FID를 지키지만 FGR이 높게 남았으며 ρ=0.8은 망각은 좋아지지만 FID 손실이 컸다. 저자들은 FM CIFAR-10에 ρ=0.6, FM MNIST에 ρ=0.4를 최적 균형으로 택했다. SiD에서는 ρ를 0.2에서 0.8로 올릴수록 Retain FID가 점진적으로 나빠졌고 FGR은 비단조적으로 변했다. ρ=0.9에서 SiD는 발산하지 않았지만 FGR이 12.05%와 9.55%로 전체 데이터 교사 수준에 근접했고 Retain FID는 10.36으로 올랐다.

다른 방법과의 비교도 제시된다. 한 단계 생성기용 UOT-Unlearn과 CIFAR-10 클래스 1, 6, 8 망각에서 비교했을 때 UOT의 Retain FID는 (9.90, 5.11, 5.88), IDU는 (2.96, 4.16, 3.22)였고, FGR은 UOT가 약 (2, 0.9, 1.5), IDU가 (0.61, 0.53, 0.54)였다. 클래스 망각 방법인 SFD와 CIFAR-10 클래스 0 망각에서 비교하면 SFD는 Retain FID 약 3.1, FGR 0.36, IDU는 3.54와 0.58로 비슷한 수준이다. 다만 IDU는 조건부 프롬프트 없이 샘플 단위로 지울 대상을 지정할 수 있다는 점이 다르다.

개발자 관점에서 이 방법은 이미 학습된 다단계 교사 모델이 있고, 지우고 싶은 데이터 샘플은 있지만 남은 데이터 전체는 없거나 쓰고 싶지 않은 상황에 맞는다. 별도 특징 추출기나 분류기 학습 없이 교사와 망각 샘플만으로 한 단계 생성기를 얻는다는 점이 실무적으로 크다. 다만 ρ 선택이 품질-망각 트레이드오프를 좌우하므로, 남은 데이터를 평가에 쓸 수 있다면 FID/FGR로 ρ를 고르고, 쓸 수 없다면 망각 데이터 비율에서 출발해 넓은 범위를 탐색하는 편이 안전하다. 클래스 라벨은 망각 집합 구성과 평가에만 쓰이고 학습 입력으로는 쓰이지 않는다.

한계도 분명하다. FM 백본에서 확장 학습을 하면 대상 클래스가 초반에 억제되다가 약 2만 번의 생성기 업데이트 이후 다시 생성 빈도가 올라가는 현상이 나타났고, 그 시점은 ρ에 따라 달라진다. 그래서 FM에서는 ρ와 체크포인트를 함께 골라야 한다. SiD 학습 구간에서는 이런 역전이 관찰되지 않았고, 저자들은 이를 IDU 목적함수가 아니라 현재 FM 구현의 한계로 본다. 또한 ρ 선택에 남은 데이터 기반 평가를 사용했으며, 남은 데이터가 없을 때는 교사 샘플에서 분류나 수동 선택으로 평가용 데이터를 만들 수 있다고 밝힌다.