💭

Stable Diffusionからの概念消去⑬：DoCo（論文）

に公開
機械学習
論文
Stable Diffusion
拡散モデル
tech

Unlearning Concepts in Diffusion Model via Concept Domain Correction and Concept Preserving Gradient (AAAI2025)

AAAI2025のスケジュールを見て採択論文を眺めていると発見したので紹介します. ちなみに2025/3/1の12:30からポスター発表です (スケジュールを見ると土曜日の部分にあります. paper IDは8293です).

書籍情報

AAAI2025はproceedingsがまだなのでarXivのものにしておきます. 断りのない限りは以下の文献から図表を引用します.

Yongliang Wu, Shiji Zhou, Mingzhuo Yang, Lianzhe Wang, Heng Chang, Wenbo Zhu, Xinting Hu, Xiao Zhou, and Xu Yang. Unlearning concepts in diffusion model via concept domain correction and concept preserving gradient, 2024.

関連リンク

概要

machine unlearningの手法を概念消去に適用する際に, 2つの問題があります.

  1. 汎化性 (generalizability): 限られたサンプルで完璧にunlearnする必要があります.
  2. 概念保護 (concept preservation): 他の概念には影響を与えてはいけません.

既存手法はこの2つについてうまく扱えていません. 汎化性については, 例えばtarget conceptとanchor conceptの距離を最小化することや, cross-attentionのvalueとkeyを最小化するといった手法がありますが効果は限定的です (と述べられていますが根拠不明です). 概念保護については正則化が取り入れられていますが, 概念保護に固執するあまり, unlearnとの間のconflictを見落としているようです (ここでのconflictについて説明はないですが, unlearnとerasureのlossによって計算される勾配の向きが大きく異なることを表すようです).

提案手法であるDoCo (Domain Correction)は, Membership Inference Attack (MIA)による敵対的学習を導入している手法です. 提案手法では概念が直接的に記述されていないOut-of-Distribution (OOD) promptsでも効果を発揮すると述べられています.

OOD promptsについて

Out-of-Distribution (OOD) promptsですが, Implicit Promptという名前が一応あります. 以下のPosition PaperがarXivに投稿されたのは2024年の3月で, この論文は5月なので先行研究の調査不足と考えられます.

https://proceedings.mlr.press/v235/yang24o.html

概念保護に関しては, concept-preserving gradient approachを採用します. これは, unlearning gradientと正則化のconflictを防ぎ, 概念消去と概念保護を両立します.

著者らによると貢献は3つにまとめられ,

  • 敵対的振る舞いによる分布分類器を採用し, 概念消去を行う
  • unlearning gradientとretraining gradientを両立する手法を導入する
  • 実験では他の概念への影響が軽微で概念消去でき, 他の手法に比べてOODに強いことを確認

手法

notationは一般的なものばかりが使われているので, この論文特有のもののみ述べますが, 他は説明なしで登場します.

問題設定

事前学習済みモデル M_{\mathrm{init}} から特定の概念 c^* を消すことを考えます. Unlearningのアルゴリズム \mathcal{U} が M_{\mathrm{init}} から特定の概念 c^* を消すと仮定すると, M_{\mathrm{init}} のパラメータは \hat{M} に置き換えられます. この操作を

\mathcal{U}(M_{\mathrm{init}}, c^*)=\hat{M}

と書くことにします. 既存研究 (ここではESD, AC, SPM)は c^* とanchor concept c を一致させようとします.

\min_{\hat{M}}\ \mathrm{dis}(\hat{M}(c^*), M_{\mathrm{init}}(c))=0

この手法は訓練時に用いるpromptに強く依存しており, 汎化性能の欠如を招きます. 理想的には, 分布を一致させることがいいです.

\min_{\hat{M}}\ \mathrm{dis}(P(\hat{M}(c^*)), P(M_{\mathrm{init}}(c)))

同時に, original anchorと他の概念の分布は変わらないことが望まれるので

P(\hat{M}(c^*))\sim P(M_{\mathrm{init}}(c))

となります. すると,

\mathrm{dis}(P(\hat{M}(c^*)), P(M_{\mathrm{init}}(c)))\approx0

とも書けます.

Concept Domain Correction

\displaystyle\min_{\hat{M}}\ \mathrm{dis}(P(\hat{M}(c^*)), P(M_{\mathrm{init}}(c))) を最適化するために, 著者らは敵対的学習に基づくConcept Domain Correction frameworkを提案しています.

