Variational autoencoder — ELBO, reparameterization আর latent space

১ · Autoencoder থেকে VAE — লুকানো কারণের মডেল

পর্ব ৭৫ (Autoencoder — compression, denoising আর representation)-এ একটি ছবিকে ছোট কোডে চেপে আবার ফিরিয়ে আনা হয়েছে। সেখানে একটা সমস্যাও দেখা যায়: কোডের জায়গায় এলোমেলো সংখ্যা বসিয়ে decoder চালালে বেশিরভাগ সময় অর্থহীন ছবি আসে। কারণ সাধারণ autoencoder কোডগুলোকে কোথায় রাখবে, তার কোনো নিয়ম নেই; train ছবির কোডগুলো ছড়িয়ে-ছিটিয়ে থাকে, মাঝখানের ফাঁকা জায়গা decoder কখনো দেখেনি। Variational autoencoder (VAE) এই কোড-জায়গাকে একটি নির্দিষ্ট বণ্টনে বাঁধে, যাতে সেখান থেকে নমুনা তুলে নতুন ছবি বানানো যায়।

VAE-র গল্পটা দুই লাইনের। প্রথমে একটি latent variable z তোলা হয় prior p(z) = N(০, I) থেকে, অর্থাৎ ১৬টি স্বাধীন সংখ্যা, প্রতিটির গড় ০ আর variance ১. তারপর decoder নামের একটি neural network z থেকে ৭৮৪টি pixel-এর প্রতিটির জন্য একটি probability দেয়, আর প্রতিটি pixel আলাদা টসে ০ বা ১ হয়: p_θ(x | z). এখানে θ হলো decoder-এর weight. z যে জায়গায় থাকে, তাকে বলে latent space. একটি ছবির মোট probability পেতে সব সম্ভাব্য z-এর ওপর যোগ (আসলে integral) করতে হয়:

p_θ(x) = ∫ p_θ(x | z) p(z) dz

এই সংখ্যাকে বলে evidence বা marginal likelihood. এর log-likelihood বাড়ানোই হতো স্বাভাবিক লক্ষ্য, পর্ব ৭৪ (Generative modeling-এর মানচিত্র)-এর maximum likelihood. কিন্তু ১৬ মাত্রার ওপর একটি neural network-এর integral-এর কোনো বন্ধ সূত্র নেই। নমুনা তুলে গড় নেওয়ার (Monte Carlo) সরল চেষ্টাও প্রায় ব্যর্থ হয়: prior থেকে তোলা বেশিরভাগ z একটি নির্দিষ্ট “৭”-এর ছবির জন্য p(x | z) প্রায় শূন্য দেয়, যেন খড়ের গাদায় সুঁই খোঁজা। সেকশন ৫-এ এটি মাপা হয়েছে: MNIST-এর ১,০০০টি test ছবিতে prior থেকে ১,০০০টি z তুলে গড় নিলে log p(x)-এর অনুমান দাঁড়ায় −১৮২.৩২ nats, অথচ encoder-কে পথপ্রদর্শক করলে একই ১,০০০টি নমুনায় −৭৪.৪৫ — প্রায় ১০৮ nats-এর পার্থক্য, পুরোটাই খারাপ অনুমানের দোষ।

উল্টো দিকের প্রশ্নটিও আটকে যায়। Bayes-এর সূত্রে posterior p(z | x) = p(x | z) p(z) / p(x), আর হরে সেই একই p(x). তাই VAE একটি দ্বিতীয় network আনে, encoder, যা প্রতিটি x-এর জন্য একটি সহজ বণ্টন দেয়: q_φ(z | x) = N(μ_φ(x), diag σ_φ²(x)). এখানে φ হলো encoder-এর weight, আর diag মানে প্রতিটি মাত্রা আলাদা, তাদের মধ্যে সম্পর্ক ধরা হয় না। প্রতিটি ছবির জন্য আলাদা করে optimisation না চালিয়ে একটি network সব ছবির জন্য এক ধাপে উত্তর দেয় — একে বলে amortized inference.

২ · ELBO — log p(x) = ELBO + KL

মেশিন লার্নিং কোর্সের পর্ব ৪৮-এ variational inference-এর সাধারণ রূপে এই অভেদ দেখেছেন। এখানে একই হিসাব VAE-র চোখে, একটি ছবি x ধরে। মূল কৌশল এক লাইনের: log p(x) z-এর ওপর নির্ভর করে না, তাই যেকোনো q(z | x)-এর ওপর তার গড় নিলেও সে একই থাকে। তারপর Bayes দিয়ে p(x) = p(x, z) / p(z | x) লিখে ভেতরে q দিয়ে গুণ আর ভাগ:

log p(x) = E_q[log p(x, z) − log q(z | x)] + E_q[log q(z | x) − log p(z | x)]

