xiangze's sparse blog

機械学習、ベイズ統計、コンピュータビジョンと関連する数学について

diffusersにSGHMC(Stochastic Gradient Hamiltonian Monte Carlo)を入れる話

ブラックホールの衝突(midjourney)

画像生成AI Advent Calendar 2022 - Adventar12/17日分の記事です。

いろいろな拡散モデルの手法が実装されているdiffusers huggingface.co にSGHMC(Stochastic Gradient Hamiltonian Monte Carlo)という手法を実装し評価しようとしています。 実装は github.com にあります(feature_SGHMC branch)。 追加モジュールはscheduling_sghmc.pyで pipeline_stable_diffusion.pyにも変更点があります。

SGHMCのアルゴリズム

SGHMCのアルゴリズム
通常の変数θに対する確率的勾配降下法

\theta_{i+1} \leftarrow \nabla U(\theta_i) +N(0,B)

(Uはポテンシャルで拡散モデルでは \nabla U(\theta) はscore functionに相当する。Nは正規分布ノイズ,Bは分散)に対して 運動量rを慣性として入れることで広い範囲のサンプリングを可能としているのが特徴です(相空間を等エネルギー面に沿ってぐるぐる回るハミルトン系の要素があると効率的にサンプリングできるという直感的理解になるかと思います。 Tokyo.stanの感想 - xiangze's sparse blog) 摩擦項を入れることで定常分布となることが保証されています(Linkの論文とpdf参照)。 統計モデリングの計算で用いられるHMCとの違いは方程式に沿って進めたステップを選択する処理(Metropolis-Hastings)が存在しないところにあります。*1 (ハミモン、ゲットだぜ！ - xiangze's sparse blog)。

diffusers ソースの構造

diffusers 以下に

├── docker
├── docs
├── examples
├── scripts
├── setup.cfg
├── setup.py
├── src
├── tests
└── utils

ソースコード本体src/diffusersの他にテストコード、ドキュメント、ライブラリ生成用のコードなどがあります。 src/diffusers以下の主なディレクトリは

├── __init__.py
├── commands
├── experimental
├── models
├── pipelines
├── schedulers
└── utils

となっており

piplines以下のpipline_*.pyというファイルには拡散モデルとその周辺にあるEncoder,Decoder、Unetなど組み合わせが記述されています。 https://huggingface.co/docs/diffusers/index#diffusers-pipelines

例えばpipeline_stable_diffusion.pyにはlatent variableを用いたstable diffusionの手法が実装されています。

models以下にはVAEなど拡散モデル以外の部分のコードがあります。これらの組み合わせをpipelineでは記述しています。

schedulers以下にはscheduling_*.pyという名前で拡散モデルの本体が実装されておりstep,step_predなどの関数が1時間ステップの動作なのでその外側で変数を保持しています。 scheduling_ddpm_flax.pyなどpytorchだけでなくJAX/Flaxの実装もあります。

今回はimportのしやすさからsrcに下記のようなテストスクリプト、ipynbなどを作って作業しています。

追加点

schedulerに運動量に当たる変数pを追加し、方程式 \partial x / \partial t= M^{-1} p dt

\partial p/ \partial t = - \nabla U(x)dt - \gamma pdt + N(0,var)

を既存のschedulerを改造して実装します。Mは質量に相当しますがここでは1に固定しました。

\partial p/\partial t=0 のときに追加前の式

\partial x / \partial t= - \nabla U(x)dt + N(0,var)

(overdamped)の場合と同じようになるように係数γを決定します。すると外力- \nabla U(x)がない場合は p(t)=exp(-\gamma t) となり解は指数的に減衰するのでγは摩擦係数と解釈できます。ここで問題になるのがノイズの分散との関係です。 定常分布が

[tex: exp(-H(x,p))=exp(-1/2 pT M^{-1}p+U(x))]

がランジュバン方程式の解になるようにすると分散は \sqrt{\gamma} となります。

ただし拡散モデルではノイズの分散はステップによって変化していく(温度Tが変化して分散は \sqrt{2k_B T \gamma } となる)ので摩擦係数との関係は一定にはならない(はず)です*2。

