Zenn
💭

Stable Diffusionからの概念消去㉓：EraseAnything（論文）

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

EraseAnything: Enabling Concept Erasure in Rectified Flow Transformers

CVPR関連の論文も大体確認できました (正確には増えれていないものもありますが, スコープがNSFWオンリーだったりするので触れていないだけです). さて, 最近ICMLの採択可否が出たようですが, 今回はICML2025で採択されたものを確認します. とは言ってもまだ公式にaccepted listが出ているわけではない(2025/05/04現在)ので, Twitterとかで調べて著者がacceptされたと述べているものになります.

書籍情報

断りのない限り, 図表は以下の論文からの引用です.

Daiheng Gao, Shilin Lu, Shaw Walters, Wenbo Zhou, Jiaming Chu, Jie Zhang, Bang Zhang, Mengxi Jia, Jian Zhao, Zhaoxin Fan, and Weiming Zhang. Eraseanything: Enabling concept erasure in rectified flow transformers, 2025.

関連リンク

導入

これまでは, 「Stable Diffusionからの概念消去」というタイトルをつけていましたが, それはStable Diffusion, 特にStable Diffusion 1.4を対象にした研究群だからです. ところが, 周知の通り, Stable Diffusion以外にも画像生成モデルはたくさんあります. 特に最近はFluxも使われています. Fluxが登場したのはおよそ1年前ですが, まだFluxからの概念消去の研究例はありませんでした. 今回はFluxからの概念消去になります. とは言いつつStable Diffusion 3も同じような感じです.

FluxやStable Diffusion 3

FluxやStable Diffusion 3はこれまでのStable Diffusion 1.4とは何が違うのかという点について簡単に触れておきます. Stable DiffusionはLatent Diffusion Modelsをベースにしたモデルです. 潜在空間において拡散モデルを学習しています. ところが, この拡散モデルの部分に時間がかかることが知られています. それは拡散モデルの定式化によるもの (逆過程では微小時間変化に限り遷移がガウス分布に従います)なので仕方がないのですが, 実用上は高速化したいです. これまでのメインアプローチは拡散モデルを別の側面から眺め, スケジューラで頑張るという形です. それとは別に新しい学習パラダイムも探求されていました. それがflow macthing系統です. flow macthing自体は目新しいものでもないですが, 最近は可逆でなくてもいい雰囲気なので用いられている印象です. FluxやStable Diffusion 3で用いられているのはflow-basedな手法のなかでもRectified Flowというものです.

拡散モデルを確率微分方程式の観点から見たときに,


Score-Based Generative Modeling through Stochastic Differential Equationsより引用
https://openreview.net/forum?id=PxTIG12RRHS

のような解軌道を辿って生成が行われているとされています. SDEの定式化にはランダム性があるのでその分不安定な (ジグザグな)軌道を描いていることがわかります. Rectified Flowのアプローチではこの軌道をまっすぐにすることを考えます. 時間幅はわかっている (t\in[0, 1])ので, 直線軌道にするには進むべき方向 (上の図だと直線の傾き)が分かればOKです. 詳しい部分は省略しますが, そのようにして少ないステップ数でも高品質な生成を実現します.

既存の概念消去の課題

これまでの概念消去の手法はそのままFluxに適用すると失敗してしまいます. 理由として考えられるものはいくつかありますが, ここではいくつかの要素を見てみます.

cross-attention

ひとつに明示的なcross attentionの不在がありそうです. Stable Diffusion 1.4では一般的なcross attentionが用いられています. すなわち \mathrm{Attention}(Q, K, V) のうち, K=V です. また, Q は (時刻依存の)画像オンリーの情報, K, V はテキストオンリーの情報と分けられます.　ところがFluxやStable Diffusion 3では話が変わります. まずはStable Diffusion 3のアーキテクチャを見てみます.


Scaling Rectified Flow Transformers for High-Resolution Image Synthesisより引用

y はさておき c, x の流れを追います. (b)のMM-DiT blockを見ると, Q, K, V の全てにおいて x の情報と c の情報が入っていることがわかります. 実際には y にもテキスト情報が入っていますから, 様々な部分でテキストの条件付けが行われています. また, Fluxのアーキテクチャでも,


https://www.reddit.com/r/LocalLLaMA/comments/1ekr7ji/fluxs_architecture_diagram_dont_think_theres_a/ より引用

Stable Diffusion 3と同様に情報は画像とテキストで分けられることなくミックスされます. 定式化すると

\begin{align*} &\bold{Q}=\mathtt{concat}(\bold{Q}_{text}, \bold{Q}_{pixel}, \mathtt{dim}=-1) \\ &\bold{K}=\mathtt{concat}(\bold{K}_{text}, \bold{K}_{pixel}, \mathtt{dim}=-1) \\ &\bold{W}_{\bold{attn}}=\mathtt{Softmax}(\bold{Q}\times\bold{K}) \end{align*}

です. なのでこれまでと同様なattention mapは得られなさそうな雰囲気もあります. そこで, attention mapを可視化してみます.