প্রথম গড়টির নাম ELBO (evidence lower bound), দ্বিতীয়টি KL(q(z | x) ‖ p(z | x)), আসল posterior থেকে q-এর দূরত্ব। KL divergence কখনো ঋণাত্মক হয় না (AI-এর ভিত্তি কোর্সের পর্ব ৩৮), তাই ELBO ≤ log p(x). দূরত্বটি আমরা হিসাব করতে পারি না, কারণ তাতে আসল posterior লাগে; কিন্তু ELBO পারি। এই ভাঙনের পুরো পথ একটি খেলনা উদাহরণে দেখুন, যেখানে z মাত্র দুটি মান নিতে পারে, তাই সবকিছু হাতে হিসাব করা যায়।

VAE-র loss তাই এই ELBO-র ঋণাত্মক, প্রতিটি ছবির জন্য:

loss(x) = −E_q(z|x)[log p_θ(x | z)] + KL(q_φ(z | x) ‖ N(০, I))

প্রথম পদ হলো reconstruction term: encoder-এর দেওয়া z থেকে decoder আসল ছবিটি কতটা ফিরিয়ে আনতে পারে। ০/১ pixel-এ এটি ঠিক binary cross-entropy-র যোগফল, যা পর্ব ১৩ (Deep learning-এর loss)-এ দেখেছেন। দ্বিতীয় পদ একটি নিয়ন্ত্রক: প্রতিটি ছবির q-কে prior-এর কাছে থাকতে বলে, যাতে কোডগুলো N(০, I)-এর চারপাশে ঘন হয়ে থাকে আর নমুনা তোলার সময় ফাঁকা জায়গা কম পড়ে। দুই পদের টানাটানিই VAE-র চরিত্র ঠিক করে, যা সেকশন ৭-এ মাপা হবে। আরেক পথেও এই সীমা পাওয়া যায়: log-এর ভেতরে গড় থাকলে Jensen-এর অসমতা দিয়ে log E[·] ≥ E[log ·].

৩ · Gaussian KL হাতে — নমুনা ছাড়াই

Loss-এর KL অংশে দুটিই Gaussian: q = N(μ, σ²) আর prior N(০, ১), প্রতি মাত্রায় আলাদা। KL = E_q[log q − log p] খুলে লিখলে log q-এর গড় −½ ln(২πσ²) − ½, আর log p-এর গড় −½ ln(২π) − ½(μ² + σ²), কারণ q-এর ওপর E[z²] = μ² + σ². বিয়োগ করলে ২π কেটে যায়:

KL(N(μ, σ²) ‖ N(০, ১)) = ½ (μ² + σ² − ১ − ln σ²)

প্রতিটি পদের অর্থ আছে। μ² বলে: কোড কেন্দ্র থেকে যত দূরে, দণ্ড তত বেশি। σ² − ১ − ln σ² শূন্য হয় শুধু σ = ১-এ, আর σ ছোট হলে −ln σ² দ্রুত বাড়ে: encoder যদি খুব নিশ্চিত হয়ে একটি বিন্দুতে চেপে বসতে চায়, তার দাম দিতে হয়। মাত্রাগুলো স্বাধীন বলে মোট KL হলো প্রতি মাত্রার KL-এর যোগফল। Encoder সাধারণত σ নয়, log σ² দেয়, যাতে output যেকোনো বাস্তব সংখ্যা হতে পারে আর σ সবসময় ধনাত্মক থাকে।

সূত্রটি ঠিক কি না, পাতার ভেতরেই নমুনা দিয়ে যাচাই করুন: q থেকে অনেক z তুলে log q − log p-এর গড় নিলে একই সংখ্যার কাছে আসা উচিত। সঙ্গে সেকশন ৫-এর সত্যিকারের MNIST run থেকে ১৬টি মাত্রার গড় KL বসানো আছে।

import numpy as np
np.set_printoptions(suppress=True, linewidth=56)

def kl_gauss(mu, sd):                       # KL(N(mu, sd²) ‖ N(0, 1)), প্রতি মাত্রা
    return 0.5 * (mu**2 + sd**2 - 1 - np.log(sd**2))

mu = np.array([0.8, -0.1])
sd = np.array([0.5, 0.95])
k = kl_gauss(mu, sd)
print("সূত্রে:       ", k.round(4), " মোট", k.sum().round(4))

rng = np.random.default_rng(42)
z = mu + sd * rng.standard_normal((200_000, 2))      # q থেকে নমুনা
log_q = -0.5 * np.log(2 * np.pi * sd**2) - (z - mu)**2 / (2 * sd**2)
log_p = -0.5 * np.log(2 * np.pi) - z**2 / 2
d = log_q - log_p
print("Monte Carlo:  ", d.mean(0).round(4), " মোট", d.mean(0).sum().round(4))
print("std error:    ", (d.std(0) / np.sqrt(len(d))).round(4))

# সত্যিকারের run: MNIST, z = 16, β = 1, seed 42; test-এ প্রতি মাত্রার গড় KL
kl_real = np.array([2.367, 2.238, 2.435, 2.768, 0.002, 1.837, 2.326, 1.440,
                    0.001, 1.413, 0.001, 0.001, 1.805, 1.603, 1.562, 1.985])
print("বড় থেকে ছোট:", np.sort(kl_real)[::-1])
print("মোট KL =", kl_real.sum().round(2), "nats =",
      (kl_real.sum() / np.log(2)).round(1), "bits")
