Stable Diffusionからの概念消去⑩：R.A.C.E (論文)

R.A.C.E.: Robust Adversarial Concept Erasure for Secure Text-to-Image Diffusion Model (ECCV2024)

引き続き, ECCV2024での概念消去の論文を見ていきます. とは言いつつECCV2024はまだaccepted papersのタイトルを確認できないので見つけたものから順番に, という感じです.

図表はことわりのない限りこの論文からの引用です.

書籍情報

Changhoon Kim, Kyle Min, and Yezhou Yang. R.a.c.e.: Robust adversarial concept erasure for secure text-to-image diffusion model, 2024.

関連リンク

• arXivのリンク
• 公式実装

TL;DR

貢献は著者らが論文内でまとめてくれていますので, それをそのまま日本語にします.

• 敵対的学習の枠組みを初めて取り入れ, prompt-basedな攻撃手法によって概念消去を達成した. なお, 追加モジュールはない.
• 提案手法であるRACEは計算効率の良い敵対的攻撃手法を実装しているので概念消去に組み込める
• RACEは、white/black-boxな攻撃に基づくpromptに対するT2Iモデルの頑健性を大幅に向上させた

準備

簡単にベース技術を示します. まず, Stable Diffusionです. これは主にimage autoencoderとU-Netの2つから構成されていています. image autoencoderはencoder $\mathcal{E}$ とdecoder $\mathcal{D}$ の2つがあり, 学習の枠組みとしては $\mathcal{D}(z=\mathcal{E}(x))=\hat{x}\approx x$ となるように学習します. $x$ は画像, $z$ は潜在変数です. U-Netは拡散モデルを担当します. $c=\mathcal{E}_{txt}(y)$ で表されるpromptに関する条件を基に, promptにあった画像を生成できるように学習します. 学習は潜在空間で行われ,

$$L_{SD}\coloneqq\mathbb{E}_{n\in\mathcal{N}(0, 1), z, c, t}[\|n-\varPhi_{\theta}(z_t, t, c)\|_2^2]$$

で学習を行います.

続いて, ESDです. ESDは以下の損失関数によって特定の概念を削除します.

$$L_{erase}\coloneqq \|\varPhi_\theta(z_t, t, c)-(\varPhi_{\theta^*}(z_t, t)-\eta(\varPhi_{\theta^*}(z_t, t, c)-\varPhi_{\theta^*}(z_t, t)))\|_2^2$$

ここで, 追加のデータセットは必要ない点が特徴です.

このような概念消去に対して, red tamingというものがあります. これは, 消去された概念を復活させようというものです. これを達成するには敵対的なprompt $\tilde{y}$ を作成します. white-boxなアプローチは $\varPhi_{\theta}$ の勾配を用いてpromptの最適化を行います. 逆に, black-boxなアプローチは購買は使わずにpromptを最適化します. しかし, どちらの場合も必要な計算量が大きく, 敵対的訓練の枠組みに統合することは難しいです.

Adversarial Training on Concept Erased Diffusion Models

では, どのように統合すればいいのでしょう.

モチベーション

T2Iの拡散モデルは,

$$L_{SD}\coloneqq\mathbb{E}_{n\in\mathcal{N}(0, 1), z, c, t}[\|n-\varPhi_{\theta}(z_t, t, c)\|_2^2]$$

で学習を行いますが, 実はこの損失関数は分類タスクにも用いることができます. この分類能力はモデルの予測 $p_{\theta}(x|c_i)$ と, 条件 $c_i$ の集合に対する事前分布 $p(c)$ にベイズの定理を適用することで得られます.

$$p_{\theta}(c_i|x)=\dfrac{p(c_i)p_{\theta}(x|c_i)}{\sum_j p(c_j)p_{\theta}(x|c_j)}$$

特に, 事前分布がprompt $c_i$ に対して一様に分布しているとき, すなわち $p(c_i)=\dfrac{1}{N}$ のとき, 事前項は無視できます. 拡散モデルでは $p_{\theta}(x|c_i)$ を直接計算することは困難なので $\log p_\theta(x|c_i)$ の計算と最適化のための変分下限の利用を行います. DDPMで行われた近似を導入することでprompt $c_i$ に対する事後分布を以下のように近似できます.