コードはscheduling_sghmc.pyのようになり、 pipeline_stable_diffusion.pyにも変更点があります。 本来はschedulerに対応したpipelineを用意すべきですがpipeline_stable_diffusionを改造する形になってしまっています。

pipeの呼び出し、画像生成コードは以下のようになります。 pipline定義以降もschedulerを動的に差し替えることができます。

import torch
from torch import autocast
import diffusers
import datetime

import diffusers.schedulers
from diffusers import StableDiffusionPipeline
from diffusers.schedulers import SGHMCScheduler

YOUR_TOKEN ="ここにhuggingfaceで取得したtokenを書く"

def gen_images(pipe,name:str,num_images=10,init_seed=0,num_inference_steps=50):
    for i in range(num_images):
#seed固定
        SEED=i+init_seed
        generator = torch.Generator(device=DEVICE).manual_seed(SEED) 
        #print(pipe.scheduler)
        with autocast(DEVICE): 
            image = pipe(prompt, guidance_scale=7.5, num_inference_steps=num_inference_steps,generator=generator)["images"][0] 
        # 現在時間 
        dt_now = datetime.datetime.now() 
        now = dt_now.strftime("%Y%m%d%H%M%S") 
        # ファイル名 
        file_path = name+str(SEED) + "_" + str(now)+"_"+str(num_inference_steps) + ".png"
        # ファイル保存 
        image.save(file_path)  


MODEL_ID = "CompVis/stable-diffusion-v1-4"
DEVICE = "cuda"

pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID,                                         
                                               revision="fp16", torch_dtype=torch.float16, use_auth_token=YOUR_TOKEN)
pipe.to(DEVICE)

prompt = "a dog painted by Katsuhika Hokusai"

## schedulerの変更
pipe.scheduler = SGHMCScheduler.from_config(pipe.scheduler.config,sampler_type="DDPM")
gen_images(pipe,"DDPM",1,num_inference_steps=500)

途中結果

プロンプト"a dog painted by Katsuhika Hokusai"に対して

Euler discrete ベースで摩擦係数dump_coef=1(定数)の場合

dump_coef=1(const)100 step
dump_coef=1(const)990 step
犬がどんどん薄くなっていってしまう！

noise variance=varに対して dump_coef=varvar の場合

Eucler discrete dump_coef=varvar ,100 step

DDPMベースで摩擦係数dump_coef=1(定数)の場合

3000 step

運動量の可視化

そしてstepを続けるとUnetの出力がNaNになってしまう現象が見られる。 なにか仮定が不適切なのか デバッグ中です。‥ DDPM、DDIMなど各サンプリング手法では分散のとり方を工夫しているがそれと調和するようなunderdamping手法があるのか、あるいはないのか‥

評価方法

既存のscheduler(sampler)の違いが画質に及ぼす影響はこちらが参考になります。 zenn.dev bookyakuno.com

Eulerでは淡白な印象になります。DPM++2Mは数ステップで収束するらしいのですが、濃いめの出力になるようです。

定量評価はFID(Frechet Inception Distance ,2つの確率分布μ、ν間の距離を表す量で分布

FID
と定義される量)を使いますがまだ実行できてはいません。。。

追記(2/19) Score-Based Generative Modeling through Stochastic Differential Equations の4.2 PREDICTOR-CORRECTOR SAMPLERSで既にHMCの存在は考慮されておりもしかしたらうまく行かなかったのかもしれない。 Predictor-Corrector methodsとは最初に計算した差分に対して補正を行うようなサンプリングステップのことでLeapFlogに近い(Appendix G参照)

追記(3/19)

Link

SGHMCについて

pythonについて

めちゃめちゃ助けられました。型ヒント、VScode, jupyterが三種の神器

Langevin方程式、Fokker plank方程式について

やりたかったこと

経路に沿って生成した静止画像をアニメーションとして再生したかったです。 色々な方の作品をお楽しみ下さい

www.youtube.com

*1:条件分岐がないというのはCPU,GPUの分岐予測がなくてもパイプライン的に計算が実行できるという利点があるかもしれません。前後のUnetと合わせて専用回路のアイデアになります。

*2:以下の結果がうまく行かない原因かもしれません