الاستدلال المتغير على النماذج الرسومية الاحتمالية ذات التوزيعات المشتركة

عرض على TensorFlow.org تشغيل في Google Colab عرض المصدر على جيثب تحميل دفتر

يلقي الاستدلال المتغير (VI) الاستدلال البايزي التقريبي باعتباره مشكلة تحسين ويسعى إلى توزيع لاحق "بديل" يقلل من تباعد KL مع الخلفي الحقيقي. غالبًا ما يكون VI المستند إلى التدرج أسرع من طرق MCMC ، ويتكون بشكل طبيعي مع تحسين معلمات النموذج ، ويوفر حدًا أقل على دليل النموذج الذي يمكن استخدامه مباشرة لمقارنة النماذج وتشخيص التقارب والاستدلال القابل للتكوين.

تقدم TensorFlow Probability أدوات لسرعة ، ومرونة ، وقابلة للتطوير والتي تتلاءم بشكل طبيعي مع مكدس TFP. تتيح هذه الأدوات بناء خلفية بديلة مع هياكل التغاير الناتجة عن التحولات الخطية أو التدفقات الطبيعية.

VI يمكن استخدامها لتقدير النظرية الافتراضية فترات موثوقة لمعلمات نموذج الانحدار لتقدير آثار العلاجات المختلفة أو ميزات الملاحظة على هذه النتيجة المثيرة للاهتمام. ربطت الفواصل الزمنية الموثوقة قيم المعلمة غير المرصودة باحتمال معين ، وفقًا للتوزيع اللاحق للمعامل المشروط بالبيانات المرصودة ومع افتراض التوزيع المسبق للمعامل.

في هذا Colab، ونحن لشرح كيفية استخدام VI الحصول على فترات معقولة لالمعلمات من النظرية الافتراضية الخطية نموذج الانحدار لمستويات غاز الرادون في المنازل قياس (باستخدام غيلمان وآخرون (2007) الرادون البيانات. وانظر أمثلة مشابهة في ستان). نحن لشرح كيفية TFP JointDistribution الصورة تتحد مع bijectors لبناء وتناسب نوعين من مؤخرات بديل التعبيرية:

  • توزيع عادي قياسي يتم تحويله بواسطة مصفوفة كتلة. قد تعكس المصفوفة الاستقلال بين بعض مكونات الجزء الخلفي والاعتماد من بين أمور أخرى ، مما يخفف من افتراض المجال المتوسط ​​أو التغاير الكامل اللاحق.
  • وأكثر تعقيدا، والقدرات العليا معكوس تدفق الانحدار .

يتم تدريب الخلفيات البديلة ومقارنتها بالنتائج من خط أساس خلفي بديل متوسط ​​المجال ، بالإضافة إلى عينات من الحقيقة الأرضية من هاميلتونيان مونت كارلو.

نظرة عامة على الاستدلال المتغير بايزي

لنفترض أن لدينا عملية توليدي التالي، حيث \(\theta\) يمثل المعلمات عشوائية، \(\omega\) يمثل المعلمات القطعية، و \(x_i\) هي ميزات و \(y_i\) هي القيم المستهدفة ل \(i=1,\ldots,n\) احظ نقاط البيانات: \ تبدأ {محاذاة } & \ theta \ sim r (\ Theta) && \ text {(Prior)} \ & \ text {for} i = 1 \ ldots n: \ non number \ & \ quad y_i \ sim p (Y_i | x_i، \ theta \ أوميغا) && \ النص {(احتمالية)} \ نهاية {محاذاة}

يتم بعد ذلك تمييز VI بـ: $ \ newcommand {\ E} {\ operatorname {\ mathbb {E}}} \ newcommand {\ K} {\ operatorname {\ mathbb {K}}} \ newcommand {\ defeq} {\ overset {\ tiny \ text {def}} {=}} \ DeclareMathOperator * {\ argmin} {arg \، min} $

