Kiến trúc mạng học sâu: từ nền tảng đến LLM·Chương 9 — Autoencoder và VAE: latent representation

Chế độ xem

Variational Autoencoder: latent space có cấu trúc xác suất

Bài này trả lời lỗ của autoencoder khi ta muốn sinh mẫu: mỗi đầu vào bị gán đúng một điểm trong không gian mã. Những điểm đó có thể thành những hòn đảo. Khoảng giữa các đảo không được mất mát tái tạo chăm sóc, nên giải mã một điểm rơi vào khoảng giữa có thể ra dữ liệu vô nghĩa. VAE, variational autoencoder, đổi điểm ấy thành một vùng xác suất.

Encoder, khối biến đầu vào thành mã, sẽ nêu tâm và độ rộng của vùng đó. Decoder, khối đọc mã để dựng dữ liệu, nhận một mẫu rút từ vùng ấy. Bài này không dạy cách GAN cho hai mạng học chống nhau, và không coi mọi mẫu rút từ vùng mã là một ảnh đã đẹp. Nó giải thích vì sao bước rút mẫu ngẫu nhiên chặn lan truyền ngược, cách viết lại bước đó, và vì sao mất mát phải cân việc dựng lại dữ liệu với việc giữ phân phối mã gần một phân phối đơn giản.

Đọc nhanh. Đi từ lỗ một điểm mỗi đầu vào, sang tâm và độ rộng, sang lý do không lan truyền ngược qua một lần rút mẫu, rồi sang hai thành phần mất mát và cách sinh mẫu mới. Các mục Đào sâu không nằm trên đường đó.

Một điểm mỗi đầu vào thì lấy mẫu không có chỗ đứng

Autoencoder học một mã xác định. Cùng một đầu vào, encoder trả cùng một điểm. Mất mát tái tạo chỉ nhìn những điểm encoder thực sự tạo ra cho dữ liệu huấn luyện. Nó không yêu cầu các điểm nằm gần nhau phải giải mã thành dữ liệu đổi một cách mượt, và không chỉ một phân phối để rút mã mới.

Muốn sinh, ta cần một điều mạnh hơn tái tạo. Ta cần biết mã mới nên rút từ phân phối nào, và muốn những vùng hay được rút ấy giải mã ra dữ liệu hợp lệ. Hợp lệ ở đây là một mục tiêu, không phải một việc VAE bảo đảm với mọi lần rút. VAE tiến tới mục tiêu đó bằng cách biến mã thành biến ngẫu nhiên có phân phối, thay vì một tọa độ cứng.

Encoder trả một vùng, không trả một điểm

Encoder của VAE không trả một mã duy nhất. Nó trả một phân phối phỏng đoán trên các mã có thể đi cùng đầu vào đang xét. Phân phối phỏng đoán đó được gọi là q. Đây là ước lượng do encoder đưa ra, không phải phân phối hậu nghiệm tính đóng bằng định lý Bayes.

Bộ trọng số của encoder, thứ được cập nhật khi huấn luyện, được gọi là phi.

Với mỗi đầu vào, q thường là một phân phối chuẩn chéo: các chiều của mã độc lập với nhau khi đã biết đầu vào. Tâm của phân phối đó, giá trị trung bình riêng cho đầu vào đang xét, được gọi là mu. Độ lệch chuẩn quanh tâm, cho biết một lần rút được phép nằm xa tâm đến đâu, được gọi là sigma. Phương sai là bình phương của sigma. Vì các chiều độc lập, ma trận hiệp phương sai là ma trận chéo của các phương sai ấy.

Gộp lại, phân phối encoder đưa ra cho một đầu vào có dạng sau. Ký hiệu diag⁡\operatorname{diag} nghĩa là ma trận chéo.

qϕ(z∣x)=N(μ(x),diag⁡(σ2(x)))q_\phi(z\mid x)=\mathcal{N}\left(\mu(x),\operatorname{diag}(\sigma^2(x))\right)

