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)p(x) của dữ liệu {xi}\{x_i\}. 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)\}, 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 z\boldsymbol{z}, phân bố của dữ liệu được mô tả dưới dạng điều kiện p(xz)p(x|z). Để tính marginal likelihood p(x)p(x), cũng như posterior p(zx)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 z\boldsymbol{z}. Mô hình VAE giải quyết điều này bằng cách xấp xỉ posterior p(zx)p(z|x) bằng một phân bố q(z)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}}

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

    logpθ(x)=Eθ(x)logZθ\log p_{\theta}(x)= -E_{\theta}(x) - \log Z_{\theta}

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.

    θlogpθ(x)=θEθ(x)θlogZθ\nabla_{\theta}\log p_{\theta}(x)= -\nabla_{\theta}E_{\theta}(x) - \nabla_{\theta}\log Z_{\theta}

    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\log, \exp, chain-rule và tính tuyến tính của toán tử gradient, ta có khai triển sau

    θZθ=θlogexp(Eθ(x))dx=Expθ(x)[θEθ(X)]\begin{aligned} \nabla_{\theta}Z_{\theta} &= \nabla_{\theta}\log \int\exp(-E_{\theta}(x))dx\\ &= \mathbb{E}_{x\sim p_{\theta}(x)}[-\nabla_{\theta}E_{\theta}(X)] \end{aligned}

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

    θlogpθ(x)=θEθ(x)+Expθ(x)[θEθ(X)] \nabla_{\theta}\log p_{\theta}(x)= -\nabla_{\theta}E_{\theta}(x) +\mathbb{E}_{x\sim p_{\theta}(x)}[\nabla_{\theta}E_{\theta}(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)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 xx 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 xx của log\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)

    xlogp(x)=xEθ(x)xlogZθ=xEθ(x)\nabla_x\log p(x) = -\nabla_{x}E_{\theta}(x) - \nabla_{x}\log Z_{\theta}= -\nabla_{x}E_{\theta}(x)

    Phương pháp SGLD thêm vào gradient một giá trị ngẫu nhiên để đảm bảo xx 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ệ ϵ\epsilon (có thể hiểu như learning rate khi học tham số), ta có

    xt+1xt=ϵxEθ(x)+2ϵηvới ηN(0,1)\begin{aligned} x_{t+1} - x_{t}& = -\epsilon\nabla_{x}E_{\theta}(x) +\sqrt{2\epsilon}\eta\\ \text{với}\ &\eta\sim \mathcal{N}(0,1) \end{aligned}

    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

    với BtB_t là chuyển động Brown, 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}}1^1.

    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)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(θ)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)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) để xấp xỉ đại lượng này (trong các phân bố hàm mũ, phân bố qϕq_{\phi} còn được gọi là đối ngẫu Fenchel của pθp_{\theta}). Từ bất đẳng thức Jensen ta có

    logZ(θ)=logexp(Eθ(x))dx=logexp(Eθ(x))qϕ(x)dqϕqϕ(x)logexp(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)\log\frac{\exp(-E_{\theta}(x))}{q_{\phi}(x)}dx\\ &=\mathbb{E}_{q_{\phi}(x)}[-E_{\theta}(x)] + H(q_{\phi}) \end{aligned}

    Để log\log-likelihood đạt giá trị lớn nhất, ta cần cực tiểu Eθ(x)logZθ-E_{\theta}(x) - \log Z_{\theta}, và từ công thức trên ta có chặn dưới của Z(θ)Z(\theta). Như vậy một mặt ta cần tìm θ\theta sao cho Eθ(x)logZθ-E_{\theta}(x) - \log Z_{\theta} nhỏ nhất, mặt khác ta muốn tối ưu ϕ\phi sao cho phân bố biến phân xấp xỉ logZ(θ)\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})

    Cách tiếp cận này gần giống như adversarial training trong GAN, ở đây tham số ϕ\phi của phân bố q(x)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}, 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^2, 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 TT mẫu, ta sẽ thêm TT biến ngẫu nhiên ẩn viv_i đại diện cho sự thay đổi của xx giữa mỗi bước, lúc này ta có phân bố sau

    p(x,{vi}i=1T)=exp(E(x)λ2iTvi22)Zp(x, \{v^i\}_{i=1}^T) = \frac{\exp(-E(x) - \frac{\lambda}{2}\sum^T_i||v_i||^2_2)}{Z}

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

    maxθminϕEθ(x)Eqϕ(x,{vi}i=1T)[Eθ(x)λ2iTvi22]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})

    Kí hiệu phân bố của xx{vi}i=1k\{v_i\}_{i=1}^kqk(x,{vi}i=1k)q^k(x, \{v_i\}_{i=1}^k), pHp_{\Eta} 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 xxviv^i, ta có

    qk1(xk1,{vi}i=1k1).pH(η)=qk(xk1,{vi}i=1k1,η)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)

    Ở mỗi bước thứ ii, viv_i được tính từ gradient hàm năng lượng tại xi1x_{i-1} như sau

    vi=ϵEθ(xi1)+η,ηpHi1v_i = -\epsilon\nabla E_{\theta}(x_{i-1}) +\eta\, ,\quad \eta\sim p_{\Eta_{i-1}}

    Do đó

    qk(xk1,{vi}i=1k1,η)=det(Jf)qk(xk1,{vi}i=1k1,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)

    trong đó JfJ_f 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, -\epsilon\nabla E_{\theta}(x) + v). Ma trận này có định thức bằng 11, suy ra qk(xk1,{vi}i=1k1,η)=qk(xk1,{vi}i=1k1,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)

    Sau đó xkx_k được tính từ xk1x_{k-1}vkv_k. Tương tự ta cũng có qk(xk1,{vi}i=1k1,vk)=qk(xk,{vi}i=1k1,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). 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, pHi\prod p_{\Eta_i} sẽ gần với 00. Một phương pháp tổng quát hơn2^2 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

Bình luận
Vui lòng đăng nhập để bình luận
Một số bài viết liên quan