\ تبدأ {محاذاة} - \ سجل ص ({y_i} _I ^ ن | {x_i} _I ^ ن، \ أوميغا) و \ defeq - \ تسجيل \ كثافة العمليات \ textrm {د} \ ثيتا \، ص (\ ثيتا) \ prod_i ^ np (y_i | x_i، \ theta، \ omega) && \ text {(لا يتجزأ حقًا)} \ & = - \ log \ int \ textrm {d} \ theta \، q (\ theta) \ frac {1 } {q (\ theta)} r (\ theta) \ prod_i ^ np (y_i | x_i، \ theta، \ omega) && \ text {(الضرب في 1)} \ & \ le - \ int \ textrm {d} \ ثيتا \، ف (\ ثيتا) \ سجل \ فارك {ص (\ ثيتا) \ prod_i ^ أرستها (y_i | خ ط، \ ثيتا، \ أوميغا)} {س (\ ثيتا)} && \ النص {(عدم المساواة جنسن )} \ & \ defeq \ E {س (\ ثيتا)} [- \ سجل ص (y_i | x_i، \ ثيتا، \ أوميغا)] + \ K [ف (\ ثيتا)، ص (\ ثيتا)] \ & \ defeq \text{expected negative log likelihood"} + نهاية \ النص {regularizer كوالا لمبور"} \ {محاذاة}

(من الناحية التقنية نحن على افتراض \(q\) هو مستمر على الاطلاق فيما يتعلق \(r\). انظر أيضا، عدم المساواة جنسن ).

نظرًا لأن الحد ينطبق على كل q ، فمن الواضح أنه ضيق للغاية لـ:

\[q^*,w^* = \argmin_{q \in \mathcal{Q},\omega\in\mathbb{R}^d} \left\{ \sum_i^n\E_{q(\Theta)}\left[ -\log p(y_i|x_i,\Theta, \omega) \right] + \K[q(\Theta), r(\Theta)] \right\}\]

فيما يتعلق بالمصطلحات ، نسميها

  • \(q^*\) في "الخلفية البديلة"، و،
  • \(\mathcal{Q}\) "الأسرة البديلة".

\(\omega^*\) يمثل القيم كحد أقصى احتمال المعلمات القطعية على فقدان VI. ترى هذه الدراسة لمزيد من المعلومات حول الاستدلال التغاير.

مثال: الانحدار الخطي الهرمي البايزي على قياسات الرادون

الرادون هو غاز مشع يدخل المنازل من خلال نقاط التلامس مع الأرض. وهي مادة مسرطنة هي السبب الرئيسي لسرطان الرئة لدى غير المدخنين. تختلف مستويات الرادون اختلافًا كبيرًا من منزل إلى آخر.

قامت وكالة حماية البيئة بدراسة مستويات الرادون في 80000 منزل. اثنان من المتنبئين المهمين هما:

  • الأرضية التي تم القياس عليها (الرادون أعلى في الأقبية)
  • مستوى اليورانيوم في المقاطعة (ارتباط إيجابي بمستويات الرادون)

توقع مستويات غاز الرادون في المنازل التي تم تجميعها حسب المقاطعة هي مشكلة الكلاسيكية في النظرية الافتراضية النمذجة الهرمي، الذي عرضته غيلمان وهيل (2006) . سنقوم ببناء نموذج خطي هرمي للتنبؤ بقياسات الرادون في المنازل ، حيث يكون التسلسل الهرمي هو تجميع المنازل حسب المقاطعة. نحن مهتمون بفترات زمنية موثوقة لتأثير الموقع (المقاطعة) على مستوى الرادون للمنازل في مينيسوتا. من أجل عزل هذا التأثير ، تم أيضًا تضمين تأثيرات مستوى الأرض واليورانيوم في النموذج. إضافي ، سنقوم بتضمين تأثير سياقي يتوافق مع متوسط ​​الأرضية التي تم إجراء القياس عليها ، حسب المقاطعة ، بحيث إذا كان هناك تباين بين المقاطعات في الأرضية التي تم إجراء القياسات عليها ، فلن يُعزى ذلك إلى تأثير المقاطعة.

pip3 install -q tf-nightly tfp-nightly

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_probability as tfp
import warnings

tfd = tfp.distributions
tfb = tfp.bijectors

plt.rcParams['figure.facecolor'] = '1.'

# Load the Radon dataset from `tensorflow_datasets` and filter to data from
# Minnesota.
dataset = tfds.as_numpy(
    tfds.load('radon', split='train').filter(
        lambda x: x['features']['state'] == 'MN').batch(10**9))

# Dependent variable: Radon measurements by house.
dataset = next(iter(dataset))
radon_measurement = dataset['activity'].astype(np.float32)
radon_measurement[radon_measurement <= 0.] = 0.1
log_radon = np.log(radon_measurement)