Với cùng một đầu vào, encoder mô tả cả một vùng. Mu là tâm vùng. Sigma là mức phân tán quanh tâm. Một lần rút cụ thể, ký hiệu zz, là một điểm nằm trong vùng đó, không phải toàn bộ vùng.

Kiến trúc VAE cơ bản

Hình 1. Cấu trúc VAE cơ bản. Nguồn: EugenioTL, Wikimedia Commons, CC BY-SA 4.0. Trang giấy phép: https://commons.wikimedia.org/wiki/File:VAE_Basic.png

Phương sai phải dương. Nếu mạng xuất thẳng phương sai, ta phải ép đầu ra dương, và dễ gặp rắc rối số học khi giá trị quá gần không. Cách phổ biến là để mạng xuất log của phương sai, gọi là logvar, rồi đổi lại thành độ lệch chuẩn bằng hàm mũ. Log của phương sai bằng hai lần log của độ lệch chuẩn, nên một nửa logvar chính là log của sigma.

log⁡σ2=2log⁡σ\log \sigma^2 = 2\log\sigma

Từ đó độ lệch chuẩn là:

σ=exp⁡(12log⁡σ2)\sigma = \exp\left(\frac{1}{2}\log\sigma^2\right)

Trong code, dòng tương ứng là lũy thừa của một nửa logvar.

std = torch.exp(0.5 * logvar)

Vì sao không lan truyền ngược qua một lần rút mẫu

Lan truyền ngược, hay backprop, là cách gradient của mất mát đi từ đầu ra về các trọng số. Đường đi đó cần các phép tính trơn: đổi tham số một chút thì đại lượng phía sau đổi một chút, và tỉ số ấy là đạo hàm.

Nếu giữa encoder và decoder ta rút ngẫu nhiên một mã từ phân phối chuẩn tâm mu, độ rộng sigma, bước rút không phải một phép tính trơn của mu và sigma. Kết quả phụ thuộc bộ số vừa gieo. Không có đạo hàm của lần gieo ấy theo tâm, cũng không có đạo hàm theo độ rộng. Gradient của mất mát vì vậy không đi ngược qua bước lấy mẫu để cập nhật encoder. Decoder vẫn nhận một mã, nhưng encoder không học được từ mã đó theo cách backprop thông thường.

Reparameterization trick, trick tái tham số hóa, giữ tính ngẫu nhiên nhưng đổi chỗ nó đứng. Ta vẫn muốn mã ngẫu nhiên. Ta gieo một nhiễu từ phân phối chuẩn chuẩn, không phụ thuộc trọng số nào của mạng, rồi tính mã bằng cộng và nhân. Với nhiễu đã gieo xong, mã là hàm trơn của mu và sigma, nên đạo hàm đi qua được phép cộng và phép nhân ấy, tới encoder.

Nhiễu ký hiệu là epsilon, rút từ phân phối chuẩn nhiều chiều, tâm không, hiệp phương sai là ma trận đơn vị.

ϵ∼N(0,I)\epsilon\sim\mathcal{N}(0,I)

Mã là tâm cộng độ lệch chuẩn nhân từng thành phần với nhiễu ấy. Tích ⊙\odot là nhân từng chiều, vì sigma và epsilon đều là vector.

z=μ+σ⊙ϵz=\mu+\sigma\odot\epsilon

VAE với reparameterization trick

Hình 2. Reparameterized VAE. Nguồn: EugenioTL, Wikimedia Commons, CC BY-SA 4.0. Trang giấy phép: https://commons.wikimedia.org/wiki/File:Reparameterized_Variational_Autoencoder.png

def reparameterize(mu, logvar):
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return mu + eps * std

Epsilon mỗi lần huấn luyện là một lần gieo mới. Mu và logvar là hàm của đầu vào và của trọng số encoder. Tách hai thứ đó ra là toàn bộ nội dung của trick: tính ngẫu nhiên không còn nằm trong một nút mà backprop không đi qua.

