Mô hình năng lượng (Energy-based models – EBM) và một số cách huấn luyện.

Mục tiêu của mô hình sinh là học phân bố $p(x)$ của dữ liệu {xi}{x_i}{xi​}. Cụ thể hơn, cho trước một họ các phân bố có thể học được {pθ(x)}{p_{theta}(x)}{pθ​(x)}, ta sẽ tìm phân bố có (log-)likelihood lớn nhất, hay hiểu một cách trực qua là phân bố tương thích nhất với tập dữ liệu cho trước. Việc này cũng tương đương với tìm phân bố có cross-entropy với dữ liệu nhỏ nhất. Tuy nhiên đại lượng này không phải lúc nào cũng tính được dễ dàng.

Ví dụ: Với VAE, ta có biến ngẫu nhiên ẩn $\mathbf{z}$, phân bố của dữ liệu được mô tả dưới dạng điều kiện $p(x|z)$. Để tính marginal likelihood $p(x)$, cũng như posterior $p(z|x)$ khi huấn luyện và suy luận, ta cần lấy tích phân theo biến liên tục nhiều chiều $\mathbf{z}$. Mô hình VAE giải quyết điều này bằng cách xấp xỉ posterior $p(z|x)$ bằng một phân bố $q(z)$ khác.

Một cách tiếp cận khác là mô hình phân bố dưới dạng không chuẩn hóa. Cụ thể hơn, hàm mật độ của một phân bố dương bất kì với tham số $\theta$ có thể viết được dưới dạng

pθ(x)=exp⁡(−Eθ(x))Zθp_{theta}(x) = frac{exp(-E_{theta}(x))}{Z_{theta}}
pθ​(x)=Zθ​exp(−Eθ​(x))​

với Eθ(x)E_{theta}(x)Eθ​(x) được gọi là hàm năng lượng của $x$, Zθ=∫exp⁡(−Eθ(x))dxZ_{theta}=intexp(-E_{theta}(x))dxZθ​=∫exp(−Eθ​(x))dx là hằng số chỉ phụ thuộc vào tham số $\theta$ nhằm bảo đảm $p(x)$ là hàm mật độ xác suất. Log-likelihood của phân bố này sẽ là

log⁡pθ(x)=−Eθ(x)−log⁡Zθlog p_{theta}(x)= -E_{theta}(x) – log Z_{theta}
logpθ​(x)=−Eθ​(x)−logZθ​

Phương pháp Markov chain Monte Carlo với Langevin dynamics

Ta sẽ đi tìm hàm năng lượng thỏa mãn phân bố có log-likehood lớn nhất dựa vào gradient.

∇θlog⁡pθ(x)=−∇θEθ(x)−∇θlog⁡Zθnabla_{theta}log p_{theta}(x)= -nabla_{theta}E_{theta}(x) – nabla_{theta}log Z_{theta}
∇θ​logpθ​(x)=−∇θ​Eθ​(x)−∇θ​logZθ​

Gradient của hàm năng lượng có thể tính một cách dễ dàng. Với hằng số chuẩn hóa, sử dụng đạo hàm của hàm $\log ,\exp$, chain-rule và tính tuyến tính của toán tử gradient, ta có khai triển sau

∇θZθ=∇θlog⁡∫exp⁡(−Eθ(x))dx=Ex∼pθ(x)[−∇θEθ(X)]begin{aligned}
nabla_{theta}Z_{theta} &= nabla_{theta}log intexp(-E_{theta}(x))dx\
&= mathbb{E}_{xsim p_{theta}(x)}[-nabla_{theta}E_{theta}(X)] end{aligned}∇θ​Zθ​​=∇θ​log∫exp(−Eθ​(x))dx=Ex∼pθ​(x)​[−∇θ​Eθ​(X)]​

Lúc này gradient của log-likelihood được viết lại như sau

∇θlog⁡pθ(x)=−∇θEθ(x)+Ex∼pθ(x)[∇θEθ(X)] nabla_{theta}log p_{theta}(x)= -nabla_{theta}E_{theta}(x) +mathbb{E}_{xsim p_{theta}(x)}[nabla_{theta}E_{theta}(X)] ∇θ​logpθ​(x)=−∇θ​Eθ​(x)+Ex∼pθ​(x)​[∇θ​Eθ​(X)]

