Stable Diffusionモデルの類似度計算(ASimilarity)に対する攻撃

gcem156

　Acertaintyの作者であるJosephusCheung氏が公開しているモデル類似度計算プログラムをだます方法について紹介します。

　NovelAIリーク派生であるかどうかを確認するためにちょこちょこ使われているようですが、簡単に結果をぐちゃぐちゃにできます。
　今回はSD2.1と同じ画像を生成するモデルにもかかわらず、類似度が低くなるようなモデルを作成します。Colabに全実装をあげています。

ASimilarityの仕組み

　ASimilarityでは2つのモデルでSelf Attention層の途中の $${qk^Tv }$$ を計算しコサイン類似度を出しているようです。何でこうしてるかは知りません。

攻撃手法

　今回はSelf Attention層において、トークンの次元方向の順序をばらばらにすることで、結果を改変します。 順序をばらばらにすると、$${qk^Tv }$$は変わりますが、to_outでその順序を元に戻すだけで、生成される画像が変わらないようにできます。次元方向に順序を入れ替えるのは、行列の順序入れ替えのみによって実現できます。そのため改変するのは重みファイルの数値のみで、生成コードやモデル構造等を変える必要はありません。

実装

今回は次元方向を半分にわけてひっくり返すだけの処理をやってみます。multi head attentionの分割を気にしなければいけなそうですが、この方法なら多分変わらないです。

new_state_dict = state_dict.copy()
for n in range(3,11):
    for x in ["q","k","v"]:
        to_x = state_dict[f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_{x}.weight"]
        output_size , _ = to_x.shape

        #トークンの次元方向の順番を入れ替える
        to_x = torch.cat([to_x[output_size//2:,:],to_x[:output_size//2,:]], dim = 0)
        new_state_dict[f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_{x}.weight"] = to_x

    #to_outで順番を元に戻す
    to_out = state_dict[f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_out.0.weight"]
    to_out = torch.cat([to_out[:,output_size//2:],to_out[:,:output_size//2]], dim = 1)
    new_state_dict[f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_out.0.weight"] = to_out

q, kは変えなくていいかも・・・。

結果

SD2.1でプロンプトは、a professional photograph of an astronaut riding a horse

改変前
改変後

　なんかちょっと違いますね。数学的には多分同じですが実装上は全く同じにはならないようです。ただほとんど同じような画像が生成できるようです。そして類似度は・・・

seed: 114514

base: v2-1_768-ema-pruned.ckpt [4bdfc29c]

v2-1_768-ema-pruned-pseudo.ckpt [9c87c7ad] - -11.01%

qwerty.pyの出力

マイナス11.01%ｗｗｗ。ここまで低いと逆に怪しくなっちゃいますが、順番の入れ替えを少なくすれば、ほどよい結果になるよう調整できるんじゃないでしょうか。

対策

　この手法に対する対策だけで言えば、to_outまで計算する、で終わるんですが、こういうのはいたちごっこになるだけなので根本的な解決にはなってないですね。
　また今回の手法は順番を入れ替えてるだけで、分布は変わらないので平均や分散をとってみるというのもありかもしれません。それも定数倍して元に戻すみたいなのすればよさそうですが。

おわりに

　こんな改変する人いないと思いますが、ASimilarityみたいなプログラムを利用するときは、結果だけでなく実装も注意深くみておくべきですね。