Hai lực trên cùng một mất mát

Nếu chỉ tối ưu tái tạo, ta quay về autoencoder, dù encoder có xuất thêm độ rộng. VAE thêm một lực thứ hai: phân phối mã của từng mẫu không được đi quá xa một prior đơn giản. Prior ở đây vẫn là giả định về mã trước khi thấy mẫu, thường là phân phối chuẩn chuẩn.

p(z)=N(0,I)p(z)=\mathcal{N}(0,I)

Khoảng cách giữa hai phân phối được đo bằng KL, viết tắt của Kullback–Leibler. KL là một số không âm. Nó bằng không khi hai phân phối trùng nhau, và lớn khi phân phối mã do encoder tạo ra đi xa phân phối chuẩn mà ta biết cách rút mẫu. KL trong bài này luôn là KL từ q sang prior, theo chiều đó.

Hai lực này không được tối ưu riêng lẻ. Hàm mất mát cân bằng reconstruction và KL: một bên giữ đủ thông tin để dựng lại dữ liệu, bên kia giữ phân phối mã gần prior để lúc sau còn biết rút mã từ đâu.

LVAE=Lrec+DKL(qϕ(z∣x)  ∥  p(z))\mathcal{L}_{VAE} = \mathcal{L}_{rec} + D_{KL}\left(q_\phi(z\mid x)\;\|\;p(z)\right)

Trong thực hành thường có một trọng số dương, ký hiệu β\beta, nhân vào vế KL. Trọng số bằng một thì hai vế đứng như công thức trên. Trọng số khác một là cách chủ động lệch về một trong hai lực. Họ này được gọi là beta-VAE.

Lβ-VAE=Lrec+β DKL(qϕ(z∣x)  ∥  p(z))\mathcal{L}_{\beta\text{-VAE}} = \mathcal{L}_{rec} + \beta\,D_{KL}\left(q_\phi(z\mid x)\;\|\;p(z)\right)

Vế tái tạo buộc decoder giữ đủ thông tin của đầu vào. Chỉ có vế này, bản dựng lại có thể rất tốt, trong khi không gian mã vẫn lộn xộn và khó rút ngẫu nhiên. Vế KL kéo q của từng mẫu về prior chuẩn. Vế này quá mạnh thì các mã trông rất hợp prior, nhưng encoder có thể không giữ đủ thông tin về đầu vào, và bản dựng lại kém đi. VAE vì vậy luôn đứng trong một đánh đổi giữa độ trung thực của bản dựng lại và độ đều của không gian mã.

Sinh mẫu mới và chỗ khác autoencoder

Lúc huấn luyện, mã thường được rút từ q của đúng mẫu đang tái tạo, vì decoder đang phải giải thích mẫu ấy. Lúc sinh một mẫu không gắn với một đầu vào có sẵn, ta không có xx để chạy encoder. Mã được rút thẳng từ prior.

z∼p(z)=N(0,I)z\sim p(z)=\mathcal{N}(0,I)

Rồi dữ liệu mới được lấy từ phân phối decoder gán cho mã đó. Tham số decoder vẫn là θ\theta.

xnew∼pθ(x∣z)x_{new}\sim p_\theta(x\mid z)

Trong một cài đặt đơn giản, đầu ra của decoder được dùng như mẫu hoặc như ước lượng của mẫu, tùy phân phối đầu ra đã chọn. Hai cách này không luôn trùng nhau: trung bình của một phân phối không phải là một lần rút từ phân phối ấy.

Đây là chỗ VAE khác autoencoder thường. VAE được huấn luyện để không gian mã tương thích với một prior mà ta biết cách rút. Tương thích là mục tiêu của vế KL, không phải một chứng nhận rằng mọi mẫu đều sắc và đúng nghĩa.

