看了一些 score matching 的論文後, 就在網路上找了一下有沒有範例, 然後找到了這個 repo: [toy_gradlogp]
非常清晰易懂, 重點實作了:
- Denoising Score Matching (DSM)
- Deep Energy Estimator Networks (DEEN, 我怎麼覺得跟 DSM 等價?!)
- Sliced Score Matching (SSM)
- Sliced Score Matching with Variance Reduction (SSM-VR)
- Langevin Dynamics (可以只根據 score function 取 sample)
雖然主要都是 PyTorch, 但 data pipeline 仍然使用 tensorflow
因此我就改寫了一下, 變成純 PyTorch, 並且也改成我習慣的資料夾結構和 config 使用 Hydra
改寫後的在這: [Score Matching Practicing in PyTorch]
以下接著說明兩個重點:
- Langevin Dynamics 簡介
- 怎麼把 gradient 也當成 loss?
Langevin Dynamics 簡介
Langevin dynamics 可以使用 score function, i.e. $\nabla_x\log p_{data}(x)$, 來取 sample, 取出來的 sample 具有 $p_{data}(x)$ 的分佈
而我們知道可以用 SM 來訓練一個 NN $s_\theta(x)$ 逼近 score function, 因此就可以用 $s_\theta(x)$ 來取 sample, 步驟如下:
給定一個固定的 step size $\epsilon>0$, initial value $z=\tilde{x}_0\sim\pi(x)$, 其中 $\pi(x)$ 是固定的 prior distribution, and $z_t\sim\mathcal{N}(0,I)$
當 $\epsilon\rightarrow 0$, and $T\rightarrow\infty$, 則 $\tilde{x}_T$ 等同於從 $p_{data}(x)$ 取樣!
只要我們能把 score function 用 NN 學得很好, 就可以利用 Langevin dynamics 採樣了. 疑?! 這樣不就完成了一個 generative modeling 了嗎?
沒錯, 這就是最初的 score-based generative modeling 想發開頭. 但這最 naive 的方法事實上會面臨到一些困難.
下一篇會介紹 Yang Song 使用這種方法時, 他是如何解決這些困難, 並成功鍊成生成模型.
怎麼把 gradient 也當成 loss?
其實直接參考這篇文章 retain_graph和create_graph参数 就可以了, 說得很清楚.
重點是利用 torch.autograd.grad 並且將其 create_graph 設定為 True
這邊簡單說明一下.
一般我們要得到 gradient 是對 loss 取 loss.backward(), pytorch 會在其 computational graph 裡做 backprop, 然後 tensor 如果 requires_grad 為 True 的話, 該 tensor 就會保留住 gradient.
但現在問題是, 這些 gradients 本身也是 loss 的一部分, 這該怎麼辦?
只要使用 torch.autograd.grad 並且將其 create_graph 設定為 True, pytorch 就會針對這些求 gradients 的 operations 生出其對應的 computational graph 並加進原來的圖裡.
也可以利用此方法求更高階的導數
在 SM 裡會用到是因為其 loss 包含 score function $\nabla_x\log NN(x;\theta)$, 而 score function 正是我們 energy-based model (是一個 NN) 的微分
一般我們定義 energy-based model $E(x)$ 為:
$$E(x)=-\log q(x)$$
其中 $q(x)$ is non-normalized distribution, i.e. $p(x)=q(x)/Z$ where $Z$ is partition function
不過如果我們的 NN 直接就是估計 score function (也就是不透過 energy-based model), 就不需要這麼做, 後面 Yang Song 的很多工作就直接這樣了.