# Measured uranium concentrations in surrounding soil.
uranium_measurement = dataset['features']['Uppm'].astype(np.float32)
log_uranium = np.log(uranium_measurement)

# County indicator.
county_strings = dataset['features']['county'].astype('U13')
unique_counties, county = np.unique(county_strings, return_inverse=True)
county = county.astype(np.int32)
num_counties = unique_counties.size

# Floor on which the measurement was taken.
floor_of_house = dataset['features']['floor'].astype(np.int32)

# Average floor by county (contextual effect).
county_mean_floor = []
for i in range(num_counties):
  county_mean_floor.append(floor_of_house[county == i].mean())
county_mean_floor = np.array(county_mean_floor, dtype=log_radon.dtype)
floor_by_county = county_mean_floor[county]

تم تحديد نموذج الانحدار على النحو التالي:

\(\newcommand{\Normal}{\operatorname{\sf Normal} }\)\ تبدأ {محاذاة} & \ النص {uranium_weight} \ سيم \ عادي (0، 1) \ & \ النص {county_floor_weight} \ سيم \ عادي (0، 1) \ & \ النص {ل} ي = 1 \ ldots \ text {num_counties}: \ & \ quad \ text {County_effect} _j \ sim \ Normal (0، \ sigma_c) \ & \ text {for} i = 1 \ ldots n: \ & \ quad \ mu_i = (\ و\ رباعية \ رباعية \ النص {التحيز} \ & \ رباعية \ رباعية + \ النص {تأثير مقاطعة} {\ النص {مقاطعة} _I} \ & \ رباعية \ رباعية + \ النص {log_uranium} _I \ النص مرات \ {uranium_weight } \ & \ رباعية \ رباعية + \ النص {floor_of_house} _I \ مرات \ النص {floor_weight} \ & \ رباعية \ رباعية + \ النص {floor_by مقاطعة} {\ النص {مقاطعة} _I} \ النص مرات \ {county_floor_weight}) \ & \ رباعية \ النص {log_radon} _I \ سيم \ عادي (\ mu_i، \ sigma_y) \ نهاية {محاذاة} التي \(i\) مؤشرات الملاحظات و \(\text{county}_i\) هي مقاطعة فيه \(i\)كانت عشر المراقبة تؤخذ.

نستخدم تأثيرًا عشوائيًا على مستوى المقاطعة لالتقاط التباين الجغرافي. المعلمات uranium_weight و county_floor_weight وعلى غرار احتماليا، و floor_weight وثابت bias هي حتمية. تعتبر خيارات النمذجة هذه تعسفية إلى حد كبير ، ويتم إجراؤها لغرض إظهار VI على نموذج احتمالي من التعقيد المعقول. للنقاش حول النمذجة متعددة المستويات مع الآثار الثابتة والعشوائية في TFP، وذلك باستخدام بيانات الرادون، انظر متعدد المستويات النمذجة التمهيدي و تركيب المعمم الخطي آثار مختلطة نماذج عن طريق تباين الاستدلال .

# Create variables for fixed effects.
floor_weight = tf.Variable(0.)
bias = tf.Variable(0.)

# Variables for scale parameters.
log_radon_scale = tfp.util.TransformedVariable(1., tfb.Exp())
county_effect_scale = tfp.util.TransformedVariable(1., tfb.Exp())

# Define the probabilistic graphical model as a 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') 

# Pin the observed `log_radon` values to model the un-normalized posterior.
target_model = model.experimental_pin(log_radon=log_radon)

خلفية بديلة معبرة

بعد ذلك نقوم بتقدير التوزيعات اللاحقة للتأثيرات العشوائية باستخدام VI مع نوعين مختلفين من المؤثرات الخلفية البديلة:

  • توزيع طبيعي مقيد متعدد المتغيرات ، بهيكل التغاير الناتج عن تحويل المصفوفة الكتلي.
  • A التوزيع مستوى عادي متعدد المتغيرات تتحول من قبل نماذج الانحدار الذاتى والتدفق العكسي ، والذي أدى ذلك إلى انقسام وإعادة هيكلتها لتتناسب مع دعم الخلفي.

متعدد المتغيرات طبيعي بديل خلفي

لبناء هذا البديل الخلفي ، يتم استخدام عامل تشغيل خطي قابل للتدريب للحث على الارتباط بين مكونات الجزء الخلفي.

