Reproducing Kingma & Welling (2013), Fig. 2 (MNIST VAE)#
This example reproduces Figure 2 from the paper on Auto-Encoding Variational Bayes (Kingma & Welling (2013)). The experiment trains Variational Autoencoders (VAEs) with different latent dimensionalities \(N_z\) on MNIST and tracks the Evidence Lower Bound (ELBO) as a function of the number of training samples processed.
The setup follows the configuration in the paper:
one-hidden-layer MLP encoder and decoder with Tanh activations;
hidden size \(H = 500\);
mini-batch size \(M = 100\);
Adagrad, learning rate selected from \(\{0.01, 0.02, 0.1\}\) (here:
0.02);all parameters initialized from \(\mathcal{N}(0, 0.01)\) (variance \(0.01\), i.e. standard deviation \(0.1\));
small weight decay corresponding to a prior \(\mathcal{N}(0, I)\) on all parameters, i.e. a weight decay of \(1/N\) with \(N\) training points, since the loss is a per-datapoint mean;
\(L = 1\) Monte Carlo sample for the latent variable estimator.
Configuration#
We begin with imports, global settings, and the reproducibility utilities. This includes the list of latent dimensions used for \(N_z\), hyperparameters, and device configuration.
1"""
2Reproduce the MNIST VAE experiment from Kingma & Welling (2013), Fig. 2.
3
4We train VAEs with different latent dimensionalities N_z and track the ELBO
5as a function of the number of training samples seen.
6The setup follows the original paper:
7
8- One hidden layer MLPs with Tanh activations in encoder/decoder
9- Hidden size H = 500
10- Mini-batch size M = 100
11- Adagrad, learning rate selected from {0.01, 0.02, 0.1} (we use 0.02 here)
12- All parameters initialized from N(0, 0.01) (variance 0.01, i.e. std 0.1)
13- Small weight decay corresponding to a prior N(0, I) on all parameters
14- L = 1 Monte Carlo sample for the stochastic latent variable
15"""
16
17from __future__ import annotations
18
19import random
20import time
21from pathlib import Path
22from typing import Dict, List, Tuple
23
24import matplotlib.pyplot as plt
25import numpy as np
26import torch
27import torch.nn as nn
28from matplotlib.ticker import (
29 LogFormatterMathtext,
30 LogLocator,
31 NullFormatter,
32)
33from torch.utils.data import DataLoader
34from torchvision import datasets, transforms
35
36from pyautoencoder.variational import VAE
37
38
39# ---------------- Configuration ---------------- #
40LATENT_DIMS: List[int] = [3, 5, 10, 20, 200] # N_z values (paper)
41HIDDEN_SIZE: int = 500 # hidden layer size (encoder/decoder, MNIST)
42BATCH_SIZE: int = 100 # M = 100 (paper)
43LEARNING_RATE: float = 0.02 # chosen from {0.01, 0.02, 0.1} (paper)
44MC_SAMPLES: int = 1 # L = 1 (paper)
45TARGET_TRAIN_SAMPLES: int = int(1e7) # stop after this many training samples
46EVAL_EVERY_SAMPLES: int = int(1e5) # evaluate and log every this many samples
47USE_STOCHASTIC_BINARIZATION: bool = False # binarize x ~ Bernoulli(p = p_x) if True
48
49DEVICE = torch.device("cpu")
50SEED: int = 1926
51
52
53# ---------------- Utilities ------------------- #
54def set_seed(seed: int) -> None:
55 random.seed(seed)
56 np.random.seed(seed)
57 torch.manual_seed(seed)
58
59
60@torch.no_grad()
61def init_params_paper(model: nn.Module) -> None:
62 # Paper, Sec. 5: "All parameters, both variational and generative, were initialized
63 # by random sampling from N(0, 0.01)". As in N(0, I), the second argument is the
64 # variance: std = 0.1, for weights and biases alike.
65 for param in model.parameters():
66 param.normal_(mean=0.0, std=0.1)
67
68
Key parameters:
LATENT_DIMS– list of latent sizes \(N_z\) evaluated.HIDDEN_SIZE– hidden layer size in encoder and decoder.MC_SAMPLES– number of Monte Carlo samples (\(L=1\) in the paper).set_seed()– deterministic seeding for reproducibility.
Data Loading#
The MNIST dataset is loaded through torchvision.datasets.
Optionally, inputs may be stochastically binarized as in the original
AEVB experiments.
1def make_dataloaders(
2 batch_size: int,
3 use_stochastic_binarization: bool = USE_STOCHASTIC_BINARIZATION,
4) -> Tuple[DataLoader, DataLoader]:
5
6 tfms: List[transforms.Compose | transforms.ToTensor] = [transforms.ToTensor()]
7
8 if use_stochastic_binarization:
9 class BernoulliBinarize:
10 def __call__(self, x: torch.Tensor) -> torch.Tensor:
11 # Sample x ~ Bernoulli(p = x); assumes x in [0, 1].
12 return torch.bernoulli(x)
13
14 tfms.append(BernoulliBinarize()) # type: ignore
15
16 transform = transforms.Compose(tfms)
17
18 train_dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
19 test_dataset = datasets.MNIST("./data", train=False, download=True, transform=transform)
20
21 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
22 test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, drop_last=False)
23 return train_loader, test_loader
The make_dataloaders() function returns training and test dataloaders
with appropriate preprocessing and optional binarization.
Model Definition#
Each VAE consists of:
a single-hidden-layer encoder mapping \(x \mapsto (\mu, \log \sigma^2)\): the user-defined encoder computes the Tanh hidden layer, and the VAE’s sampling layer maps it to \(\mu\) and \(\log \sigma^2\),
a single-hidden-layer decoder mapping \(z \mapsto \hat{x}\),
Tanh nonlinearities,
linear output logits interpreted by
VAE.compute_loss()under a Bernoulli likelihood.
All parameters (weights and biases) are initialized from \(\mathcal{N}(0, 0.01)\) as described in the paper, reading the second argument as the variance.
1def make_vae(latent_dim: int, hidden: int = HIDDEN_SIZE) -> VAE:
2 encoder = nn.Sequential(
3 nn.Flatten(),
4 nn.Linear(28 * 28, hidden),
5 nn.Tanh(),
6 )
7
8 decoder = nn.Sequential(
9 nn.Linear(latent_dim, hidden),
10 nn.Tanh(),
11 nn.Linear(hidden, 28 * 28),
12 nn.Unflatten(-1, (1, 28, 28)), # keep last layer linear: the Bernoulli likelihood expects logits
13 )
14
15 model = VAE(encoder=encoder, decoder=decoder, latent_dim=latent_dim)
16 model.build(input_sample=torch.randn(1, 1, 28, 28))
17 init_params_paper(model)
18 return model
The model is built explicitly via VAE.build(), which infers
dimension-dependent components from a representative input sample.
ELBO Evaluation#
We evaluate the ELBO over a dataloader by averaging the batch-wise elbo
diagnostics returned by VAE.compute_loss(), weighted by batch size.
The forward pass samples \(z \sim q(z \mid x)\) also in eval mode, so the
returned elbo is a Monte Carlo estimate of the ELBO on the evaluated data.
1@torch.no_grad()
2def average_elbo(dataloader: DataLoader, model: VAE) -> float:
3 model.eval() # forward still samples z ~ q(z|x), so this is an ELBO estimate
4 total_elbo = torch.zeros((), device=DEVICE)
5 n = 0
6
7 for x, _ in dataloader:
8 x = x.to(DEVICE)
9 out = model(x, S=MC_SAMPLES)
10 loss_info = model.compute_loss(x, out, beta=1, likelihood='bernoulli')
11 elbo_batch = loss_info.diagnostics['elbo']
12 batch_size = x.size(0)
13 total_elbo += elbo_batch * batch_size
14 n += batch_size
15
16 return total_elbo.item() / n
This routine is used during training to record both training and test ELBO values.
Training Loop#
For each latent dimension \(N_z\), the VAE is trained until a fixed
number of training samples have been processed
(TARGET_TRAIN_SAMPLES). Periodic evaluations of the ELBO on both
train and test sets are logged.
1def train_one_setting(
2 latent_dim: int,
3 train_loader: DataLoader,
4 test_loader: DataLoader,
5) -> Tuple[VAE, List[Dict[str, float]]]:
6 model = make_vae(latent_dim).to(DEVICE)
7
8 # Paper, Sec. 5: "a small weight decay term corresponding to a prior p(theta) = N(0, I)",
9 # and all parameters "were jointly stochastically optimized using the MAP criterion",
10 # so the prior is applied to encoder and decoder alike. The loss is a per-datapoint mean
11 # (the dataset objective divided by N), so the prior contributes -||params||^2 / (2N):
12 # a weight decay of 1/N.
13 optimizer = torch.optim.Adagrad(
14 model.parameters(),
15 lr=LEARNING_RATE,
16 weight_decay=1.0 / len(train_loader.dataset),
17 )
18
19 logs: List[Dict[str, float]] = []
20 samples_seen = 0
21 next_eval = EVAL_EVERY_SAMPLES
22
23 start_time = time.time()
24 while samples_seen < TARGET_TRAIN_SAMPLES:
25 for x, _ in train_loader:
26 x = x.to(DEVICE)
27 model.train()
28 optimizer.zero_grad()
29
30 out = model(x, S=MC_SAMPLES)
31 loss_info = model.compute_loss(x, out, beta=1, likelihood='bernoulli')
32 loss_info.objective.backward()
33 optimizer.step()
34
35 batch_size = x.size(0)
36 samples_seen += batch_size
37
38 if samples_seen >= next_eval:
39 train_elbo = average_elbo(train_loader, model)
40 test_elbo = average_elbo(test_loader, model)
41 logs.append(
42 {
43 "samples": float(samples_seen),
44 "train_elbo": float(train_elbo),
45 "test_elbo": float(test_elbo),
46 }
47 )
48 elapsed = time.time() - start_time
49 print(
50 f"N_z={latent_dim:3d} | "
51 f"samples={samples_seen:>9d} | "
52 f"ELBO_train={train_elbo:.2f}, "
53 f"ELBO_test={test_elbo:.2f} | "
54 f"(elapsed {elapsed:.1f}s)"
55 )
56 next_eval += EVAL_EVERY_SAMPLES
57
58 if samples_seen >= TARGET_TRAIN_SAMPLES:
59 break
60
61 return model, logs
Each training step:
Draws
MC_SAMPLES(\(L = 1\)) Monte Carlo samples from \(q(z \mid x)\),Computes the negative ELBO via
VAE.compute_loss()withlikelihood='bernoulli',Backpropagates gradients through encoder and decoder,
Records evaluation metrics periodically.
The ELBO diagnostics contain:
elbo– batch-mean Evidence Lower Bound (beta=1here),log_likelihood– batch-mean reconstruction term,kl_divergence– batch-mean KL divergence.
Plotting the Results#
After all training runs complete, the ELBO curves are plotted as a function of the number of processed samples (log-scaled), following the appearance of Fig. 2 from Kingma & Welling (2013).
1def plot_elbo_curves(all_logs: Dict[int, List[Dict[str, float]]]) -> Path:
2 num_settings = len(all_logs)
3 fig, axes = plt.subplots(1, num_settings, figsize=(3.0 * num_settings, 3.0))
4
5 # Ensure axes is iterable even when num_settings == 1
6 if num_settings == 1:
7 axes = np.array([axes])
8
9 latent_dims_sorted = sorted(all_logs.keys())
10
11 for i, nz in enumerate(latent_dims_sorted):
12 ax = axes[i]
13 xs = [entry["samples"] for entry in all_logs[nz]]
14 ys_train = [entry["train_elbo"] for entry in all_logs[nz]]
15 ys_test = [entry["test_elbo"] for entry in all_logs[nz]]
16
17 ax.plot(xs, ys_train, label="AEVB (train)", color="r")
18 ax.plot(xs, ys_test, linestyle="--", label="AEVB (test)", color="r")
19
20 ax.set_xscale("log")
21 ax.set_ylim(-150, -95)
22 ax.set_title(f"MNIST, $N_z = {nz}$")
23
24 if i == 0:
25 ax.set_xlabel("# training samples evaluated")
26 ax.set_ylabel(r"$\mathcal{L}$ (ELBO)")
27
28 # Log-scale formatting on x-axis
29 ax.xaxis.set_major_locator(LogLocator(base=10))
30 ax.xaxis.set_major_formatter(LogFormatterMathtext(base=10))
31 ax.xaxis.set_minor_locator(LogLocator(base=10, subs=range(2, 10)))
32 ax.xaxis.set_minor_formatter(NullFormatter())
33
34 axes[0].legend(loc="lower right")
35 plt.tight_layout()
36
37 out_path = Path("vae_mnist_fig2_repro.png")
38 plt.savefig(out_path, dpi=200, bbox_inches="tight")
39 plt.close(fig)
40 print(f"Saved figure -> {out_path}")
41 return out_path
The figure is saved as vae_mnist_fig2_repro.png.
Entry Point#
The main() function orchestrates the full experiment: seeding,
data loading, training VAEs for each \(N_z\), logging results, and
plotting the final curves.
1def main() -> None:
2 print("=== Reproducing Kingma & Welling (2013) Fig. 2 on MNIST ===")
3 print(f"Using device: {DEVICE}")
4
5 set_seed(SEED)
6 train_loader, test_loader = make_dataloaders(BATCH_SIZE)
7
8 all_logs: Dict[int, List[Dict[str, float]]] = {}
9 for nz in LATENT_DIMS:
10 print(f"\n--- Training VAE with N_z = {nz} ---")
11 _, logs = train_one_setting(nz, train_loader, test_loader)
12 all_logs[nz] = logs
13
14 plot_elbo_curves(all_logs)
15 print("Done.")
The script can be executed directly:
1if __name__ == "__main__":
2 main()
Full Example Script#
For convenience, here is the complete script in one block:
1"""
2Reproduce the MNIST VAE experiment from Kingma & Welling (2013), Fig. 2.
3
4We train VAEs with different latent dimensionalities N_z and track the ELBO
5as a function of the number of training samples seen.
6The setup follows the original paper:
7
8- One hidden layer MLPs with Tanh activations in encoder/decoder
9- Hidden size H = 500
10- Mini-batch size M = 100
11- Adagrad, learning rate selected from {0.01, 0.02, 0.1} (we use 0.02 here)
12- All parameters initialized from N(0, 0.01) (variance 0.01, i.e. std 0.1)
13- Small weight decay corresponding to a prior N(0, I) on all parameters
14- L = 1 Monte Carlo sample for the stochastic latent variable
15"""
16
17from __future__ import annotations
18
19import random
20import time
21from pathlib import Path
22from typing import Dict, List, Tuple
23
24import matplotlib.pyplot as plt
25import numpy as np
26import torch
27import torch.nn as nn
28from matplotlib.ticker import (
29 LogFormatterMathtext,
30 LogLocator,
31 NullFormatter,
32)
33from torch.utils.data import DataLoader
34from torchvision import datasets, transforms
35
36from pyautoencoder.variational import VAE
37
38
39# ---------------- Configuration ---------------- #
40LATENT_DIMS: List[int] = [3, 5, 10, 20, 200] # N_z values (paper)
41HIDDEN_SIZE: int = 500 # hidden layer size (encoder/decoder, MNIST)
42BATCH_SIZE: int = 100 # M = 100 (paper)
43LEARNING_RATE: float = 0.02 # chosen from {0.01, 0.02, 0.1} (paper)
44MC_SAMPLES: int = 1 # L = 1 (paper)
45TARGET_TRAIN_SAMPLES: int = int(1e7) # stop after this many training samples
46EVAL_EVERY_SAMPLES: int = int(1e5) # evaluate and log every this many samples
47USE_STOCHASTIC_BINARIZATION: bool = False # binarize x ~ Bernoulli(p = p_x) if True
48
49DEVICE = torch.device("cpu")
50SEED: int = 1926
51
52
53# ---------------- Utilities ------------------- #
54def set_seed(seed: int) -> None:
55 random.seed(seed)
56 np.random.seed(seed)
57 torch.manual_seed(seed)
58
59
60@torch.no_grad()
61def init_params_paper(model: nn.Module) -> None:
62 # Paper, Sec. 5: "All parameters, both variational and generative, were initialized
63 # by random sampling from N(0, 0.01)". As in N(0, I), the second argument is the
64 # variance: std = 0.1, for weights and biases alike.
65 for param in model.parameters():
66 param.normal_(mean=0.0, std=0.1)
67
68
69# ---------------- Data ----------------------- #
70def make_dataloaders(
71 batch_size: int,
72 use_stochastic_binarization: bool = USE_STOCHASTIC_BINARIZATION,
73) -> Tuple[DataLoader, DataLoader]:
74
75 tfms: List[transforms.Compose | transforms.ToTensor] = [transforms.ToTensor()]
76
77 if use_stochastic_binarization:
78 class BernoulliBinarize:
79 def __call__(self, x: torch.Tensor) -> torch.Tensor:
80 # Sample x ~ Bernoulli(p = x); assumes x in [0, 1].
81 return torch.bernoulli(x)
82
83 tfms.append(BernoulliBinarize()) # type: ignore
84
85 transform = transforms.Compose(tfms)
86
87 train_dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
88 test_dataset = datasets.MNIST("./data", train=False, download=True, transform=transform)
89
90 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
91 test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, drop_last=False)
92 return train_loader, test_loader
93
94
95# ---------------- Model ---------------------- #
96def make_vae(latent_dim: int, hidden: int = HIDDEN_SIZE) -> VAE:
97 encoder = nn.Sequential(
98 nn.Flatten(),
99 nn.Linear(28 * 28, hidden),
100 nn.Tanh(),
101 )
102
103 decoder = nn.Sequential(
104 nn.Linear(latent_dim, hidden),
105 nn.Tanh(),
106 nn.Linear(hidden, 28 * 28),
107 nn.Unflatten(-1, (1, 28, 28)), # keep last layer linear: the Bernoulli likelihood expects logits
108 )
109
110 model = VAE(encoder=encoder, decoder=decoder, latent_dim=latent_dim)
111 model.build(input_sample=torch.randn(1, 1, 28, 28))
112 init_params_paper(model)
113 return model
114
115# ---------------- Evaluation ----------------- #
116@torch.no_grad()
117def average_elbo(dataloader: DataLoader, model: VAE) -> float:
118 model.eval() # forward still samples z ~ q(z|x), so this is an ELBO estimate
119 total_elbo = torch.zeros((), device=DEVICE)
120 n = 0
121
122 for x, _ in dataloader:
123 x = x.to(DEVICE)
124 out = model(x, S=MC_SAMPLES)
125 loss_info = model.compute_loss(x, out, beta=1, likelihood='bernoulli')
126 elbo_batch = loss_info.diagnostics['elbo']
127 batch_size = x.size(0)
128 total_elbo += elbo_batch * batch_size
129 n += batch_size
130
131 return total_elbo.item() / n
132
133
134# ---------------- Training ------------------- #
135def train_one_setting(
136 latent_dim: int,
137 train_loader: DataLoader,
138 test_loader: DataLoader,
139) -> Tuple[VAE, List[Dict[str, float]]]:
140 model = make_vae(latent_dim).to(DEVICE)
141
142 # Paper, Sec. 5: "a small weight decay term corresponding to a prior p(theta) = N(0, I)",
143 # and all parameters "were jointly stochastically optimized using the MAP criterion",
144 # so the prior is applied to encoder and decoder alike. The loss is a per-datapoint mean
145 # (the dataset objective divided by N), so the prior contributes -||params||^2 / (2N):
146 # a weight decay of 1/N.
147 optimizer = torch.optim.Adagrad(
148 model.parameters(),
149 lr=LEARNING_RATE,
150 weight_decay=1.0 / len(train_loader.dataset),
151 )
152
153 logs: List[Dict[str, float]] = []
154 samples_seen = 0
155 next_eval = EVAL_EVERY_SAMPLES
156
157 start_time = time.time()
158 while samples_seen < TARGET_TRAIN_SAMPLES:
159 for x, _ in train_loader:
160 x = x.to(DEVICE)
161 model.train()
162 optimizer.zero_grad()
163
164 out = model(x, S=MC_SAMPLES)
165 loss_info = model.compute_loss(x, out, beta=1, likelihood='bernoulli')
166 loss_info.objective.backward()
167 optimizer.step()
168
169 batch_size = x.size(0)
170 samples_seen += batch_size
171
172 if samples_seen >= next_eval:
173 train_elbo = average_elbo(train_loader, model)
174 test_elbo = average_elbo(test_loader, model)
175 logs.append(
176 {
177 "samples": float(samples_seen),
178 "train_elbo": float(train_elbo),
179 "test_elbo": float(test_elbo),
180 }
181 )
182 elapsed = time.time() - start_time
183 print(
184 f"N_z={latent_dim:3d} | "
185 f"samples={samples_seen:>9d} | "
186 f"ELBO_train={train_elbo:.2f}, "
187 f"ELBO_test={test_elbo:.2f} | "
188 f"(elapsed {elapsed:.1f}s)"
189 )
190 next_eval += EVAL_EVERY_SAMPLES
191
192 if samples_seen >= TARGET_TRAIN_SAMPLES:
193 break
194
195 return model, logs
196
197
198# ---------------- Plotting ------------------- #
199def plot_elbo_curves(all_logs: Dict[int, List[Dict[str, float]]]) -> Path:
200 num_settings = len(all_logs)
201 fig, axes = plt.subplots(1, num_settings, figsize=(3.0 * num_settings, 3.0))
202
203 # Ensure axes is iterable even when num_settings == 1
204 if num_settings == 1:
205 axes = np.array([axes])
206
207 latent_dims_sorted = sorted(all_logs.keys())
208
209 for i, nz in enumerate(latent_dims_sorted):
210 ax = axes[i]
211 xs = [entry["samples"] for entry in all_logs[nz]]
212 ys_train = [entry["train_elbo"] for entry in all_logs[nz]]
213 ys_test = [entry["test_elbo"] for entry in all_logs[nz]]
214
215 ax.plot(xs, ys_train, label="AEVB (train)", color="r")
216 ax.plot(xs, ys_test, linestyle="--", label="AEVB (test)", color="r")
217
218 ax.set_xscale("log")
219 ax.set_ylim(-150, -95)
220 ax.set_title(f"MNIST, $N_z = {nz}$")
221
222 if i == 0:
223 ax.set_xlabel("# training samples evaluated")
224 ax.set_ylabel(r"$\mathcal{L}$ (ELBO)")
225
226 # Log-scale formatting on x-axis
227 ax.xaxis.set_major_locator(LogLocator(base=10))
228 ax.xaxis.set_major_formatter(LogFormatterMathtext(base=10))
229 ax.xaxis.set_minor_locator(LogLocator(base=10, subs=range(2, 10)))
230 ax.xaxis.set_minor_formatter(NullFormatter())
231
232 axes[0].legend(loc="lower right")
233 plt.tight_layout()
234
235 out_path = Path("vae_mnist_fig2_repro.png")
236 plt.savefig(out_path, dpi=200, bbox_inches="tight")
237 plt.close(fig)
238 print(f"Saved figure -> {out_path}")
239 return out_path
240
241
242# ---------------- Main ---------------------- #
243def main() -> None:
244 print("=== Reproducing Kingma & Welling (2013) Fig. 2 on MNIST ===")
245 print(f"Using device: {DEVICE}")
246
247 set_seed(SEED)
248 train_loader, test_loader = make_dataloaders(BATCH_SIZE)
249
250 all_logs: Dict[int, List[Dict[str, float]]] = {}
251 for nz in LATENT_DIMS:
252 print(f"\n--- Training VAE with N_z = {nz} ---")
253 _, logs = train_one_setting(nz, train_loader, test_loader)
254 all_logs[nz] = logs
255
256 plot_elbo_curves(all_logs)
257 print("Done.")
258
259if __name__ == "__main__":
260 main()