Giá trị kì vọng của gradient có thể xấp xỉ bằng cách lấy mẫu từ phân bố $p(x)$ thông qua Markov chain Monte Carlo (MCMC). Với phân bố nhiều chiều, cách tiếp cận thường được sử dụng là Stochastic gradient Langevin dynamic (SGLD). Một cách trực quan, ta khởi tạo $x$ ngẫu nhiên, sau đó ‘di chuyển’ x để sao cho mô phỏng phân bố tốt nhất. Việc thay đổi này được thực hiện bằng cách tính gradient tại $x$ của $\log$ xác suất. Ở bước này, ta có thể tận dụng cách mô hình của EBM để bỏ qua hằng số chuẩn hóa (lưu ý trong quá trình lấy mẫu $\theta$ không đổi)

∇xlog⁡p(x)=−∇xEθ(x)−∇xlog⁡Zθ=−∇xEθ(x)nabla_xlog p(x) = -nabla_{x}E_{theta}(x) – nabla_{x}log Z_{theta}= -nabla_{x}E_{theta}(x)
∇x​logp(x)=−∇x​Eθ​(x)−∇x​logZθ​=−∇x​Eθ​(x)

Phương pháp SGLD thêm vào gradient một giá trị ngẫu nhiên để đảm bảo $x$ là phân bố thực sự chứ không chỉ hội tụ để giá trị MAP. Cụ thể hơn, với tỉ lệ $ϵ$ (có thể hiểu như learning rate khi học tham số), ta có

xt+1−xt=−ϵ∇xEθ(x)+2ϵηvới η∼N(0,1)begin{aligned}
x_{t+1} – x_{t}& = -epsilonnabla_{x}E_{theta}(x) +sqrt{2epsilon}eta\
text{với} &etasim mathcal{N}(0,1)
end{aligned}xt+1​−xt​với ​=−ϵ∇x​Eθ​(x)+2ϵ​ηη∼N(0,1)​

Ví dụ minh họa cho phân bố chuẩn 2 chiều

import matplotlib.pyplot as plt
import numpy as np
mus = np.random.rand(2)
sigmas = np.array([[2, 1], [1, 2]], dtype = float)
pre = np.linalg.inv(sigmas)
def grad(point, mean, pre):
    #gradient hàm năng lượng của phân bố chuẩn
    return - np.matmul(point - mean, pre) 
x = np.random.rand(2)
lr = 0.3
arr = [np.copy(x)]
step = 1000
for i in range(step):
    x_grad = grad(x, mus, pre)
    noise = np.random.normal(0, 1, size = 2)
    x += lr * x_grad + noise * np.sqrt(2 * lr) # bỏ nhiễu để thấy sự khác nhau giữa nghiệm MAP và SGLD
    arr.append(np.copy(x))
plt.scatter(*zip(*arr))
plt.show()

Cách làm này có thể xem như xấp xỉ của phương trình vi phân ngẫu nhiên Langevin

dXt=−∇Eθ(Xt)dt+2dBtdX_t = -nabla E_{theta}(X_t)dt + sqrt2 dB_t
dXt​=−∇Eθ​(Xt​)dt+2​dBt​

với BtB_tBt​ là chuyển động Brownian, vi phân của nó được mô hình thông qua phân bố chuẩn. Phương trình này có nghiệm duy nhất là phân bố p(x)=exp⁡(−Eθ(x))Zθp(x)=frac{exp(-E_{theta}(x))}{Z_{theta}}p(x)=Zθ​exp(−Eθ​(x))​1^11.

Bên cạnh việc sử dụng để huấn luyện mô hình, việc lấy mẫu trực tiếp từ $p(x)$ còn được áp dụng để sinh dữ liệu từ phân bố đã được học.

Adversarial training

Như ta đã thấy ở trên, vấn đề của việc học và suy luận với EBM là tính toán cho hằng số chuẩn hóa $Z(\theta )$. Phương pháp MCMC giải quyết điều này bằng việc tìm cách lấy mẫu từ phân bố $p(x)$. Một cách tiếp cận khác là dùng một phân bố biến phân qϕ(x)q_{phi}(x)qϕ​(x) để xấp xỉ đại lượng này (trong các phân bố hàm mũ, phân bố qϕq_{phi}qϕ​ còn được gọi là đối ngẫu Fenchel của pθp_{theta}pθ​). Từ bất đẳng thức Jensen ta có