# Determine the `event_shape` of the posterior, and calculate the size of each
# `event_shape` component. These determine the sizes of the components of the
# underlying standard Normal distribution, and the dimensions of the blocks in
# the blockwise matrix transformation.
event_shape = target_model.event_shape_tensor()
flat_event_shape = tf.nest.flatten(event_shape)
flat_event_size = tf.nest.map_structure(tf.reduce_prod, flat_event_shape)

# The `event_space_bijector` maps unconstrained values (in R^n) to the support
# of the prior -- we'll need this at the end to constrain Multivariate Normal
# samples to the prior's support.
event_space_bijector = target_model.experimental_default_event_space_bijector()

بناء JointDistribution مع المكونات العادية القياسية قيمتها ناقلات، مع أحجام تحددها مكونات السابقة المقابلة. يجب أن تكون المكونات ذات قيمة متجهة بحيث يمكن تحويلها بواسطة المشغل الخطي.

base_standard_dist = tfd.JointDistributionSequential(
      [tfd.Sample(tfd.Normal(0., 1.), s) for s in flat_event_size])

بناء مشغل خطي مثلث منخفض قابل للتدريب. سنقوم بتطبيقه على التوزيع الطبيعي القياسي لتنفيذ تحويل مصفوفة (قابل للتدريب) وإحداث بنية الارتباط للخلف.

داخل المشغل blockwise الخطية، ويمثل كتلة كاملة مصفوفة التغاير للتدريب الكامل بين مكونات الخلفية، في حين أن كتلة من الأصفار (أو None ) يعبر عن الاستقلال. تكون الكتل الموجودة على القطر إما مصفوفات مثلثية أو قطرية منخفضة ، بحيث تمثل بنية الكتلة بأكملها مصفوفة ذات مثلث منخفض.

تطبيق هذا bijector على التوزيع الأساسي ينتج عنه توزيع طبيعي متعدد المتغيرات بمتوسط ​​0 وتغاير (Cholesky-Factored) يساوي مصفوفة الكتلة المثلثية السفلية.

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 bijector للسماح سيلة لاتخاذ قيم غير صفرية.

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 each component to match the prior, using a nested structure of
# `Reshape` bijectors wrapped in `JointMap` to form a multipart bijector.
reshape_bijector = tfb.JointMap(
    tf.nest.map_structure(tfb.Reshape, flat_event_shape))

# Restructure the flat list of components to match the prior's structure
unflatten_bijector = tfb.Restructure(
        tf.nest.pack_sequence_as(
            event_shape, range(len(flat_event_shape))))

الآن ، ضع كل ذلك معًا - اربط المحرضين القابل للتدريب معًا وقم بتطبيقهم على التوزيع الطبيعي القياسي الأساسي لإنشاء الجزء الخلفي البديل.

surrogate_posterior = tfd.TransformedDistribution(
    base_standard_dist,
    bijector = tfb.Chain(  # Note that the chained bijectors are applied in reverse order
        [
         event_space_bijector,  # constrain the surrogate to the support of the prior
         unflatten_bijector,  # pack the reshaped components into the `event_shape` structure of the posterior
         reshape_bijector,  # reshape the vector-valued components to match the shapes of the posterior components
         loc_bijector,  # allow for nonzero mean
         scale_bijector  # apply the block matrix transformation to the standard Normal distribution
         ]))

تدريب البديل البديل العادي متعدد المتغيرات.

optimizer = tf.optimizers.Adam(learning_rate=1e-2)
mvn_loss = tfp.vi.fit_surrogate_posterior(
    target_model.unnormalized_log_prob,
    surrogate_posterior,
    optimizer=optimizer,
    num_steps=10**4,
    sample_size=16,
    jit_compile=True)

mvn_samples = surrogate_posterior.sample(1000)
mvn_final_elbo = tf.reduce_mean(
    target_model.unnormalized_log_prob(*mvn_samples)
    - surrogate_posterior.log_prob(mvn_samples))

print('Multivariate Normal surrogate posterior ELBO: {}'.format(mvn_final_elbo))

plt.plot(mvn_loss)
plt.xlabel('Training step')
_ = plt.ylabel('Loss value')

Multivariate Normal surrogate posterior ELBO: -1065.705322265625

بي إن جي

نظرًا لأن البديل البديل المدرب هو توزيع TFP ، فيمكننا أخذ عينات منه ومعالجتها لإنتاج فترات خلفية موثوقة للمعلمات.