Thuộc tínhAutoencoderVAE
Đầu ra encoderMột vector mãTham số của một phân phối, thường là mu và logvar
MãXác địnhNgẫu nhiên, qua một lần rút
Mất mátTái tạoTái tạo cộng KL
Prior rõ để rút mẫuKhông bắt buộcThường là phân phối chuẩn chuẩn
Sinh mẫu mớiKhông có chỗ rút mã tự nhiênRút mã từ prior rồi giải mã
Bản dựng lạiThường sắc hơnCó thể kém hơn vì có vế KL
Không gian mãCó thể thủng những vùng trốngBị ép gần prior hơn
Việc điển hìnhBiểu diễn, khử nhiễu, bất thườngMô hình biến ẩn có phân phối

Hai mô hình tối ưu hai việc khác nhau. Bản dựng lại sát từng mẫu thuộc autoencoder. Một cách rút mã mới từ phân phối đã biết thuộc VAE, kèm đánh đổi ở bản dựng lại.

Đào sâu: dạng đóng của KL và cận dưới

Với q là chuẩn chéo và prior là chuẩn chuẩn, KL có dạng đóng. Không cần lấy mẫu để ước lượng số này. Mỗi chiều đóng góp một hạng từ tâm và từ phương sai của chiều đó.

DKL(q∥p)=12∑j(μj2+σj2−log⁡σj2−1)D_{KL}(q\|p) = \frac{1}{2}\sum_j \left( \mu_j^2+\sigma_j^2-\log\sigma_j^2-1 \right)

Nếu code giữ logvar, tức log của phương sai, cùng một đại lượng viết thành tổng quen thuộc trong cài đặt:

DKL(q∥p)=−12∑j(1+logvarj−μj2−exp⁡(logvarj))D_{KL}(q\|p) = -\frac{1}{2}\sum_j \left( 1+\text{logvar}_j-\mu_j^2-\exp(\text{logvar}_j) \right)
kl = -0.5 * torch.sum(
    1 + logvar - mu.pow(2) - logvar.exp(),
    dim=1
).mean()

Ở mức bài này, không cần tự chứng minh hai dòng ấy trùng nhau. Điều cần giữ là chúng đo độ lệch giữa phân phối mã do encoder tạo ra và phân phối chuẩn chuẩn, đúng chiều KL đã nói ở trên.

Paper Auto-Encoding Variational Bayes tối ưu một cận dưới của log xác suất dữ liệu, không tối ưu trực tiếp chính log xác suất ấy, vì chính nó khó tính. Cận dưới này cân bằng reconstruction và KL: hạng kỳ vọng thưởng decoder khi mã giải thích được dữ liệu, hạng KL phạt khi q đi xa prior. Cận dưới đó được gọi là ELBO. ELBO là evidence lower bound, tức cận dưới của bằng chứng, hiểu ở đây là cận dưới cho log xác suất của dữ liệu đã quan sát. ELBO không phải chính log xác suất ấy. Ta muốn tăng ELBO. Khi đổi sang mất mát để giảm, dấu đổi thành tái tạo cộng KL, đúng dạng đã viết ở trên.

log⁡pθ(x)≥Eqϕ(z∣x)[log⁡pθ(x∣z)]⏟reconstruction / data fit−DKL(qϕ(z∣x)∥p(z))⏟regularize latent distribution\log p_\theta(x) \ge \underbrace{ \mathbb{E}_{q_\phi(z\mid x)} \left[ \log p_\theta(x\mid z) \right] }_{\text{reconstruction / data fit}} - \underbrace{ D_{KL} \left(q_\phi(z\mid x)\|p(z)\right) }_{\text{regularize latent distribution}}

Vế phải là ELBO. Bất đẳng thức nói log xác suất của dữ liệu ít nhất bằng đại lượng này. Mất mát thực hành là mặt sau của việc tăng cận dưới ấy.

Đào sâu: beta, sụp hậu nghiệm, và những chỗ hay hiểu sai