log⁡Z(θ)=log⁡∫exp⁡(−Eθ(x))dx=log⁡∫exp⁡(−Eθ(x))qϕ(x)dqϕ≥∫qϕ(x)log⁡exp⁡(−Eθ(x))qϕ(x)dx=Eqϕ(x)[−Eθ(x)]+H(qϕ)begin{aligned}
log Z(theta) &= log int exp(-E_{theta}(x))dx\
&= log int frac{exp(-E_{theta}(x))}{q_{phi}(x)}dq_{phi}\
&geq int q_{phi}(x)logfrac{exp(-E_{theta}(x))}{q_{phi}(x)}dx\
&=mathbb{E}_{q_{phi}(x)}[-E_{theta}(x)] + H(q_{phi})
end{aligned}logZ(θ)​=log∫exp(−Eθ​(x))dx=log∫qϕ​(x)exp(−Eθ​(x))​dqϕ​≥∫qϕ​(x)logqϕ​(x)exp(−Eθ​(x))​dx=Eqϕ​(x)​[−Eθ​(x)]+H(qϕ​)​

Để $\log$-likelihood đạt giá trị lớn nhất, ta cần cực tiểu −Eθ(x)−log⁡Zθ-E_{theta}(x) – log Z_{theta}−Eθ​(x)−logZθ​, và từ công thức trên ta có chặn dưới của $Z(\theta )$. Như vậy một mặt ta cần tìm $\theta$ sao cho −Eθ(x)−log⁡Zθ-E_{theta}(x) – log Z_{theta}−Eθ​(x)−logZθ​ nhỏ nhất, mặt khác ta muốn tối ưu $\varphi$ sao cho phân bố biến phân xấp xỉ $\log Z(\theta )$ tốt nhất. Tổng hợp lại ta có muốn hàm mục tiêu sau

max⁡θmin⁡ϕ−Eθ(x)−Eqϕ(x)[−Eθ(x)]−H(qϕ) max_{theta}min_{phi} -E_{theta}(x) – mathbb{E}_{q_{phi}(x)}[-E_{theta}(x)] – H(q_{phi})
θmax​ϕmin​−Eθ​(x)−Eqϕ​(x)​[−Eθ​(x)]−H(qϕ​)

Cách tiếp cận này gần giống như adversarial training trong GAN, ở đây tham số $\varphi$ của phân bố $q(x)$ đóng vai trò như tham số của mô hình sinh (không phải mô hình sinh của EBM mà ta cần), tham số $\theta$ đóng vai trò như discriminator khi kiểm tra phân bố của mô hình sinh thông qua log-likelihood.

Một vấn đề của phương pháp này nằm ở việc tính entropy của qϕq_{phi}qϕ​, do ta cần phải biết chính xác giá trị của hàm mật độ. Có một vài cách tiếp cận bằng việc dùng một phân bố biến phân nữa để xấp xỉ.

Một cách giải quyết khác là mô tả quá trình lấy mẫu dưới dạng flow2^22, lúc này hàm mật độ sẽ được tính thông qua hàm mật độ tại thời điểm trước. Ví dụ với phương pháp SGLD bên trên, giả sử cần lấy $T$ mẫu, ta sẽ thêm $T$ biến ngẫu nhiên ẩn viv_ivi​ đại diện cho sự thay đổi của $x$ giữa mỗi bước, lúc này ta có phân bố sau

p(x,{vi}i=1T)=exp⁡(−E(x)−λ2∑iT∣∣vi∣∣22)Zp(x, {v^i}_{i=1}^T) = frac{exp(-E(x) – frac{lambda}{2}sum^T_i||v_i||^2_2)}{Z}

p(x,{vi}i=1T​)=Zexp(−E(x)−2λ​∑iT​∣∣vi​∣∣22​)​

Likelihood lúc này là của phân bố trên cả $x$ và viv_ivi​, tuy nhiên do $x$ và viv_ivi​ độc lập nên tối ưu đại lượng này tương đương với cực đại $\log$-likelihood của phân bố marginal trên $x$. Hàm mục tiêu lúc này trở thành