إظهار مؤامرات مربع وشعيرات أقل من 50٪ و 95٪ فترات ذات مصداقية للتأثير مقاطعة من اثنين من أكبر المقاطعات وأوزان الانحدار على قياسات اليورانيوم التربة والكلمة تعني مقاطعة. تشير الفواصل الزمنية الموثوقة اللاحقة لتأثيرات المقاطعة إلى أن الموقع في مقاطعة سانت لويس مرتبط بمستويات أقل من الرادون ، بعد حساب المتغيرات الأخرى ، وأن تأثير الموقع في مقاطعة هينيبين قريب من الحياد.

تظهر الفترات الخلفية الموثوقة على أوزان الانحدار أن المستويات الأعلى من اليورانيوم في التربة ترتبط بمستويات أعلى من الرادون ، وتميل المقاطعات التي تم فيها أخذ القياسات في الطوابق العليا (على الأرجح لأن المنزل لم يكن به قبو) إلى الحصول على مستويات أعلى من الرادون ، والتي يمكن أن تتعلق بخصائص التربة وتأثيرها على نوع الهياكل المبنية.

المعامل (القطعي) للأرضية سالب ، مما يشير إلى أن الطوابق السفلية بها مستويات أعلى من الرادون ، كما هو متوقع.

st_louis_co = 69  # Index of St. Louis, the county with the most observations.
hennepin_co = 25  # Index of Hennepin, with the second-most observations.

def pack_samples(samples):
  return {'County effect (St. Louis)': samples.county_effect[..., st_louis_co],
          'County effect (Hennepin)': samples.county_effect[..., hennepin_co],
          'Uranium weight': samples.uranium_weight,
          'Floor-by-county weight': samples.county_floor_weight}

def plot_boxplot(posterior_samples):
  fig, axes = plt.subplots(1, 4, figsize=(16, 4))

  # Invert the results dict for easier plotting.
  k = list(posterior_samples.values())[0].keys()
  plot_results = {
      v: {p: posterior_samples[p][v] for p in posterior_samples} for v in k}
  for i, (var, var_results) in enumerate(plot_results.items()):
    sns.boxplot(data=list(var_results.values()), ax=axes[i],
                width=0.18*len(var_results), whis=(2.5, 97.5))
    # axes[i].boxplot(list(var_results.values()), whis=(2.5, 97.5))
    axes[i].title.set_text(var)
    fs = 10 if len(var_results) < 4 else 8
    axes[i].set_xticklabels(list(var_results.keys()), fontsize=fs)

results = {'Multivariate Normal': pack_samples(mvn_samples)}

print('Bias is: {:.2f}'.format(bias.numpy()))
print('Floor fixed effect is: {:.2f}'.format(floor_weight.numpy()))
plot_boxplot(results)

Bias is: 1.40
Floor fixed effect is: -0.72

بي إن جي

عكس تدفق بديل ذاتي الانحدار الخلفي

تدفقات الانحدار التلقائي العكسية (IAFs) هي تدفقات تطبيع تستخدم الشبكات العصبية لالتقاط التبعيات المعقدة وغير الخطية بين مكونات التوزيع. بعد ذلك ، قمنا ببناء نموذج خلفي بديل لـ IAF لمعرفة ما إذا كان هذا النموذج ذو السعة العالية والأكثر مرونة يتفوق على المتغير العادي المقيد متعدد المتغيرات.

# Build a standard Normal with a vector `event_shape`, with length equal to the
# total number of degrees of freedom in the posterior.
base_distribution = tfd.Sample(
    tfd.Normal(0., 1.), sample_shape=[tf.reduce_sum(flat_event_size)])

# Apply an IAF to the base distribution.
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)
]

# Split the base distribution's `event_shape` into components that are equal
# in size to the prior's components.
split = tfb.Split(flat_event_size)

# Chain these bijectors and apply them to the standard Normal base distribution
# to build the surrogate posterior. `event_space_bijector`,
# `unflatten_bijector`, and `reshape_bijector` are the same as in the
# multivariate Normal surrogate posterior.
iaf_surrogate_posterior = tfd.TransformedDistribution(
    base_distribution,
    bijector=tfb.Chain([
         event_space_bijector,  # constrain the surrogate to the support of the prior
         unflatten_bijector,  # pack the reshaped components into the `event_shape` structure of the prior
         reshape_bijector,  # reshape the vector-valued components to match the shapes of the prior components
         split] +  # Split the samples into components of the same size as the prior components
         iaf_bijectors  # Apply a flow model to the Tensor-valued standard Normal distribution
         ))