$$p_{\theta}(c_i|x)=\dfrac{\exp\{-\mathbb{E}_{z, n, t}[\|n-\varPhi_{\theta}(z_t, t, c_i)\|_2^2]\}}{\sum_j\exp\{-\mathbb{E}_{z, n, t}[\|n-\varPhi_{\theta}(z_t, t, c_j)\|_2^2]\}}$$

ここから, 拡散モデルを用いた画像分類では次の式を計算することで画像 $x$ のクラスを推定します.

$$\argmin_{c}\sum_{t=1}^T[\|n-\varPhi_{\theta}(z_t, t, c)\|^2]$$

ここで, 単一のタイムステップ $t^*$ を用いて計算を行った場合でも画像 $x$ を分類することができるところは大事な点です.

Single-Timestep Adversarial Attacks on Erased T2I Models

先ほどの観点から, SDの損失関数を分類機構として再構築します. Textual Inversionを, 画像 $x$ を再生成するためのtext embedding $c$ を最適化することを目的とした敵対的攻撃の一種として見ることができます.

ここで, Textual Inversionは全てのタイムステップにわたる最適化が必要なので計算量は大きいです. そのため, 課題は「敵対的text embeddingは単一のタイムステップで特定できるのか？」となります.

概念消去の目的は, 消去された概念の画像 $\tilde{x}$ の再生成を促進する可能性のある対象概念のembedding $c$ および, その近接のembeddingを無効化することにあります. この目的達成のために, 概念消去済みの拡散モデル $\varPhi_{\theta}$ から $\tilde{x}$ を生成するために, 敵対的摂動 $\delta$ を導入します. この $\delta$ は以下の最適化によって決定します.

$$\argmin_{\|\delta\|_{\infty}\leq\varepsilon}\|n-\varPhi_{\theta}(\tilde{z_t}, t, c+\delta)\|_2^2$$

$\tilde{z}=\mathcal{E}(\tilde{x})$ で, $\varepsilon$ は微小な数, $c=\mathcal{E}_{txt}(y)$ です. この最適化にはProjected Gradient Descent (PGD) を用います. なお, 実験的に単一の $t^*$ で効果が出ることを論文では確認しています. 概念によって異なりますが, $t^*$ におけるAttack Success Rate (ASR) は例えばNudityでは $t^*=500$ 以降で増加します.

Adversarial Training for T2I Concept Erasure

さて, 単一のタイムステップでの敵対的攻撃が可能であることがわかりましたが, これが概念消去において頑健性を向上させるのかは不明です. ここでは, 「敵対的攻撃は概念消去メカニズムの耐性を向上させることができるのか？」について調査します.

RACEでは既存手法とは異なり, 対象概念のembeddingだけでなく隣接概念のembeddingも消去することを目指します. これは隣接概念によって意図せず消した概念が復活することを防ぐためです. 敵対的攻撃をESDに組み込みます.

$$L_{RACE}\coloneqq \|\varPhi_\theta(z_t, t, c+\delta)-(\varPhi_{\theta^*}(z_t, t)-\eta(\varPhi_{\theta^*}(z_t, t, c)-\varPhi_{\theta^*}(z_t, t)))\|_2^2$$

本来なら実験を行って確認するところですが, それは実験セクションに譲ります.

ここで, アルゴリズムの全体像を確認します. 二重のfor-loopになっていますが, 内側が敵対的の部分です.

実験

実験設定

基本的には既存研究と同様の評価を行います.

RACEをESDに組み込んで評価を行います. 先ほどのアルゴリズムにおいて $\varepsilon=0.1, \alpha = \dfrac{\varepsilon}{4}$ を用います. ESDのメインとなるモデル更新においては, 学習率1e-5を用いてAdamで最適化します.

頑健性をテストするために, white/black box approachで敵対的攻撃を行います. 初期評価ではI2P Red Taming Promptを用います. Black-boxのシナリオではCLIPを用いて敵対的promptを生成するPEZという手法を使用します. White-boxのシナリオではSDの勾配を用いてpromptを生成するP4DやUnlearnDiffを用います.

評価指標ですが, ASRを用います. 特に, 基盤モデルではなくドメイン固有のものを用います. 例えばstyleではImageNetで訓練され, WikiArtでfine-tuningされたViT-base, NudityはNudeNet, objectはImageNetで訓練されたResNet-50を使用します. ASRは