ところがtext embeddingとattention mapには関係がありそうな結果となりました. 実験的には \bold{W}_{\bold{attn}} で関連づけられているようです.

text encoder

概念消去では無関係な概念の保持が求められます. 既存手法を適用することは自然ですが, 大きな問題として用いているtext encoderが異なるというものがあります. FluxではT5を用いています. CLIP text encoderは単語レベルの埋め込みや類似度計算に適していますがT5は文レベルの埋め込みを目的としており, 単語レベルの類似度を捉えられません.

実際に, nudeという単語に対する類似度を, T5のtokenizerの語彙の中で見てみます.

結果は微妙です. girlはまだしも, leanやdeerがnudeと結びつくかと言われると違うように思えます (少なくとも著者は). さらに, T5の埋め込みはCLIPの 77\times768 に比べると 256\times 4096 と大きく, 敵対的なプロンプトの計算は高コストです.

ロバストな消去

ここまで既存手法ではできなさそうな理由を述べてきましたが, 実は単に消すだけなら非常に簡単です. prompt内のキーワードのトークンインデックスを特定すれば, \bold{W}_{\bold{attn}}\in[24, 1280, 1280] で対応する部分を消すことで実現できるからです. ただ, この手法ではちょっと文字を変えただけで概念を復活させることができます.

上図の例ではsoccerをsoccerrsにしたり, NikeをNikkeにするだけで概念の復活が確認できます. すなわち, 単なるattention mapの削除ではダメということがわかります.

提案手法

提案手法はattention mapの操作ではなく学習ベースでの手法となります. 例えばESDなんかは学習ベースの手法です. 学習は2段階に分かれており, それぞれ順番に見ていきます.

Lower-Level: Concept Erasure

手法は単純でLoRAを用いてFluxを更新します. 消去したい概念のデータセットを D_{un} とします. ESD同様に, 以下の損失を用います.

\mathcal{L}_{esd}=\mathbb{E}\left[v_{\theta_{o}+\Delta\theta}(x_t, c_{un}, t)-\eta\|v_{\theta_o}(x_t, c_{un}, t)-v_{\theta_o}(x_t, \emptyset, t)\|_2^2\right]

v_{\theta_o} はoriginalのパラメータ, \Delta\theta はLoRAの重みです. v を用いるのは見慣れないですが, 著者らはFlow Matchingにおける速度 (velocity)を表しており, v-predictionと概念的に等価であると述べています. 先ほどcross attetionの話をしましたが, 入力プロンプトにおいてtarget conceptに割り当てられたattention weight F_{idx}^{un}=\bold{W}_{\bold{attn}}[:, :,idx] を弱めることでさらに効果的な消去を目指します.

\mathcal{L}_{attn}=\sum_{idx=start}^{end}F_{idx}^{un}

ただ, 一般的な文章だとこのidxが固定されてしまいます. そのため, 順序を並び替えて対応します. Fluxは文章の中の順序を並び替えても同一内容の画像を生成することができるため, このような対処で適当に添字を動かします.

Upper-Level: Irrelevant Concept Preservation

毎度お馴染み他の概念を維持する手法です. ここではtarget conceptと無関係な概念 c_{ir} を同時に含むpromptを用意して画像 I_f を6-10枚ほど生成します. これらの画像を使ってLoRAを訓練します.

\mathcal{L}_{lora}=\mathbb{E}\left[\|v-v_{\theta+\Delta\theta}(u_t, c, t)\|_2^2\right]

です. ここで, v=x_T-u_{pix} で, u_{pix} は I_f をVAEで潜在表現に圧縮したものです. u_t は u_t=(1-t)u_{pix}+tx_T です. この手法はobjectのような局所的な概念には有効ですがstyleのような大域的な概念に対しては不十分です. そこで, attention mapに基づいた対照学習を行います. この手法では I_{f} を用意する必要はなく, LLMに頑張ってもらって D_{ir} を生成します. まず, GPT-4oで c_{ir}\in D_{ir} を生成します. また, 効率性の観点からNLTKを用いてtarget conceptの類義語 (syn)を生成します. 例えばnudeであればnakeなどが得られます. K=3 個の概念をここでは選びます. 次に, x_T を固定し, nudeをnakeに置き換えた c_{ir}^i (i=\{1, 2, 3\}) でdenoisingを行います. denoisingの過程でattention feature F^{un}, F^{syn}, F^{ir}=\{F^{k_1}, \ldots, F^{k_K}\} が得られます.

目標は F^{un} を動的に変化する F^{ir} に整合させつつも, F^{syn} からは離れるようにすることです. これはReverse Self Contrastive loss (RSC)によって行います.

\mathcal{L}_{rsc}=\log\left(\dfrac{\sum_{i=0}^K\exp\left(\dfrac{F^{un}\cdot F^{k_i}}{\tau}\right)}{\exp\left(\dfrac{F^{syn}\cdot F^{k_i}}{\tau}\right)}\right)