تدريب اللاحق IAF بديل.

optimizer=tf.optimizers.Adam(learning_rate=1e-2)
iaf_loss = tfp.vi.fit_surrogate_posterior(
  target_model.unnormalized_log_prob,
  iaf_surrogate_posterior,
  optimizer=optimizer,
  num_steps=10**4,
  sample_size=4,
  jit_compile=True)

iaf_samples = iaf_surrogate_posterior.sample(1000)
iaf_final_elbo = tf.reduce_mean(
    target_model.unnormalized_log_prob(*iaf_samples)
    - iaf_surrogate_posterior.log_prob(iaf_samples))
print('IAF surrogate posterior ELBO: {}'.format(iaf_final_elbo))

plt.plot(iaf_loss)
plt.xlabel('Training step')
_ = plt.ylabel('Loss value')

IAF surrogate posterior ELBO: -1065.3663330078125

بي إن جي

تظهر الفترات الزمنية الموثوقة للجزء الخلفي البديل IAF مشابهة لتلك الخاصة بالمتغير العادي متعدد المتغيرات المقيد.

results['IAF'] = pack_samples(iaf_samples)
plot_boxplot(results)

بي إن جي

خط الأساس: متوسط ​​المجال الخلفي الخلفي

غالبًا ما يُفترض أن الخلفيات البديلة السادسة هي توزيعات ميدانية (مستقلة) عادية ، بوسائل وتباينات قابلة للتدريب ، مقيدة بدعم السابق مع تحول حيوي. نحدد الخلفية البديلة ذات المجال المتوسط ​​بالإضافة إلى الخلفيتين البديلتين الأكثر تعبيراً ، باستخدام نفس الصيغة العامة مثل البديل العادي البديل متعدد المتغيرات.

# A block-diagonal linear operator, in which each block is a diagonal operator,
# transforms the standard Normal base distribution to produce a mean-field
# surrogate posterior.
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))
mean_field_scale = tfb.ScaleMatvecLinearOperatorBlock(block_diag_linop)

mean_field_loc = 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))

mean_field_surrogate_posterior = tfd.TransformedDistribution(
    base_standard_dist,
    bijector = tfb.Chain(  # Note that the chained bijectors are applied in reverse order
        [
         event_space_bijector,  # constrain the surrogate to the support of the prior
         unflatten_bijector,  # pack the reshaped components into the `event_shape` structure of the posterior
         reshape_bijector, # reshape the vector-valued components to match the shapes of the posterior components
         mean_field_loc,   # allow for nonzero mean
         mean_field_scale  # apply the block matrix transformation to the standard Normal distribution
         ]))

optimizer=tf.optimizers.Adam(learning_rate=1e-2)
mean_field_loss = tfp.vi.fit_surrogate_posterior(
    target_model.unnormalized_log_prob,
    mean_field_surrogate_posterior,
    optimizer=optimizer,
    num_steps=10**4,
    sample_size=16,
    jit_compile=True)

mean_field_samples = mean_field_surrogate_posterior.sample(1000)
mean_field_final_elbo = tf.reduce_mean(
    target_model.unnormalized_log_prob(*mean_field_samples)
    - mean_field_surrogate_posterior.log_prob(mean_field_samples))
print('Mean-field surrogate posterior ELBO: {}'.format(mean_field_final_elbo))

plt.plot(mean_field_loss)
plt.xlabel('Training step')
_ = plt.ylabel('Loss value')

Mean-field surrogate posterior ELBO: -1065.7652587890625

بي إن جي

في هذه الحالة ، يعطي البديل البديل للحقل المتوسط ​​نتائج مماثلة للخلفيات البديلة الأكثر تعبيرًا ، مما يشير إلى أن هذا النموذج الأبسط قد يكون مناسبًا لمهمة الاستدلال.

results['Mean Field'] = pack_samples(mean_field_samples)
plot_boxplot(results)

بي إن جي

الحقيقة الأساسية: هاميلتوني مونتي كارلو (HMC)

نحن نستخدم HMC لتوليد عينات "الحقيقة الأساسية" من الخلف الحقيقي ، للمقارنة مع نتائج البدائل البديلة.

num_chains = 8
num_leapfrog_steps = 3
step_size = 0.4
num_steps=20000

