文 / Emily Fertig、Joshua V. Dillon、Wynn Vonnegut、Dave Moore 和 TensorFlow Probability 團隊

本文將介紹 TensorFlow Probability 中新的聯合分布變分推斷工具，展示如何在回歸模型中使用這些工具估計權重的貝葉斯可信區間，

TensorFlow Probability
https://tensorflow.google.cn/probability/

概述

變分推斷 (Variational Inference, VI) 將近似貝葉斯推斷轉換為優化問題，并尋求一個“代理 (Surrogate)”后驗分布，使 KL 與真后驗散度最小化，基于梯度的 VI 通常比 MCMC 方法更快，與模型引數的優化自然組合，并提供模型證據的下限，可直接用于模型比較、收斂診斷和可組合推斷，

TensorFlow Probability (TFP) 提供了快速、靈活和可擴展的 VI 工具，自然融入 TFP 堆疊，這些工具可以通過線性變換或歸一化流引起的協方差結構構建代理后驗，

VI 可用于估計回歸模型引數的貝葉斯可信區間 (Credible Intervals)，以估計各種處理或觀察到的特征對目標結果的影響，可信區間根據引數以觀察資料為條件的后驗分布，給定引數的先驗分布的假設，以一定的概率約束未觀察到的引數的值，

本文演示了如何使用 VI 獲得貝葉斯線性回歸模型引數的可信區間，該模型用于測量家庭中的氡水平（使用 Gelman et al.'s (2007) Radon dataset；參見 Stan 中的類似示例），我們演示了 TFP JointDistributions 如何與雙射器結合以建立和擬合兩種型別的表達性代理后驗：

由塊矩陣變換的標準正態分布，矩陣可以反映后驗的某些分量之間的獨立性和其他分量之間的依賴性，放寬均值場或全方差后驗的假設，
一個更復雜、容量更大的逆自回歸流 (Inverse Autoregressive Flow) ，

Gelman et al.'s (2007) Radon dataset
http://www.stat.columbia.edu/~gelman/arm/
類似示例
https://mc-stan.org/users/documentation/case-studies/radon.html#Correlations-among-levels
逆自回歸流
https://arxiv.org/abs/1606.04934

訓練代理后驗，并與均值場代理后驗基線結果進行比較，下圖顯示了三個 VI 代理后驗得到的四個模型引數的可信區間，以及用于比較的漢密爾頓蒙特卡洛 (Hamiltonian Monte Carlo) 演算法，

您可以在此 Google Colab 中跟進和了解所有詳細資訊，

Google Colab
https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Variational_Inference_and_Joint_Distributions.ipynb

示例：氡測量的貝葉斯分層線性回歸

氡是一種放射性氣體，通過地面的接觸點進入房屋，這種致癌物質是非吸煙者患肺癌的主要原因，不同家庭的氡水平差異很大，

EPA 對 8 萬間房屋的氡水平進行了研究，兩個重要的預測因素是：

測量的樓層（地下室的氡含量較高）
縣級鈾水平（與氡水平正相關）

預測按縣分組的房屋的氡水平是貝葉斯分層建模中的一個經典問題，由Gelman 和 Hill (2006) 提出，我們感興趣的是位置（縣，Countyi）對明尼蘇達州房屋氡水平影響的可信區間，為了隔離這種影響，模型中還包括樓層和鈾水平的影響，此外，我們還將按縣納入一個與測量的平均樓層相對應的背景效應，這樣如果測量的樓層在各縣之間存在差異，就不會歸因于縣效應，

Gelman 和 Hill (2006)
http://www.stat.columbia.edu/~gelman/arm/

回歸模型具體如下：

其中 i 為觀測值，countyi 即縣的第 i 個觀測值，

我們使用 county-level 隨機效應捕捉地理差異，引數 uranium_weight 和 county_floor_weight 經過概率建模，floor_weight 和常數 bias 是確定性的，這些建模選擇在很大程度上是任意的，目的是為了在一個合理復雜的概率模型上證明 VI，關于使用氡資料集在 TFP 中通過固定和隨機效應進行多層建模的更深入討論，請參閱多級建模 Primer 和使用變分推斷擬合泛化的線性混合效應模型，

多級建模 Primer
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Multilevel_Modeling_Primer.ipynb
使用變分推斷擬合泛化的線性混合效應模型
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Linear_Mixed_Effects_Model_Variational_Inference.ipynb

Github 上提供了此示例的完整代碼，

Github
https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Variational_Inference_and_Joint_Distributions.ipynb

變數為確定性引數和正態分布尺度引數定義，后者被約束為正數，

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors


floor_weight = tf.Variable(0.)
bias = tf.Variable(0.)
log_radon_scale = tfp.util.TransformedVariable(1., tfb.Exp())
county_effect_scale = tfp.util.TransformedVariable(1., tfb.Exp())

我們將回歸的概率圖模型指定為 TFP JointDistribution