Hệ số beta lệch cán cân. Beta nhỏ hơn một nghiêng về tái tạo. Beta quanh một là dạng không nhân thêm trọng số. Beta lớn hơn một ép mã mạnh hơn về prior. Lực ép ấy đôi khi giúp tách yếu tố, và bản dựng lại có thể xấu đi. Không có một beta đúng cho mọi dữ liệu. Chọn theo việc cần: bản dựng lại, hay mã dễ rút.

Một thất bại có tên posterior collapse, sụp hậu nghiệm: encoder học q gần như trùng prior với gần như mọi đầu vào.

q(z∣x)≈p(z)q(z\mid x)\approx p(z)

Lúc đó mã mang rất ít thông tin về đầu vào. Decoder quá mạnh có thể bỏ qua mã mà vẫn ra đầu ra tương đối. Dấu hiệu thường gặp: KL rất gần không, và đổi mã thì đầu ra ít đổi có nghĩa. Gần không ở đây là encoder bỏ cuộc, không phải mọi mẫu đã được giải thích tốt.

Các hướng thường thử là tăng dần trọng số KL, giảm sức decoder, miễn một phần KL nhỏ, hoặc chỉnh beta. KL quá yếu thì mã có thể vẫn lộn xộn. KL quá mạnh thì mô hình có thể bỏ qua mã.

Nội suy nhìn cấu trúc mã bằng đoạn thẳng giữa hai mã. Với tt đi từ 0 đến 1:

z(t)=(1−t)zA+tzBz(t)=(1-t)z_A+t z_B

Giải mã từng điểm trên đoạn. Không gian đủ mượt thì đầu ra chuyển dần, thay vì ra rác ở giữa. Nội suy đẹp chưa có nghĩa mỗi chiều mã là một thuộc tính gọi tên được. Việc tách yếu tố là mục tiêu riêng.

Vế tái tạo trong ELBO đến từ âm log xác suất của dữ liệu khi mã đã biết, nên phân phối đầu ra phải khớp dữ liệu. Pixel nhị phân hay đi với Bernoulli. Số thực đơn giản hay đi với Gaussian. Ảnh phức tạp có thể cần hỗn hợp logistic hoặc decoder riêng. BCE và MSE đủ cho bài này. Chúng không phải likelihood mặc định của mọi bài toán.

Encoder xuất mu và sigma riêng cho từng đầu vào, không phải một cặp số dùng chung cả tập. Mẫu hoàn toàn mới rút mã từ prior. KL gần không có thể là sụp hậu nghiệm, tức mã ít thông tin. VAE không có nhiệm vụ tái tạo sắc hơn autoencoder.

mu, logvar = encoder(x)
z = reparameterize(mu, logvar)
x_hat = decoder(z)

recon = reconstruction_loss(x_hat, x)
kl = kl_divergence(mu, logvar)
loss = recon + beta * kl

optimizer.zero_grad()
loss.backward()
optimizer.step()

Nguồn đã dùng

Nếu chỉ đọc bài này

Autoencoder gán mỗi đầu vào một điểm, nên khoảng giữa các điểm không phải nơi ta biết cách rút mẫu. VAE để encoder nêu một vùng: tâm là mu, độ rộng là sigma, thường tính từ log của phương sai để khỏi xuất một số âm. Một lần rút ngẫu nhiên không cho backprop đi qua, vì lần gieo không có đạo hàm theo tâm và theo độ rộng. Trick viết mã thành tâm cộng nhiễu đã gieo nhân với độ rộng, nên đường đạo hàm trở lại. Mất mát cân reconstruction và KL. Tăng cận dưới ELBO là mặt kia của cùng việc đó. Sinh mẫu mới thì rút mã từ prior chuẩn, rồi đưa qua decoder. Bản dựng lại có thể kém autoencoder hơn, và KL gần không có thể là dấu encoder đã bỏ mã.

Bài kế sang GAN, nơi hai mạng học bằng đối kháng và mục tiêu không còn là tái tạo từng mẫu đã có.