print("KL < 0.01 এমন মাত্রা:", np.flatnonzero(kl_real < 0.01))

৪ · Reparameterization trick — এলোমেলোতার ভেতর দিয়ে gradient

Reconstruction পদে একটি গড় আছে: E_q(z|x)[log p(x | z)]. এটি হুবহু হিসাব করা যায় না, তাই প্রতিটি ছবির জন্য q থেকে একটি z তুলে তার মান দিয়ে কাজ চালানো হয়। সমস্যা হলো, “একটি নমুনা তোলা” কোনো মসৃণ function নয়। Encoder যদি μ একটু বাড়ায়, তোলা z কতটা বদলায়? প্রশ্নটাই অস্পষ্ট, কারণ নমুনা প্রতিবার নতুন এলোমেলো সংখ্যা। ফলে backpropagation (পর্ব ৪, হাতে-কলমে backpropagation) encoder পর্যন্ত পৌঁছানোর পথ পায় না।

সমাধানটি চোখে পড়ার মতো সরল। N(μ, σ²) থেকে নমুনা তোলা আর আগে N(০, ১) থেকে ε তুলে z = μ + σ·ε বানানো — বণ্টনের দিক থেকে হুবহু একই। কিন্তু দ্বিতীয় রূপে এলোমেলোতা পুরোটা ε-তে, যা কোনো parameter-এর ওপর নির্ভর করে না; আর z হয়ে যায় μ আর σ-এর একটি সাধারণ, অবকলনযোগ্য function: ∂z/∂μ = ১, ∂z/∂σ = ε. একে বলে reparameterization trick. মেশিন লার্নিং কোর্সের পর্ব ৪৮-এ black-box VI-তে এর সাধারণ রূপ দেখেছেন; নিচের ছবিতে VAE-র ভেতরে তার জায়গা।

এবার একটি ছোট উদাহরণে gradient-এর পথ সংখ্যা দিয়ে। Decoder-এর loss-কে সরল রাখতে ধরা হয়েছে L(z) = (z − ১)², যেন decoder চায় z = ১ হোক। Encoder-এর output μ = ০.৮, σ = ০.৫, আর তোলা ε = −১.২.

Reparameterization ছাড়াও gradient পাওয়ার একটি পুরোনো পথ আছে: score-function estimator (REINFORCE নামেও পরিচিত), যা ∇_μ E[f(z)] = E[f(z) · ∇_μ log q(z)] ব্যবহার করে। এতে f-এর অবকলন লাগে না, তাই discrete z-এও চলে, কিন্তু সাধারণত অনেক বেশি ওঠানামা করে। নিচের ব্লকে দুটিকে একই সমস্যায় পাশাপাশি রাখা হয়েছে।

import numpy as np

rng = np.random.default_rng(0)
mu, sd = 0.8, 0.5
f = lambda z: (z - 1.0) ** 2              # খেলনা decoder-loss
# হুবহু: E[f] = (mu - 1)^2 + sd^2
print("হুবহু gradient:  μ", round(2 * (mu - 1), 3), " σ", 2 * sd)

eps = -1.2                                # stepper-এর একটি ε
z = mu + sd * eps
print("একটি ε = -1.2:   z", round(z, 2), " μ", round(2 * (z - 1), 3),
      " σ", round(2 * (z - 1) * eps, 3))

for n in (10, 1000, 100_000):
    e = rng.standard_normal(n)
    z = mu + sd * e
    rep_mu = 2 * (z - 1)                  # pathwise: df/dz · dz/dμ
    rep_sd = 2 * (z - 1) * e              # pathwise: df/dz · dz/dσ
    sf_mu = f(z) * (z - mu) / sd**2       # score-function (REINFORCE)
    print(f"n={n:>6}  reparam μ {rep_mu.mean():+.3f} σ {rep_sd.mean():+.3f}"
          f" | score μ {sf_mu.mean():+.3f}")

e = rng.standard_normal(100_000)
z = mu + sd * e
print("প্রতি নমুনার std (μ-এর gradient):")
print("  reparam", round((2 * (z - 1)).std(), 3),
      "  score", round((f(z) * (z - mu) / sd**2).std(), 3))

৫ · সত্যিকারের run — MNIST আর Fashion-MNIST-এ VAE

এবার আসল আকারে, পর্ব ৭৪-এর হুবহু একই শর্তে, যাতে সংখ্যাগুলো সেখানকার মডেলের পাশে বসানো যায়। বাজারঘরের পণ্যের ছবির বদলে Fashion-MNIST (Zalando Research, MIT licence): জামা, জুতা, ব্যাগের ২৮×২৮ ধূসর ছবি; এগুলো আসল দোকানের ছবি নয়, একটি পরিচিত গবেষণা-ডেটাসেট। পাশে MNIST-এর হাতের লেখা সংখ্যা (এর licence সাধারণত CC BY-SA 3.0 বলা হয়; নিজের কাজে ব্যবহারের আগে উৎস দেখে নিন)। প্রতিটি pixel ১২৮ সীমায় ০/১; train-এর প্রথম ৫৫,০০০টি ছবিতে শেখা, শেষ ৫,০০০টি val, আর test-এর ১০,০০০টি ছবি সবশেষে একবার মাপা।