@tfd.JointDistributionCoroutineAutoBatched
def model():
 uranium_weight = yield tfd.Normal(0., scale=1., name='uranium_weight')
 county_floor_weight = yield tfd.Normal(
     0., scale=1., name='county_floor_weight')
 county_effect = yield tfd.Sample(
     tfd.Normal(0., scale=county_effect_scale),
     sample_shape=[num_counties], name='county_effect')
 yield tfd.Normal(
     loc=(log_uranium * uranium_weight
          + floor_of_house * floor_weight
          + floor_by_county * county_floor_weight
          + tf.gather(county_effect, county, axis=-1)
          + bias),
     scale=log_radon_scale[..., tf.newaxis],
     name='log_radon')

我們將 log_radon 固定在觀測到的氡資料上，以對未歸一化的后驗進行建模，

target_model = model.experimental_pin(log_radon=log_radon)

貝葉斯變分推斷快速摘要

假設我們有以下生成程序，其中 𝜃 代表隨機引數（回歸模型中的 uranium_weight、county_floor_weight, 和 county_effect），ω 代表確定性引數（floor_weight、log_radon_scale、county_effect_scale 和 bias），x𝑖 為特征（log_uranium、floor_of_house 和 floor_by_county）；𝑦𝑖 為目標值 (log_radon)，𝑖 = 1...n 個觀測資料點：

VI 的特征為：

（從技術上看，我們假設 q 相對于 r 是絕對連續的，另請參閱琴生不等式）

由于該約束對所有 q 都成立，顯然最嚴格的是：

關于術語，我們稱

q* 為“代理后驗”，
Q 為“代理族”，

ω* 表示 VI 損失上確定性引數的最大似然值，有關變分推斷的更多資訊，請參見此調研，

調研
https://arxiv.org/abs/1601.00670

表達性代理后驗

接下來我們用 VI 估計引數的后驗分布，有兩種不同型別的代理后驗：

受約束的多元正態分布，其協方差結構由分塊矩陣變換誘導，
由逆自回歸流變換的多元標準正態分布，隨后進行拆分和重構以匹配后驗的支持，

多元正態代理后驗

為了構建該代理后驗，使用可訓練的線性算子誘導后驗分量之間的相關性，

首先構建一個具有向量值的標準正態分量的基礎分布，其大小等于相應的先驗分量的大小，這些分量具有向量值，因此可以被線性算子變換，

flat_event_size = tf.nest.map_structure(
     tf.reduce_prod,
     tf.nest.flatten(target_model.event_shape_tensor()))
base_standard_dist = tfd.JointDistributionSequential(
     [tfd.Sample(tfd.Normal(loc=0., scale=1.), s)
      for s in flat_event_size])

對此分布應用可訓練的分塊下三角線性算子誘導后驗的相關性，在線性算子中，可訓練的全矩陣塊代表后驗兩個分量之間的全協方差，而零（或 None）的塊表示獨立性，對角線上的塊是下三角矩陣或對角線矩陣，因此整個塊結構代表下三角矩陣，

將此雙射器應用于基礎分布會導致均值為 0、（Cholesky 分解的）協方差等于下三角塊矩陣的多元正態分布，

operators = (
  (tf.linalg.LinearOperatorDiag,),  # Variance of uranium weight (scalar).
  (tf.linalg.LinearOperatorFullMatrix,  # Covariance between uranium and floor-by-county weights.
   tf.linalg.LinearOperatorDiag),  # Variance of floor-by-county weight (scalar).
  (None,  # Independence between uranium weight and county effects.
   None,  # Independence between floor-by-county and county effects.
   tf.linalg.LinearOperatorDiag)  # Independence among the 85 county effects.
)
block_tril_linop = (
  tfp.experimental.vi.util.build_trainable_linear_operator_block(
      operators, flat_event_size))
scale_bijector = tfb.ScaleMatvecLinearOperatorBlock(block_tril_linop)

最后，應用可訓練的 Shift 雙射器允許平均值取非零值，

loc_bijector = tfb.JointMap(
   tf.nest.map_structure(
       lambda s: tfb.Shift(
           tf.Variable(tf.random.uniform(
               (s,), minval=-2., maxval=2., dtype=tf.float32))),
       flat_event_size))

以比例尺和位置雙射器對標準正態分布進行變換所得的多元正態分布必須進行重塑和重構才能匹配先驗，并最終制約于先驗的支持，

reshape_bijector = tfb.JointMap(
   tf.nest.map_structure(tfb.Reshape, flat_event_shape))
unflatten_bijector = tfb.Restructure(
       tf.nest.pack_sequence_as(
           event_shape, range(len(flat_event_shape))))
event_space_bijector = target_model.experimental_default_event_space_bijector()

接下來將其整合 - 將可訓練的雙射器鏈接在一起，并應用于基礎標準正態分布以構建代理后驗，

surrogate_posterior = tfd.TransformedDistribution(
   base_standard_dist,
   bijector = tfb.Chain([  # Chained bijectors are applied in reverse order.
        event_space_bijector,  # Constrain to the support of the prior.
        unflatten_bijector,  # Pack components into the event_shape structure.
        reshape_bijector,  # Reshape the vector-valued components.
        loc_bijector,  # Allow for nonzero mean.
        scale_bijector  # Apply the block matrix transformation.
      ]))