となります. \tau は温度パラメータでハイパーパラメータとなりますが, 著者らは様々なケースでのテスト結果から \tau=0.07 としたようです. 余談ですが, この損失関数だけ見ると新しい感じがしますがやってることはただの勾配上昇です.

アルゴリズムのまとめ

最終的には以下の最適化問題に帰着されます.

\min\mathcal{L}_{lora+rsc}(\Delta^*\theta; D_{ir}) \quad \mathrm{s.t.}\ \Delta^*\theta=\min\mathcal{L_{esd+attn}}(\Delta\theta; D_{un})

実験

Flux.1 [dev]に対して適用します. 特に断りのない限りはflow matching Euler samplerを28 stepsで使います. また, LoRAモジュールはAdamWを1000 stepsで更新します. low-levelの学習率は0.001, upper-levelの学習率は0.0005です. ESDのguidanceは \eta=1 です. LoRAは \bold{Q} と \bold{K} の部分集合である \mathtt{add}\_\mathtt{q}\_\mathtt{proj}, \mathtt{add}\_\mathtt{k}\_\mathtt{proj} に取り付けます. これはtext-related parameterです. 実験はNVIDIA A100 GPU, batch size 1で行われています.

NSFW

I2P datasetでの結果をみます. 既存研究と評価条件は同じなので特に述べません. 結果をみます.

最初の方で既存手法がFluxに適用できない理由を述べていたのにFluxに適用するのか...という感じはしますが, なぜか消去性能はUCEに劣っています. その辺りの分析はないです.

Entity (object), Abstraction (style), Relationship

Entity (object), Abstraction (style), Relationshipの3カテゴリの消去を試みます. Relationshipというのは例えば「AとBが握手している」の「握手している」の部分になります. まずは定性的な結果を見てみます.

なぜ古めの手法であるACと比較しているのかは分かりませんが, 少なくともACで消せていない概念も消せていることがわかります. また, 複数概念の消去もでき,

のような結果となります. 続いて定量評価を確認します. こちらはMACEなどでも使われているACC_eなどを使いますが, 分類器としてCLIPを使っている点に注意が必要です.

MACEではoriginalモデルの結果も載っていたのですが, ここでは載っていません. これでは表からわかるようにCAとの比較で数値がいいことしか読み取れず, 実際にどの程度の消去性能かは不明です. 例えばMACEの論文では, AirplaneのACC_eは96.06です. FluxはStable Diffusion 1.4より性能が良いことを勘案するとoriginalモデルのACC_eはさらに高いと考えられます. このことを考えると14.8が12.5になったのは絶対的な優位性はなさそうに思えてきます (例えば実行時間を比較して提案手法の方が明らかに長い訓練時間を要する場合は「CAでよくないか」という話になってしまうと思います).

User Study

この手の研究では珍しく, 人による評価を行っています. 評価項目は5つです.

  1. Erasing Cleanliness: target conceptを含むプロンプトでtarget conceptが出現しないこと
  2. Irrelevant Preservation: 他の概念が適切に出現すること
  3. Imaging Quality: 説明はありませんが生成画像の品質でしょう
  4. Prompt Adherence: これも説明はありませんがtext-image alignmentのことでしょう
  5. Output Diversity: これも説明はありませんが生成画像の多様性のことと考えられます

1, 2と3-5では聞き方が違います. 1, 2では以下のUIで行います.

赤が消した概念, 緑がその他の概念に関する問いです.

3-5では以下のUIで行います.

どちらの場合も5が最高評価です. 参加者はartistではない人を20人集めて行います. それぞれについて200枚の評価を行い, 結果を測定します. では, 結果を確認します.

全体的に良い性能であることがわかります. これもOriginalモデルの評価があると3-5の結果の受け止め方がわかりやすいので不親切だなと感じます.

思ったこと

  • 実験のセクションでは記述が雑な印象を受けます. 少なくとも読み進めるのに必要十分な情報は提供されていないように思えます. 例えば複数概念の消去結果を示していますが, 「同時の消去」か「順番の消去」かは不明ですし, 後者の場合LoRAは別々に用意してマージするなどの手法も考えられます. Limitationでは同時消去が示唆されています. 揚げ足を取るような感じになってしまいますが, Entity, Abstraction, Relationshipの定量評価では生成枚数も不明で判断に困ります (relationshipがCLIPで判定できるのかもよくわかっていないです).
  • 一方で画像生成のスタンダードはFluxやStable Diffusion 3.5に移っているのでFluxに着目した研究はとても良いと思います. Stable Diffusion 3.5でもやればよりよいように思えます.
  • LoRAを取り付けるタイプの概念消去はblack-boxで見た時に消去されているように振る舞うだけで実際に消去されているわけではない (LoRAを外せば元のモデル)ので, 難しいところです.

参考文献

  • Daiheng Gao, Shilin Lu, Shaw Walters, Wenbo Zhou, Jiaming Chu, Jie Zhang, Bang Zhang, Mengxi Jia, Jian Zhao, Zhaoxin Fan, and Weiming Zhang. Eraseanything: Enabling concept erasure in rectified flow transformers, 2025.

Discussion