Stochastic Unlearning with Knowledge Preserving Loss

Abstract

Recently, machine unlearning has received considerable attention, in the context of responsible artificial intelligence and privacy regulations. This technical report introduces novel machine unlearning methods, such as stochastic re-initialization, knowledge preserving loss, gaussian noise, and forget-remember cycle. We present successful unlearning results, validated in the NeurIPS 2023 Machine Unlearning Challenge, accompanied by visualizations of logit distributions and several interim experiments.

Overview

Forget set에 접근하지 않고 retain set만으로, 재학습 없이 특정 데이터의 영향을 효과적으로 제거하는 machine unlearning 방법을 제안한다.

1. Stochastic re-initialization — 모델 레이어를 랜덤으로 선택하여 재초기화함으로써, 특정 데이터에 대한 기억을 확률적으로 파괴한다.
2. Knowledge preserving — 원본 모델의 출력을 MSE loss로 재현하여 retain set에 대한 성능을 유지한다.
3. Forget-remember cycles — Forgetting과 remembering phase를 3-4회 반복하여, 점진적으로 re-initialization 비율을 높이면서 retain set 성능을 유지한다.

전체 아키텍처. (a) Retain set 이미지에 gaussian noise 추가, (b) 랜덤 레이어 stochastic re-initialization, (c) 원본 모델과의 knowledge preserving loss 계산, (d) (a)-(c)를 $n$ cycle 반복.

Method

전체 파이프라인은 forgetting phase와 remembering phase를 반복하는 cycle 구조다.

Stochastic Re-initialization

모델의 일부 레이어를 랜덤으로 선택하여 re-initialize 한 뒤 retain set으로 fine-tune한다. Fisher information matrix (FIM) 기반의 고정 선택보다 랜덤 선택이 더 효과적이었는데, 대회 메트릭이 512회 독립 실행의 분포를 평가하기 때문에 충분한 randomness가 필요했기 때문이다.

FC layer와 projection-shortcut layer는 클래스/해상도 정보를 담고 있어 re-initialization 대상에서 제외했다.

Knowledge Preserving Loss

Unlearning의 정의에 따르면, unlearn된 모델은 retrain된 모델과 유사해야 한다. Retain set에 대해서는 원본 모델의 출력을 재현하도록 MSE loss로 학습한다:

Gaussian Noise Augmentation

Retain set 이미지에 gaussian noise (\(\sigma=0.01\) )를 추가하여 randomness를 확보하면서도 robust한 knowledge preserving 효과를 얻는다.

Forget-Remember Cycles

Re-initialization 비율을 단순히 높이면 모델이 너무 많이 잊어버린다. 대신 3~4 cycle로 forgetting과 remembering을 반복하면 점진적으로 re-initialization 비율을 높이면서도 retain set 성능을 유지할 수 있다.

Results

Quantitative

Qualitative

Forget set과 retain set에서의 logit 분포를 비교하면, 우리 방법이 fine-tuning 대비 retrained model의 분포에 훨씬 가깝다.

(a) Fine-tune

(b) Ours

Forget set의 logit 분포. Fine-tune 대비 우리 방법이 retrained model의 분포에 훨씬 가깝다.

(a) Fine-tune

(b) Ours

Retain set의 logit 분포. 우리 방법이 retain set에서도 retrained model과 유사한 출력을 유지한다.

BibTeX

@misc{ahn2023stochastic,
  author    = {Ahn, Jaesin and Lee, Chaehyeon and Bae, Jun-Hyun
               and Yim, Junho and Jung, Heechul},
  title     = {Stochastic Unlearning with Knowledge Preserving Loss},
  year      = {2023},
  note      = {NeurIPS 2023 Machine Unlearning Challenge,
               8th place out of 1,188 teams},
  url       = {https://www.kaggle.com/competitions/neurips-2023-machine-unlearning/writeups/forget-9th-place-solution-forget-set-free-approach}
}