{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.12" }, "title": "LAB07｜一个图像 VAE 怎样同时学重建与可用起点", "course": "Lecture04", "evidence_status": "saved figures from executed standalone NumPy experiment", "execution_verification": "Every code block matches a successful Deepnote run snapshot by ID and hash. Per-cell metadata distinguishes imported cloud outputs from retained local baseline outputs.", "visual_verification": "All student-facing figures inspected; latent titles corrected by re-rendering saved arrays without retraining.", "lecture04": { "source_notebook_id": "6567b1168915486c8de4ab9cfea54a8e", "saved_results_origin": "Static baseline figures retained; code outputs carry per-cell run provenance." }, "deepnote_readback": { "notebook_id": "6567b1168915486c8de4ab9cfea54a8e", "updated_at": "2026-09-30T02:27:56.599Z", "block_count": 14, "all_source_blocks_match": true } }, "cells": [ { "cell_type": "markdown", "id": "lab07-reading-00", "metadata": { "deepnote_block_id": "86b59ab14bca47fb91069a3e84905782", "cloud_content_sha256": "sha256:32a69b7ea9d78da66564c69e215b9a5c21d139f32323d932776eba6deb564ff2" }, "source": "# LAB07｜一个图像 VAE 怎样同时学重建与可用起点\n\n你在课堂上已经看过两条路径：把一张图像编码后再解码，以及直接抽一个潜变量再解码。本实验把这两条路径放进**同一个真正训练的小网络**。先看保存结果，再读它学了什么；最后只改一个条件，重新运行。这里使用公开的 UCI 8×8 手写数字，由 scikit-learn 自带数据提供，**不是 MNIST**。所有展示均来自本次 NumPy 训练，没有用生成式插画替换模型输出。\n\n本册承接 N15–N20 与 CLS03。LAB02 用可解的高斯例子解释 q、后验与 ELBO；本册则让 1,437 张训练图像共用一个编码器和一个解码器，检查学到的参数是否能处理 360 张未用于更新的测试图像。\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-01", "metadata": { "deepnote_block_id": "d97c246eea494ce38f87a844d4ff5537", "cloud_content_sha256": "sha256:77c7706e50fc0dc9ffebf29d6f6f3288287d7faeb7a3f448c8bf9788bf95e992" }, "source": "## 1. 先看学习究竟改变了什么\n\n![固定十张测试数字，原始灰度、二值输入、初始化输出，以及β=0与β=1训练后的解码概率图。各列测试输入不变。](https://codingai-lec04.pages.dev/assets/experiments/vae/reconstruction.png)\n\n从上往下读一列。第一行是公开数据原来的灰度图，第二行是本实验真正输入的二值图：原像素值大于 8 记为 1，否则记为 0。第三行是随机初始化模型的输出。最后两行分别来自 β=0、β=1 的训练。先找数字 0 的空心区域，再看数字 3、8 是否仍容易分辨：学习已经让输出产生笔画结构，但两维潜变量和小模型也丢失了细节。它没有把全部测试数字都清楚重建出来。\n\n这张图把编码器输出的均值 μ(x) 送入解码器，用来做稳定的重建对照。灰度值是每个二值像素取 1 的概率，不是二值像素本身；`decoder(mu)` 也不等于对所有后验 z 求平均的 `E_q[decoder(z)]`。训练时并没有固定取 μ，而是随机抽取 z。\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-02", "metadata": { "deepnote_block_id": "0774747598614c4dba3637c8d448e1cc", "cloud_content_sha256": "sha256:fecc6b6608e0b99573819424b544ebc087c6d846afa7b03206dc0dc9492f0eb5" }, "source": "## 2. 已经能重建，为什么还不能随便取一个 z？\n\n![初始化与两组训练模型使用完全相同的十个标准正态潜变量。每列起点z固定，只比较解码器训练前后的输出概率图。](https://codingai-lec04.pages.dev/assets/experiments/vae/prior_same_seed.png)\n\n这一次没有输入数字图像。我们先从 N(0,I) 取出固定的 25 个二维向量，图中展示前十个，再把同一组 z 交给初始化、β=0 和 β=1 三个解码器。重建时起点来自“这张图被编码到哪里”；现在起点来自预先约定的标准正态分布。两种起点没有理由天然一致。\n\n这正是本册的模拟学生追问：“我已经能把输入重建出来，为什么还不能随便取一个 z？”先别用“因为 KL”结束回答。请看下一组图：彩色点是 360 张测试图的后验均值，颜色仅用于显示原数字标签，标签没有进入训练损失；黑叉是刚才那 25 个标准正态起点。虚线圆只是半径 2 的尺度参照，不是所有先验概率的边界。两幅左侧散点图按各自数据范围定轴；同一个半径 2 圆在屏幕上大小不同，先读轴刻度再比较区域，不能把像素距离当成潜坐标距离。\n\n![β=0后验均值与固定先验起点的位置；右边是[-3,3]平方区域的均匀潜空间网格解码结果。](https://codingai-lec04.pages.dev/assets/experiments/vae/latent_beta0.png)\n\n![β=1后验均值与同一组先验起点；右边仍是相同坐标范围的潜空间网格。](https://codingai-lec04.pages.dev/assets/experiments/vae/latent_beta1.png)\n\nβ=0 时，编码器把图像送到的区域明显移出了标准正态起点常见的尺度：后验均值距原点的平均距离约 **12.32**。模型能够在那些区域支持重建，却没有被要求照顾 N(0,I) 经常选到的区域。β=1 时，平均距离约为 **1.28**，两类起点的尺度更接近。因而先验入口的差异有了可观察的原因，不只是一个正则项名称。\n\n右边的网格把 z₁、z₂ 均匀扫过 [-3,3]，每个格子显示一个解码概率图。它帮助我们看解码器在这块区域如何变化，**不是从高斯先验独立采样出来的图阵**。左图也只显示后验均值；完整 q 还有方差，不能把均值散点当成整个聚合后验分布。\n\nβ=1 的先验输出仍有明显模糊和重复，并没有证明十类数字已均匀覆盖。这里观察到的是重建与先验入口之间的一次真实权衡，不是“大 β 的图片一定更好”。\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-03", "metadata": { "deepnote_block_id": "60204947b1fe476bb0fdfc771d97e79e", "cloud_content_sha256": "sha256:4e23d30071b97c94e8fb59c9526ee896f48f3ea2956c9b223daf1c4ef88e8b77" }, "source": "## 3. 两组训练只改变 β\n\n固定数据划分、网络、初始参数、每个 epoch 的样本顺序、重参数化噪声、优化器、训练次数。每组都训练 150 个 epoch，每批最多 64 张图，共 3,450 次更新。β=0、β=1 各采用一个预先固定的随机设置，没有搜索超参数，也没有按输出好坏换种子；另在独立进程执行 Notebook 代码复核，同设置得到完全相同的模型与观察数组。编码器是 64→64→(μ₂,logσ²₂)，解码器是 2→64→64，两边共享参数总数为 8,772。\n\n| 固定测试集上的量 | 初始化 | β=0，训练后 | β=1，训练后 |\n|---|---:|---:|---:|\n| 重建负对数似然，nats/图 | 44.875 | 16.061 | 18.638 |\n| KL(q(z\\|x)\\|N(0,I))，nats/图 | 0.454 | 103.375 | 2.258 |\n| 后验均值距原点的平均距离 | — | 12.323 | 1.285 |\n| 后验标准差的平均值 | — | 0.0098 | 0.3496 |\n\n重建项是 64 个像素求和，再对图像和 8 组固定后验噪声取平均；KL 是两个潜变量维度求和，再对图像平均。这样 β 的尺度才有明确含义。β=0 的重建项更低，但 KL 大幅增加；β=1 为兼顾先验入口付出了一些重建代价。不要直接比较两个不同权重目标的总损失并宣布一个更好。\n\n![相同测试集和固定蒙特卡洛噪声下的重建负对数似然与KL曲线。β=0的重建下降并不阻止后验远离标准先验。](https://codingai-lec04.pages.dev/assets/experiments/vae/learning_curves.png)\n\n本次 CPU 上，每组训练连同周期性测试评估约 1.6 秒；图像绘制与文件导出不计入这两个时间。运行时间只记录本次环境，不承诺其他机器同速。\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-04", "metadata": { "deepnote_block_id": "1996a1cd0a414a69ba6a05e9f1b2ff76", "cloud_content_sha256": "sha256:d5115f79e6d7f53d3431ca6af6714d67a86c0d1bb8d4cea29677b9ca60762781" }, "source": "## 4. 网络究竟学什么，随机性怎样进入梯度？\n\n编码器读取二值图像 x，给出两维 μ(x) 和 logσ²(x)，定义近似后验 qφ(z|x)。解码器读取 z，输出 64 个概率 πθ(z)，定义条件独立的 Bernoulli 观测模型。标准先验 p(z)=N(0,I) 固定，没有待学习参数。网络学的是两组共享权重 φ、θ，不是为每张图分别存一套解码器。\n\n训练时先抽 ε∼N(0,I)，再计算 z=μ+σ⊙ε。一次训练步把这次 ε 看作固定输入，损失便能沿 z 对 μ、σ 求导。σ 并未被绕过去：对 logσ² 的重建梯度包含 `(dL/dz) * epsilon * 0.5 * sigma`。这就是本实现的重参数化路径。\n\n我们最小化的批次目标为：\n\n$$J_\\beta = \\frac{1}{B}\\sum_i\\left[-\\sum_{j=1}^{64}\\log p_\\theta(x_{ij}|z_i)+\\beta\\,\\frac12\\sum_{k=1}^{2}(\\mu_{ik}^2+\\sigma_{ik}^2-1-\\log\\sigma_{ik}^2)\\right].$$\n\n每个样本每次更新使用一个后验 z。β=1 时，它是标准负 ELBO 的蒙特卡洛估计；β=0 则只保留随机编码下的重建目标。**β=0 仍然抽取 z，所以不能直接把它改称确定性 AE。** 本次它学到的 σ 很小，但“很小”不等于实现中完全没有抽样。\n\n用一个像素核对损失：若真实二值像素为 1，模型预测它为 1 的概率是 0.8，这个像素贡献 −log(0.8)≈0.223 nats；若预测 0.2，则贡献约 1.609 nats。模型靠重复更新共享参数，把这类损失反馈到解码器、z，再到编码器。\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-05", "metadata": { "deepnote_block_id": "d46f38703d524b5e9c4dfee8924c342a", "cloud_content_sha256": "sha256:e7c63b772268e439137334a46806c1f6ea57a15599819be79e0d74b7cda495b1" }, "source": "## 5. 后验抽样、插值和先验抽样不是同一操作\n\n![β=1时，数字0、3、6、9各固定一张输入，每行从该输入的近似后验抽八个z，显示解码概率图。](https://codingai-lec04.pages.dev/assets/experiments/vae/posterior_beta1.png)\n\n每一行只固定一张输入图。编码器先给出这一张图的 μ、σ；改变 ε 会得到同一 q(z|x) 下不同的 z，再产生不同概率图。因此它仍然借助输入图像，不是无输入的先验生成。观察一行内部能否保留同一数字的笔画，再与其他行比较；本次有些数字仍模糊，不能把每张输出都当作识别正确的样本。\n\n![β=1下，从测试数字3的后验均值到数字8的后验均值作11个等距点，并显示解码概率图。](https://codingai-lec04.pages.dev/assets/experiments/vae/interpolation_beta1.png)\n\n插值则先编码两张端点图，再计算 z(α)=(1−α)μ(x₃)+αμ(x₈)。中间点是我们指定的路径，不是后验抽样，也不是先验抽样。平滑变化说明这个解码器沿这条路径连续响应；它不保证中间每个图都属于清楚可辨的数字，也不证明潜空间的任意方向都有相同语义。\n\n| 操作 | 已经给定什么 | 改变或抽取什么 | 本图显示什么 |\n|---|---|---|---|\n| 稳定重建 | 一张 x | 取 z=μ(x) | decoder(μ) 的概率图 |\n| 后验抽样 | 一张 x 与 qφ(z\\|x) | ε，继而 z | 不同 z 的解码概率图 |\n| 潜变量插值 | 两张端点图的 μ | 路径位置 α | 指定路径上的概率图 |\n| 先验生成 | 固定 p(z) | z∼N(0,I) | 无输入图像的概率图 |\n| 潜空间网格 | 坐标范围与格点 | 按网格扫坐标 | 均匀格点的概率图 |\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-06", "metadata": { "deepnote_block_id": "5e0d7c5688b64df487478cc7cb939fd9", "cloud_content_sha256": "sha256:5f07b290f25f2161ae9662018d07b2ab9d9cab819b4b819b0e9f581af7677f99" }, "source": "## 6. 概率图之后，还能再抽一次观测\n\n![同一批β=1的先验潜变量，上行是解码器输出的Bernoulli概率，下行在对应概率下逐像素抽取0或1。](https://codingai-lec04.pages.dev/assets/experiments/vae/mean_vs_observation.png)\n\n上行一个灰色像素的值若为 0.7，它表示“这个像素取 1 的概率为 0.7”。下行用同一张概率图和固定均匀随机数 u，按 `u < probability` 得到真正的二值观测。下行更粗糙是这层观测抽样的直接结果，不能说它用了另一个 z 或另一个解码器。\n\n所以“从 VAE 生成图像”至少要交代显示的是哪一层。这里常用概率图，是为了稳定比较 z 和模型变化；若声称从整个生成模型抽样，就还包含 Bernoulli 观测这一步。\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-07", "metadata": { "deepnote_block_id": "2b3fe0af3cd842ecbb69891341616d19", "cloud_content_sha256": "sha256:9596d220c539427faac4f9a91054142cd1806208b9056be8d0656b661726d71e" }, "source": "## 7. 读实现，然后在独立进程中重算\n\n完整实现放在同目录 `lab07_image_vae.py`，只需要 NumPy、SciPy、scikit-learn、Matplotlib 与 threadpoolctl。数据由 scikit-learn 自带，无须在线下载；没有其他 Notebook 留下的隐藏变量。运行 `python lab07_image_vae.py --out results` 会从初始化开始训练两组模型，保存数据划分、初始和最终权重、固定测试输入、逐期曲线、原始观察数组与全部图片。\n\n`forward` 对像素的求和、对批次的平均和 KL 梯度都显式写出。运行前会固定 ε，对 β=0 与 β=1 两条路径、每个参数张量的一些元素作中心有限差分核验；本次最大缩放误差见 `results/gradient_check.json`。这项核验验证梯度实现的局部一致性，不替代生成效果评价。\n\n保存结果来自一次独立 Python 进程执行；本册的 Notebook 包含同一份实现与独立运行入口。浏览保存图不需要启动训练，点击运行才会重新学习参数。参数、原始数组、图像与软件版本都保存在同一个结果目录。\n\n\n\n**在 Deepnote 重算时怎样读新结果。** 从下方实现开始按顺序运行；运行格会直接显示本次目录中的重建、共同先验起点和两组潜空间图。先确认两组共用输入与起点，再分别比较重建项、KL 与先验图。上方图仍是最初保存的基线；修改 β 后，运行格显示的是你的新结果。\n\n\n**云端复核（2026-09-30）。** 本册在 Deepnote 的独立内核按自身步骤完整运行。150 epoch 的两组 β 对照完整执行，重建项、KL 和平均潜半径与上面的基线在显示精度内一致；新结果图由运行格直接显示。上方原始保存图继续标为本地基线，两次测量分别记录；大图在插件的运行快照预览中可能因尺寸被省略，这不代替学生在同册查看自己的输出。\n" }, { "cell_type": "markdown", "id": "lab07-reading-08", "metadata": { "deepnote_block_id": "3cd454fa9d8b4abaa4fb2b2c069bebfd", "cloud_content_sha256": "sha256:d7f0764f0867032aec659c61048ae66c448654bd44e78f218483b6cca0703337" }, "source": "## 8. 只改一个变量，留下可解释的观察\n\n先保留这次结果。把 `CONFIG[\"betas\"]` 改为 `[0.0, 0.5, 1.0]`，其余设置不动，并输出到另一个目录。先写下你预期改变的两项：重建项、后验与先验的相对位置。再看新结果有没有支持这个预期。图像会按新的 β 列表生成对应行，新增组另存权重和原始观察。本次保存证据只覆盖 β=0 与 1，没有运行 β=0.5。\n\n更直接的单变量练习是运行 `python lab07_image_vae.py --out shorter_training --epochs 50`，保持两组 β 不变，只改变训练预算。比较的是“150→50 个 epoch”这一干预，不能把较短训练结果当作新的 β 对照。观察完成后，用一段话回答：哪个量真的被学习改变了？哪条生成路径借助输入图像？重建更好是否同时改善了标准先验入口？\n\n自检：如果只展示 `decoder(mu)`，能否说已经检验从 p(z) 生成？不能，因为它的起点仍来自测试输入。若看到先验图阵里十张图相似，能否直接宣布完整模型坍塌？也不能；先记录固定十个起点下的重复现象，再扩大事先定义的观察，区分两维瓶颈、小数据、训练预算与分布覆盖问题。\n\n" }, { "cell_type": "markdown", "id": "lab07-reading-09", "metadata": { "deepnote_block_id": "757f65a806524d8daaa3c8120052f2ac", "cloud_content_sha256": "sha256:f1df5c43c749f622a63cc5846e4a0a093779ff46776e2a8f452fb3e9d74f2e38" }, "source": "## 9. 从结果回到机制\n\n本实验让同一组共享参数实际连接重建与无输入生成，并给出了一个可以指着图解释的断点：**重建路径常去的潜空间区域，不必等于先验生成经常访问的区域。** KL 权重改变了这两种需求的权衡，代价和收益都需要分别观察。请回到 CLS03 的 q 与先验部分，或在 LAB02 的可解例子中计算一次 KL，再回来检查你算的量在这张图里对应什么。\n\n原理核对使用 [Kingma 与 Welling《Auto-Encoding Variational Bayes》](https://arxiv.org/html/1312.6114v11) 的重参数化、Gaussian q、Bernoulli decoder 与 Gaussian KL；数据定义使用 [scikit-learn 官方 load_digits 文档](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_digits.html)。本实验是课程自己的轻量实现，不是论文实验指标的复现，也不是用于评价真实高清图像生成的基准。\n" }, { "cell_type": "markdown", "id": "lab07-implementation-note", "metadata": { "deepnote_block_id": "ae9e81c6d4084e528bce2b3b5eac805e", "cloud_content_sha256": "sha256:32676b77ef40483d51be539a9b0fb8574a31fa62bff23b933c147c6f0079a7fc" }, "source": "## 完整可运行实现\n下面的代码就是随附 `.py` 的同源实现。在新内核中从这里开始运行，下一格会重训；不会读取其他册的变量。上面的图是已经保存的真实实验结果。\n" }, { "cell_type": "code", "id": "lab07-implementation", "metadata": { "jupyter": { "source_hidden": true }, "deepnote_block_id": "ddfd9d2a9fbc4262b1db810652f6d2a4", "output_provenance": { "origin": "Deepnote executed snapshot", "run_id": "3fe20b5a-7faa-4e24-aafe-5fc4b1b95a00", "content_hash": "sha256:f704cb6651f3090a74b324fb7f5aa517b25c3311cd14a0c22a6d9ed94bb82777" }, "verified_cloud_execution": { "run_id": "3fe20b5a-7faa-4e24-aafe-5fc4b1b95a00", "content_hash": "sha256:f704cb6651f3090a74b324fb7f5aa517b25c3311cd14a0c22a6d9ed94bb82777", "output_preview_omitted": false, "completed_at": "2026-09-30T02:25:07.392Z" }, "cloud_content_sha256": "sha256:f704cb6651f3090a74b324fb7f5aa517b25c3311cd14a0c22a6d9ed94bb82777" }, "execution_count": null, "source": "\"\"\"LAB07: genuine small image VAE training, NumPy-only backpropagation.\n\nRun: python lab07_image_vae.py --out results\nRequires numpy, scipy, scikit-learn, matplotlib, threadpoolctl. No download, GPU or hidden state.\nUses sklearn's bundled 8x8 UCI digits, NOT MNIST. Pixel > 8 -> 1.\nPrimary references:\nhttps://arxiv.org/html/1312.6114v11 (sections 2-3, appendices B/C)\nhttps://scikit-learn.org/stable/modules/generated/sklearn.datasets.load_digits.html\n\"\"\"\nimport os\nfor name in (\"OPENBLAS_NUM_THREADS\", \"OMP_NUM_THREADS\", \"MKL_NUM_THREADS\"):\n os.environ[name] = \"1\"\nimport argparse\nimport copy\nimport csv\nimport hashlib\nimport json\nfrom pathlib import Path\nimport platform\nimport time\nimport numpy as np\nfrom threadpoolctl import threadpool_limits\nimport scipy\nfrom scipy.special import expit\nimport sklearn\nfrom sklearn.datasets import load_digits\nfrom sklearn.model_selection import train_test_split\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\n\nCONFIG = dict(data=\"sklearn.datasets.load_digits (UCI 8x8 digits)\",\n binarization=\"pixel > 8\", train_fraction=0.8,\n split_seed=7301, init_seed=7302, train_seed=7303,\n evaluation_seed=7304, visual_seed=7305,\n architecture=\"64-tanh64-(mu2,logvar2); z2-tanh64-logits64\",\n latent_dim=2, hidden_dim=64, epochs=150, batch_size=64,\n learning_rate=0.001, betas=[0.0, 1.0],\n reconstruction=\"Bernoulli negative log likelihood: sum 64 pixels, mean batch\",\n kl=\"diagonal Gaussian to N(0,I): sum 2 latent dimensions, mean batch\",\n train_mc_samples=1, evaluation_mc_samples=8,\n optimizer=\"Adam beta1=0.9 beta2=0.999 eps=1e-8\",\n dtype=\"float64\", version=\"lab07-v1\")\n\n\ndef initialize(seed):\n rng = np.random.default_rng(seed)\n p = {}\n for name, a, b in [(\"e\", 64, 64), (\"mu\", 64, 2), (\"lv\", 64, 2),\n (\"d\", 2, 64), (\"out\", 64, 64)]:\n p[name+\"W\"] = rng.normal(0, np.sqrt(2/(a+b)), (a, b))\n p[name+\"b\"] = np.zeros(b)\n return p\n\n\ndef encode(p, x):\n h = np.tanh(x @ p[\"eW\"] + p[\"eb\"])\n return h @ p[\"muW\"] + p[\"mub\"], h @ p[\"lvW\"] + p[\"lvb\"]\n\n\ndef decode(p, z):\n return expit(np.tanh(z @ p[\"dW\"] + p[\"db\"]) @ p[\"outW\"] + p[\"outb\"])\n\n\ndef forward(p, x, eps, beta, gradient=False):\n h = np.tanh(x @ p[\"eW\"] + p[\"eb\"])\n mu, lv = h @ p[\"muW\"] + p[\"mub\"], h @ p[\"lvW\"] + p[\"lvb\"]\n sd = np.exp(0.5 * lv)\n z = mu + sd * eps # reparameterization, no detached mu/sd\n hd = np.tanh(z @ p[\"dW\"] + p[\"db\"])\n logits = hd @ p[\"outW\"] + p[\"outb\"]\n prob = expit(logits)\n recon = np.mean(np.sum(np.logaddexp(0, logits) - x * logits, axis=1))\n kl = np.mean(0.5 * np.sum(mu**2 + np.exp(lv) - 1 - lv, axis=1))\n result = {\"objective\": float(recon + beta * kl), \"reconstruction_nll\": float(recon),\n \"kl\": float(kl), \"negative_elbo\": float(recon + kl)}\n if not gradient:\n return result\n n = len(x)\n g = {}\n dl = (prob-x)/n # pixel SUM; batch MEAN\n g[\"outW\"], g[\"outb\"] = hd.T @ dl, dl.sum(0)\n dhd = (dl @ p[\"outW\"].T) * (1-hd**2)\n g[\"dW\"], g[\"db\"] = z.T @ dhd, dhd.sum(0)\n dz = dhd @ p[\"dW\"].T\n dmu = dz + beta * mu/n\n dlv = dz * eps * 0.5 * sd + beta * 0.5 * (np.exp(lv)-1)/n\n g[\"muW\"], g[\"mub\"] = h.T @ dmu, dmu.sum(0)\n g[\"lvW\"], g[\"lvb\"] = h.T @ dlv, dlv.sum(0)\n dh = (dmu @ p[\"muW\"].T + dlv @ p[\"lvW\"].T) * (1-h**2)\n g[\"eW\"], g[\"eb\"] = x.T @ dh, dh.sum(0)\n return result, g\n\n\ndef gradient_check(p, x):\n \"\"\"Finite differences with fixed epsilon; check both beta paths, every tensor.\"\"\"\n rng = np.random.default_rng(7306)\n eps = rng.normal(size=(len(x), 2))\n checks = []\n for beta in [0.0, 1.0]:\n _, analytic = forward(p, x, eps, beta, True)\n for name, tensor in p.items():\n for flat in rng.choice(tensor.size, min(4, tensor.size), replace=False):\n index = np.unravel_index(flat, tensor.shape)\n original = tensor[index]\n tensor[index] = original+1e-5\n up = forward(p, x, eps, beta)[\"objective\"]\n tensor[index] = original-1e-5\n down = forward(p, x, eps, beta)[\"objective\"]\n tensor[index] = original\n numerical = (up-down)/2e-5\n actual = analytic[name][index]\n scaled = abs(numerical-actual)/max(1, abs(numerical), abs(actual))\n checks.append({\"beta\": beta, \"parameter\": name, \"index\": list(index),\n \"analytic\": float(actual), \"finite_difference\": float(numerical),\n \"scaled_error\": float(scaled)})\n assert max(c[\"scaled_error\"] for c in checks) < 1e-6\n return checks\n\n\ndef evaluate(p, x, eps_samples, beta):\n records = [forward(p, x, eps, beta) for eps in eps_samples]\n return {key: float(np.mean([v[key] for v in records])) for key in records[0]}\n\n\ndef train(initial, xtrain, xtest, test_eps, beta, cfg):\n p = copy.deepcopy(initial)\n m, v = ({k: np.zeros_like(a) for k, a in p.items()} for _ in range(2))\n rng = np.random.default_rng(cfg[\"train_seed\"])\n log = [{\"epoch\": 0, **evaluate(p, xtest, test_eps, beta)}]\n tick, step = time.perf_counter(), 0\n for epoch in range(1, cfg[\"epochs\"]+1):\n permutation = rng.permutation(len(xtrain))\n for start in range(0, len(xtrain), cfg[\"batch_size\"]):\n batch = xtrain[permutation[start:start+cfg[\"batch_size\"]]]\n eps = rng.normal(size=(len(batch), cfg[\"latent_dim\"]))\n _, grad = forward(p, batch, eps, beta, True)\n step += 1\n for k in p:\n m[k] = .9*m[k] + .1*grad[k]\n v[k] = .999*v[k] + .001*grad[k]**2\n p[k] -= cfg[\"learning_rate\"] * (m[k]/(1-.9**step))/(np.sqrt(v[k]/(1-.999**step))+1e-8)\n if epoch % 5 == 0 or epoch == cfg[\"epochs\"]:\n log.append({\"epoch\": epoch, **evaluate(p, xtest, test_eps, beta)})\n return p, log, time.perf_counter()-tick, step\n\n\ndef image_rows(rows, labels, target, titles=None, title=\"\", scale=1):\n nrows, ncols = len(rows), len(rows[0])\n fig, axes = plt.subplots(nrows, ncols, figsize=(ncols*1.15*scale, nrows*1.25*scale),\n squeeze=False)\n for i, row in enumerate(rows):\n for j, value in enumerate(row):\n ax = axes[i,j]\n ax.imshow(value.reshape(8,8), cmap=\"gray\", vmin=0, vmax=1, interpolation=\"nearest\")\n ax.set_xticks([]); ax.set_yticks([])\n if j == 0:\n ax.set_ylabel(labels[i], fontsize=9)\n if i == 0 and titles:\n ax.set_title(titles[j], fontsize=9)\n fig.suptitle(title, fontsize=12)\n fig.tight_layout()\n fig.savefig(target, dpi=160, bbox_inches=\"tight\")\n plt.close(fig)\n\n\ndef json_dump(path, value):\n path.write_text(json.dumps(value, ensure_ascii=False, indent=2, default=lambda x: int(x) if isinstance(x,np.integer) else str(x)))\n\n\n@threadpool_limits.wrap(limits=1, user_api='blas')\ndef run(out, epochs=None):\n out = Path(out); out.mkdir(parents=True, exist_ok=True)\n cfg = copy.deepcopy(CONFIG)\n if epochs is not None: cfg[\"epochs\"] = epochs\n d = load_digits()\n # Fixed binary observations make the product Bernoulli likelihood explicit.\n x = (d.data > 8).astype(np.float64)\n train_id, test_id = train_test_split(np.arange(len(x)), test_size=.2,\n stratify=d.target, random_state=cfg[\"split_seed\"])\n xtrain, xtest = x[train_id], x[test_id]\n cfg.update(train_size=len(train_id), test_size=len(test_id), parameters=0)\n initial = initialize(cfg[\"init_seed\"])\n cfg[\"parameters\"] = sum(a.size for a in initial.values())\n json_dump(out/\"config.json\", cfg)\n checks = gradient_check(initial, xtrain[:5])\n json_dump(out/\"gradient_check.json\", {\"passed\": True, \"checks\": checks,\n \"max_scaled_error\": max(c[\"scaled_error\"] for c in checks)})\n np.savez_compressed(out/\"data.npz\", gray=d.data/16, binary=x, labels=d.target,\n train_indices=train_id, test_indices=test_id)\n (out/\"DATASET_DESCRIPTION.txt\").write_text(d.DESCR)\n test_eps = np.random.default_rng(cfg[\"evaluation_seed\"]).normal(\n size=(cfg[\"evaluation_mc_samples\"], len(xtest), cfg[\"latent_dim\"]))\n visual_rng = np.random.default_rng(cfg[\"visual_seed\"])\n fixed_prior = visual_rng.normal(size=(25,2))\n fixed_uniform = visual_rng.uniform(size=(25,64))\n selected = np.array([np.flatnonzero(d.target[test_id] == digit)[0] for digit in range(10)])\n selected_x = xtest[selected]\n post_eps = visual_rng.normal(size=(4,8,2))\n np.savez_compressed(out/\"fixed_inputs.npz\", prior=fixed_prior, observation_uniform=fixed_uniform,\n selected_global_indices=test_id[selected], posterior_epsilon=post_eps, test_epsilon=test_eps)\n np.savez_compressed(out/\"weights_initial.npz\", **initial)\n models, histories, summaries = {}, {}, {}\n for beta in cfg[\"betas\"]:\n tag = f\"beta{beta:g}\"\n p, log, seconds, steps = train(initial, xtrain, xtest, test_eps, beta, cfg)\n models[tag], histories[tag] = p, log\n np.savez_compressed(out/f\"weights_{tag}.npz\", **p)\n mu, lv = encode(p, xtest)\n summaries[tag] = {\"beta\": beta, \"elapsed_training_seconds\": seconds, \"updates\": steps,\n \"initial_test\": log[0], \"final_test\": log[-1],\n \"posterior_mean_radius_mean\": float(np.linalg.norm(mu,axis=1).mean()),\n \"posterior_std_mean\": float(np.exp(lv/2).mean())}\n print(json.dumps(summaries[tag]), flush=True)\n json_dump(out/\"summary.json\", {\"config\":cfg,\"runs\":summaries,\n \"environment\":{\"python\":platform.python_version(),\"numpy\":np.__version__,\n \"scipy\":scipy.__version__,\"sklearn\":sklearn.__version__,\"matplotlib\":matplotlib.__version__},\n \"gradient_check_max_scaled_error\":max(c[\"scaled_error\"] for c in checks)})\n with (out/\"learning_curves.csv\").open(\"w\", newline=\"\") as f:\n writer = csv.DictWriter(f,fieldnames=[\"beta\",\"epoch\",\"objective\",\"reconstruction_nll\",\"kl\",\"negative_elbo\"])\n writer.writeheader()\n for tag, log in histories.items():\n for row in log: writer.writerow({\"beta\": summaries[tag][\"beta\"],**row})\n # Fixed test input; using z=mu is deterministic reconstruction, not E_q[decoder(z)].\n recon_rows = [d.data[test_id[selected]]/16, selected_x, decode(initial, encode(initial,selected_x)[0])]\n for p in models.values(): recon_rows.append(decode(p, encode(p,selected_x)[0]))\n model_labels = [f\"beta={b:g}\" for b in cfg[\"betas\"]]\n image_rows(recon_rows,[\"original\\ngray\",\"binary\\ninput\",\"initial\\ndecoder(mu)\"]+[label+\"\\ndecoder(mu)\" for label in model_labels],\n out/\"reconstruction.png\",[str(i) for i in range(10)],\"Same held-out digits, before and after learning\")\n # Exactly the same z vectors in all rows; second figure uses same observation uniforms.\n prior_rows = [decode(p,fixed_prior) for p in [initial,*models.values()]]\n image_rows([r[:10] for r in prior_rows],[\"initial\",*model_labels],out/\"prior_same_seed.png\",\n title=\"Same ten z ~ N(0,I): decoder Bernoulli probabilities\")\n observation_tag = \"beta1\" if \"beta1\" in models else list(models)[-1]\n p = models[observation_tag]\n mean = decode(p,fixed_prior)\n image_rows([mean[:10],(fixed_uniform different posterior z -> mean images\")\n ends = encode(p,selected_x[[3,8]])[0]\n alpha = np.linspace(0,1,11)\n interp = (1-alpha[:,None])*ends[0]+alpha[:,None]*ends[1]\n image_rows([decode(p,interp)],[\"linear z\\npath\"],out/f\"interpolation_{tag}.png\",\n [f\"{a:.1f}\" for a in alpha],title=f\"{tag}: interpolation from digit 3 to digit 8 (not prior draws)\")\n grid_axis = np.linspace(-3,3,15)\n gx,gy = np.meshgrid(grid_axis,grid_axis[::-1])\n grid = decode(p,np.column_stack([gx.ravel(),gy.ravel()]))\n mosaic = grid.reshape(15,15,8,8).transpose(0,2,1,3).reshape(120,120)\n fig,axs=plt.subplots(1,2,figsize=(11,5),constrained_layout=True)\n axs[0].scatter(means[:,0],means[:,1],c=d.target[test_id],cmap=\"tab10\",s=13,alpha=.65)\n angle=np.linspace(0,2*np.pi,200)\n axs[0].plot(2*np.cos(angle),2*np.sin(angle),color=\"black\",ls=\"--\",lw=1,label=\"radius 2 reference\")\n axs[0].scatter(fixed_prior[:,0],fixed_prior[:,1],marker=\"x\",s=35,c=\"black\",label=\"25 fixed prior draws\")\n axs[0].set(xlabel=\"z1\",ylabel=\"z2\")\n axs[0].set_title(f\"{tag}: held-out posterior means\",fontsize=11,pad=10)\n axs[0].legend(fontsize=7);axs[0].set_aspect(\"equal\",adjustable=\"box\")\n axs[1].imshow(mosaic,cmap=\"gray\",vmin=0,vmax=1,extent=(-3,3,-3,3),interpolation=\"nearest\")\n axs[1].set(xlabel=\"z1\",ylabel=\"z2\")\n axs[1].set_title(\"Uniform grid in [-3,3]^2 (not prior samples)\",fontsize=11,pad=10)\n fig.savefig(out/f\"latent_{tag}.png\",dpi=160,bbox_inches=\"tight\");plt.close(fig)\n np.savez_compressed(out/f\"observations_{tag}.npz\",test_mu=means,test_logvar=logs,\n selected_reconstruction=decode(p,encode(p,selected_x)[0]),prior_mean=decode(p,fixed_prior),\n prior_observation=(fixed_uniform