直感的なアプローチとしては, M_{\mathrm{init}} をfine-tuningしますが, これはgeneratorとして振る舞います. そして, discriminator D を学習します. D が M_{\mathrm{init}}(c^*) と \hat{M}(c) を見分けられないように学習します. この辺りは完全にGANです. マルコフ連鎖を用いる拡散モデルの都合から, 直接画像を変えていくのではなく潜在空間内でノイズをalignmentすることを考えます. これによって, 識別器は画像の真贋判定をするのではなくノイズのテキスト状態を予測することになります (anchor or targetを予測します). これを表したのが下図です.

損失関数は

\mathcal{L}_{u}=\min_{\varepsilon_{\theta}}\max_{D}V(\varepsilon_{\theta}, D)=\mathbb{E}_{x\sim p(x|c)}[\log D(x_{t-1}(c))]+\mathbb{E}_{x\sim p(x|c^*)}[\log(1-D(x_{t-1}(c^*)))]

となります. もはや普通のGANです.

Concept Preserving Gradient

P(\hat{M}(c^*))\sim P(M_{\mathrm{init}}(c)) を最適化することを考えます. 既存手法 (AC)では, preservation setを考えて

\mathcal{L}_r=\|\varepsilon-\varepsilon_{\theta}(x_t, c, t)\|_2^2

を考えていました. これでも実現できますが, unlearningとretrainingのconflictは残ります. ここではeffective gradient surgeryを取り入れます.

https://proceedings.neurips.cc/paper/2020/hash/3fe78a8acf5fda99de95303940a2420c-Abstract.html

図にすると

のようになります. unlearningとretrainingの目的関数をそれぞれ \bold{G}_u, \bold{G}_r とします. 訓練の過程では2つのシナリオが考えられます.

  1. \bold{G}_u, \bold{G}_r のなす角が 90^\circ 以下の場合: \bold{G}_u にaligenmentさせ, \bold{G}=\bold{G}_u とします.
  2. \bold{G}_u, \bold{G}_r のなす角が 90^\circ を超えるの場合: conflictが発生します. この場合, \bold{G}_u を \bold{G}_r の方向に射影します.

定式化すると

\bold{G}=\begin{cases} \bold{G}_u & \bold{G}_u\cdot\bold{G}_r\geq 0 \\ \bold{G}_u-\lambda\dfrac{\bold{G}_u\cdot\bold{G}_r}{\|\bold{G}_r\|^2}\bold{G}_r & \mathrm{otherwise} \end{cases}

です. \lambda はハイパーパラメータですがデフォルトは1とします.

実験

batch size 8で2000回更新します. 学習率は6e-6とし, 最初の1000回は識別器の更新のみwarm-up strategyを採用します. 識別器にはPatchGANを用います. 実験は全て2枚のA100 GPUで行われます. 評価指標にはCLIP Score, CLIP Accuracy, FIDを用います. celebrityの消去のみIdentity Score Matching (ISM)を用います. これは生成された顔画像の埋め込みとclean image setの顔の埋め込みの平均のcosine distanceを計算したものです.

では結果を見ます.

論文には何に対して評価指標を計算するかが書かれていませんが, SPMの研究に準拠しているものと思われます. 定量結果を上の表で見ると, 太字であるbest performanceのほとんどが提案手法であることがわかります. これは, 概念消去と概念保護を両立できていることを表すと主張していますが, 疑わしい部分もあります.

私の考え

例えば他の概念におけるFIDが低いことは他の概念の生成結果がoriginalにしか強いことを意味していますが, これではOODの評価はできなくなってしまいます. OODでは間接的に表現するので他の画像が出る場合も往々にしてあります. それに対しては維持できていないといけないので, FIDは低くなりますが, 消した概念に関する場合はFIDが高くなければならないので, 必ずしも低ければいいというわけではないはずです (もちろん, 今回の設定では低ければいいように思えますし私自身もそう思いますが, 厳密にはそこで用いたpromptがOODでないことを保証する必要があると思います. 1000枚生成して何枚以上再現できたらOODとみなすみたいな基準が欲しいですね.). ちなみに私自身も最近OODの実験をしていて, 初期ノイズによって同じpromptでも再現できる場合とできない場合が半々くらいだなと感じています.

続いて, 生成結果を見てみます. まずはobjectの消去です.