নকশা: encoder ৭৮৪ → ৫১২ → ২৫৬ (ReLU), তারপর দুটি আলাদা layer থেকে μ আর log σ²; decoder উল্টো পথে ২৫৬ → ৫১২ → ৭৮৪টি logit. Latent ২ আর ১৬ মাত্রা। Adam (lr ০.০০১), batch ১২৮, ৫০ epoch; val-এর −ELBO সবচেয়ে কম যে epoch-এ, সেই weight রাখা হয়েছে। Python 3.12.13, torch 2.15.0 (CPU build), numpy 2.5.4; সব CPU-তে, ৩টি thread-এ (লেখার সময়ের version), অন্য কাজের সঙ্গে ভাগ করা যন্ত্রে। মূল অংশ (vae.py থেকে ছাঁটা ও সংক্ষেপিত; আসল ফাইলে encoder আলাদা method-এ):

def forward(self, x):                                 # x: (batch, 784), 0/1
    h = self.enc(x)
    mu, logvar = self.mu(h), self.logvar(h)
    eps = torch.randn_like(mu)
    z = mu + torch.exp(0.5 * logvar) * eps            # reparameterisation
    logits = self.dec(z)
    rec = F.binary_cross_entropy_with_logits(logits, x, reduction="none").sum(1)
    kl = 0.5 * (mu.pow(2) + logvar.exp() - 1 - logvar).sum(1)
    return rec, kl                                    # প্রতি ছবি, nats

rec, kl = model(xb)
loss = (rec + beta * kl).mean()                       # batch-এর গড়
opt.zero_grad(); loss.backward(); opt.step()

খেয়াল করুন, rec আর kl দুটোই প্রতি ছবিতে যোগফল (sum), তারপর batch-এ গড়। একটিকে pixel-গড় আর অন্যটিকে যোগফল নিলে KL-এর ওজন অজান্তে ৭৮৪ গুণ বদলে যায় — AI-এর লেখা VAE কোডের একটি চেনা ভুল। MNIST, ১৬ মাত্রা, seed ৪২-এর লগ (ছাঁটা; প্রতি লাইনে val-এর reconstruction, KL, −ELBO আর তার bits/dim):

$ python vae.py mnist 16 1 42 50        # data z_dim beta seed epochs
mnist_z16_b1_s42 params=1079600 torch=2.15.0+cpu
epoch beta_eff val_rec val_kl val_negelbo val_bpd sec
  1 1.00   113.09  15.41   128.50 0.2365     2
  5 1.00    72.02  20.81    92.83 0.1708    12
 10 1.00    65.12  22.27    87.39 0.1608    29
 20 1.00    60.89  23.17    84.06 0.1547    57
 30 1.00    59.26  23.61    82.87 0.1525    86
 40 1.00    58.32  24.31    82.63 0.1521   134
 50 1.00    58.40  23.85    82.25 0.1514   171
{"data": "mnist", "z": 16, "beta": 1.0, "seed": 42, "warmup": 0, "epochs": 50, "best_epoch": 50, "params": 1079600, "test_rec": 58.2763265625, "test_kl": 23.7836359375, "test_negelbo": 82.05996250000001, "test_elbo_bpd": 0.1510044655029137, "test_iwae200_nats": 75.53347778320312, "test_iwae200_bpd": 0.13899460434913635, "iwae_s": 15.8, "active_units": 12, "fit_s": 170.9, …}
var_mu 1.038 0.980 0.989 1.031 0.001 1.004 1.047 0.907 0.001 0.893 0.001 0.001 0.975 0.931 0.902 0.974
kl_per_dim 2.37 2.24 2.43 2.77 0.00 1.84 2.33 1.44 0.00 1.41 0.00 0.00 1.80 1.60 1.56 1.99

শেষ লাইনের test_iwae200 হলো importance sampling-এ log p(x)-এর অনুমান। ELBO নিম্নসীমা মাত্র, তাই পর্ব ৭৪-এর হুবহু likelihood-এর পাশে বসাতে আরেকটু ভালো অনুমান দরকার। কৌশলটি importance sampling: q(z | x) থেকে K টি z তুলে p(x, z) / q(z | x) অনুপাতের গড় নেওয়া, যার log আবার একটি নিম্নসীমা, তবে K বাড়লে log p(x)-এর দিকে ওঠে (IWAE সীমা)। Prior থেকে নমুনা তুললে (সেকশন ১-এর “খড়ের গাদা”) কী হয়, সেটিও পাশে মাপা হয়েছে, প্রথম ১,০০০টি test ছবিতে:

সত্যিকারের run, seed ৪২-এর মডেল, test-এর প্রথম ১,০০০টি ছবি; সংখ্যা বড় মানে ভালো। K = ১-এর সারিটি নিজেই এলোমেলো: একই মডেলে আরেকবার মাপায় MNIST-এ −৮২.০৯ আর Fashion-এ −১২০.০৮ এসেছিল। বন্ধনীতে CPU-তে ১,০০০টি ছবির সময়। K বাড়ার সঙ্গে অনুমান একঘেয়ে ভাবে ওঠে, কিন্তু ধীর হয়ে আসে; K = ১,০০০-ও একটি নিম্নসীমা, তাই আসল log p(x) এর চেয়ে সামান্য বেশি হতে পারে।