flat_event_shape = tf.nest.flatten(target_model.event_shape)
enum_components = list(range(len(flat_event_shape)))
bijector = tfb.Restructure(
    enum_components,
    tf.nest.pack_sequence_as(target_model.event_shape, enum_components))(
        target_model.experimental_default_event_space_bijector())

current_state = bijector(
    tf.nest.map_structure(
        lambda e: tf.zeros([num_chains] + list(e), dtype=tf.float32),
    target_model.event_shape))

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=target_model.unnormalized_log_prob,
    num_leapfrog_steps=num_leapfrog_steps,
    step_size=[tf.fill(s.shape, step_size) for s in current_state])

hmc = tfp.mcmc.TransformedTransitionKernel(
    hmc, bijector)
hmc = tfp.mcmc.DualAveragingStepSizeAdaptation(
    hmc,
    num_adaptation_steps=int(num_steps // 2 * 0.8),
    target_accept_prob=0.9)

chain, is_accepted = tf.function(
    lambda current_state: tfp.mcmc.sample_chain(
        current_state=current_state,
        kernel=hmc,
        num_results=num_steps // 2,
        num_burnin_steps=num_steps // 2,
        trace_fn=lambda _, pkr:
        (pkr.inner_results.inner_results.is_accepted),
        ),
    autograph=False,
    jit_compile=True)(current_state)

accept_rate = tf.reduce_mean(tf.cast(is_accepted, tf.float32))
ess = tf.nest.map_structure(
    lambda c: tfp.mcmc.effective_sample_size(
        c,
        cross_chain_dims=1,
        filter_beyond_positive_pairs=True),
    chain)

r_hat = tf.nest.map_structure(tfp.mcmc.potential_scale_reduction, chain)
hmc_samples = pack_samples(
    tf.nest.pack_sequence_as(target_model.event_shape, chain))
print('Acceptance rate is {}'.format(accept_rate))

Acceptance rate is 0.9008625149726868

ارسم آثارًا للعينة للتحقق من صحة نتائج HMC.

def plot_traces(var_name, samples):
  fig, axes = plt.subplots(1, 2, figsize=(14, 1.5), sharex='col', sharey='col')
  for chain in range(num_chains):
    s = samples.numpy()[:, chain]
    axes[0].plot(s, alpha=0.7)
    sns.kdeplot(s, ax=axes[1], shade=False)
    axes[0].title.set_text("'{}' trace".format(var_name))
    axes[1].title.set_text("'{}' distribution".format(var_name))
    axes[0].set_xlabel('Iteration')

warnings.filterwarnings('ignore')
for var, var_samples in hmc_samples.items():
  plot_traces(var, var_samples)

بي إن جي

بي إن جي

بي إن جي

بي إن جي

أنتجت جميع الأجزاء الخلفية الثلاثة البديلة فترات زمنية موثوقة تشبه بصريًا عينات HMC ، على الرغم من أنها في بعض الأحيان غير مشتتة بسبب تأثير خسارة ELBO ، كما هو شائع في VI.

results['HMC'] = hmc_samples
plot_boxplot(results)

بي إن جي

نتائج إضافية

وظائف التآمر

الدليل السفلي منضم (ELBO)

IAF ، إلى حد بعيد ، البديل الأكبر والأكثر مرونة ، يتقارب مع أعلى دليل منخفض (ELBO).

plot_loss_and_elbo()

بي إن جي

العينات الخلفية

عينات من كل خلفي بديل ، مقارنة مع عينات الحقيقة الأرضية لمؤسسة حمد الطبية (تصور مختلف للعينات الموضحة في مخططات الصندوق).

plot_kdes()

بي إن جي

استنتاج

في Colab ، قمنا ببناء خلفيات بديلة بديلة VI باستخدام توزيعات مشتركة و bijectors متعددة الأجزاء ، وقمنا بتلائمها لتقدير فترات موثوقة للأوزان في نموذج الانحدار على مجموعة بيانات الرادون. بالنسبة لهذا النموذج البسيط ، يتم تنفيذ المؤثرات الخلفية البديلة الأكثر تعبيراً بشكل مشابه للخلفية البديلة ذات المجال المتوسط. ومع ذلك ، يمكن استخدام الأدوات التي أظهرناها لبناء مجموعة واسعة من الخلفيات البديلة المرنة المناسبة لنماذج أكثر تعقيدًا.