
wd14-taggerについて色々
どうも、cellaです。エンジニアではないです。素人です。
LoRA/FPFTいずれにしても、Stable diffusionのファインチューンをするなら避けては通れないタグ付け問題。これをいかに質良くつけるかでモデルの性能に大きく差が出ることはトレーナーなら周知の事実と思います。
よく使用されているのはSmilingWolf氏のwd-tagger-v2/v3だと思います。ただしこれには大きな弱点があり、
Tags with less than 600 images were filtered out.
datasetにそのタグを含む画像が600枚以上含まれていないと、タグ自体が学習に使用されていません(タグ自体が抜かれて学習されているのか、学習してからMLPを抜いているのかはnot mentioned)。
これは非常に具合が悪く、というのも、ファインチューンで得たい概念というのは大抵それほどのタグ数が収載されていないことが多く、ルーチンで既存のwd-taggerでタグ付け→学習するだけでは全く意味がないという現実があります。基本的に手作業でつけることになります。
v1.0: P=R: threshold = 0.5296, F1 = 0.4772
またSmilingWolf氏のtaggerの性能は実は決していいものではなく、F1 scoreも0.48と据え置きです(ただし、F1 score自体はあまりアテにしない方がいいです)。使用感として規定のthresholdだとそこそこ誤タグもあるかな、という印象は使っている方なら感じられているのではないでしょうか。
FPFTは荷が重いので…
当然改良したいという欲が出てくるわけですが、SmilingWolf氏のGithubリポジトリにJAX-CVという公式のファインチューンリポジトリがありますが、こちらはFull-parameter Finetuneを前提としたものかつ、jax/fluxというおそらくSD使いには使い慣れないリポジトリで書かれているコードなので、手を出しづらいと思います。
そこで、wd-taggerにLoRAを当ててしまおう、という魂胆です。すでにある取り組みとしては、plat氏がsiglip-tagger-test-3というモデルを開発していますが、これはImage encoderとしてのgoogle/siglip-so400m-patch14-384に線形層を加え、分類器としたもので、siglip部分自体のファインチューンは行われていません。
自分は、すでにあるwdtaggerの各ブロックに対してLoRAを適応することで、新たなデータ範囲への推論性能が得られないかと考えました。tagutlでは、attentionブロック各層にLoRAを適応し、かつheadの線形層をトレーニング可能とすることで、低コストで新規のデータセットでのファインチューンができるように設計しています。Rank 64のLoRAでVRAM 48GbでBs = 16、VRAM 12GbでもBs = 4で動作することを確認しています。
結果
これを使って30万枚程度のデータセットで学習した結果です。

おおむね結果は良好で、F1 0.8以上がコンスタントに出る結果になっています(計算量の都合で、Validation Datasetからランダムに100枚で施行)。
ただ、学習途中のデータを直接見る限りFalse Positiveでpredが高いタグはほとんどが実際にはデータセット作成時にタグ付けを忘れた(か、他のタグで十分に説明できるため抜かれた)ものであることが多くありました。これはある程度予想できた結果です。

作品名は誤タグだが、ほかはキャラの特徴を考えるとおおむね正しい
したがって実際のデータでは、F1 scoreを必ずしも高くとる必要はない、というのが基本的な考え方になります。データ上のFalse Positiveをある程度許容して、True Positiveを拾える範囲を広げる(このグラフであれば、0.5程度までthresholdを下げる)のが妥当です。
Predとは一体何か?
上のグラフの横軸の値がpredですが、これを「そのカテゴリに含まれる確率」と言う人が散見されます。Encoderの出力に対して、SigmoidあるいはSoftmax関数を通すことで0-1の区間に正規化されることを便宜上確率とみなす、というような説明はしばしばNeural networkの本を読んでいると登場しますが、果たしてこれを確率と言ってしまっていいのか、と思います。
上のグラフのFalse Positiveに青いエラーバーを表示していますが、これはValidation Datasetで実際にnegativeであった画像のpred値の分布を表示しています(95%CIがどうみても怪しいですが、少なくとも平均は信頼できると思います)。
たとえば"Commentary Request"のような画像の内容によらないタグは、うまく推論できるようにはならないはずなので、negativeのpred値は高いままなかなか下がりません。このようなタグに対して「この画像のCommentary Request該当率は70%です」というのはかなり乱暴な議論であることは想像に容易いと思います。
この点の数学的な導出については自分もちゃんと取り組んだことないので、時間に余裕ができたら考えてみようかと思っています。

これはHIVEなどで表示されるResultでも同様のことが言えます。「便宜上確率と呼んでいる」程度の理解に留めておかず、ネタ抜きで実際に〇〇%でAIだ!AIじゃない!のような言説をする方は率直に申し上げてリテラシーが低いと思います。
また本当に確率と考えることができたとしても、多数のタグを同時に推論しているため、多重性(多重検定の問題)を考える必要があるので、これもまた理解していないと誤った結論を導いてしまいます。
統計学的な話は検索すればいくらでも出てくると思いますので割愛。
今後
さてモデルですが現在、5万程度までタグ数を拡張し、細かく調整を行っていっています。完成したtagger自体を今後公開するかどうかはわかりませんが、自分のdiscordサーバー内でちまちま開発を継続しています。愚痴を聞きながらテスターになってくれる方はXのDMまで。
以下は開発課程に興味のある人向けのマニアックなメモです。大したことは書いていないので、開発を応援していただける方だけ見ていっていただればと思います(モデル自体は2025.03.19現在公開していません)。