এবার মূল টেবিল: পুরো ১০,০০০টি test ছবিতে, K = ২০০, পর্ব ৭৪-এর ফলের পাশে।

সব সংখ্যা এই যন্ত্রের সত্যিকারের run, একটি ভাগে, β = ১. z = ১৬-এর ঘরে seed ৪২, ৪৩, ৪৪-এর পরিসর; z = ২ শুধু seed ৪২. Test ১০,০০০টি ছবি, IWAE-এর K = ২০০ (পুরো test-এ প্রায় ১১–১৯ সেকেন্ড)। দুই নকশাতেই parameter প্রায় ১০.৭–১০.৮ লাখ (z = ১৬-এ ১০,৭৯,৬০০)। সময় CPU-তে, শুধু মাত্রা হিসেবে পড়ুন।

• পর্ব ৭৪-এর পাশে MNIST: স্বাধীন Bernoulli ০.৩৭৮৮, FVSBN ০.১৪০১, MADE ০.১২৭৩–০.১২৭৬ bits/dim. ১৬ মাত্রার VAE-র importance-sampled অনুমান ০.১৩৯০–০.১৩৯৩: FVSBN-এর প্রায় সমান, MADE-এর চেয়ে প্রায় ০.০১১–০.০১২ bits/dim (প্রতি ছবিতে প্রায় ৬.২–৬.৫ nats) খারাপ। শুধু ELBO (০.১৫১) বসালে ব্যবধান দ্বিগুণ দেখাত — তাই ELBO দিয়ে তুলনা অন্যায্য।
• Fashion-MNIST-এ: স্বাধীন Bernoulli ০.৭০৫০, PCA + GMM ০.৫৭৯০ (ঊর্ধ্বসীমা), FVSBN ০.২০৩১–০.২০৩২, MADE ০.১৯৩৫–০.১৯৩৬. VAE ০.২১৩২–০.২১৩৭: প্রথম দুটির চেয়ে অনেক ভালো, দুটি autoregressive মডেলের চেয়েই খারাপ। স্বাধীন pixel-এর decoder ছোট কিনারা আর নকশার খুঁটিনাটি ধরতে পারে না, যা autoregressive মডেল প্রতিবেশী pixel দেখে ধরে।
• বিনিময়ে কী পাওয়া গেল: নমুনা এক forward pass-এ (MADE-এ ৭৮৪টি ধাপ), আর ১৬টি সংখ্যার একটি ব্যবহারযোগ্য latent কোড (সেকশন ৬)।
• Seed আর মাত্রা: তিন seed-এর পার্থক্য ০.০০০৬ bits/dim-এর মধ্যে। ২ মাত্রায় VAE অনেক খারাপ (MNIST ০.২৪১), কারণ দুটি সংখ্যায় একটি ছবির বেশি তথ্য ধরা যায় না; তবু তা স্বাধীন Bernoulli-র চেয়ে ভালো। MNIST-এর একটি run-এর সেরা epoch শেষটি (৫০), তাই আরও epoch-এ হয়তো সামান্য উন্নতি হতো; পরীক্ষা করা হয়নি। প্রবন্ধের এলোমেলোভাবে ০/১ করা MNIST-এর সংখ্যার সঙ্গে এগুলো মেলানো যাবে না।

Epoch ধরে reconstruction আর KL কীভাবে বদলায়, আর ২-মাত্রার VAE ছবিগুলোকে latent-এ কোথায় বসায়, নিচের ব্লকে সত্যিকারের লগ আর encoder output থেকে আঁকা হয়েছে (val অংশ; MNIST-এর ২-D মানচিত্রে ২০০টি এলোমেলো test ছবির μ)।

import pandas as pd
import matplotlib.pyplot as plt

c = pd.read_csv("vae_curves.csv")         # val, প্রতি epoch, nats
m = pd.read_csv("map2d.csv")              # MNIST z=2: ২০০টি test ছবির μ
c["negelbo"] = c.rec + c.kl
for run, g in c.groupby("run", sort=False):
    b = g.loc[g.negelbo.idxmin()]
    print(f"{run:12s} সেরা epoch {int(b.epoch):2d}: rec {b.rec:6.2f}"
          f"  KL {b.kl:5.2f}  -ELBO {b.negelbo:6.2f}")
print("epoch 1-এ KL:", c[c.epoch == 1].set_index("run").kl.to_dict())
print("২-D μ-এর std:", m[["z1", "z2"]].std().round(2).tolist())
print("শ্রেণিভিত্তিক গড় μ (০, ১, ৭):")
print(m.groupby("label")[["z1", "z2"]].mean().loc[[0, 1, 7]].round(2))