訓練多元正態代理后驗，

optimizer = tf.optimizers.Adam(learning_rate=1e-2)
@tf.function(jit_compile=True)
def run_vi():
 return tfp.vi.fit_surrogate_posterior(
   target_model.unnormalized_log_prob,
   surrogate_posterior,
   optimizer=optimizer,
   num_steps=10**4,
   sample_size=16)
mvn_loss = run_vi()
mvn_samples = surrogate_posterior.sample(1000)

由于訓練后的代理后驗是 TFP 分布，因此我們可以從中獲取樣本并進行處理，得出引數的后驗可信區間，

以下箱須圖顯示了兩個最大縣的縣效應的 50% 和 95% 可信區間，以及各個縣對土壤鈾測量值和平均樓層的回歸權重，縣效應的后驗可信區間表明，在考慮其他變數之后，圣路易斯縣的地理位置與較低的氡水平相關，亨內平縣的地理位置效應接近中性，

回歸權重的后驗可信區間表明，較高的土壤鈾水平與較高的氡水平相關，在較高樓層進行測量的縣（很可能是因為房屋沒有地下室）往往有較高的氡水平，這可能與土壤特性及其對建筑結構型別的影響有關，

樓層的（確定性）系數為 -0.7，表明較低樓層的氡水平較高，符合預期，

逆自回歸流代理后驗

逆自回歸流 (IAF) 是一種正態化流，使用神經網路捕獲分布中各分量之間復雜的非線性依賴關系，接下來，我們構建 IAF 代理后驗，探明此容量更大、更靈活的模型是否優于受約束的多元正態，

首先建立具有矢量事件形狀的標準正態分布，其長度等于后驗自由度的總數，

base_distribution = tfd.Sample(

   tfd.Normal(loc=0., scale=1.),
   sample_shape=[tf.reduce_sum(flat_event_size)])

可訓練的 IAF 變換正態分布，

num_iafs = 2
iaf_bijectors = [
   tfb.Invert(tfb.MaskedAutoregressiveFlow(
       shift_and_log_scale_fn=tfb.AutoregressiveNetwork(
           params=2,
           hidden_units=[256, 256],
           activation='relu')))
   for _ in range(num_iafs)
]

IAF 雙射器與其他雙射器鏈接，建立具有與先驗相同的事件形狀和支持的代理后驗，

iaf_surrogate_posterior = tfd.TransformedDistribution(
   base_distribution,
   bijector=tfb.Chain([
        event_space_bijector,  # Constrain to the support of the prior.
        unflatten_bijector,  # Pack components into the event_shape structure.
        reshape_bijector,  # Reshape the vector-valued components.
        tfb.Split(flat_event_size),  # Split into parts, same size as prior.
   ] + iaf_bijectors))  # Apply a flow model.

與多元正態代理后驗相同，IAF 代理后驗也使用 tfp.vi.fit_surrogate_posterior 訓練，IAF 代理后驗的可信區間與受約束的多元正態的可信區間相似，

均值場代理后驗

VI 代理后驗通常被假定為均值場（獨立）正態分布，具有可訓練的均值和方差，通過雙射變換約束到先驗支持，除了兩個更具表達性的代理后驗之外，我們還使用一個與多元正態代理后驗相同的通用公式來定義均值場代理后驗，我們不使用分塊下三角線性算子，而是使用分塊對角線性算子，其中每個塊都是對角線的：

operators = (
  tf.linalg.LinearOperatorDiag,
  tf.linalg.LinearOperatorDiag,
  tf.linalg.LinearOperatorDiag,
)
block_diag_linop = (
   tfp.experimental.vi.util.build_trainable_linear_operator_block(
       operators, flat_event_size))

在這種情況下，均值場代理后驗給出了與更具表達性的代理后驗相似的結果，表明這種更簡單的模型可能足以完成推斷任務，作為“基本事實”，我們還使用漢密爾頓蒙特卡洛演算法采樣（有關完整示例，請參見 Colab），這三個代理后驗產生的可信區間在視覺上均與 HMC 樣本相似，盡管有時由于 ELBO 損失的影響而分散不足，這常見于 VI，

Colab
https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Variational_Inference_and_Joint_Distributions.ipynb

結論

在這篇文章中，我們使用聯合分布和多部分雙射器構建了 VI 代理后驗，并將其擬合以估計氡資料集回歸模型中權重的可信區間，對于此簡單模型，更具表達性的代理后驗的表現似乎與均值場代理后驗相似，不過我們展示的工具可構建各種靈活的代理后驗，適用于更復雜的模型，

前往 TFP 首頁查看我們的代碼、檔案和更多示例，

TFP 首頁
https://tensorflow.google.cn/probability/

想了解 TensorFlow 更多實用工具，請點擊下方鏈接，前往 TensorFlow 官網，在機器學習的道路上，我們共同進步，

https://tensorflow.google.cn

轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/275734.html

標籤：AI

上一篇：答應我，這次搞懂 I/O 多路復用！

下一篇：Gradle專案的jar發布到私有倉庫

TensorFlow Probability 聯合分布變分推斷工具，估計權重的貝葉斯可信區間更簡單