左の例では, ESDが他の概念を維持できていないことがわかります. 例えば右のケースでは, a clear blue sky aboveとありますが, 再現できていません (正直これはノイズの選び方によるのではと思います.). さらに, 左の例では複雑なプロンプトに対応できていないことがわかります. ConAbl (CA)はそもそも消せてないです. SPMは提案手法よりもcolorやtextureが維持できていると主張していますがどこのことかはよくわかりません.

次にstyleの結果を見ます.

ESDは他のスタイルに大きく影響があることがわかります. SPMに関しても, 私は絵画については素人なのでstyleは変わらないように見えますが, 著者らによると美的な部分で魅力が失われているようです (これは定量化すべきだと思います. SPMはstyleのFIDに関して結構いい感じの性能を出しているのが表からわかるので主観的な判断はよくない気がします. それこそuser studyを入れるべき事項でしょう). そしてConAblに関しては言及がありません.

再び定量評価に戻ってcelebritiesに関して見てみます.

論文には「提案手法が最もいい結果」ということ以外の言及はないので, 個人的な意見を述べます. 消す概念以外のISMについてはSDより悪いことが予想されますがいい結果もあります (10^{-4}とかなら気にしませんが 10^{-2}のオーダーです.). これは素直に解釈すれば「unlearnで性能が上がることもある」または「評価指標が適当ではない」のどちらかですが, 後者の方が適当な感じがします. 前者ならもっと詳細な説明が欲しいです.

最後の比較結果としてI2Pでの結果を確認します. NudeNetで検知します.

SD 1.4と比較してどれだけ減ったかが示されています. ESDとConAblよりはいいですが, SPMには負けています. 著者らはfemale breastとbuttocksのようなsensitive body partsではSPMよりいいと主張しています. 確かにそうですが, それについてもっと詳しい考察が欲しいなと思いました.

OOD prompt

実験のセクションでは全く触れられていないOOD promptですが, 一応生成例のみ示されています (示されているだけでrefされて言及とかは全くないです).

左がVan Goghを消した場合, 右がDogを消した場合です. どちらも関連するpromptであるStarry Nightやhuskyに対して概念生成を防ぐことに成功しています. しかし, このようなOOD promptに対しては敵対的学習を行った手法と比較するのが筋といえます. ちょうど同じ日にarXivに投稿されたAdvUnlearnが敵対的学習を行っている手法です.

https://openreview.net/forum?id=dkpmfIydrF

arXivのversion 1で実験されていないのは仕方ないと思いますが, version 2は12月更新で, AAAIの投稿期限は8月なので十分実験を行って反映することはできるんじゃないかなと思います (しかもこの論文は引用されています).

思ったこと

  • まず, 論文が短いです. AAAIは7ページまでなのでそれは仕方ないのですが, それはAAAIも認めていて, Supplementary Materialを提出することが可能です. そこで更なる実験などを行うべきではないかと思います.
  • GANは学習が不安定なことが知られています. そのため, 様々な学習安定策 (例えばTTURなど)が採用されたりします. ここではそのようなことをしなくても上手くいったということでしょうか.
  • 毎回denoiseしてdiscriminatorで判定はちょっと計算量が多い気がします. 例えばProjectedGANのように事前学習するとかはないのかなと思いました.
  • A100が2枚いるというのは結構なGPUを要求する気がしますがどうでしょう. A100だと80GBだと思いますが, これが2枚というのはStable Diffusion 1.4の概念消去では過剰な要求に見えます (多分2つのU-Netと1つのDiscriminatorをbackwardしても80GBもいらないと思います).

参考文献

  • Yongliang Wu, Shiji Zhou, Mingzhuo Yang, Lianzhe Wang, Heng Chang, Wenbo Zhu, Xinting Hu, Xiao Zhou, and Xu Yang. Unlearning concepts in diffusion model via concept domain correction and concept preserving gradient, 2024.
  • Yue Yang, Yuqi Lin, Hong Liu, Wenqi Shao, Runjian Chen, Hailong Shang, Yu Wang, Yu Qiao, Kaipeng Zhang, and Ping Luo. Position: Towards implicit prompt for text-to-image models. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 56235–56250. PMLR, 21–27 Jul 2024.
  • Yimeng Zhang, Xin Chen, Jinghan Jia, Yihua Zhang, Chongyu Fan, Jiancheng Liu, Mingyi Hong, Ke Ding, and Sijia Liu. Defensive unlearning with adversarial training for robust concept erasure in diffusion models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.

Discussion