fig, ax = plt.subplots(1, 3, figsize=(11, 3.4))
for run, g in c.groupby("run", sort=False):
    ax[0].plot(g.epoch, g.rec, label=run)
    ax[1].plot(g.epoch, g.kl, label=run)
ax[0].set_title("val reconstruction (nats)")
ax[1].set_title("val KL (nats)")
for a in ax[:2]:
    a.set_xlabel("epoch"); a.legend(fontsize=8)
sc = ax[2].scatter(m.z1, m.z2, c=m.label, cmap="tab10", s=8)
ax[2].set_title("MNIST z=2: encoder mean")
ax[2].set_xlabel("z1"); ax[2].set_ylabel("z2")
fig.colorbar(sc, ax=ax[2], ticks=range(10))
fig.subplots_adjust(wspace=0.4)
plt.show()

৬ · Latent space — নমুনা, interpolation আর মানচিত্র

নমুনা: prior N(০, I) থেকে z তুলে decoder চালানো — encoder লাগে না। Interpolation: দুটি ছবির μ-এর মাঝখানের সরলরেখায় কয়েকটি বিন্দু নিয়ে decode করা; KL পদ latent-কে ফাঁকহীন রাখে বলে মাঝখানের বিন্দুগুলোও সাধারণত বিশ্বাসযোগ্য ছবি দেয়। মানচিত্র: ২-মাত্রার latent-এ একটি নিয়মিত grid decode করে দেখা, কোন অঞ্চলে কী থাকে।

নিচের ব্লক সত্যিকারের run-এর ছবি আঁকে (seed ৪২-এর মডেল, সেরা val epoch-এর weight): প্রথম সারিতে MNIST-এর ১৬-মাত্রার VAE থেকে ৮টি নমুনা, দ্বিতীয় সারিতে Fashion-MNIST-এর, তৃতীয় সারিতে MNIST-এর একটি “১” থেকে একটি “৭”-এ interpolation (test-এর প্রথম ১ আর প্রথম ৭), আর নিচে MNIST-এর ২-মাত্রার VAE-তে z₁, z₂ ∈ {−২, −১, ০, ১, ২}-এর ৫×৫ grid. প্রতিটি ছবি decoder-এর pixel-probability, ১৬ ধাপে (৪ bit) চেপে base64-এ পাতায় রাখা; তাই এগুলো নমুনার “গড় ছবি”, ০/১ টস করা ছবি নয়।

import base64
import numpy as np
import matplotlib.pyplot as plt

rows = {}
for line in open("vae_imgs.txt"):
    name, data = line.strip().split(":")
    imgs = []
    for b in data.split(","):
        v = np.frombuffer(base64.b64decode(b), np.uint8)
        imgs.append(np.stack([v >> 4, v & 15], 1).reshape(28, 28) / 15)
    rows[name] = np.array(imgs)            # 4-bit → pixel probability

fig, axes = plt.subplots(8, 8, figsize=(8, 8.4))
for ax in axes.ravel():
    ax.axis("off")
for r, name in enumerate(["mnist_samples", "fashion_samples", "mnist_interp"]):
    for c in range(8):
        axes[r, c].imshow(rows[name][c], cmap="gray_r", vmin=0, vmax=1)
    print(f"{name:16s} গড় pixel probability {rows[name].mean():.3f}")