$$\mathrm{ASR} = \dfrac{1}{N} \sum_{i=1}^{N} \mathbb{1}(f(SD(\tilde{y_i})) = \tilde{y_i})$$

で計算します. ここで, $\tilde{y_i}$ は敵対的prompt, $f$ は分類器, $N$ はpromptの数です. 消去後の品質評価にはMS-COCOから5000枚の画像を生成して計算されるFIDとCLIP Scoreを用います.

結果

定性評価を行うのですが, ここではNudityの評価は省略しますので, 結果の方は論文を参照ください.

White-box手法であるUnlearnDiffと比較すると, RACEの方がより優れた敵対的手法と言えそうな結果になっています. ここで重要なのは, RACEの方法論的な利点がprompt生成のために外部の画像やpromptに依存しない点で, 既存手法とは一線を画している点です.

続いて, 定量評価を行います. RACEはASRを最大で30%ほど低下させることに成功しています.

上表からもわかるように, NudityでもASRを30%以上低下させていることがわかります. どのようにしてASRが改善しているかを明らかにするために消去対象のカテゴリや要素を分析します. NudeNetを用いてI2PからNudity promptsに対してoriginal SD, ESD, UnlearnDiff, RACEが生成したものを列挙します.

ESDはかなりの消去に成功していますが, UnlearnDiff攻撃によって再生成が可能になっていることがわかります. これとは対照的にRACEはUnlearnDiff攻撃に対しても頑健性を維持しています.

ただし, 表のFIDの部分から分かるように, 頑健性と引き換えに生成品質が悪化していることがわかります. いわゆるトレードオフです. その原因については考察がされていますが, あくまで考察レベルなのでここでは取り上げません. なお, 後述する工夫によってトレードオフの解消を試みます.

Disentanglement

概念消去ではいくつか達成すべき事項がありますが, そのひとつに他の概念への影響が軽微であることが挙げられます. 生成例を見てみます. たしかに, 定性的には他の概念への影響はあまりなさそうです.

定量評価も行います. 5000のpromptを作成して画像生成を行います. 前述したResNet-50を使用しtop-1 accuracyを測定します.

RACEはESDと比較して消去能力が向上していることがわかります. しかしながら, 他の概念の維持能力は低下しているようです. これについて著者らはRACEが対象概念の $\varepsilon$ 近傍を狙う方法であり, 近接概念に影響を与える可能性があるためとしています. また, 先ほどの定性評価から他の概念への視覚的影響を最小限に抑えつつ概念消去ができると述べています. 個人的にはこのロジックはやや意味不明です. というのも, 定性評価で示された概念が近接概念でない可能性があるからです. この主張をするには各クラスにおいてのaccuracyを測定する必要があり, Acc. Othersでまとめていいものではないと思います.

Discussion

先ほど頑健性と品質がトレードオフになるという話をしました. これに対して著者らは以下の式のように損失関数に正則化項を加えて実験を行いました.

$$L_{RACE+Reg}\coloneqq L_{RACE}+\lambda\|\theta-\theta^*\|_1,$$

ここで, $\theta$ はRACE, $\theta^*$ はoriginal SDのパラメータ, $\lambda$ は正則化の強さでここでは0.1を使います.

ESDと比較して頑健性を維持しつつ品質を改善していて, トレードオフの部分的解消に成功していると述べていますが, 少し微妙な結果に思えます.

この結果に出てくるRACE+keywordsですが, 対象概念が同義語で現れることを考慮したものです. 例えばNudityの場合, nude, nsfw, bareなどがあります. この手法は確かに頑健性を向上させますが, 却ってトレードオフを強めています.

思ったこと

• 敵対的学習をうまく組み込むのは新しくていいと思いました. 既存手法だと時間がかかりますが, 特定のタイムステップだけでうまくいくのも非常に魅力的です.
• 結果は微妙な気がします. 特に, 学習を行うことが前提なのでclosed-formで更新するような手法には使えない気がします
• 論文とは関係ないですが, ESDの公式実装の上に実装されているので論文の実装の仕方がかなり古いことが気になりました. ESDはdiffusersの実装も公開されているのでそっちをベースにした方が研究的にはやりやすいと思います.

参考文献

• Changhoon Kim, Kyle Min, and Yezhou Yang. R.a.c.e.: Robust adversarial concept erasure for secure text-to-image diffusion model, 2024.