現(xiàn)指南：用全批量 Hamiltonian Monte Carlo 探究貝葉斯神經(jīng)網(wǎng)絡(luò)后驗(yàn)的真實(shí)形態(tài))
bnn_hmc 復(fù)現(xiàn)指南用全批量 Hamiltonian Monte Carlo 探究貝葉斯神經(jīng)網(wǎng)絡(luò)后驗(yàn)的真實(shí)形態(tài)【免費(fèi)下載鏈接】google-researchGoogle Research項(xiàng)目地址: https://gitcode.com/gh_mirrors/go/google-research本文以 Google Research 倉庫中的 bnn_hmc/README.md 為骨架結(jié)合其 JAX 源碼完整講解如何復(fù)現(xiàn)論文What Are Bayesian Neural Network Posteriors Really Like?arXiv:2104.14421作者 Pavel Izmailov、Sharad Vikram、Matthew D. Hoffman 與 Andrew Gordon Wilson中的全部實(shí)驗(yàn)。讀完本文你將掌握bnn_hmc項(xiàng)目的環(huán)境搭建、HMC / SGD / Deep Ensembles / SGMCMC / MFVI 五類訓(xùn)練腳本的完整命令行用法、關(guān)鍵超參數(shù)的含義與推薦取值以及后驗(yàn)密度可視化工具的使用方式并能對照源碼理解每一步計(jì)算的底層實(shí)現(xiàn)。一、項(xiàng)目背景論文的核心問題與主要發(fā)現(xiàn)該項(xiàng)目使用**全批量full-batchHamiltonian Monte CarloHMC**作為近乎精確的貝葉斯推理工具去回答貝葉斯深度學(xué)習(xí)中的一系列基礎(chǔ)問題貝葉斯神經(jīng)網(wǎng)絡(luò)BNN的后驗(yàn)到底是什么樣的它能否帶來實(shí)際性能收益常用的近似方法深度集成、SGMCMC、變分推斷與真實(shí)后驗(yàn)有多接近論文在實(shí)驗(yàn)中得到的主要結(jié)論如下也是本項(xiàng)目所有腳本設(shè)計(jì)意圖的集中體現(xiàn)BNN 相比標(biāo)準(zhǔn)訓(xùn)練和深度集成可以獲得顯著的性能提升單條長 HMC 鏈即可提供與多條短鏈相當(dāng)?shù)暮篁?yàn)表示即長鏈采樣效率優(yōu)于多條短鏈并聯(lián)與近期一些研究相反論文發(fā)現(xiàn)后驗(yàn)降溫posterior tempering并非達(dá)到近似最優(yōu)性能所必需幾乎沒有“冷后驗(yàn)cold posterior”效應(yīng)的證據(jù)——論文進(jìn)一步指出該效應(yīng)在很大程度上是數(shù)據(jù)增強(qiáng)data augmentation的產(chǎn)物貝葉斯模型平均BMA的性能對先驗(yàn)尺度prior scale的選擇比較魯棒并且對角高斯先驗(yàn)、高斯混合先驗(yàn)與邏輯斯蒂先驗(yàn)下的結(jié)果較為接近BNN 在分布漂移domain shift下泛化表現(xiàn)令人意外地差深度集成與 SGMCMC 等更廉價(jià)的替代方案雖然能取得不錯(cuò)的泛化但它們給出的預(yù)測分布與 HMC 明顯不同值得注意的是深度集成的預(yù)測分布與 HMC 的接近程度與標(biāo)準(zhǔn) SGLD 相當(dāng)且明顯優(yōu)于標(biāo)準(zhǔn)變分推斷。倉庫以 JAX 實(shí)現(xiàn)了上述全部實(shí)驗(yàn)代碼并附帶 bibtex 引用信息見 bnn_hmc/README.md引用該論文時(shí)可直接使用article{izmailov2021bayesian, title{What Are Bayesian Neural Network Posteriors Really Like?}, author{Izmailov, Pavel and Vikram, Sharad and Hoffman, Matthew D and Wilson, Andrew Gordon}, journal{arXiv preprint arXiv:2104.14421}, year{2021} }二、環(huán)境搭建requirements.txt 與 pip 安裝倉庫提供了 bnn_hmc/requirements.txt其頭部注釋寫明可以直接用來創(chuàng)建 conda 環(huán)境conda create --name env --file requirements.txt從該文件內(nèi)容看這是一個(gè)基于Python 3.8.8、Linux-64平臺(tái)鎖定的環(huán)境清單核心依賴包括jax0.2.12、jaxlib0.1.65cuda112CUDA 11.2 版本、dm-haiku0.0.5.dev0、optax0.0.6、tensorflow2.4.1、tensorflow-datasets4.2.0、numpy1.19.5、scipy1.6.3、tabulate0.8.9、tensorboard2.5.0等?？梢姶a運(yùn)行同時(shí)依賴 JAX 生態(tài)數(shù)值計(jì)算、Haiku網(wǎng)絡(luò)定義與 TensorFlow數(shù)據(jù)加載與日志。如果使用 pip 手動(dòng)安裝README 給出的步驟為pip install tensorflow pip install --upgrade pip pip install --upgrade jax jaxlib0.1.65cuda112 -f \ https://storage.googleapis.com/jax-releases/jax_releases.html pip install githttps://github.com/deepmind/dm-haiku pip install tensorflow_datasets pip install tabulate pip install optax其中 JAX 的安裝方式與硬件強(qiáng)相關(guān)示例中的jaxlib0.1.65cuda112對應(yīng) CUDA 11.2 環(huán)境安裝其他硬件上的 JAX 請以 JAX 官方倉庫的最新安裝說明為準(zhǔn)。三、文件結(jié)構(gòu)總覽倉庫目錄結(jié)構(gòu)在 README 中給出了清晰說明詳見 bnn_hmc/README.md實(shí)際目錄與 README 描述一致可歸納為三層核心算法層bnn_hmc/corehmc.pyHamiltonian Monte Carlo 算法實(shí)現(xiàn)sgmcmc.py以 optax optimizer 形式實(shí)現(xiàn)的 SGMCMC 方法SGLD / SGHMCvi.py均值場變分推斷Mean Field VI實(shí)現(xiàn)。工具層bnn_hmc/utils訓(xùn)練腳本共用的功能模塊包括train_utils.py訓(xùn)練輪次與更新規(guī)則、models.py實(shí)驗(yàn)用網(wǎng)絡(luò)結(jié)構(gòu)、losses.py先驗(yàn)與似然函數(shù)、data_utils.py數(shù)據(jù)加載與預(yù)處理、optim_utils.py優(yōu)化器與學(xué)習(xí)率調(diào)度、ensemble_utils.py預(yù)測集成、metrics.py評(píng)估指標(biāo)、cmd_args_utils.py公共命令行參數(shù)、script_utils.py訓(xùn)練腳本公共功能、checkpoint_utils.py檢查點(diǎn)保存/加載、logging_utils.py日志打印、precision_utils.py數(shù)值精度控制、tree_utils.pypytree 通用操作。訓(xùn)練腳本層bnn_hmc 根目錄run_hmc.pyHMC 訓(xùn)練腳本run_sgd.pySGD 訓(xùn)練腳本深度集成即多次以不同種子運(yùn)行它run_sgmcmc.pySGMCMC 訓(xùn)練腳本run_vi.pyMFVI 訓(xùn)練腳本make_posterior_surface_plot.py后驗(yàn)密度可視化腳本。此外 bnn_hmc/notebooks 目錄還提供了若干交互式 Notebook如synthetic_regression_inference.ipynb、mcmc_gaussian_test.ipynb可用于小規(guī)模驗(yàn)證 MCMC 推斷流程。四、公共命令行參數(shù)所有訓(xùn)練腳本共用的基礎(chǔ)配置所有訓(xùn)練腳本run_hmc.py、run_sgd.py、run_sgmcmc.py、run_vi.py共享由 utils/cmd_args_utils.py 中add_common_flags()定義的一組參數(shù)參數(shù)說明seed隨機(jī)種子源碼默認(rèn)0dir訓(xùn)練目錄用于保存檢查點(diǎn)與 TensorBoard 日志必填dataset_name數(shù)據(jù)集名稱例如cifar10、cifar100、imdbUCI 數(shù)據(jù)集寫作UCI 數(shù)據(jù)集名_隨機(jī)種子的形式例如yacht_2其中種子決定訓(xùn)練/測試劃分subset_train_to從數(shù)據(jù)集使用的樣本數(shù)默認(rèn)使用全部數(shù)據(jù)model_name神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)名稱例如lenet、resnet20_frn_swish、cnn_lstm、mlp_regression_smallweight_decay權(quán)重衰減對貝葉斯方法權(quán)重衰減決定了先驗(yàn)方差prior_var 1 / weight_decaytemperature后驗(yàn)溫度默認(rèn)1init_checkpoint用于初始化鏈/模型的檢查點(diǎn)路徑可選tabulate_freq表格表頭打印頻率use_float64使用 float64 精度在 TPU 及部分 GPU 上不可用默認(rèn)使用float32另外源碼中還包含--tpu_ip參數(shù)默認(rèn)None用于指定 Cloud TPU 的內(nèi)部 IP配合 utils/train_utils.py 中set_up_jax()的jax_backend_target配置可將計(jì)算運(yùn)行在 TPU 上use_float64則通過jax_enable_x64開關(guān)開啟。需要特別說明的是weight_decay與先驗(yàn)的關(guān)系在 utils/losses.py 的make_gaussian_log_prior()中高斯先驗(yàn)的對數(shù)密度為-0.5 * weight_decay * ||params||2 常數(shù)項(xiàng)再除以溫度temperature因此weight_decay越大意味著先驗(yàn)方差越小、先驗(yàn)越強(qiáng)。這也是論文中討論“先驗(yàn)尺度魯棒性”實(shí)驗(yàn)的操作入口。五、運(yùn)行 HMCrun_hmc.py5.1 HMC 專屬參數(shù)除公共參數(shù)外HMC 訓(xùn)練腳本 run_hmc.py 還定義以下參數(shù)參數(shù)說明源碼默認(rèn)值step_sizeHMC 步長1e-4trajectory_lenHMC 軌跡長度1e-3num_iterationsHMC 迭代總次數(shù)1000max_num_leapfrog_steps允許的最大 leapfrog 步數(shù)作為 sanity check應(yīng)大于trajectory_len / step_size10000num_burn_in_iterationsburn-in 迭代次數(shù)0--no_mh若設(shè)置忽略 Metropolis–Hastings 校正False其中 leapfrog 步數(shù)在訓(xùn)練循環(huán)中由ceil(trajectory_len / step_size)計(jì)算得出見 utils/train_utils.py 的make_hmc_update()若超過max_num_leapfrog_steps會(huì)直接斷言報(bào)錯(cuò)。5.2 示例命令I(lǐng)MDB 上的 CNN-LSTM論文在 IMDB 情感分類任務(wù)上使用 CNN-LSTM 模型cnn_lstm并以不同溫度運(yùn)行 HMC 驗(yàn)證“冷后驗(yàn)”問題。README 給出的四組命令完整復(fù)刻如下運(yùn)行環(huán)境為 8 張 NVIDIA Tesla V-100 GPU# Temperature 1 python3 run_hmc.py --seed1 --weight_decay40. --temperature1. \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size1e-5 --trajectory_len0.24 \ --max_num_leapfrog_steps30000 # Temperature 0.3 python3 run_hmc.py --seed1 --weight_decay40. --temperature0.3 \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size3e-6 --trajectory_len0.136 \ --max_num_leapfrog_steps46000 # Temperature 0.1 python3 run_hmc.py --seed1 --weight_decay40. --temperature0.1 \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size1e-6 --trajectory_len0.078 \ --max_num_leapfrog_steps90000 # Temperature 0.03 python3 run_hmc.py --seed1 --weight_decay40. --temperature0.03 \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size1e-6 --trajectory_len0.043 \ --max_num_leapfrog_steps45000可以觀察到溫度越低所需步長越小、leapfrog 步數(shù)越多——這是因?yàn)榻禍厥沟煤篁?yàn)更“尖”需要更精細(xì)的積分步長。IMDB 任務(wù)使用--use_float64以獲得數(shù)值穩(wěn)定的接受率計(jì)算。5.3 示例命令MNIST 160 樣本子集上的 MLP論文中還有一個(gè)輕量示例——在 MNIST 隨機(jī)取 160 個(gè)樣本的子集上訓(xùn)練 MLP 分類器mlp_classification單張 GPU 即可運(yùn)行適合快速驗(yàn)證整個(gè) HMC 流程python3 run_hmc.py --seed0 --weight_decay1. --temperature1. \ --dirruns/hmc/mnist_subset160 --dataset_namemnist \ --model_namemlp_classification --step_size3.e-5 --trajectory_len1.5 \ --num_iterations100 --max_num_leapfrog_steps50000 \ --num_burn_in_iterations10 --subset_train_to160注意這里通過--subset_train_to160限制訓(xùn)練樣本數(shù)并設(shè)置了 10 次 burn-in 迭代。5.4 HMC 的源碼級(jí)原理HMC 更新邏輯在 core/hmc.py 中實(shí)現(xiàn)核心流程清晰可循leapfrog 積分make_leapfrog對動(dòng)量做半步更新 → 更新參數(shù) → 重新計(jì)算梯度 → 再對動(dòng)量做半步更新整個(gè)過程用jax.lax.fori_loop串起n_leapfrog步接受概率make_accept_prob基于動(dòng)能差、對數(shù)似然差與對數(shù)先驗(yàn)差計(jì)算能量差得到min(1, exp(energy_diff))代碼注釋特別指出將似然與先驗(yàn)分開返回是為了在 float32 下更精確地計(jì)算接受/拒絕步驟中的先驗(yàn)密度比自適應(yīng)步長adapt_step_size按step_size * exp(speed * (accept_prob - target_accept_rate))調(diào)整步長使接受率逼近目標(biāo)值MH 校正make_adaptive_hmc_updateburn-in 之外默認(rèn)開啟 Metropolis–Hastings 校正--no_mh可關(guān)閉。在 run_hmc.py 的訓(xùn)練主循環(huán)中每輪迭代都會(huì)評(píng)估 train/test 指標(biāo)、將被接受的樣本或--no_mh時(shí)的所有樣本累加進(jìn)ensemble_predictions即 BMA 集成、保存檢查點(diǎn)并把a(bǔ)ccept_prob、num_ensembled、當(dāng)前step_size等寫入 TensorBoard。訓(xùn)練目錄名由模型、權(quán)重衰減、步長、軌跡長度、burn-in、MH 開關(guān)、溫度與種子自動(dòng)拼接而成便于對比實(shí)驗(yàn)歸檔。需要說明的是README 明確指出論文中 CIFAR-10 上的 HMC 實(shí)驗(yàn)是在包含 512 個(gè) TPU 設(shè)備的 TPU pod 上、使用后續(xù)將發(fā)布的修改版代碼運(yùn)行的——即本倉庫開箱即用的代碼主要面向 GPU 場景CIFAR 級(jí)全量 HMC 需要特殊規(guī)模的硬件。六、運(yùn)行 SGD 與深度集成run_sgd.pySGD 腳本 run_sgd.py 使用cosine 學(xué)習(xí)率調(diào)度其專屬參數(shù)如下由 utils/cmd_args_utils.py 的add_sgd_flags()定義參數(shù)說明源碼默認(rèn)值init_step_size初始 SGD 步長1e-6num_epochsSGD 總 epoch 數(shù)300batch_size批大小80eval_freq評(píng)估頻率epoch10save_freq檢查點(diǎn)保存頻率epoch50momentum_decaySGD 動(dòng)量衰減0.9README 給出的示例ResNet-20-FRNresnet20_frn_swish在 CIFAR-10 / CIFAR-100 上訓(xùn)練 500 個(gè) epoch并使用--subset_train_to40960限制樣本數(shù)python3 run_sgd.py --seed1 --weight_decay10 --dirruns/sgd/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish \ --init_step_size3e-7 --num_epochs500 --eval_freq10 --batch_size80 \ --save_freq500 --subset_train_to40960 python3 run_sgd.py --seed1 --weight_decay10 --dirruns/sgd/cifar100/ \ --dataset_namecifar100 --model_nameresnet20_frn_swish \ --init_step_size1e-6 --num_epochs500 --eval_freq10 --batch_size80 \ --save_freq500 --subset_train_to40960IMDB 上的 CNN-LSTMpython3 run_sgd.py --seed1 --weight_decay3. --dirruns/sgd/imdb/ \ --dataset_nameimdb --model_namecnn_lstm --init_step_size3e-7 \ --num_epochs500 --eval_freq10 --batch_size80 --save_freq500深度集成的訓(xùn)練方式README 明確說明只需用不同隨機(jī)種子多次運(yùn)行run_sgd.py將得到的多個(gè) SGD 解進(jìn)行預(yù)測集成即可。論文比較了深度集成與 HMC 的預(yù)測分布差異得出“深度集成與 HMC 的接近程度和標(biāo)準(zhǔn) SGLD 相當(dāng)、且優(yōu)于標(biāo)準(zhǔn)變分推斷”的結(jié)論。從實(shí)現(xiàn)上看SGD 訓(xùn)練過程位于 utils/train_utils.py 的make_sgd_train_epoch()數(shù)據(jù)被切分為num_batches份每個(gè) epoch 內(nèi)用jax.lax.scan串行執(zhí)行各 batch 的梯度更新并由jax.lax.psum完成跨設(shè)備梯度聚合體現(xiàn) JAX 多設(shè)備如多 GPU并行訓(xùn)練能力。七、運(yùn)行 SGMCMCrun_sgmcmc.pySGMCMC 腳本 run_sgmcmc.py 共享 SGD 的全部參數(shù)并額外引入以下參數(shù)參數(shù)說明源碼默認(rèn)值preconditioner預(yù)處理器選擇None或RMSpropNonestep_size_schedule步長調(diào)度constant或cyclical。constant表示先做num_burnin_epochs個(gè) epoch 的 cosine burn-in隨后步長固定為final_step_sizecyclical表示先做固定步長的 burn-in再進(jìn)入 cosine 循環(huán)調(diào)度constantnum_burnin_epochs達(dá)到最終學(xué)習(xí)率前的 epoch 數(shù)300final_step_size最終步長僅用于constant調(diào)度默認(rèn)取init_step_sizestep_size_cycle_length_epochs循環(huán)周期長度epoch僅用于cyclical調(diào)度50save_all_ensembled保存所有被集成的網(wǎng)絡(luò)Falseensemble_freq集成 iterates 的頻率epoch10CIFAR-10 上 ResNet-20-FRN 的四種 SGMCMC 變體示例# SGLD python3 run_sgmcmc.py --seed1 --weight_decay5. --dirruns/sgmcmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size1e-6 \ --final_step_size1e-6 --num_epochs10000 --num_burnin_epochs1000 \ --eval_freq10 --batch_size80 --save_freq10 --momentum0. \ --subset_train_to40960 # SGHMC python3 run_sgmcmc.py --seed1 --weight_decay5 --dirruns/sgmcmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size3e-7 \ --final_step_size3e-7 --num_epochs10000 --num_burnin_epochs1000 \ --eval_freq10 --batch_size80 --save_freq10 --subset_train_to40960 \ --momentum0.9 # SGHMC-CLR python3 run_sgmcmc.py --seed1 --weight_decay5 --dirruns/sgmcmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size3e-7 \ --num_epochs10000 --num_burnin_epochs1000 --step_size_schedulecyclical \ --step_size_cycle_length_epochs50 --ensemble_freq50 --eval_freq10 \ --batch_size80 --save_freq1000 --subset_train_to40960 \ --preconditionerNone --momentum0.95 --eval_freq10 --save_all_ensembled # SGHMC-CLR-Prec python3 run_sgmcmc.py --seed1 --weight_decay5 --dirruns/sghmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size3e-5 \ --num_epochs10000 --num_burnin_epochs1000 --step_size_schedulecyclical \ --step_size_cycle_length_epochs50 --ensemble_freq50 --eval_freq10 \ --batch_size80 --save_freq50 --subset_train_to40960 \ --preconditionerRMSprop --momentum0.95 --eval_freq10 --save_all_ensembled四個(gè)變體對應(yīng)論文中比較的 SGMCMC 方法族標(biāo)準(zhǔn) SGLDmomentum0、SGHMC動(dòng)量0.9、SGHMC 循環(huán)學(xué)習(xí)率CLRstep_size_schedulecyclical且momentum0.95以及在此基礎(chǔ)上疊加 RMSprop 預(yù)處理preconditionerRMSprop。注意 SGLD 與 SGHMC 使用了不同的初始/最終步長1e-6與3e-7這與各自方法的噪聲特性有關(guān)。SGMCMC 的實(shí)現(xiàn)位于 core/sgmcmc.py以 optaxGradientTransformation形式提供sgld_gradient_update()實(shí)現(xiàn) SGLD當(dāng)momentum_decay0時(shí)為經(jīng)典 SGLDWelling Teh, ICML 2011否則退化為欠阻尼 SGLD即 SGHMCChen et al., ICML 2014。更新規(guī)則中梯度項(xiàng)乘lr_sqrt高斯噪聲項(xiàng)的標(biāo)準(zhǔn)差為sqrt(2 * (1 - momentum_decay))再疊加動(dòng)量項(xiàng)momentum_decay * m預(yù)處理通過Preconditioner抽象實(shí)現(xiàn)None對應(yīng)IdentityPreconditionerState恒等變換RMSprop對應(yīng)get_rmsprop_preconditioner()其基于梯度二階矩估計(jì)默認(rèn)衰減因子0.99、eps1e-7對噪聲與更新量分別做M^{-1/2}與M^{-1}變換即論文中 SGHMC-CLR-Prec 的“Prec”來源。步長調(diào)度在 run_sgmcmc.py 的get_lr_schedule()中構(gòu)建constant調(diào)度調(diào)用make_constant_lr_schedule_with_cosine_burnin(init_step_size, final_step_size, burnin_steps)cyclical調(diào)度調(diào)用make_cyclical_cosine_lr_schedule_with_const_burnin(...)周期由step_size_cycle_length_epochs換算為步數(shù)。訓(xùn)練過程中ensemble_freq決定每隔多少個(gè) epoch 將當(dāng)前 iterate 的預(yù)測累加進(jìn)集成save_all_ensembled則同時(shí)把被集成的網(wǎng)絡(luò)參數(shù)落盤。八、運(yùn)行 MFVIrun_vi.py均值場變分推斷腳本 run_vi.py 同樣共享 SGD 參數(shù)并額外定義參數(shù)說明源碼默認(rèn)值optimizer優(yōu)化器選擇SGD或Adam源碼默認(rèn)AdamREADME 描述為默認(rèn) SGD請以倉庫代碼為準(zhǔn)vi_sigma_initMFVI 中權(quán)重標(biāo)準(zhǔn)差σ的初始值1e-3vi_ensemble_sizeVI 評(píng)估時(shí)采樣的集成規(guī)模20mean_init_checkpoint用于初始化 MFVI 均值mean的 SGD 檢查點(diǎn)None示例命令# ResNet-20-FRN on CIFAR-10 or CIFAR-100 python3 run_vi.py --seed11 --weight_decay5. --dirruns/vi/cifar100/ \ --dataset_name[cifar10 | cifar100] --model_nameresnet20_frn_swish \ --init_step_size1e-4 --num_epochs300 --eval_freq10 --batch_size80 \ --save_freq300 --subset_train_to40960 --optimizerAdam \ --vi_sigma_init0.01 --temperature1. --vi_ensemble_size20 \ --mean_init_checkpointpath-to-sgd-solution # CNN-LSTM on IMDB python3 run_vi.py --seed11 --weight_decay5. --dirruns/vi/imdb/ \ --dataset_nameimdb --model_namecnn_lstm --init_step_size1e-4 \ --num_epochs500 --eval_freq10 --batch_size80 --save_freq200 \ --optimizerAdam --vi_sigma_init0.01 --temperature1. --vi_ensemble_size20 \ --mean_init_checkpointpath-to-sgd-solution其中--mean_init_checkpoint通常指向前面用run_sgd.py訓(xùn)練出的解即論文中的“用 SGD 解初始化 MFVI 均值”策略。MFVI 的實(shí)現(xiàn)在 core/vi.pyget_mfvi_model_fn()將每個(gè)權(quán)重參數(shù)化為{mean, inv_softplus_std}兩個(gè)量均值直接用原參數(shù)復(fù)制標(biāo)準(zhǔn)差則存其 inverse-softplus 形式sigma softplus(inv_softplus_std)保證標(biāo)準(zhǔn)差恒為正sample_parms_fn用重參數(shù)化技巧m n * sn 為標(biāo)準(zhǔn)正態(tài)噪聲從變分分布采樣make_kl_with_gaussian_prior(weight_decay, temperature)實(shí)現(xiàn) ELBO 中的先驗(yàn) KL 項(xiàng)對每個(gè)參數(shù)計(jì)算log(σ_prior/σ_vi) (σ_vi2 μ_vi2)/(2σ_prior2) - 1/2其中σ_prior sqrt(1/weight_decay)并以temperature作為 KL 項(xiàng)權(quán)重評(píng)估時(shí)vi_ensemble_predict_fn從變分后驗(yàn)采樣vi_ensemble_size次并集成預(yù)測同時(shí)run_vi.py還會(huì)將每個(gè)參數(shù)的 softplus 標(biāo)準(zhǔn)差以直方圖形式寫入 TensorBoardMFVI/param_stds便于觀察各層不確定性的分布。九、可視化后驗(yàn)密度make_posterior_surface_plot.py腳本 make_posterior_surface_plot.pyREADME 中筆誤寫作makemake_posterior_surface_plot.py實(shí)際文件名以此處為準(zhǔn)可以在由三個(gè)檢查點(diǎn)張成的二維平面上可視化后驗(yàn)的 log 密度、log 似然與 log 先驗(yàn)復(fù)現(xiàn)論文中的后驗(yàn)曲面圖。參數(shù)如下參數(shù)說明源碼默認(rèn)值limit_bottom曲面可視化在垂直方向底部的范圍-0.25limit_top曲面可視化在垂直方向頂部的范圍1.25limit_left曲面可視化在水平方向左側(cè)的范圍-0.25limit_right曲面可視化在水平方向右側(cè)的范圍1.25grid_size每個(gè)方向上的網(wǎng)格點(diǎn)數(shù)20checkpoint1第一個(gè)檢查點(diǎn)路徑—checkpoint2第二個(gè)檢查點(diǎn)路徑—checkpoint3第三個(gè)檢查點(diǎn)路徑—示例IMDB 上的 CNN-LSTMpython3 make_posterior_surface_plot.py --weight_decay40 --temperature1. \ --dirruns/surface_plots/imdb/ --model_namecnn_lstm --dataset_nameimdb \ --checkpoint1ckpt1 --checkpoint2ckpt2 --checkpoint3ckpt3 --limit_bottom-0.75 --limit_left-0.75 --limit_right1.75 --limit_top1.75 \ --grid_size50從源碼看該腳本以三個(gè)檢查點(diǎn)構(gòu)造平面坐標(biāo)系每點(diǎn)對應(yīng)平面上的一個(gè)二維坐標(biāo)在grid_size × grid_size網(wǎng)格上逐點(diǎn)計(jì)算 log 后驗(yàn)密度、log 似然與 log 先驗(yàn)依賴 utils/losses.py 與 utils/models.py 提供的前向計(jì)算并用 matplotlib 輸出可視化結(jié)果是理解“BNN 后驗(yàn)為何不像高斯”“冷后驗(yàn)現(xiàn)象由何而來”等論文論點(diǎn)的直觀工具。十、小結(jié)如何組織一次完整的復(fù)現(xiàn)實(shí)驗(yàn)綜合以上內(nèi)容一次完整的論文實(shí)驗(yàn)復(fù)現(xiàn)可按如下流程組織按 第二節(jié) 搭建 conda/pip 環(huán)境用run_sgd.py在不同種子下訓(xùn)練若干 SGD 解得到基線精度與深度集成用run_hmc.py在對應(yīng)任務(wù)上運(yùn)行全批量 HMC建議先用 5.3 節(jié) 的 MNIST 160 樣本示例在小規(guī)模上驗(yàn)證流程得到近乎精確的后驗(yàn)與 BMA 集成結(jié)果用run_sgmcmc.py分別運(yùn)行 SGLD / SGHMC / SGHMC-CLR / SGHMC-CLR-Prec與 HMC 對比預(yù)測分布的接近程度用run_vi.py以 SGD 解初始化均值運(yùn)行 MFVI作為變分推斷基線用make_posterior_surface_plot.py對若干代表性檢查點(diǎn)生成后驗(yàn)曲面圖直觀檢驗(yàn)論文關(guān)于先驗(yàn)尺度、溫度與后驗(yàn)形態(tài)的結(jié)論。所有腳本都將檢查點(diǎn)與 TensorBoard 日志寫入--dir指定目錄可通過 TensorBoard 對比各方法的 accuracy、NLL、ECE分類任務(wù)或 RMSE / NLL回歸任務(wù)等指標(biāo)以及test/ens_*前綴的集成指標(biāo)。值得留意的是論文涉及溫度對比實(shí)驗(yàn)時(shí)如 HMC 的四個(gè)溫度檔位需保持weight_decay一致而僅改變temperature并在解讀結(jié)果時(shí)記住該倉庫代碼默認(rèn)不使用數(shù)據(jù)增強(qiáng)而論文指出的“冷后驗(yàn)效應(yīng)很大程度上是數(shù)據(jù)增強(qiáng)的產(chǎn)物”這一結(jié)論正是建立在這樣的實(shí)驗(yàn)設(shè)置之上的?！久赓M(fèi)下載鏈接】google-researchGoogle Research項(xiàng)目地址: https://gitcode.com/gh_mirrors/go/google-research創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 bnn_hmc 復(fù)現(xiàn)指南：用全批量 Hamiltonian Monte Carlo 探究貝葉斯神經(jīng)網(wǎng)絡(luò)后驗(yàn)的真實(shí)形態(tài)
bnn_hmc 復(fù)現(xiàn)指南用全批量 Hamiltonian Monte Carlo 探究貝葉斯神經(jīng)網(wǎng)絡(luò)后驗(yàn)的真實(shí)形態(tài)【免費(fèi)下載鏈接】google-researchGoogle Research項(xiàng)目地址: https://gitcode.com/gh_mirrors/go/google-research本文以 Google Research 倉庫中的 bnn_hmc/README.md 為骨架結(jié)合其 JAX 源碼完整講解如何復(fù)現(xiàn)論文What Are Bayesian Neural Network Posteriors Really Like?arXiv:2104.14421作者 Pavel Izmailov、Sharad Vikram、Matthew D. Hoffman 與 Andrew Gordon Wilson中的全部實(shí)驗(yàn)。讀完本文你將掌握bnn_hmc項(xiàng)目的環(huán)境搭建、HMC / SGD / Deep Ensembles / SGMCMC / MFVI 五類訓(xùn)練腳本的完整命令行用法、關(guān)鍵超參數(shù)的含義與推薦取值以及后驗(yàn)密度可視化工具的使用方式并能對照源碼理解每一步計(jì)算的底層實(shí)現(xiàn)。一、項(xiàng)目背景論文的核心問題與主要發(fā)現(xiàn)該項(xiàng)目使用**全批量full-batchHamiltonian Monte CarloHMC**作為近乎精確的貝葉斯推理工具去回答貝葉斯深度學(xué)習(xí)中的一系列基礎(chǔ)問題貝葉斯神經(jīng)網(wǎng)絡(luò)BNN的后驗(yàn)到底是什么樣的它能否帶來實(shí)際性能收益常用的近似方法深度集成、SGMCMC、變分推斷與真實(shí)后驗(yàn)有多接近論文在實(shí)驗(yàn)中得到的主要結(jié)論如下也是本項(xiàng)目所有腳本設(shè)計(jì)意圖的集中體現(xiàn)BNN 相比標(biāo)準(zhǔn)訓(xùn)練和深度集成可以獲得顯著的性能提升單條長 HMC 鏈即可提供與多條短鏈相當(dāng)?shù)暮篁?yàn)表示即長鏈采樣效率優(yōu)于多條短鏈并聯(lián)與近期一些研究相反論文發(fā)現(xiàn)后驗(yàn)降溫posterior tempering并非達(dá)到近似最優(yōu)性能所必需幾乎沒有“冷后驗(yàn)cold posterior”效應(yīng)的證據(jù)——論文進(jìn)一步指出該效應(yīng)在很大程度上是數(shù)據(jù)增強(qiáng)data augmentation的產(chǎn)物貝葉斯模型平均BMA的性能對先驗(yàn)尺度prior scale的選擇比較魯棒并且對角高斯先驗(yàn)、高斯混合先驗(yàn)與邏輯斯蒂先驗(yàn)下的結(jié)果較為接近BNN 在分布漂移domain shift下泛化表現(xiàn)令人意外地差深度集成與 SGMCMC 等更廉價(jià)的替代方案雖然能取得不錯(cuò)的泛化但它們給出的預(yù)測分布與 HMC 明顯不同值得注意的是深度集成的預(yù)測分布與 HMC 的接近程度與標(biāo)準(zhǔn) SGLD 相當(dāng)且明顯優(yōu)于標(biāo)準(zhǔn)變分推斷。倉庫以 JAX 實(shí)現(xiàn)了上述全部實(shí)驗(yàn)代碼并附帶 bibtex 引用信息見 bnn_hmc/README.md引用該論文時(shí)可直接使用article{izmailov2021bayesian, title{What Are Bayesian Neural Network Posteriors Really Like?}, author{Izmailov, Pavel and Vikram, Sharad and Hoffman, Matthew D and Wilson, Andrew Gordon}, journal{arXiv preprint arXiv:2104.14421}, year{2021} }二、環(huán)境搭建requirements.txt 與 pip 安裝倉庫提供了 bnn_hmc/requirements.txt其頭部注釋寫明可以直接用來創(chuàng)建 conda 環(huán)境conda create --name env --file requirements.txt從該文件內(nèi)容看這是一個(gè)基于Python 3.8.8、Linux-64平臺(tái)鎖定的環(huán)境清單核心依賴包括jax0.2.12、jaxlib0.1.65cuda112CUDA 11.2 版本、dm-haiku0.0.5.dev0、optax0.0.6、tensorflow2.4.1、tensorflow-datasets4.2.0、numpy1.19.5、scipy1.6.3、tabulate0.8.9、tensorboard2.5.0等?？梢姶a運(yùn)行同時(shí)依賴 JAX 生態(tài)數(shù)值計(jì)算、Haiku網(wǎng)絡(luò)定義與 TensorFlow數(shù)據(jù)加載與日志。如果使用 pip 手動(dòng)安裝README 給出的步驟為pip install tensorflow pip install --upgrade pip pip install --upgrade jax jaxlib0.1.65cuda112 -f \ https://storage.googleapis.com/jax-releases/jax_releases.html pip install githttps://github.com/deepmind/dm-haiku pip install tensorflow_datasets pip install tabulate pip install optax其中 JAX 的安裝方式與硬件強(qiáng)相關(guān)示例中的jaxlib0.1.65cuda112對應(yīng) CUDA 11.2 環(huán)境安裝其他硬件上的 JAX 請以 JAX 官方倉庫的最新安裝說明為準(zhǔn)。三、文件結(jié)構(gòu)總覽倉庫目錄結(jié)構(gòu)在 README 中給出了清晰說明詳見 bnn_hmc/README.md實(shí)際目錄與 README 描述一致可歸納為三層核心算法層bnn_hmc/corehmc.pyHamiltonian Monte Carlo 算法實(shí)現(xiàn)sgmcmc.py以 optax optimizer 形式實(shí)現(xiàn)的 SGMCMC 方法SGLD / SGHMCvi.py均值場變分推斷Mean Field VI實(shí)現(xiàn)。工具層bnn_hmc/utils訓(xùn)練腳本共用的功能模塊包括train_utils.py訓(xùn)練輪次與更新規(guī)則、models.py實(shí)驗(yàn)用網(wǎng)絡(luò)結(jié)構(gòu)、losses.py先驗(yàn)與似然函數(shù)、data_utils.py數(shù)據(jù)加載與預(yù)處理、optim_utils.py優(yōu)化器與學(xué)習(xí)率調(diào)度、ensemble_utils.py預(yù)測集成、metrics.py評(píng)估指標(biāo)、cmd_args_utils.py公共命令行參數(shù)、script_utils.py訓(xùn)練腳本公共功能、checkpoint_utils.py檢查點(diǎn)保存/加載、logging_utils.py日志打印、precision_utils.py數(shù)值精度控制、tree_utils.pypytree 通用操作。訓(xùn)練腳本層bnn_hmc 根目錄run_hmc.pyHMC 訓(xùn)練腳本run_sgd.pySGD 訓(xùn)練腳本深度集成即多次以不同種子運(yùn)行它run_sgmcmc.pySGMCMC 訓(xùn)練腳本run_vi.pyMFVI 訓(xùn)練腳本make_posterior_surface_plot.py后驗(yàn)密度可視化腳本。此外 bnn_hmc/notebooks 目錄還提供了若干交互式 Notebook如synthetic_regression_inference.ipynb、mcmc_gaussian_test.ipynb可用于小規(guī)模驗(yàn)證 MCMC 推斷流程。四、公共命令行參數(shù)所有訓(xùn)練腳本共用的基礎(chǔ)配置所有訓(xùn)練腳本run_hmc.py、run_sgd.py、run_sgmcmc.py、run_vi.py共享由 utils/cmd_args_utils.py 中add_common_flags()定義的一組參數(shù)參數(shù)說明seed隨機(jī)種子源碼默認(rèn)0dir訓(xùn)練目錄用于保存檢查點(diǎn)與 TensorBoard 日志必填dataset_name數(shù)據(jù)集名稱例如cifar10、cifar100、imdbUCI 數(shù)據(jù)集寫作UCI 數(shù)據(jù)集名_隨機(jī)種子的形式例如yacht_2其中種子決定訓(xùn)練/測試劃分subset_train_to從數(shù)據(jù)集使用的樣本數(shù)默認(rèn)使用全部數(shù)據(jù)model_name神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)名稱例如lenet、resnet20_frn_swish、cnn_lstm、mlp_regression_smallweight_decay權(quán)重衰減對貝葉斯方法權(quán)重衰減決定了先驗(yàn)方差prior_var 1 / weight_decaytemperature后驗(yàn)溫度默認(rèn)1init_checkpoint用于初始化鏈/模型的檢查點(diǎn)路徑可選tabulate_freq表格表頭打印頻率use_float64使用 float64 精度在 TPU 及部分 GPU 上不可用默認(rèn)使用float32另外源碼中還包含--tpu_ip參數(shù)默認(rèn)None用于指定 Cloud TPU 的內(nèi)部 IP配合 utils/train_utils.py 中set_up_jax()的jax_backend_target配置可將計(jì)算運(yùn)行在 TPU 上use_float64則通過jax_enable_x64開關(guān)開啟。需要特別說明的是weight_decay與先驗(yàn)的關(guān)系在 utils/losses.py 的make_gaussian_log_prior()中高斯先驗(yàn)的對數(shù)密度為-0.5 * weight_decay * ||params||2 常數(shù)項(xiàng)再除以溫度temperature因此weight_decay越大意味著先驗(yàn)方差越小、先驗(yàn)越強(qiáng)。這也是論文中討論“先驗(yàn)尺度魯棒性”實(shí)驗(yàn)的操作入口。五、運(yùn)行 HMCrun_hmc.py5.1 HMC 專屬參數(shù)除公共參數(shù)外HMC 訓(xùn)練腳本 run_hmc.py 還定義以下參數(shù)參數(shù)說明源碼默認(rèn)值step_sizeHMC 步長1e-4trajectory_lenHMC 軌跡長度1e-3num_iterationsHMC 迭代總次數(shù)1000max_num_leapfrog_steps允許的最大 leapfrog 步數(shù)作為 sanity check應(yīng)大于trajectory_len / step_size10000num_burn_in_iterationsburn-in 迭代次數(shù)0--no_mh若設(shè)置忽略 Metropolis–Hastings 校正False其中 leapfrog 步數(shù)在訓(xùn)練循環(huán)中由ceil(trajectory_len / step_size)計(jì)算得出見 utils/train_utils.py 的make_hmc_update()若超過max_num_leapfrog_steps會(huì)直接斷言報(bào)錯(cuò)。5.2 示例命令I(lǐng)MDB 上的 CNN-LSTM論文在 IMDB 情感分類任務(wù)上使用 CNN-LSTM 模型cnn_lstm并以不同溫度運(yùn)行 HMC 驗(yàn)證“冷后驗(yàn)”問題。README 給出的四組命令完整復(fù)刻如下運(yùn)行環(huán)境為 8 張 NVIDIA Tesla V-100 GPU# Temperature 1 python3 run_hmc.py --seed1 --weight_decay40. --temperature1. \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size1e-5 --trajectory_len0.24 \ --max_num_leapfrog_steps30000 # Temperature 0.3 python3 run_hmc.py --seed1 --weight_decay40. --temperature0.3 \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size3e-6 --trajectory_len0.136 \ --max_num_leapfrog_steps46000 # Temperature 0.1 python3 run_hmc.py --seed1 --weight_decay40. --temperature0.1 \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size1e-6 --trajectory_len0.078 \ --max_num_leapfrog_steps90000 # Temperature 0.03 python3 run_hmc.py --seed1 --weight_decay40. --temperature0.03 \ --dirruns/hmc/imdb/ --dataset_nameimdb --model_namecnn_lstm \ --use_float64 --step_size1e-6 --trajectory_len0.043 \ --max_num_leapfrog_steps45000可以觀察到溫度越低所需步長越小、leapfrog 步數(shù)越多——這是因?yàn)榻禍厥沟煤篁?yàn)更“尖”需要更精細(xì)的積分步長。IMDB 任務(wù)使用--use_float64以獲得數(shù)值穩(wěn)定的接受率計(jì)算。5.3 示例命令MNIST 160 樣本子集上的 MLP論文中還有一個(gè)輕量示例——在 MNIST 隨機(jī)取 160 個(gè)樣本的子集上訓(xùn)練 MLP 分類器mlp_classification單張 GPU 即可運(yùn)行適合快速驗(yàn)證整個(gè) HMC 流程python3 run_hmc.py --seed0 --weight_decay1. --temperature1. \ --dirruns/hmc/mnist_subset160 --dataset_namemnist \ --model_namemlp_classification --step_size3.e-5 --trajectory_len1.5 \ --num_iterations100 --max_num_leapfrog_steps50000 \ --num_burn_in_iterations10 --subset_train_to160注意這里通過--subset_train_to160限制訓(xùn)練樣本數(shù)并設(shè)置了 10 次 burn-in 迭代。5.4 HMC 的源碼級(jí)原理HMC 更新邏輯在 core/hmc.py 中實(shí)現(xiàn)核心流程清晰可循leapfrog 積分make_leapfrog對動(dòng)量做半步更新 → 更新參數(shù) → 重新計(jì)算梯度 → 再對動(dòng)量做半步更新整個(gè)過程用jax.lax.fori_loop串起n_leapfrog步接受概率make_accept_prob基于動(dòng)能差、對數(shù)似然差與對數(shù)先驗(yàn)差計(jì)算能量差得到min(1, exp(energy_diff))代碼注釋特別指出將似然與先驗(yàn)分開返回是為了在 float32 下更精確地計(jì)算接受/拒絕步驟中的先驗(yàn)密度比自適應(yīng)步長adapt_step_size按step_size * exp(speed * (accept_prob - target_accept_rate))調(diào)整步長使接受率逼近目標(biāo)值MH 校正make_adaptive_hmc_updateburn-in 之外默認(rèn)開啟 Metropolis–Hastings 校正--no_mh可關(guān)閉。在 run_hmc.py 的訓(xùn)練主循環(huán)中每輪迭代都會(huì)評(píng)估 train/test 指標(biāo)、將被接受的樣本或--no_mh時(shí)的所有樣本累加進(jìn)ensemble_predictions即 BMA 集成、保存檢查點(diǎn)并把a(bǔ)ccept_prob、num_ensembled、當(dāng)前step_size等寫入 TensorBoard。訓(xùn)練目錄名由模型、權(quán)重衰減、步長、軌跡長度、burn-in、MH 開關(guān)、溫度與種子自動(dòng)拼接而成便于對比實(shí)驗(yàn)歸檔。需要說明的是README 明確指出論文中 CIFAR-10 上的 HMC 實(shí)驗(yàn)是在包含 512 個(gè) TPU 設(shè)備的 TPU pod 上、使用后續(xù)將發(fā)布的修改版代碼運(yùn)行的——即本倉庫開箱即用的代碼主要面向 GPU 場景CIFAR 級(jí)全量 HMC 需要特殊規(guī)模的硬件。六、運(yùn)行 SGD 與深度集成run_sgd.pySGD 腳本 run_sgd.py 使用cosine 學(xué)習(xí)率調(diào)度其專屬參數(shù)如下由 utils/cmd_args_utils.py 的add_sgd_flags()定義參數(shù)說明源碼默認(rèn)值init_step_size初始 SGD 步長1e-6num_epochsSGD 總 epoch 數(shù)300batch_size批大小80eval_freq評(píng)估頻率epoch10save_freq檢查點(diǎn)保存頻率epoch50momentum_decaySGD 動(dòng)量衰減0.9README 給出的示例ResNet-20-FRNresnet20_frn_swish在 CIFAR-10 / CIFAR-100 上訓(xùn)練 500 個(gè) epoch并使用--subset_train_to40960限制樣本數(shù)python3 run_sgd.py --seed1 --weight_decay10 --dirruns/sgd/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish \ --init_step_size3e-7 --num_epochs500 --eval_freq10 --batch_size80 \ --save_freq500 --subset_train_to40960 python3 run_sgd.py --seed1 --weight_decay10 --dirruns/sgd/cifar100/ \ --dataset_namecifar100 --model_nameresnet20_frn_swish \ --init_step_size1e-6 --num_epochs500 --eval_freq10 --batch_size80 \ --save_freq500 --subset_train_to40960IMDB 上的 CNN-LSTMpython3 run_sgd.py --seed1 --weight_decay3. --dirruns/sgd/imdb/ \ --dataset_nameimdb --model_namecnn_lstm --init_step_size3e-7 \ --num_epochs500 --eval_freq10 --batch_size80 --save_freq500深度集成的訓(xùn)練方式README 明確說明只需用不同隨機(jī)種子多次運(yùn)行run_sgd.py將得到的多個(gè) SGD 解進(jìn)行預(yù)測集成即可。論文比較了深度集成與 HMC 的預(yù)測分布差異得出“深度集成與 HMC 的接近程度和標(biāo)準(zhǔn) SGLD 相當(dāng)、且優(yōu)于標(biāo)準(zhǔn)變分推斷”的結(jié)論。從實(shí)現(xiàn)上看SGD 訓(xùn)練過程位于 utils/train_utils.py 的make_sgd_train_epoch()數(shù)據(jù)被切分為num_batches份每個(gè) epoch 內(nèi)用jax.lax.scan串行執(zhí)行各 batch 的梯度更新并由jax.lax.psum完成跨設(shè)備梯度聚合體現(xiàn) JAX 多設(shè)備如多 GPU并行訓(xùn)練能力。七、運(yùn)行 SGMCMCrun_sgmcmc.pySGMCMC 腳本 run_sgmcmc.py 共享 SGD 的全部參數(shù)并額外引入以下參數(shù)參數(shù)說明源碼默認(rèn)值preconditioner預(yù)處理器選擇None或RMSpropNonestep_size_schedule步長調(diào)度constant或cyclical。constant表示先做num_burnin_epochs個(gè) epoch 的 cosine burn-in隨后步長固定為final_step_sizecyclical表示先做固定步長的 burn-in再進(jìn)入 cosine 循環(huán)調(diào)度constantnum_burnin_epochs達(dá)到最終學(xué)習(xí)率前的 epoch 數(shù)300final_step_size最終步長僅用于constant調(diào)度默認(rèn)取init_step_sizestep_size_cycle_length_epochs循環(huán)周期長度epoch僅用于cyclical調(diào)度50save_all_ensembled保存所有被集成的網(wǎng)絡(luò)Falseensemble_freq集成 iterates 的頻率epoch10CIFAR-10 上 ResNet-20-FRN 的四種 SGMCMC 變體示例# SGLD python3 run_sgmcmc.py --seed1 --weight_decay5. --dirruns/sgmcmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size1e-6 \ --final_step_size1e-6 --num_epochs10000 --num_burnin_epochs1000 \ --eval_freq10 --batch_size80 --save_freq10 --momentum0. \ --subset_train_to40960 # SGHMC python3 run_sgmcmc.py --seed1 --weight_decay5 --dirruns/sgmcmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size3e-7 \ --final_step_size3e-7 --num_epochs10000 --num_burnin_epochs1000 \ --eval_freq10 --batch_size80 --save_freq10 --subset_train_to40960 \ --momentum0.9 # SGHMC-CLR python3 run_sgmcmc.py --seed1 --weight_decay5 --dirruns/sgmcmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size3e-7 \ --num_epochs10000 --num_burnin_epochs1000 --step_size_schedulecyclical \ --step_size_cycle_length_epochs50 --ensemble_freq50 --eval_freq10 \ --batch_size80 --save_freq1000 --subset_train_to40960 \ --preconditionerNone --momentum0.95 --eval_freq10 --save_all_ensembled # SGHMC-CLR-Prec python3 run_sgmcmc.py --seed1 --weight_decay5 --dirruns/sghmc/cifar10/ \ --dataset_namecifar10 --model_nameresnet20_frn_swish --init_step_size3e-5 \ --num_epochs10000 --num_burnin_epochs1000 --step_size_schedulecyclical \ --step_size_cycle_length_epochs50 --ensemble_freq50 --eval_freq10 \ --batch_size80 --save_freq50 --subset_train_to40960 \ --preconditionerRMSprop --momentum0.95 --eval_freq10 --save_all_ensembled四個(gè)變體對應(yīng)論文中比較的 SGMCMC 方法族標(biāo)準(zhǔn) SGLDmomentum0、SGHMC動(dòng)量0.9、SGHMC 循環(huán)學(xué)習(xí)率CLRstep_size_schedulecyclical且momentum0.95以及在此基礎(chǔ)上疊加 RMSprop 預(yù)處理preconditionerRMSprop。注意 SGLD 與 SGHMC 使用了不同的初始/最終步長1e-6與3e-7這與各自方法的噪聲特性有關(guān)。SGMCMC 的實(shí)現(xiàn)位于 core/sgmcmc.py以 optaxGradientTransformation形式提供sgld_gradient_update()實(shí)現(xiàn) SGLD當(dāng)momentum_decay0時(shí)為經(jīng)典 SGLDWelling Teh, ICML 2011否則退化為欠阻尼 SGLD即 SGHMCChen et al., ICML 2014。更新規(guī)則中梯度項(xiàng)乘lr_sqrt高斯噪聲項(xiàng)的標(biāo)準(zhǔn)差為sqrt(2 * (1 - momentum_decay))再疊加動(dòng)量項(xiàng)momentum_decay * m預(yù)處理通過Preconditioner抽象實(shí)現(xiàn)None對應(yīng)IdentityPreconditionerState恒等變換RMSprop對應(yīng)get_rmsprop_preconditioner()其基于梯度二階矩估計(jì)默認(rèn)衰減因子0.99、eps1e-7對噪聲與更新量分別做M^{-1/2}與M^{-1}變換即論文中 SGHMC-CLR-Prec 的“Prec”來源。步長調(diào)度在 run_sgmcmc.py 的get_lr_schedule()中構(gòu)建constant調(diào)度調(diào)用make_constant_lr_schedule_with_cosine_burnin(init_step_size, final_step_size, burnin_steps)cyclical調(diào)度調(diào)用make_cyclical_cosine_lr_schedule_with_const_burnin(...)周期由step_size_cycle_length_epochs換算為步數(shù)。訓(xùn)練過程中ensemble_freq決定每隔多少個(gè) epoch 將當(dāng)前 iterate 的預(yù)測累加進(jìn)集成save_all_ensembled則同時(shí)把被集成的網(wǎng)絡(luò)參數(shù)落盤。八、運(yùn)行 MFVIrun_vi.py均值場變分推斷腳本 run_vi.py 同樣共享 SGD 參數(shù)并額外定義參數(shù)說明源碼默認(rèn)值optimizer優(yōu)化器選擇SGD或Adam源碼默認(rèn)AdamREADME 描述為默認(rèn) SGD請以倉庫代碼為準(zhǔn)vi_sigma_initMFVI 中權(quán)重標(biāo)準(zhǔn)差σ的初始值1e-3vi_ensemble_sizeVI 評(píng)估時(shí)采樣的集成規(guī)模20mean_init_checkpoint用于初始化 MFVI 均值mean的 SGD 檢查點(diǎn)None示例命令# ResNet-20-FRN on CIFAR-10 or CIFAR-100 python3 run_vi.py --seed11 --weight_decay5. --dirruns/vi/cifar100/ \ --dataset_name[cifar10 | cifar100] --model_nameresnet20_frn_swish \ --init_step_size1e-4 --num_epochs300 --eval_freq10 --batch_size80 \ --save_freq300 --subset_train_to40960 --optimizerAdam \ --vi_sigma_init0.01 --temperature1. --vi_ensemble_size20 \ --mean_init_checkpointpath-to-sgd-solution # CNN-LSTM on IMDB python3 run_vi.py --seed11 --weight_decay5. --dirruns/vi/imdb/ \ --dataset_nameimdb --model_namecnn_lstm --init_step_size1e-4 \ --num_epochs500 --eval_freq10 --batch_size80 --save_freq200 \ --optimizerAdam --vi_sigma_init0.01 --temperature1. --vi_ensemble_size20 \ --mean_init_checkpointpath-to-sgd-solution其中--mean_init_checkpoint通常指向前面用run_sgd.py訓(xùn)練出的解即論文中的“用 SGD 解初始化 MFVI 均值”策略。MFVI 的實(shí)現(xiàn)在 core/vi.pyget_mfvi_model_fn()將每個(gè)權(quán)重參數(shù)化為{mean, inv_softplus_std}兩個(gè)量均值直接用原參數(shù)復(fù)制標(biāo)準(zhǔn)差則存其 inverse-softplus 形式sigma softplus(inv_softplus_std)保證標(biāo)準(zhǔn)差恒為正sample_parms_fn用重參數(shù)化技巧m n * sn 為標(biāo)準(zhǔn)正態(tài)噪聲從變分分布采樣make_kl_with_gaussian_prior(weight_decay, temperature)實(shí)現(xiàn) ELBO 中的先驗(yàn) KL 項(xiàng)對每個(gè)參數(shù)計(jì)算log(σ_prior/σ_vi) (σ_vi2 μ_vi2)/(2σ_prior2) - 1/2其中σ_prior sqrt(1/weight_decay)并以temperature作為 KL 項(xiàng)權(quán)重評(píng)估時(shí)vi_ensemble_predict_fn從變分后驗(yàn)采樣vi_ensemble_size次并集成預(yù)測同時(shí)run_vi.py還會(huì)將每個(gè)參數(shù)的 softplus 標(biāo)準(zhǔn)差以直方圖形式寫入 TensorBoardMFVI/param_stds便于觀察各層不確定性的分布。九、可視化后驗(yàn)密度make_posterior_surface_plot.py腳本 make_posterior_surface_plot.pyREADME 中筆誤寫作makemake_posterior_surface_plot.py實(shí)際文件名以此處為準(zhǔn)可以在由三個(gè)檢查點(diǎn)張成的二維平面上可視化后驗(yàn)的 log 密度、log 似然與 log 先驗(yàn)復(fù)現(xiàn)論文中的后驗(yàn)曲面圖。參數(shù)如下參數(shù)說明源碼默認(rèn)值limit_bottom曲面可視化在垂直方向底部的范圍-0.25limit_top曲面可視化在垂直方向頂部的范圍1.25limit_left曲面可視化在水平方向左側(cè)的范圍-0.25limit_right曲面可視化在水平方向右側(cè)的范圍1.25grid_size每個(gè)方向上的網(wǎng)格點(diǎn)數(shù)20checkpoint1第一個(gè)檢查點(diǎn)路徑—checkpoint2第二個(gè)檢查點(diǎn)路徑—checkpoint3第三個(gè)檢查點(diǎn)路徑—示例IMDB 上的 CNN-LSTMpython3 make_posterior_surface_plot.py --weight_decay40 --temperature1. \ --dirruns/surface_plots/imdb/ --model_namecnn_lstm --dataset_nameimdb \ --checkpoint1ckpt1 --checkpoint2ckpt2 --checkpoint3ckpt3 --limit_bottom-0.75 --limit_left-0.75 --limit_right1.75 --limit_top1.75 \ --grid_size50從源碼看該腳本以三個(gè)檢查點(diǎn)構(gòu)造平面坐標(biāo)系每點(diǎn)對應(yīng)平面上的一個(gè)二維坐標(biāo)在grid_size × grid_size網(wǎng)格上逐點(diǎn)計(jì)算 log 后驗(yàn)密度、log 似然與 log 先驗(yàn)依賴 utils/losses.py 與 utils/models.py 提供的前向計(jì)算并用 matplotlib 輸出可視化結(jié)果是理解“BNN 后驗(yàn)為何不像高斯”“冷后驗(yàn)現(xiàn)象由何而來”等論文論點(diǎn)的直觀工具。十、小結(jié)如何組織一次完整的復(fù)現(xiàn)實(shí)驗(yàn)綜合以上內(nèi)容一次完整的論文實(shí)驗(yàn)復(fù)現(xiàn)可按如下流程組織按 第二節(jié) 搭建 conda/pip 環(huán)境用run_sgd.py在不同種子下訓(xùn)練若干 SGD 解得到基線精度與深度集成用run_hmc.py在對應(yīng)任務(wù)上運(yùn)行全批量 HMC建議先用 5.3 節(jié) 的 MNIST 160 樣本示例在小規(guī)模上驗(yàn)證流程得到近乎精確的后驗(yàn)與 BMA 集成結(jié)果用run_sgmcmc.py分別運(yùn)行 SGLD / SGHMC / SGHMC-CLR / SGHMC-CLR-Prec與 HMC 對比預(yù)測分布的接近程度用run_vi.py以 SGD 解初始化均值運(yùn)行 MFVI作為變分推斷基線用make_posterior_surface_plot.py對若干代表性檢查點(diǎn)生成后驗(yàn)曲面圖直觀檢驗(yàn)論文關(guān)于先驗(yàn)尺度、溫度與后驗(yàn)形態(tài)的結(jié)論。所有腳本都將檢查點(diǎn)與 TensorBoard 日志寫入--dir指定目錄可通過 TensorBoard 對比各方法的 accuracy、NLL、ECE分類任務(wù)或 RMSE / NLL回歸任務(wù)等指標(biāo)以及test/ens_*前綴的集成指標(biāo)。值得留意的是論文涉及溫度對比實(shí)驗(yàn)時(shí)如 HMC 的四個(gè)溫度檔位需保持weight_decay一致而僅改變temperature并在解讀結(jié)果時(shí)記住該倉庫代碼默認(rèn)不使用數(shù)據(jù)增強(qiáng)而論文指出的“冷后驗(yàn)效應(yīng)很大程度上是數(shù)據(jù)增強(qiáng)的產(chǎn)物”這一結(jié)論正是建立在這樣的實(shí)驗(yàn)設(shè)置之上的?！久赓M(fèi)下載鏈接】google-researchGoogle Research項(xiàng)目地址: https://gitcode.com/gh_mirrors/go/google-research創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考