g = rows["mnist_grid"]                     # z1, z2 ∈ {-2, -1, 0, 1, 2}
for i in range(25):
    axes[3 + i // 5, i % 5].imshow(g[i], cmap="gray_r", vmin=0, vmax=1)
axes[3, 6].text(0, 0.5, "5x5 grid\nof z=2\nlatent\n(top: z2=-2)", fontsize=9)
print("grid-এর প্রতিটি ছবির ‘কালি’ (সারি = z2, কলাম = z1):")
print(g.reshape(5, 5, -1).mean(2).round(2))
fig.subplots_adjust(wspace=0.05, hspace=0.05)
plt.show()

৭ · β-VAE, posterior collapse আর KL annealing

সেকশন ২-এর loss-এ KL-কে একটি সংখ্যা β দিয়ে গুণ করলে পাওয়া যায় β-VAE: loss = rec + β·KL. β = ১ হলে এটি ঠিক −ELBO. β > ১ মানে latent-এ তথ্য রাখার দাম বেশি; β < ১ মানে কম। কতগুলো latent মাত্রা আসলে কাজ করছে, তা মাপার প্রচলিত উপায় active units: test ছবিগুলোর ওপর E_q[z_u] (অর্থাৎ μ_u)-এর variance ০.০১-এর বেশি হলে মাত্রা u সক্রিয় (Burda ও সহকর্মী, ২০১৬)। যে মাত্রায় μ সব ছবিতে প্রায় একই, encoder সেখানে কোনো ছবির খবর রাখছে না।

সত্যিকারের run, z = ১৬, seed ৪২, ৫০ epoch, test ১০,০০০টি ছবি; rec, KL আর −ELBO nats-এ। −ELBO সবসময় β = ১ ধরে হিসাব করা (rec + KL), তাই সারিগুলো তুলনাযোগ্য। Active মাত্রা গোনা হয়েছে test-এ μ-এর variance > ০.০১ নিয়মে; সক্রিয় মাত্রাগুলোর variance ছিল প্রায় ০.৩৭ থেকে ১.২৫, নিষ্ক্রিয়গুলোর ০.০০১৪-এর বেশি নয়, মাঝামাঝি প্রায় কিছু নেই।

ছবিটা দুই ডেটাসেটেই একই। β বাড়ালে KL কমে (MNIST-এ ৩৩.৬৭ → ২৩.৭৮ → ১০.৭৪), reconstruction খারাপ হয় (৪৯.৬১ → ৫৮.২৮ → ৮৭.৪৩), আর active মাত্রা কমে (১৫ → ১২ → ৮)। আসল ELBO আর likelihood-এর সেরা β = ১-এ, যেমনটা হওয়ার কথা, কারণ β = ১-ই ঠিক ELBO বাড়ায়। β = ০.৫ প্রায় সমান, কিন্তু সামান্য খারাপ: সে z-এ বেশি তথ্য রেখে reconstruction ভালো করেছে, কিন্তু তার দাম KL-এ বেশি দিয়েছে। β = ৪ সবচেয়ে কম তথ্য রাখে: ১৬টির মধ্যে মাত্র ৮টি মাত্রা কাজ করে। ছোট, “পরিষ্কার” latent চাইলে এটি কাজে লাগতে পারে, কিন্তু likelihood-এ প্রায় ০.০২ bits/dim-এর দাম দিতে হয়। খেয়াল করুন, β = ১-এও ১৬টির মধ্যে ৪–৫টি মাত্রা নিষ্ক্রিয় ছিল, সব seed-এই।

নিষ্ক্রিয় মাত্রার চরম রূপের নাম posterior collapse: প্রায় সব মাত্রায়, প্রায় সব ছবির জন্য q(z | x) prior-এর সমান হয়ে যায়, আর decoder z-কে উপেক্ষা করে। আমাদের MLP decoder z ছাড়া চলতে পারে না, তাই collapse আংশিক থেকেছে। কিন্তু decoder যদি নিজেই আগের pixel বা শব্দ দেখে অনুমান করতে পারে (autoregressive decoder), z ছাড়াই reconstruction ভালো হয় আর KL-এর দাম এড়াতে মডেল z ফেলে দিতে পারে — ভাবনার প্রশ্ন ২-এ এর ইতিহাস।

একটি প্রচলিত প্রতিকার KL annealing (KL warm-up নামেও পরিচিত): শুরুতে β = ০ থেকে ধীরে ধীরে ১-এ তোলা, যাতে encoder আগে কাজের তথ্য রাখতে শেখে, তারপর KL-এর চাপ আসে। আরেকটি free bits: প্রতিটি মাত্রায় একটি ন্যূনতম KL পর্যন্ত দণ্ড না দেওয়া। MNIST-এ, ১৬ মাত্রায়, প্রথম ১০ epoch-এ β-কে ০.১ থেকে ১-এ রৈখিকভাবে তুলে একটি run চালানো হয়েছে:

$ python vae.py mnist 16 1 42 50 10     # শেষে warm-up epoch
mnist_z16_b1_s42_w10 params=1079600 torch=2.15.0+cpu
epoch beta_eff val_rec val_kl val_negelbo val_bpd sec
  1 0.10    94.05  47.19   141.24 0.2599     2
  2 0.20    71.25  42.01   113.26 0.2084     5
  5 0.50    59.93  32.80    92.73 0.1706    12
 10 1.00    58.79  27.29    86.09 0.1584    25
 20 1.00    56.10  27.23    83.33 0.1533    51
 46 1.00    53.82  27.21    81.03 0.1491   121
 50 1.00    53.59  27.72    81.31 0.1496   133
{"data": "mnist", "z": 16, "beta": 1.0, "seed": 42, "warmup": 10, "epochs": 50, "best_epoch": 46, …, "test_rec": 53.32564375, "test_kl": 27.162065625, "test_negelbo": 80.487709375, "test_elbo_bpd": 0.14811124893855182, "test_iwae200_nats": 73.34008026123047, "test_iwae200_bpd": 0.13495837152004242, …, "active_units": 16, …}
var_mu 0.979 0.874 0.941 1.082 0.937 1.073 1.053 0.976 0.845 1.020 1.030 0.974 0.843 0.901 0.907 0.849

প্রথম epoch-এ β মাত্র ০.১, তাই KL ৪৭.১৯ পর্যন্ত উঠেছে (warm-up ছাড়া একই epoch-এ ছিল ১৫.৪১): encoder শুরুতেই অনেক তথ্য রাখতে শিখেছে। β = ১-এ পৌঁছানোর পরে KL কমে প্রায় ২৭-এ থিতু হয়, কিন্তু কোনো মাত্রা নিষ্ক্রিয় হয়নি — ১৬টির ১৬টিই সক্রিয়, warm-up ছাড়া তিন seed-এ ছিল ১২টি করে। Test-এ −ELBO ৮০.৪৯ (warm-up ছাড়া ৮১.৭৫–৮২.০৬), importance-sampled অনুমান ০.১৩৫০ bits/dim (০.১৩৯০–০.১৩৯৩-এর জায়গায়)। তবে এটি একটি seed-এর একটি run; সিদ্ধান্তের আগে কয়েকটি seed আর warm-up-এর দৈর্ঘ্য বদলে দেখা দরকার। আর MNIST-এর মতো সহজ ডেটায় MLP decoder-এ সমস্যাটা ছোট; শক্তিশালী autoregressive decoder-এ (text-এ যেমন) প্রভাব অনেক বড় হতে পারে।

৮ · VAE-র পরিবার — conditional VAE আর VQ-VAE

মূল VAE থেকে অনেক রূপ এসেছে। দুটি এই মডিউলের জন্য জরুরি।

VQ-VAE (van den Oord, Vinyals আর Kavukcuoglu, ২০১৭; প্রবন্ধের সারাংশ অনুযায়ী) সাধারণ VAE থেকে দুই জায়গায় আলাদা: কোড discrete, আর prior স্থির N(০, I) নয়, পরে আলাদা করে শেখা (যেমন কোডগুলোর ওপর একটি autoregressive মডেল)। “সবচেয়ে কাছেরটি বাছা” কাজটির অবকলন নেই, তাই training-এ সাধারণত straight-through estimator ব্যবহার হয়: forward-এ quantized vector, backward-এ gradient যেন quantization ছিলই না এমনভাবে সোজা encoder-এ পাঠানো — পর্ব ৮ (PyTorch autograd)-এ round করার উদাহরণে এই কৌশল দেখেছেন। সঙ্গে দুটি বাড়তি loss পদ codebook vector আর encoder output-কে পরস্পরের কাছে টানে। প্রবন্ধের সারাংশে এটাও বলা আছে যে এই নকশা শক্তিশালী autoregressive decoder-এর সঙ্গেও posterior collapse এড়াতে সাহায্য করে।

কেন এটি পর্ব ৮০ (Latent diffusion আর guidance)-এর সেতু? আজকের অনেক ছবি-বানানো ব্যবস্থা দুই ধাপে কাজ করে: একটি autoencoder ছবিকে অনেক ছোট latent-এ চাপে, আর আরেকটি মডেল সেই latent-এ নমুনা তোলা শেখে। Rombach ও সহকর্মীদের latent diffusion প্রবন্ধের সারাংশ বলে, তাঁরা diffusion চালিয়েছেন “pretrained autoencoder”-এর latent space-এ; সেই autoencoder-এর নিয়ন্ত্রক হিসেবে একটি হালকা KL পদ বা vector quantization ব্যবহারের কথা প্রবন্ধে আছে বলে জানা যায়। বিস্তারিত পর্ব ৮০-তে।

৯ · AI-এর যুগে — VAE-র কোড আর derivation যাচাই

AI সহকারী একটি VAE কয়েক সেকেন্ডে লিখে দেয়, আর সেই কোড প্রায়ই চলে — যা বিপদের অংশ, কারণ এখানকার ভুলগুলো চুপচাপ ভুল সংখ্যা দেয়, কোনো error নয়। চেনা ভুলগুলো: KL-এর চিহ্ন উল্টো বা ½ বাদ; log σ²-কে σ ভেবে বসানো; reconstruction-কে pixel-এর ওপর গড় (mean) আর KL-কে যোগফল (sum) ধরা, ফলে KL-এর ওজন অজান্তে ৭৮৪ গুণ বেড়ে যাওয়া; sigmoid-এর পরে আবার binary_cross_entropy_with_logits; test-এ model.eval() আর torch.no_grad() ভুলে যাওয়া; ELBO-কে সরাসরি log-likelihood বলে প্রকাশ করা; একটি seed-এর ফলকে সাধারণ সত্য বলা; আর derivation-এ বিশ্বাসযোগ্য কিন্তু ভুল একটি ধাপ বা বানানো উদ্ধৃতি। সমাধান: version-এর অফিসিয়াল docs, ছোট batch-এ যাচাই, আর প্রতিটি সূত্র খেলনা উদাহরণে হাতে মেলানো।

AI-এর লেখা VAE হাতে পেলে — যাচাইয়ের চেকলিস্ট

• KL-এর লাইনটি একটি সংখ্যায় হাতে মিলিয়েছি, আর সব মাত্রায় যোগ হচ্ছে কি না দেখেছি
• Reconstruction আর KL দুটোই প্রতি ছবিতে sum, তারপর batch-এ mean — একই এককে আছে
• একটি ছোট batch-এ overfit করিয়ে reconstruction নামতে দেখেছি, আর shape ও dtype মিলিয়েছি
• Test-এ model.eval() আর torch.no_grad() আছে; ELBO আর importance-sampled likelihood আলাদা নামে লিখেছি
• Active latent মাত্রা গুনেছি, যাতে posterior collapse চোখ এড়িয়ে না যায়
• স্বাধীন Bernoulli-র মতো baseline আর কয়েকটি seed-এর সঙ্গে তুলনা করেছি
• ব্যক্তিগত ডেটা বা key কোনো prompt-এ দিইনি