max⁡θmin⁡ϕ−Eθ(x)−Eqϕ(x,{vi}i=1T)[−Eθ(x)−λ2∑iT∣∣vi∣∣22]−H(qϕ) max_{theta}min_{phi} -E_{theta}(x) – mathbb{E}_{q_{phi}(x,{v^i}_{i=1}^T)}[-E_{theta}(x)- frac{lambda}{2}sum^T_i||v_i||^2_2] – H(q_{phi})
θmax​ϕmin​−Eθ​(x)−Eqϕ​(x,{vi}i=1T​)​[−Eθ​(x)−2λ​i∑T​∣∣vi​∣∣22​]−H(qϕ​)

Kí hiệu phân bố của $x$ và {vi}i=1k{v_i}_{i=1}^k{vi​}i=1k​ là qk(x,{vi}i=1k)q^k(x, {v_i}_{i=1}^k)qk(x,{vi​}i=1k​), pHp_{Eta}pH​ là phân bố của nhiễu thêm vào. Do nhiễu $\eta$ tại mỗi bước độc lập với $x$ và viv^ivi, ta có

qk−1(xk−1,{vi}i=1k−1).pH(η)=qk(xk−1,{vi}i=1k−1,η)q^{k-1}(x_{k-1}, {v_i}_{i=1}^{k-1}).p_{Eta}(eta)=q^{k}(x_{k-1}, {v_i}_{i=1}^{k-1}, eta)
qk−1(xk−1​,{vi​}i=1k−1​).pH​(η)=qk(xk−1​,{vi​}i=1k−1​,η)

Ở mỗi bước thứ $i$, viv_ivi​ được tính từ gradient hàm năng lượng tại xi−1x_{i-1}xi−1​ như sau

vi=−ϵ∇Eθ(xi−1)+η ,η∼pHi−1v_i = -epsilonnabla E_{theta}(x_{i-1}) +eta, ,quad etasim p_{Eta_{i-1}}
vi​=−ϵ∇Eθ​(xi−1​)+η,η∼pHi−1​​

Do đó

qk(xk−1,{vi}i=1k−1,η)=∣det⁡(Jf)∣qk(xk−1,{vi}i=1k−1,vk)q^{k}(x_{k-1}, {v_i}_{i=1}^{k-1}, eta)= |det(J_f)| q^{k}(x_{k-1}, {v_i}_{i=1}^{k-1}, v_k)
qk(xk−1​,{vi​}i=1k−1​,η)=∣det(Jf​)∣qk(xk−1​,{vi​}i=1k−1​,vk​)

trong đó JfJ_fJf​ là ma trận Jacobian của ánh xạ f:(x,V,v)↦(x,V,−ϵ∇Eθ(x)+v)f:(x, V, v) mapsto (x, V, -epsilonnabla E_{theta}(x) + v)f:(x,V,v)↦(x,V,−ϵ∇Eθ​(x)+v). Ma trận này có định thức bằng $1$, suy ra qk(xk−1,{vi}i=1k−1,η)=qk(xk−1,{vi}i=1k−1,vk)q^{k}(x_{k-1}, {v_i}_{i=1}^{k-1}, eta)= q^{k}(x_{k-1}, {v_i}_{i=1}^{k-1}, v_k)qk(xk−1​,{vi​}i=1k−1​,η)=qk(xk−1​,{vi​}i=1k−1​,vk​)

Sau đó xkx_kxk​ được tính từ xk−1x_{k-1}xk−1​ và vkv_kvk​. Tương tự ta cũng có qk(xk−1,{vi}i=1k−1,vk)=qk(xk,{vi}i=1k−1,vk)q^{k}(x_{k-1}, {v_i}_{i=1}^{k-1}, v_k)=q^{k}(x_{k}, {v_i}_{i=1}^{k-1}, v_k)qk(xk−1​,{vi​}i=1k−1​,vk​)=qk(xk​,{vi​}i=1k−1​,vk​). Như vậy ta có thể tính được hàm mật độ thông qua hàm mật độ của các phân bố khởi tạo.

Tuy nhiên phương pháp này có nhược điểm khi dùng số bước lớn, ∏pHiprod p_{Eta_i}∏pHi​​ sẽ gần với $0$. Một phương pháp tổng quát hơn2^22 sử dụng Hamilton dynamic có tính chất hàm mật độ không đổi sau mỗi bước, giúp việc tính toán đơn giản hơn.

Tham khảo

How to Train Your Energy-Based Models