יש שאלה? התחבר לקהילה בפורום הביקור של TensorFlow

מחקר הסתברות של TensorFlow: אומדן משתנות

צפה ב- TensorFlow.org הפעל בגוגל קולאב צפה במקור ב- GitHub הורד מחברת

כתבתי מחברת זו כמקרה למחקר על הסתברות TensorFlow. הבעיה שבחרתי לפתור היא הערכת מטריצת משתנות לדגימות של 2-D ממוצע משתנה אקראי גאוסית. לבעיה יש כמה תכונות נחמדות:

  • אם אנו משתמשים ב- Wishart הפוך לפני המשתנות (גישה נפוצה), לבעיה יש פתרון אנליטי, כך שנוכל לבדוק את התוצאות שלנו.
  • הבעיה כוללת דגימה של פרמטר מוגבל, מה שמוסיף מורכבות מעניינת.
  • הפיתרון הכי פשוט הוא לא המהיר ביותר, ולכן יש לבצע עבודת אופטימיזציה.

החלטתי לכתוב את החוויות שלי תוך כדי. לקח לי זמן לעטוף את הראש סביב הנקודות הטובות יותר של TFP, אז המחברת הזו מתחילה בצורה די פשוטה ואז עובדת בהדרגה לתכונות TFP מסובכות יותר. נתקלתי בהרבה בעיות לאורך הדרך, וניסיתי לתפוס את התהליכים שעזרו לי לזהות אותם וגם את הפתרונות לעקיפת הבעיה. ניסיתי לכלול המון פרטים (כולל הרבה בדיקות כדי לוודא שצעדים בודדים נכונים).

למה ללמוד TensorFlow הסתברות?

מצאתי ש- TensorFlow הסתברות מושך את הפרויקט שלי מכמה סיבות:

  • ההסתברות של TensorFlow מאפשרת לכם ליצור טיפוס לפתח מודלים מורכבים באופן אינטראקטיבי במחברת. אתה יכול לפרק את הקוד שלך לחתיכות קטנות שתוכל לבדוק באופן אינטראקטיבי ועם בדיקות יחידה.
  • ברגע שתהיה מוכן להגדיל, תוכל לנצל את כל התשתית שיש לנו כדי להפוך את TensorFlow לרוץ על מספר מעבדים מותאמים במספר מכונות.
  • לבסוף, בעוד שאני מאוד אוהב את סטן, אני מתקשה למדי לבצע ניפוי באגים. אתה צריך לכתוב את כל קוד הדוגמנות שלך בשפה עצמאית שיש לה מעט מאוד כלים המאפשרים לך לתקוע את הקוד שלך, לבדוק מצבי ביניים וכן הלאה.

החיסרון הוא שההסתברות של TensorFlow היא הרבה יותר חדשה מסטאן ומ- PyMC3, ולכן התיעוד הוא עבודה מתמשכת, ויש הרבה פונקציונליות שעוד לא נבנתה. לשמחתי, מצאתי שהבסיס של TFP מוצק, והוא תוכנן בצורה מודולרית המאפשרת להרחיב את הפונקציונליות שלו בצורה די פשוטה. במחברת זו, בנוסף לפתרון המקרה, אראה כמה דרכים להרחיב את ה- TFP.

למי זה מיועד

אני מניח שקוראים מגיעים למחברת זו עם כמה תנאים מוקדמים חשובים. אתה צריך:

נסיון ראשון

הנה הניסיון הראשון שלי לבעיה. ספוילר: הפיתרון שלי לא עובד, וזה ייקח כמה ניסיונות כדי לסדר את הדברים! התהליך אמנם לוקח זמן מה, אך כל ניסיון למטה היה שימושי ללמוד חלק חדש ב- TFP.

הערה אחת: TFP לא מיישמת כעת את התפלגות ה- Wishart ההפוכה (נראה בסוף כיצד לגלגל את Wishart ההפוך שלנו), אז במקום זאת אשנה את הבעיה לזו של אומדן מטריצת דיוק באמצעות Wishart קודם.

import collections
import math
import os
import time

import numpy as np
import pandas as pd
import scipy
import scipy.stats
import matplotlib.pyplot as plt

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

שלב 1: התכנס את התצפיות

הנתונים שלי כאן הם סינתטיים, כך שזה ייראה קצת יותר מסודר מאשר דוגמה בעולם האמיתי. עם זאת, אין שום סיבה שלא תוכל ליצור נתונים סינתטיים משלך.

טיפ : לאחר שהחלטת על צורת המודל שלך, תוכל לבחור כמה ערכי פרמטרים ולהשתמש במודל שבחרת כדי ליצור נתונים סינתטיים. כבדיקת שפיות של היישום שלך, תוכל לוודא שההערכות שלך כוללות את הערכים האמיתיים של הפרמטרים שבחרת. כדי להפוך את מחזור איתור הבאגים / בדיקות שלך למהיר יותר, תוכל לשקול גרסה פשוטה של ​​המודל שלך (למשל להשתמש בפחות מימדים או פחות דוגמאות).

טיפ: הכי קל לעבוד עם התצפיות שלך כמערכי NumPy. דבר חשוב שצריך לשים לב אליו הוא ש- NumPy כברירת מחדל משתמש ב- float64, בעוד ש- TensorFlow כברירת מחדל משתמש ב- float32.

באופן כללי, פעולות TensorFlow מעוניינות שכל הטיעונים יהיו באותו סוג, ועליך לבצע השלכת נתונים מפורשת בכדי לשנות סוגים. אם אתה משתמש בתצפיות float64, יהיה עליך להוסיף הרבה פעולות יציקה. NumPy, לעומת זאת, ידאג ללהק אוטומטית. לפיכך, הרבה יותר קל להמיר את הנתונים המנודלים שלך ל- float32 מאשר לאלץ את TensorFlow להשתמש ב- float64.

בחר כמה ערכי פרמטרים

# We're assuming 2-D data with a known true mean of (0, 0)
true_mean = np.zeros([2], dtype=np.float32)
# We'll make the 2 coordinates correlated
true_cor = np.array([[1.0, 0.9], [0.9, 1.0]], dtype=np.float32)
# And we'll give the 2 coordinates different variances
true_var = np.array([4.0, 1.0], dtype=np.float32)
# Combine the variances and correlations into a covariance matrix
true_cov = np.expand_dims(np.sqrt(true_var), axis=1).dot(
    np.expand_dims(np.sqrt(true_var), axis=1).T) * true_cor
# We'll be working with precision matrices, so we'll go ahead and compute the
# true precision matrix here
true_precision = np.linalg.inv(true_cov)
# Here's our resulting covariance matrix
print(true_cov)
# Verify that it's positive definite, since np.random.multivariate_normal
# complains about it not being positive definite for some reason.
# (Note that I'll be including a lot of sanity checking code in this notebook -
# it's a *huge* help for debugging)
print('eigenvalues: ', np.linalg.eigvals(true_cov))
[[4.  1.8]
 [1.8 1. ]]
eigenvalues:  [4.843075   0.15692513]

צור כמה תצפיות סינתטיות

שים לב ש- TensorFlow Probability משתמש במוסכמה שהממד (ים) הראשוניים של הנתונים שלך מייצגים מדדי מדגם, והמימד / ים הסופיים של הנתונים שלך מייצגים את הממדיות של הדגימות שלך.

כאן אנו רוצים 100 דוגמאות, שכל אחת מהן היא וקטור באורך 2. ניצור מערך my_data עם צורה (100, 2). my_data[i, :] הוא המדגם $ i $ th, והוא וקטור באורך 2.

(זכור לגרום ל- my_data להיות מסוג float32!)

# Set the seed so the results are reproducible.
np.random.seed(123)

# Now generate some observations of our random variable.
# (Note that I'm suppressing a bunch of spurious about the covariance matrix
# not being positive semidefinite via check_valid='ignore' because it really is
# positive definite!)
my_data = np.random.multivariate_normal(
    mean=true_mean, cov=true_cov, size=100,
    check_valid='ignore').astype(np.float32)
my_data.shape
(100, 2)

שפיות לבדוק את התצפיות

מקור פוטנציאלי אחד של באגים הוא לבלגן את הנתונים הסינתטיים שלך! בואו נעשה כמה בדיקות פשוטות.

# Do a scatter plot of the observations to make sure they look like what we
# expect (higher variance on the x-axis, y values strongly correlated with x)
plt.scatter(my_data[:, 0], my_data[:, 1], alpha=0.75)
plt.show()

png

print('mean of observations:', np.mean(my_data, axis=0))
print('true mean:', true_mean)
mean of observations: [-0.24009615 -0.16638893]
true mean: [0. 0.]
print('covariance of observations:\n', np.cov(my_data, rowvar=False))
print('true covariance:\n', true_cov)
covariance of observations:
 [[3.95307734 1.68718486]
 [1.68718486 0.94910269]]
true covariance:
 [[4.  1.8]
 [1.8 1. ]]

בסדר, הדוגמאות שלנו נראות סבירות. השלב הבא.

שלב 2: יישם את פונקציית הסבירות ב- NumPy

הדבר העיקרי שנצטרך לכתוב כדי לבצע את דגימת ה- MCMC שלנו בהסתברות TF הוא פונקציית סבירות יומן. באופן כללי זה קצת יותר מסובך לכתוב TF מאשר NumPy, ולכן אני מוצא מועיל לבצע יישום ראשוני ב- NumPy. אני מתכוון לפצל את פונקציית הסבירות לשתי חלקים, פונקציית סבירות נתונים המתאימה ל- $ P (נתונים | פרמטרים) $ ופונקציה של סבירות קודמת המתאימה ל- $ P (פרמטרים) $.

שים לב שפונקציות NumPy אלה לא צריכות להיות מותאמות במיוחד / וקטוריות שכן המטרה היא רק ליצור כמה ערכים לבדיקה. תקינות היא השיקול המרכזי!

ראשית נבצע את חתיכת הסבירות ביומן הנתונים. זה די פשוט. הדבר היחיד שיש לזכור הוא שאנחנו נעבוד עם מטריצות מדויקות, ולכן נפרמט בהתאם.

def log_lik_data_numpy(precision, data):
  # np.linalg.inv is a really inefficient way to get the covariance matrix, but
  # remember we don't care about speed here
  cov = np.linalg.inv(precision)
  rv = scipy.stats.multivariate_normal(true_mean, cov)
  return np.sum(rv.logpdf(data))

# test case: compute the log likelihood of the data given the true parameters
log_lik_data_numpy(true_precision, my_data)
-280.81822950593767

אנו נשתמש ב- Wishart לפני המטריצה ​​המדויקת מכיוון שיש פיתרון אנליטי לגב האחורי (ראה טבלה שימושית של ויקיפדיה של תקדימים מצומדים ).

להפצת Wishart יש שני פרמטרים:

  • מספר דרגות החופש (שכותרתו $ \ nu $ בוויקיפדיה)
  • מטריצת קנה מידה (שכותרתה $ V $ בוויקיפדיה)

הממוצע להפצת Wishart עם פרמטרים $ \ nu, V $ הוא $ E [W] = \ nu V $, והשונות היא $ \ text {Var} (W_ {ij}) = \ nu (v_ {ij} ^ 2 + v_ {ii} v_ {jj}) $

אינטואיציה שימושית כלשהי: ניתן ליצור מדגם Wishart על ידי יצירת $ \ nu $ עצמאי שואב $ x_1 \ ldots x _ {\ nu} $ ממשתנה אקראי נורמלי רב משתני עם ממוצע 0 ומשתנה $ V $ ואז יוצר את הסכום $ W = \ sum_ {i = 1} ^ {\ nu} x_i x_i ^ T $.

אם אתה מגדיל מחדש את דוגמאות ה- Wishart על ידי חלוקתן ב- $ \ nu $, תקבל את מטריצת הדו-משתנות לדוגמה של $ x_i $. מטריצת מדגם משתנות לדוגמא זו צריכה להיות נוטה לכיוון $ V $ כאשר $ \ nu $ עולה. כאשר $ \ nu $ קטן, יש הרבה וריאציות במטריצת המשתנות המדוגמת, כך שערכים קטנים של $ \ nu $ תואמים לקדימות חלשות יותר וערכים גדולים של $ \ nu $ תואמים לקדימות חזקות יותר. שים לב ש- $ \ nu $ חייב להיות גדול לפחות כמימד החלל שאתה דגימה, אחרת תיצור מטריצות יחיד.

נשתמש ב- $ \ nu = 3 $ כך שיהיה לנו חלש קודם, וניקח $ V = \ frac {1} {\ nu} I $ אשר ימשוך את אומדן השונות שלנו לעבר הזהות (כזכור שהממוצע הוא $ \ nu V $).

PRIOR_DF = 3
PRIOR_SCALE = np.eye(2, dtype=np.float32) / PRIOR_DF

def log_lik_prior_numpy(precision):
  rv = scipy.stats.wishart(df=PRIOR_DF, scale=PRIOR_SCALE)
  return rv.logpdf(precision)

# test case: compute the prior for the true parameters
log_lik_prior_numpy(true_precision)
-9.103606346649766

התפלגות Wishart היא הצמידה הקודמת לאמידת מטריצת הדיוק של נורמל רב משתני עם ממוצע ידוע $ \ mu $.

נניח שהפרמטרים הקודמים של Wishart הם $ \ nu, V $ ושיש לנו תצפיות של $ n $ על הנורמלי הרב-משתני שלנו, $ x_1, \ ldots, x_n $. הפרמטרים האחוריים הם $ n + \ nu, \ left (V ^ {- 1} + \ sum_ {i = 1} ^ n (x_i- \ mu) (x_i- \ mu) ^ T \ right) ^ {- 1 } $.

n = my_data.shape[0]
nu_prior = PRIOR_DF
v_prior = PRIOR_SCALE
nu_posterior = nu_prior + n
v_posterior = np.linalg.inv(np.linalg.inv(v_prior) + my_data.T.dot(my_data))
posterior_mean = nu_posterior * v_posterior
v_post_diag = np.expand_dims(np.diag(v_posterior), axis=1)
posterior_sd = np.sqrt(nu_posterior *
                       (v_posterior ** 2.0 + v_post_diag.dot(v_post_diag.T)))

עלילה מהירה של האחוריים והערכים האמיתיים. שימו לב שהאחוריים קרובים לאחורי המדגם אך הם הצטמצמו מעט לעבר הזהות. שים לב גם שהערכים האמיתיים רחוקים למדי ממצב האחורי - ככל הנראה הסיבה לכך היא שלפני כן אין התאמה טובה מאוד לנתונים שלנו. בבעיה אמיתית ככל הנראה נשתפר טוב יותר עם משהו כמו ווישארט הפוך בקנה מידה לפני השונות (ראו, למשל, הפרשנות של אנדרו גלמן בנושא), אבל אז לא יהיה לנו אחורי אנליטי נחמד.

sample_precision = np.linalg.inv(np.cov(my_data, rowvar=False, bias=False))
fig, axes = plt.subplots(2, 2)
fig.set_size_inches(10, 10)
for i in range(2):
  for j in range(2):
    ax = axes[i, j]
    loc = posterior_mean[i, j]
    scale = posterior_sd[i, j]
    xmin = loc - 3.0 * scale
    xmax = loc + 3.0 * scale
    x = np.linspace(xmin, xmax, 1000)
    y = scipy.stats.norm.pdf(x, loc=loc, scale=scale)
    ax.plot(x, y)
    ax.axvline(true_precision[i, j], color='red', label='True precision')
    ax.axvline(sample_precision[i, j], color='red', linestyle=':', label='Sample precision')
    ax.set_title('precision[%d, %d]' % (i, j))
plt.legend()
plt.show()

png

שלב 3: יישם את פונקציית הסבירות ב- TensorFlow

ספוילר: הניסיון הראשון שלנו לא יעבוד; נדבר על למה להלן.

טיפ : השתמש במצב להוט של TensorFlow בעת פיתוח פונקציות הסבירות שלך. מצב להוט גורם ל TF להתנהג יותר כמו NumPy - הכל מבוצע באופן מיידי, כך שתוכל לבצע ניפוי באגים באופן אינטראקטיבי במקום שתצטרך להשתמש ב- Session.run() . ראה את ההערות כאן .

ראשוני: שיעורי הפצה

ל- TFP יש אוסף של מחלקות הפצה בהן נשתמש כדי ליצור את ההסתברויות ביומן שלנו. דבר אחד שיש לציין הוא שהשיעורים הללו עובדים עם עשרות דוגמאות ולא רק דוגמאות בודדות - זה מאפשר וקטורציה ומהירות מהירות קשורות.

הפצה יכולה לעבוד עם טנסור של דגימות בשתי דרכים שונות. הפשוטה ביותר להמחיש את שתי הדרכים הללו באמצעות דוגמה קונקרטית הכוללת התפלגות עם פרמטר סקלרי יחיד. אשתמש בהפצת פואסון , הכוללת פרמטר rate .

  • אם אנו יוצרים Poisson עם ערך יחיד לפרמטר rate , קריאה לשיטת ה- sample() מחזירה ערך יחיד. ערך זה נקרא event , ובמקרה זה האירועים כולם סקלריים.
  • אם אנו יוצרים Poisson עם טנסור ערכים עבור פרמטר rate , קריאה לשיטת ה- sample() מחזירה כעת ערכים מרובים, אחד לכל ערך במתח הקצב. האובייקט פועל כאוסף של פויסונים עצמאיים, שלכל אחד מהם קצב משלו, וכל אחד מהערכים שהוחזרו על ידי קריאה sample() תואם את אחד מאותם פויסונים. אוסף זה של אירועים עצמאיים אך לא מופצים זהה נקרא batch .
  • השיטה sample() לוקחת פרמטר sample_shape אשר ברירת מחדל היא tuple ריק. העברת ערך שאינו ריק עבור sample_shape מביאה לדגימת החזרת אצוות מרובות. אוסף קבוצות זה נקרא sample .

שיטת log_prob() של הפצה צורכת נתונים באופן המקביל לאופן שבו sample() מייצר אותם. log_prob() מחזיר הסתברויות לדוגמאות, כלומר עבור קבוצות אירועים מרובות ובלתי תלויות.

  • אם יש לנו אובייקט ה- Poisson שלנו שנוצר rate סקלרי, כל אצווה היא סקלרית, ואם נעביר טנזור של דגימות, נוציא טנסור באותו גודל של הסתברויות יומן.
  • אם יש לנו אובייקט ה- Poisson שלנו שנוצר עם טנסור של צורה T של ערכי rate , כל אצווה היא טנזור של צורה T אם נעביר טנסור של דגימות של צורה D, T, נקבל טנסור של הסתברויות יומן של צורה D, T.

להלן כמה דוגמאות הממחישות מקרים אלה. ראה מחברת זו לקבלת הדרכה מפורטת יותר על אירועים, קבוצות וצורות.

# case 1: get log probabilities for a vector of iid draws from a single
# normal distribution
norm1 = tfd.Normal(loc=0., scale=1.)
probs1 = norm1.log_prob(tf.constant([1., 0.5, 0.]))

# case 2: get log probabilities for a vector of independent draws from
# multiple normal distributions with different parameters.  Note the vector
# values for loc and scale in the Normal constructor.
norm2 = tfd.Normal(loc=[0., 2., 4.], scale=[1., 1., 1.])
probs2 = norm2.log_prob(tf.constant([1., 0.5, 0.]))

print('iid draws from a single normal:', probs1.numpy())
print('draws from a batch of normals:', probs2.numpy())
iid draws from a single normal: [-1.4189385 -1.0439385 -0.9189385]
draws from a batch of normals: [-1.4189385 -2.0439386 -8.918939 ]

סבירות יומן נתונים

ראשית נבצע את הפונקציה של סבירות יומן הנתונים.

VALIDATE_ARGS = True
ALLOW_NAN_STATS = False

הבדל מרכזי אחד ממקרה NumPy הוא שתפקוד הסבירות שלנו של TensorFlow יצטרך לטפל בקטורים של מטריצות דיוק ולא רק במטריצות בודדות. וקטורים של פרמטרים ישמשו כאשר אנו מדגמים ממספר שרשראות.

ניצור אובייקט הפצה שעובד עם קבוצה של מטריצות דיוק (כלומר מטריצה ​​אחת לכל שרשרת).

בעת חישוב ההסתברויות של יומני הנתונים שלנו, נצטרך לשכפל את הנתונים שלנו באופן זהה לפרמטרים שלנו כך שיהיה עותק אחד לכל משתנה אצווה. צורת הנתונים המשוכפלים שלנו תצטרך להיות כדלקמן:

[sample shape, batch shape, event shape]

במקרה שלנו, צורת האירוע היא 2 (מכיוון שאנחנו עובדים עם גאוסים דו-ממדיים). צורת המדגם היא 100, מכיוון שיש לנו 100 דגימות. צורת האצווה תהיה רק ​​מספר המטריצות המדויקות איתן אנו עובדים. זה בזבזני לשכפל את הנתונים בכל פעם שאנחנו קוראים לפונקציית הסבירות, ולכן נשכפל את הנתונים מראש ונעביר בגרסה המשוכפלת.

שים לב שמדובר ביישום לא יעיל: MultivariateNormalFullCovariance הוא יקר ביחס לכמה אלטרנטיבות שנדבר עליהן בחלק האופטימיזציה.

def log_lik_data(precisions, replicated_data):
  n = tf.shape(precisions)[0]  # number of precision matrices
  # We're estimating a precision matrix; we have to invert to get log
  # probabilities.  Cholesky inversion should be relatively efficient,
  # but as we'll see later, it's even better if we can avoid doing the Cholesky
  # decomposition altogether.
  precisions_cholesky = tf.linalg.cholesky(precisions)
  covariances = tf.linalg.cholesky_solve(
      precisions_cholesky, tf.linalg.eye(2, batch_shape=[n]))
  rv_data = tfd.MultivariateNormalFullCovariance(
      loc=tf.zeros([n, 2]),
      covariance_matrix=covariances,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)

  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# For our test, we'll use a tensor of 2 precision matrices.
# We'll need to replicate our data for the likelihood function.
# Remember, TFP wants the data to be structured so that the sample dimensions
# are first (100 here), then the batch dimensions (2 here because we have 2
# precision matrices), then the event dimensions (2 because we have 2-D
# Gaussian data).  We'll need to add a middle dimension for the batch using
# expand_dims, and then we'll need to create 2 replicates in this new dimension
# using tile.
n = 2
replicated_data = np.tile(np.expand_dims(my_data, axis=1), reps=[1, 2, 1])
print(replicated_data.shape)
(100, 2, 2)

טיפ: דבר אחד שמצאתי שהוא מועיל ביותר הוא לכתוב מעט בדיקות שפיות של פונקציות TensorFlow שלי. זה ממש קל לבלגן את הווקטורציה ב- TF, כך שיש פונקציות פשוטות יותר של NumPy מסביב זו דרך נהדרת לאמת את פלט ה- TF. תחשוב על אלה כעל בדיקות יחידה קטנות.

# check against the numpy implementation
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
n = precisions.shape[0]
lik_tf = log_lik_data(precisions, replicated_data=replicated_data).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.8182

סיכוי יומן קודם

הקדם יותר קל יותר מכיוון שאיננו צריכים לדאוג לשכפול נתונים.

@tf.function(autograph=False)
def log_lik_prior(precisions):
  rv_precision = tfd.WishartTriL(
      df=PRIOR_DF,
      scale_tril=tf.linalg.cholesky(PRIOR_SCALE),
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)
  return rv_precision.log_prob(precisions)
# check against the numpy implementation
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
n = precisions.shape[0]
lik_tf = log_lik_prior(precisions).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]))
  print('tensorflow:', lik_tf[i])
0
numpy: -2.2351873809649625
tensorflow: -2.2351875
1
numpy: -9.103606346649766
tensorflow: -9.103608

בנה את פונקציית הסבירות היומית המשותפת

פונקציית הסבירות של יומן הנתונים לעיל תלויה בתצפיות שלנו, אך לדגימה לא יהיו כאלה. אנו יכולים להיפטר מהתלות מבלי להשתמש במשתנה גלובלי באמצעות [סגירה] (https://en.wikipedia.org/wiki/Closure_ (מחשב_תכנות). הסגירות כוללות פונקציה חיצונית הבונה סביבה המכילה משתנים הנחוצים על ידי פונקציה פנימית.

def get_log_lik(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik(precision):
    return log_lik_data(precision, replicated_data) + log_lik_prior(precision)

  return _log_lik

שלב 4: מדגם

בסדר, הגיע הזמן לדגום! כדי לפשט את העניין פשוט נשתמש בשרשרת אחת ונשתמש במטריצת הזהות כנקודת ההתחלה. נעשה דברים בזהירות רבה יותר בהמשך.

שוב, זה לא יעבוד - נקבל חריג.

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  init_precision = tf.expand_dims(tf.eye(2), axis=0)

  # Use expand_dims because we want to pass in a tensor of starting values
  log_lik_fn = get_log_lik(my_data, n_chains=1)

  # we'll just do a few steps here
  num_results = 10
  num_burnin_steps = 10
  states = tfp.mcmc.sample_chain(
     num_results=num_results,
     num_burnin_steps=num_burnin_steps,
     current_state=[
         init_precision,
     ],
     kernel=tfp.mcmc.HamiltonianMonteCarlo(
         target_log_prob_fn=log_lik_fn,
         step_size=0.1,
         num_leapfrog_steps=3),
     trace_fn=None,
     seed=123)
  return states

try:
  states = sample()
except Exception as e:
  # shorten the giant stack trace
  lines = str(e).split('\n')
  print('\n'.join(lines[:5]+['...']+lines[-3:]))
Cholesky decomposition was not successful. The input might not be valid.
     [[{ {node mcmc_sample_chain/trace_scan/while/body/_79/smart_for_loop/while/body/_371/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/body/_537/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/StatefulPartitionedCall/Cholesky} }]] [Op:__inference_sample_2849]

Function call stack:
sample
...
Function call stack:
sample

זיהוי הבעיה

InvalidArgumentError (see above for traceback): Cholesky decomposition was not successful. The input might not be valid. זה לא מועיל במיוחד. בואו נראה אם ​​נוכל לברר עוד על מה שקרה.

  • נדפיס את הפרמטרים עבור כל שלב כדי שנוכל לראות את הערך שעבורו הדברים נכשלים
  • נוסיף כמה קביעות כדי להגן מפני בעיות ספציפיות.

קביעות הן מסובכות מכיוון שהן פעולות של TensorFlow, ועלינו לדאוג שהן יבוצעו ולא יעברו אופטימיזציה מהגרף. כדאי לקרוא סקירה זו של ניפוי באגים של TensorFlow אם אינך מכיר קביעות TF. אתה יכול לכפות במפורש טענות לביצוע באמצעות tf.control_dependencies (ראה את ההערות בקוד למטה).

לפונקציית Print המקורית של TensorFlow יש התנהגות זהה לטענות - זו פעולה, ועליך להקפיד על ביצועה. Print גורמת לכאבי ראש נוספים כשאנחנו עובדים במחברת: הפלט שלה נשלח ל stderr , ו- stderr לא מוצג בתא. נשתמש בטריק כאן: במקום להשתמש ב- tf.Print , ניצור פעולת הדפסה משלנו TensorFlow באמצעות tf.pyfunc . כמו בטענות, עלינו לוודא שהשיטה שלנו מבוצעת.

def get_log_lik_verbose(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  def _log_lik(precisions):
    # An internal method we'll make into a TensorFlow operation via tf.py_func
    def _print_precisions(precisions):
      print('precisions:\n', precisions)
      return False  # operations must return something!
    # Turn our method into a TensorFlow operation
    print_op = tf.compat.v1.py_func(_print_precisions, [precisions], tf.bool)

    # Assertions are also operations, and some care needs to be taken to ensure
    # that they're executed
    assert_op = tf.assert_equal(
        precisions, tf.linalg.matrix_transpose(precisions),
        message='not symmetrical', summarize=4, name='symmetry_check')

    # The control_dependencies statement forces its arguments to be executed
    # before subsequent operations
    with tf.control_dependencies([print_op, assert_op]):
      return (log_lik_data(precisions, replicated_data) +
              log_lik_prior(precisions))

  return _log_lik
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  init_precision = tf.eye(2)[tf.newaxis, ...]
  log_lik_fn = get_log_lik_verbose(my_data)
  # we'll just do a few steps here
  num_results = 10
  num_burnin_steps = 10
  states = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          init_precision,
      ],
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=log_lik_fn,
          step_size=0.1,
          num_leapfrog_steps=3),
      trace_fn=None,
      seed=123)

try:
  states = sample()
except Exception as e:
  # shorten the giant stack trace
  lines = str(e).split('\n')
  print('\n'.join(lines[:5]+['...']+lines[-3:]))
precisions:
 [[[1. 0.]
  [0. 1.]]]
precisions:
 [[[1. 0.]
  [0. 1.]]]
precisions:
 [[[ 0.24315196 -0.2761638 ]
  [-0.33882257  0.8622    ]]]
 assertion failed: [not symmetrical] [Condition x == y did not hold element-wise:] [x (leapfrog_integrate_one_step/add:0) = ] [[[0.243151963 -0.276163787][-0.338822573 0.8622]]] [y (leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/matrix_transpose/transpose:0) = ] [[[0.243151963 -0.338822573][-0.276163787 0.8622]]]
     [[{ {node mcmc_sample_chain/trace_scan/while/body/_96/smart_for_loop/while/body/_381/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/body/_503/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/symmetry_check_1/Assert/AssertGuard/else/_577/Assert} }]] [Op:__inference_sample_4837]

Function call stack:
sample
...
Function call stack:
sample

מדוע זה נכשל

ערך הפרמטר החדש הראשון שמנסה הדגימה הוא מטריצה ​​א-סימטרית. זה גורם לפירוק Cholesky להיכשל, מכיוון שהוא מוגדר רק למטריצות סימטריות (וחיוביות).

הבעיה כאן היא שהפרמטר שמעניין אותנו הוא מטריצת דיוק, ומטריצות דיוק חייבות להיות אמיתיות, סימטריות וחיוביות. הדוגם אינו יודע שום דבר על אילוץ זה (למעט אפשרי דרך הדרגתיות), כך שייתכן לחלוטין שהדוגם יציע ערך לא חוקי, מה שמוביל לחריג, במיוחד אם גודל הצעדים גדול.

בעזרת הדגימה של המילטון מונטה קרלו, אנו עשויים להיות מסוגלים לעקוף את הבעיה על ידי שימוש בגודל צעד קטן מאוד, מכיוון ששיפוע צריך להרחיק את הפרמטרים מאזורים לא חוקיים, אך גודל צעדים קטן פירושו התכנסות איטית. עם דוגמנית מטרופוליס-הייסטינגס, שאינה יודעת דבר על שיפועים, נגזר עלינו.

גרסה 2: בדיקה חוזרת לפרמטרים לא מוגבלים

יש פיתרון פשוט לבעיה לעיל: אנו יכולים לשנות את המודל שלנו כך שהפרמטרים החדשים כבר לא יהיו אילוצים אלה. TFP מספקת סט כלים שימושי - bijectors - לשם כך בדיוק.

רפרמטריזציה עם מצמידים

מטריצת הדיוק שלנו חייבת להיות אמיתית וסימטרית; אנו רוצים פרמטריזציה חלופית שאין בה אילוצים אלה. נקודת התחלה היא פקטוריזציה Cholesky של מטריצת הדיוק. גורמי ה- Cholesky עדיין מוגבלים - הם משולשים תחתונים, והאלמנטים האלכסוניים שלהם חייבים להיות חיוביים. עם זאת, אם ניקח את יומן האלכסונים של הגורם Cholesky, היומנים כבר אינם מוגבלים להיות חיוביים, ואז אם אנו משטחים את החלק המשולש התחתון לווקטור 1-D, אין לנו עוד את האילוץ המשולש התחתון. . התוצאה במקרה שלנו תהיה וקטור באורך 3 ללא אילוצים.

( במדריך סטן יש פרק נהדר לשימוש בתמורות כדי להסיר סוגים שונים של אילוצים בפרמטרים.)

לשינוי מתאם זה אין השפעה מועטה על פונקציית הסבירות ביומן הנתונים שלנו - עלינו רק להפוך את השינוי שלנו כדי להחזיר את מטריצת הדיוק - אך ההשפעה על הקודם מורכבת יותר. ציינו שההסתברות למטריצת דיוק נתונה ניתנת על ידי התפלגות Wishart; מה ההסתברות למטריצה ​​שהפכה אותנו?

נזכיר שאם אנו מיישמים פונקציה מונוטונית $ g $ על משתנה אקראי 1-D $ X $, $ Y = g (X) $, הצפיפות עבור $ Y $ ניתנת על ידי

$$ f_Y(y) = | \frac{d}{dy}(g^{-1}(y)) | f_X(g^{-1}(y)) $$

הנגזרת של $ g ^ {- 1} $ מחזיקה את האופן שבו $ g $ משנה נפחים מקומיים. עבור משתנים אקראיים בעלי ממד גבוה יותר, הגורם המתקן הוא הערך המוחלט של הקובע של היעקוביאן של $ g ^ {- 1} $ (ראה כאן ).

נצטרך להוסיף יעקוביאנית של ההופך ההפוך לפונקציית הסבירות הקודמת שלנו. לשמחתנו, שיעור ה- Bijector של TFP יכול לטפל בזה עבורנו.

מחלקת Bijector משמשת לייצוג פונקציות חלקות הפיכות המשמשות לשינוי משתנים בפונקציות צפיפות ההסתברות. לכולנים של Bijectors יש שיטת forward() המבצעת טרנספורמציה, שיטה inverse() אותה, ושיטות forward_log_det_jacobian() ו- inverse_log_det_jacobian() המספקות את התיקונים היעקוביים שאנו זקוקים לנו כאשר אנו משנים קובץ PDF מחדש.

TFP מספק אוסף של תכשיטים שימושיים שנוכל לשלב באמצעות קומפוזיציה באמצעות מפעיל Chain ליצירת טרנספורמציות די מורכבות. במקרה שלנו נרכיב את שלושת המקשרים הבאים (הפעולות בשרשרת מבוצעות מימין לשמאל):

  1. השלב הראשון בתמורה שלנו הוא לבצע פקטורציה של Cholesky על מטריצת הדיוק. אין לכך כיתת ביג'קטור; עם זאת, המקשרCholeskyOuterProduct לוקח תוצר של 2 גורמים Cholesky. אנו יכולים להשתמש ההופכי של אותה פעולה באמצעות Invert המפעיל.
  2. השלב הבא הוא לקחת את יומן האלמנטים האלכסוניים של הגורם Cholesky. אנו משיגים זאת באמצעות המצורף TransformDiagonal וההפוך של המצורף Exp .
  3. לבסוף אנו משטחים את החלק המשולש התחתון של המטריצה ​​לווקטור באמצעות ההפך של FillTriangular FillTriangular.
# Our transform has 3 stages that we chain together via composition:
precision_to_unconstrained = tfb.Chain([
    # step 3: flatten the lower triangular portion of the matrix
    tfb.Invert(tfb.FillTriangular(validate_args=VALIDATE_ARGS)),
    # step 2: take the log of the diagonals    
    tfb.TransformDiagonal(tfb.Invert(tfb.Exp(validate_args=VALIDATE_ARGS))),
    # step 1: decompose the precision matrix into its Cholesky factors
    tfb.Invert(tfb.CholeskyOuterProduct(validate_args=VALIDATE_ARGS)),
])
# sanity checks
m = tf.constant([[1., 2.], [2., 8.]])
m_fwd = precision_to_unconstrained.forward(m)
m_inv = precision_to_unconstrained.inverse(m_fwd)

# bijectors handle tensors of values, too!
m2 = tf.stack([m, tf.eye(2)])
m2_fwd = precision_to_unconstrained.forward(m2)
m2_inv = precision_to_unconstrained.inverse(m2_fwd)

print('single input:')
print('m:\n', m.numpy())
print('precision_to_unconstrained(m):\n', m_fwd.numpy())
print('inverse(precision_to_unconstrained(m)):\n', m_inv.numpy())
print()

print('tensor of inputs:')
print('m2:\n', m2.numpy())
print('precision_to_unconstrained(m2):\n', m2_fwd.numpy())
print('inverse(precision_to_unconstrained(m2)):\n', m2_inv.numpy())
single input:
m:
 [[1. 2.]
 [2. 8.]]
precision_to_unconstrained(m):
 [0.6931472 2.        0.       ]
inverse(precision_to_unconstrained(m)):
 [[1. 2.]
 [2. 8.]]

tensor of inputs:
m2:
 [[[1. 2.]
  [2. 8.]]

 [[1. 0.]
  [0. 1.]]]
precision_to_unconstrained(m2):
 [[0.6931472 2.        0.       ]
 [0.        0.        0.       ]]
inverse(precision_to_unconstrained(m2)):
 [[[1. 2.]
  [2. 8.]]

 [[1. 0.]
  [0. 1.]]]

המחלקה TransformedDistribution אוטומטית את תהליך החלת bijector להפצה וביצוע התיקון log_prob() הדרוש ל- log_prob() . הקודם החדש שלנו הופך להיות:

def log_lik_prior_transformed(transformed_precisions):
  rv_precision = tfd.TransformedDistribution(
      tfd.WishartTriL(
          df=PRIOR_DF,
          scale_tril=tf.linalg.cholesky(PRIOR_SCALE),
          validate_args=VALIDATE_ARGS,
          allow_nan_stats=ALLOW_NAN_STATS),
      bijector=precision_to_unconstrained,
      validate_args=VALIDATE_ARGS)
  return rv_precision.log_prob(transformed_precisions)
# Check against the numpy implementation.  Note that when comparing, we need
# to add in the Jacobian correction.
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
transformed_precisions = precision_to_unconstrained.forward(precisions)
lik_tf = log_lik_prior_transformed(transformed_precisions).numpy()
corrections = precision_to_unconstrained.inverse_log_det_jacobian(
    transformed_precisions, event_ndims=1).numpy()
n = precisions.shape[0]

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]) + corrections[i])
  print('tensorflow:', lik_tf[i])
0
numpy: -0.8488930160357633
tensorflow: -0.84889317
1
numpy: -7.305657151741624
tensorflow: -7.305659

אנחנו רק צריכים להפוך את השינוי על פי הסבירות שלנו ביומן הנתונים:

precision = precision_to_unconstrained.inverse(transformed_precision)

מכיוון שאנו רוצים למעשה פקטוריזציה של Cholesky של מטריצת הדיוק, יהיה יעיל יותר לעשות כאן רק הפוך חלקי. עם זאת, נעזוב את האופטימיזציה למועד מאוחר יותר ונשאיר את החלק ההפוך כתרגיל עבור הקורא.

def log_lik_data_transformed(transformed_precisions, replicated_data):
  # We recover the precision matrix by inverting our bijector.  This is
  # inefficient since we really want the Cholesky decomposition of the
  # precision matrix, and the bijector has that in hand during the inversion,
  # but we'll worry about efficiency later.
  n = tf.shape(transformed_precisions)[0]
  precisions = precision_to_unconstrained.inverse(transformed_precisions)
  precisions_cholesky = tf.linalg.cholesky(precisions)
  covariances = tf.linalg.cholesky_solve(
      precisions_cholesky, tf.linalg.eye(2, batch_shape=[n]))
  rv_data = tfd.MultivariateNormalFullCovariance(
      loc=tf.zeros([n, 2]),
      covariance_matrix=covariances,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)

  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# sanity check
precisions = np.stack([np.eye(2, dtype=np.float32), true_precision])
transformed_precisions = precision_to_unconstrained.forward(precisions)
lik_tf = log_lik_data_transformed(
    transformed_precisions, replicated_data).numpy()

for i in range(precisions.shape[0]):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.8182

שוב אנו עוטפים את הפונקציות החדשות שלנו בסגירה.

def get_log_lik_transformed(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik_transformed(transformed_precisions):
    return (log_lik_data_transformed(transformed_precisions, replicated_data) +
            log_lik_prior_transformed(transformed_precisions))

  return _log_lik_transformed
# make sure everything runs
log_lik_fn = get_log_lik_transformed(my_data)
m = tf.eye(2)[tf.newaxis, ...]
lik = log_lik_fn(precision_to_unconstrained.forward(m)).numpy()
print(lik)
[-431.5611]

דְגִימָה

עכשיו, כשאנחנו לא צריכים לדאוג שהמדגם שלנו יתפוצץ בגלל ערכי פרמטר לא חוקיים, בואו ניצור כמה דוגמאות אמיתיות.

הדוגם עובד עם הגרסה הלא מוגבלת של הפרמטרים שלנו, לכן עלינו להפוך את הערך ההתחלתי שלנו לגרסה הלא מוגבלת שלו. הדוגמאות שאנו מייצרים יהיו גם כולן בצורה בלתי מוגבלת, ולכן עלינו לשנות אותן בחזרה. תכשיטים הם וקטוריים, כך שקל לעשות זאת.

# We'll choose a proper random initial value this time
np.random.seed(123)
initial_value_cholesky = np.array(
    [[0.5 + np.random.uniform(), 0.0],
     [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
    dtype=np.float32)
initial_value =  initial_value_cholesky.dot(
  initial_value_cholesky.T)[np.newaxis, ...]

# The sampler works with unconstrained values, so we'll transform our initial
# value
initial_value_transformed = precision_to_unconstrained.forward(
  initial_value).numpy()
# Sample!
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_transformed(my_data, n_chains=1)

  num_results = 1000
  num_burnin_steps = 1000

  states, is_accepted = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          initial_value_transformed,
      ],
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=log_lik_fn,
          step_size=0.1,
          num_leapfrog_steps=3),
      trace_fn=lambda _, pkr: pkr.is_accepted,
      seed=123)
  # transform samples back to their constrained form
  precision_samples = [precision_to_unconstrained.inverse(s) for s in states]
  return states, precision_samples, is_accepted

states, precision_samples, is_accepted = sample()

בואו נשווה את ממוצע תפוקת הדוגם לממוצע האחורי האנליטי!

print('True posterior mean:\n', posterior_mean)
print('Sample mean:\n', np.mean(np.reshape(precision_samples, [-1, 2, 2]), axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Sample mean:
 [[ 1.4315274  -0.25587553]
 [-0.25587553  0.5740424 ]]

אנחנו כבר רחוק! בואו להבין מדוע. ראשית בואו נסתכל על הדוגמאות שלנו.

np.reshape(precision_samples, [-1, 2, 2])
array([[[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       ...,

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]],

       [[ 1.4315385, -0.2558777],
        [-0.2558777,  0.5740494]]], dtype=float32)

אה אה - נראה שלכולם יש אותו ערך. בואו להבין מדוע.

המשתנה kernel_results_ הוא שם tuple שנותן מידע על הדוגם בכל מדינה. השדה is_accepted הוא המפתח כאן.

# Look at the acceptance for the last 100 samples
print(np.squeeze(is_accepted)[-100:])
print('Fraction of samples accepted:', np.mean(np.squeeze(is_accepted)))
[False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False False False False False False False False False
 False False False False]
Fraction of samples accepted: 0.0

כל הדגימות שלנו נדחו! יש להניח שגודל המדרגות שלנו היה גדול מדי. בחרתי stepsize=0.1 באופן שרירותי בלבד.

גרסה 3: דגימה בגודל צעד אדפטיבי

מכיוון שדגימה עם הבחירה השרירותית שלי בגודל הצעד נכשלה, יש לנו כמה פריטי סדר יום:

  1. ליישם גודל צעד אדפטיבי, ו
  2. לבצע כמה בדיקות התכנסות.

יש קוד דוגמה נחמד ב tensorflow_probability/python/mcmc/hmc.py ליישום גדלי צעדים אדפטיביים. התאמתי את זה בהמשך.

שים לב שיש משפט sess.run() נפרד לכל שלב. זה באמת מועיל לניפוי באגים, מכיוון שהוא מאפשר לנו להוסיף מעט אבחונים לפי שלב במידת הצורך. לדוגמא, אנו יכולים להראות התקדמות מצטברת, זמן כל שלב וכו '.

טיפ: אחת הדרכים הנפוצות ככל הנראה לבלגן את הדגימה שלך היא שהגרף שלך יגדל בלולאה. (הסיבה לסיום הגרף לפני הפעלת ההפעלה היא למנוע בעיות כאלה בלבד). אם לא השתמשת ב- finalize (), בדיקת ניפוי שגיאות שימושית אם הקוד שלך מאט לסריקה היא להדפיס את הגרף. גודל בכל שלב דרך len(mygraph.get_operations()) - אם האורך גדל, כנראה שאתה עושה משהו רע.

אנחנו הולכים להפעיל כאן 3 רשתות עצמאיות. עריכת השוואות בין הרשתות תעזור לנו לבדוק את ההתכנסות.

# The number of chains is determined by the shape of the initial values.
# Here we'll generate 3 chains, so we'll need a tensor of 3 initial values.
N_CHAINS = 3

np.random.seed(123)

initial_values = []
for i in range(N_CHAINS):
  initial_value_cholesky = np.array(
      [[0.5 + np.random.uniform(), 0.0],
       [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
      dtype=np.float32)
  initial_values.append(initial_value_cholesky.dot(initial_value_cholesky.T))
initial_values = np.stack(initial_values)

initial_values_transformed = precision_to_unconstrained.forward(
  initial_values).numpy()
@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_transformed(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=hmc,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states, is_accepted = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values_transformed,
      kernel=adapted_kernel,
      trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
      parallel_iterations=1)
  # transform samples back to their constrained form
  precision_samples = precision_to_unconstrained.inverse(states)
  return states, precision_samples, is_accepted

states, precision_samples, is_accepted = sample()

בדיקה מהירה: שיעור הקבלה שלנו במהלך הדגימה שלנו קרוב ליעד שלנו 0.651.

print(np.mean(is_accepted))
0.6190666666666667

אפילו יותר טוב, ממוצע הדוגמה שלנו וסטיית התקן קרובים למה שאנחנו מצפים מהפתרון האנליטי.

precision_samples_reshaped = np.reshape(precision_samples, [-1, 2, 2])
print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.96426415 -1.6519215 ]
 [-1.6519215   3.8614824 ]]
print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13622096 0.25235635]
 [0.25235635 0.5394968 ]]

בודק התכנסות

באופן כללי לא יהיה לנו פתרון אנליטי לבדוק, ולכן נצטרך לוודא שהדוגם התכנס. בדיקה סטנדרטית אחת היא נתון Gelman-Rubin $ \ hat {R} $, הדורש מספר שרשראות דגימה. $ \ hat {R} $ מודד את מידת השונות (של האמצעים) בין שרשראות עולה על מה שניתן היה לצפות אם השרשראות היו מופצות זהה. ערכים של $ \ hat {R} $ קרוב ל -1 משמשים לציון התכנסות משוערת. ראה את המקור לפרטים.

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0038308 1.0005717]
 [1.0005717 1.0006068]]

ביקורת מודל

אם לא היה לנו פיתרון אנליטי, זה הזמן לעשות ביקורת אמיתית על המודל.

להלן מספר היסטוגרמות מהירות של רכיבי המדגם ביחס לאמת הקרקעית שלנו (באדום). שים לב שהדוגמאות הצטמצמו מערכי המטריצות המדויקות של המדגם לקראת מטריצת הזהות לפני כן.

0319 אאד 280

png

כמה פיזורי זוגות של רכיבים מדויקים מראים שבגלל מבנה המתאם של האחורי, הערכים האחוריים האמיתיים אינם כה סבירים כפי שהם נראים מהשוליים לעיל.

fig, axes = plt.subplots(4, 4)
fig.set_size_inches(12, 12)
for i1 in range(2):
  for j1 in range(2):
    index1 = 2 * i1 + j1
    for i2 in range(2):
      for j2 in range(2):
        index2 = 2 * i2 + j2
        ax = axes[index1, index2]
        ax.scatter(precision_samples_reshaped[:, i1, j1],
                   precision_samples_reshaped[:, i2, j2], alpha=0.1)
        ax.axvline(true_precision[i1, j1], color='red')
        ax.axhline(true_precision[i2, j2], color='red')
        ax.axvline(sample_precision[i1, j1], color='red', linestyle=':')
        ax.axhline(sample_precision[i2, j2], color='red', linestyle=':')
        ax.set_title('(%d, %d) vs (%d, %d)' % (i1, j1, i2, j2))
plt.tight_layout()
plt.show()

png

גרסה 4: דגימה פשוטה יותר של פרמטרים מוגבלים

Bijectors ביצעו דגימה של המטריצה ​​המדויקת בצורה פשוטה, אך הייתה כמות לא מבוטלת של המרה ידנית אל הייצוג הלא מוגבל וממנו. יש דרך קלה יותר!

TransformedTransitionKernel

TransformedTransitionKernel מפשט את התהליך הזה. זה עוטף את הדוגם שלך ומטפל בכל ההמרות. נדרשת כטיעון רשימה של מצמידים שממפים ערכי פרמטר בלתי מוגבלים לערכים מוגבלים. אז הנה אנחנו צריכים את ההופכי של precision_to_unconstrained bijector השתמשנו לעיל. נוכל פשוט להשתמש ב- tfb.Invert(precision_to_unconstrained) , אך זה כרוך בנטילת היפוך של היפוך (TensorFlow אינו מספיק חכם כדי לפשט את tf.Invert(tf.Invert()) ל- tf.Identity()) , אז במקום זאת פשוט אכתוב מצרף חדש.

מגביל תכשיט

# The bijector we need for the TransformedTransitionKernel is the inverse of
# the one we used above
unconstrained_to_precision = tfb.Chain([
    # step 3: take the product of Cholesky factors
    tfb.CholeskyOuterProduct(validate_args=VALIDATE_ARGS),
    # step 2: exponentiate the diagonals    
    tfb.TransformDiagonal(tfb.Exp(validate_args=VALIDATE_ARGS)),
    # step 1: map a vector to a lower triangular matrix
    tfb.FillTriangular(validate_args=VALIDATE_ARGS),
])
# quick sanity check
m = [[1., 2.], [2., 8.]]
m_inv = unconstrained_to_precision.inverse(m).numpy()
m_fwd = unconstrained_to_precision.forward(m_inv).numpy()

print('m:\n', m)
print('unconstrained_to_precision.inverse(m):\n', m_inv)
print('forward(unconstrained_to_precision.inverse(m)):\n', m_fwd)
m:
 [[1.0, 2.0], [2.0, 8.0]]
unconstrained_to_precision.inverse(m):
 [0.6931472 2.        0.       ]
forward(unconstrained_to_precision.inverse(m)):
 [[1. 2.]
 [2. 8.]]

דגימה עם TransformedTransitionKernel

עם TransformedTransitionKernel , אנחנו כבר לא צריכים לבצע טרנספורמציות ידניות של הפרמטרים שלנו. הערכים הראשוניים שלנו והדוגמאות שלנו הם כולם מטריצות מדויקות; אנחנו רק צריכים להעביר את הקישוריות הלא מגבילות שלנו לליבה וזה דואג לכל התמורות.

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  ttk = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=hmc, bijector=unconstrained_to_precision)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=ttk,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values,
      kernel=adapted_kernel,
      trace_fn=None,
      parallel_iterations=1)
  # transform samples back to their constrained form
  return states

precision_samples  = sample()

בודק התכנסות

בדיקת ההתכנסות $ \ hat {R} $ נראית טוב!

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0013582 1.0019467]
 [1.0019467 1.0011805]]

השוואה מול האחורי האנליטי

שוב בואו נבדוק מול האחורי האנליטי.

# The output samples have shape [n_steps, n_chains, 2, 2]
# Flatten them to [n_steps * n_chains, 2, 2] via reshape:
precision_samples_reshaped = np.reshape(precision_samples, [-1, 2, 2])
print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.96687526 -1.6552585 ]
 [-1.6552585   3.867676  ]]
print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13329624 0.24913791]
 [0.24913791 0.53983927]]

אופטימיזציות

עכשיו שיש לנו דברים שרצים מקצה לקצה, בואו נעשה גרסה אופטימלית יותר. מהירות לא משנה יותר מדי לדוגמא זו, אך ברגע שמטריצות גדלות, כמה אופטימיזציות יעשו הבדל גדול.

שיפור מהיר אחד גדול שאנחנו יכולים לעשות הוא לבצע שינויי פרמטרים מבחינת הפירוק של Cholesky. הסיבה היא שתפקוד הסבירות שלנו דורש גם את המשתנות וגם את המטריצות המדויקות. היפוך מטריקס הוא יקר ($ O (n ^ 3) $ למטריצה ​​$ n \ פעמים n $), ואם אנו מבצעים פרמטרציה במונחים של המשתנות או של מטריצת הדיוק, עלינו לבצע היפוך כדי להשיג את האחר.

כזכור, ניתן לפרק מטריצה ​​סימטרית אמיתית, חיובית-מובהקת $ M $ למוצר של הצורה $ M = LL ^ T $ כאשר המטריצה ​​$ L $ משולשת תחתונה ובעלת אלכסונים חיוביים. בהתחשב בפירוק Cholesky של $ M $, אנו יכולים להשיג ביעילות רבה יותר הן $ M $ (תוצר של מטריצה ​​משולשת תחתונה עליונה) והן $ M ^ {- 1} $ (באמצעות החלפת גב). הפקטוריזציה של Cholesky עצמה אינה זולה לחישוב, אך אם אנו מפרמטים במונחים של גורמי Cholesky, עלינו לחשב רק את הפקטוריזציה של Choleksy של ערכי הפרמטרים הראשוניים.

שימוש בפירוק Cholesky של מטריצת השונות

ל- TFP יש גרסה של ההתפלגות הנורמלית הרב-משתנית, MultivariateNormalTriL , שעוברת פרמטרים במונחים של הגורם Cholesky של מטריצת השונות. אז אם היינו מבצעים פרמטרים במונחים של הגורם Cholesky של מטריצת השונות, היינו יכולים לחשב את סבירות יומן הנתונים ביעילות. האתגר הוא בחישוב הסבירות הקודמת ביומן ביעילות דומה.

אם הייתה לנו גרסה של הפצת Wishart ההפוכה שעבדה עם גורמי Cholesky של דגימות, היינו מסודרים. אוי ואבוי, אנחנו לא. (הצוות יקבל בברכה הגשות קוד!) כחלופה, אנו יכולים להשתמש בגרסת הפצת Wishart שעובדת עם גורמי Cholesky של דוגמאות יחד עם שרשרת bijectors.

כרגע חסרים לנו כמה מקבצים מניות כדי להפוך את העניינים ליעילים באמת, אבל אני רוצה להראות את התהליך כתרגיל והמחשה מועילה לעוצמתם של המצורפים של TFP.

הפצה של Wishart הפועלת על גורמי Cholesky

להפצת Wishart יש דגל שימושי, input_output_cholesky , המציין כי מטריצות הקלט והפלט צריכות להיות גורמי Cholesky. זה יותר יעיל ומספרי לעבוד עם הגורמים Cholesky מאשר מטריצות מלאות, ולכן זה רצוי. נקודה חשובה בנוגע לסמנטיקה של הדגל: זוהי רק אינדיקציה לכך שייצוג הקלט והפלט להפצה צריך להשתנות - זה לא מצביע על תיקון מלא מחדש של ההתפלגות, אשר יהיה כרוך בתיקון יעקובי ל log_prob() פוּנקצִיָה. אנחנו באמת רוצים לעשות את הפירמטריזציה המלאה הזו, אז נבנה את ההפצה שלנו.

# An optimized Wishart distribution that has been transformed to operate on
# Cholesky factors instead of full matrices.  Note that we gain a modest
# additional speedup by specifying the Cholesky factor of the scale matrix
# (i.e. by passing in the scale_tril parameter instead of scale).

class CholeskyWishart(tfd.TransformedDistribution):
  """Wishart distribution reparameterized to use Cholesky factors."""
  def __init__(self,
      df,
      scale_tril,
      validate_args=False,
      allow_nan_stats=True,
      name='CholeskyWishart'):
    # Wishart has a bunch of methods that we want to support but not
    # implement.  We'll subclass TransformedDistribution here to take care of
    # those.  We'll override the few for which speed is critical and implement
    # them with a separate Wishart for which input_output_cholesky=True
    super(CholeskyWishart, self).__init__(
        distribution=tfd.WishartTriL(
            df=df,
            scale_tril=scale_tril,
            input_output_cholesky=False,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats),
        bijector=tfb.Invert(tfb.CholeskyOuterProduct()),
        validate_args=validate_args,
        name=name
    )
    # Here's the Cholesky distribution we'll use for log_prob() and sample()
    self.cholesky = tfd.WishartTriL(
        df=df,
        scale_tril=scale_tril,
        input_output_cholesky=True,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats)

  def _log_prob(self, x):
    return (self.cholesky.log_prob(x) +
            self.bijector.inverse_log_det_jacobian(x, event_ndims=2))

  def _sample_n(self, n, seed=None):
    return self.cholesky._sample_n(n, seed)
# some checks
PRIOR_SCALE_CHOLESKY = np.linalg.cholesky(PRIOR_SCALE)

@tf.function(autograph=False)
def compute_log_prob(m):
  w_transformed = tfd.TransformedDistribution(
      tfd.WishartTriL(df=PRIOR_DF, scale_tril=PRIOR_SCALE_CHOLESKY),
      bijector=tfb.Invert(tfb.CholeskyOuterProduct()))
  w_optimized = CholeskyWishart(
      df=PRIOR_DF, scale_tril=PRIOR_SCALE_CHOLESKY)
  log_prob_transformed = w_transformed.log_prob(m)
  log_prob_optimized = w_optimized.log_prob(m)
  return log_prob_transformed, log_prob_optimized

for matrix in [np.eye(2, dtype=np.float32),
               np.array([[1., 0.], [2., 8.]], dtype=np.float32)]:
  log_prob_transformed, log_prob_optimized = [
      t.numpy() for t in compute_log_prob(matrix)]
  print('Transformed Wishart:', log_prob_transformed)
  print('Optimized Wishart', log_prob_optimized)
Transformed Wishart: -0.84889317
Optimized Wishart -0.84889317
Transformed Wishart: -99.269455
Optimized Wishart -99.269455

בניית התפלגות Wishart הפוכה

יש לנו את מטריצת השונות המשותפת שלנו $ C $ מפורקת ל- $ C = LL ^ T $ כאשר $ L $ נמוך משולש ובעל אלכסון חיובי. אנו רוצים לדעת את ההסתברות ל- $ L $ בהתחשב בכך ש- $ C \ sim W ^ {- 1} (\ nu, V) $ כאשר $ W ^ {- 1} $ היא התפלגות ה- Wishart ההפוכה.

להתפלגות Wishart ההפוכה יש את התכונה שאם $ C \ sim W ^ {- 1} (\ nu, V) $, אז מטריצת הדיוק $ C ^ {- 1} \ sim W (\ nu, V ^ {- 1 }) $. כך שנוכל לקבל את ההסתברות של $ L $ באמצעות TransformedDistribution שלוקח כפרמטרים את התפלגות Wishart ו- bijector הממפה את הגורם Cholesky של מטריצת דיוק לפקטור Cholesky של ההופכי שלו.

דרך פשוטה (אך לא יעילה במיוחד) להגיע מהגורם Cholesky של $ C ^ {- 1} $ ל- $ L $ היא להפוך את גורם Cholesky על ידי פתרון אחורי, ואז ליצור את מטריצת השונות מגורמים הפוכים אלה ואז עושה פקטוריזציה של Cholesky.

תן לפירוק Cholesky של $ C ^ {- 1} = MM ^ T $. $ M $ הוא משולש נמוך יותר, כך שנוכל להפוך אותו באמצעות MatrixInverseTriL MatrixInverseTriL.

גיבוש $ C $ מ $ M ^ {- 1} $ הוא קצת מסובך: $ C = (MM ^ T) ^ {- 1} = M ^ {- T} M ^ {- 1} = M ^ {- T } (M ^ {- T}) ^ T $. $ M $ הוא משולש תחתון, כך $ M ^ {- 1} $ יהיה גם משולש תחתון, ו- $ M ^ {- T} $ יהיה משולש עליון. המקשר CholeskyOuterProduct() פועל רק עם מטריצות משולשות תחתונות, ולכן איננו יכולים להשתמש בו כדי ליצור $ C $ מ $ M ^ {- T} $. הדרך לעקיפת הבעיה שלנו היא שרשרת מצמידים שמאפשרים את השורות והעמודות של מטריצה.

למרבה המזל ההיגיון הזה נעטף בתוכנת CholeskyToInvCholesky !

שילוב של כל החלקים

# verify that the bijector works
m = np.array([[1., 0.], [2., 8.]], dtype=np.float32)
c_inv = m.dot(m.T)
c = np.linalg.inv(c_inv)
c_chol = np.linalg.cholesky(c)
wishart_cholesky_to_iw_cholesky = tfb.CholeskyToInvCholesky()
w_fwd = wishart_cholesky_to_iw_cholesky.forward(m).numpy()

print('numpy =\n', c_chol)
print('bijector =\n', w_fwd)
numpy =
 [[ 1.0307764   0.        ]
 [-0.03031695  0.12126781]]
bijector =
 [[ 1.0307764   0.        ]
 [-0.03031695  0.12126781]]

ההפצה הסופית שלנו

ה- Wishart ההפוך שלנו הפועל על גורמי Cholesky הוא כדלקמן:

inverse_wishart_cholesky = tfd.TransformedDistribution(
    distribution=CholeskyWishart(
        df=PRIOR_DF,
        scale_tril=np.linalg.cholesky(np.linalg.inv(PRIOR_SCALE))),
    bijector=tfb.CholeskyToInvCholesky())

יש לנו את Wishart ההפוך שלנו, אבל זה די איטי מכיוון שאנחנו צריכים לעשות פירוק של Cholesky במכשול. בואו נחזור לפרמטריזציה של מטריצה ​​מדויקת ונראה מה אנחנו יכולים לעשות שם לצורך אופטימיזציה.

גרסה סופית (!): שימוש בפירוק Cholesky של מטריצת הדיוק

גישה חלופית היא לעבוד עם גורמי Cholesky של מטריצת הדיוק. כאן קל לחשב את פונקציית הסבירות הקודמת, אך פונקציית הסבירות ביומן הנתונים דורשת עבודה רבה יותר מכיוון של- TFP אין גרסה של הנורמל הרב-משתני שעוברת פרמטרים על ידי דיוק.

סיכוי יומן קודם מותאם

אנו משתמשים בהפצת CholeskyWishart שבנו לעיל כדי לבנות את ה- prior.

# Our new prior.
PRIOR_SCALE_CHOLESKY = np.linalg.cholesky(PRIOR_SCALE)

def log_lik_prior_cholesky(precisions_cholesky):
  rv_precision = CholeskyWishart(
      df=PRIOR_DF,
      scale_tril=PRIOR_SCALE_CHOLESKY,
      validate_args=VALIDATE_ARGS,
      allow_nan_stats=ALLOW_NAN_STATS)
  return rv_precision.log_prob(precisions_cholesky)
# Check against the slower TF implementation and the NumPy implementation.
# Note that when comparing to NumPy, we need to add in the Jacobian correction.
precisions = [np.eye(2, dtype=np.float32),
              true_precision]
precisions_cholesky = np.stack([np.linalg.cholesky(m) for m in precisions])
precisions = np.stack(precisions)
lik_tf = log_lik_prior_cholesky(precisions_cholesky).numpy()
lik_tf_slow = tfd.TransformedDistribution(
    distribution=tfd.WishartTriL(
        df=PRIOR_DF, scale_tril=tf.linalg.cholesky(PRIOR_SCALE)),
    bijector=tfb.Invert(tfb.CholeskyOuterProduct())).log_prob(
    precisions_cholesky).numpy()
corrections = tfb.Invert(tfb.CholeskyOuterProduct()).inverse_log_det_jacobian(
    precisions_cholesky, event_ndims=2).numpy()
n = precisions.shape[0]

for i in range(n):
  print(i)
  print('numpy:', log_lik_prior_numpy(precisions[i]) + corrections[i])
  print('tensorflow slow:', lik_tf_slow[i])
  print('tensorflow fast:', lik_tf[i])
0
numpy: -0.8488930160357633
tensorflow slow: -0.84889317
tensorflow fast: -0.84889317
1
numpy: -7.442875031036973
tensorflow slow: -7.442877
tensorflow fast: -7.442876

סיכוי יומן נתונים מותאם

אנו יכולים להשתמש במצורפים של TFP כדי לבנות גרסה משלנו לנורמלי הרב-משתני. הנה הרעיון המרכזי:

נניח שיש לי וקטור עמודה $ X $ שהאלמנטים שלו הם דוגמאות של $ N (0, 1) $. יש לנו $ \ text {mean} (X) = 0 $ ו- $ \ text {cov} (X) = I $

עכשיו תן ל- $ Y = AX + b $. יש לנו $ \ text {mean} (Y) = b $ ו- $ \ text {cov} (Y) = AA ^ T $

מכאן שנוכל ליצור וקטורים עם $ b $ ממוצע ומשתנה $ C $ באמצעות טרנספורמציית affine $ Ax + b $ לווקטורים של דגימות רגילות iid בתנאי $ AA ^ T = C $. לפירוק Cholesky של $ C $ יש את הנכס הרצוי. עם זאת, ישנם פתרונות אחרים.

תן ל- $ P = C ^ {- 1} $ ותן לפירוק Cholesky של $ P $ להיות $ B $, כלומר $ BB ^ T = P $. עַכשָׁיו

$ P ^ {- 1} = (BB ^ T) ^ {- 1} = B ^ {- T} B ^ {- 1} = B ^ {- T} (B ^ {- T}) ^ T $

אז דרך נוספת להשיג את הממוצע וההשתנות הרצויים שלנו היא להשתמש בטרנספורמציה affine $ Y = B ^ {- T} X + b $.

הגישה שלנו (באדיבות מחברת זו ):

  1. השתמש tfd.Independent() לשלב קבוצה של 1-D Normal משתנה אקראי לתוך משתנה אקראי רב-ממדי אחת. הפרמטר reinterpreted_batch_ndims עבור Independent() מציין את מספר מידות האצווה שיש לפרש מחדש כמידות אירוע. In our case we create a 1-D batch of length 2 that we transform into a 1-D event of length 2, so reinterpreted_batch_ndims=1 .
  2. Apply a bijector to add the desired covariance: tfb.Invert(tfb.ScaleMatvecTriL(scale_tril=precision_cholesky, adjoint=True)) . Note that above we're multiplying our iid normal random variables by the transpose of the inverse of the Cholesky factor of the precision matrix $(B^{-T}X)$. The tfb.Invert takes care of inverting $B$, and the adjoint=True flag performs the transpose.
  3. Apply a bijector to add the desired offset: tfb.Shift(shift=shift) Note that we have to do the shift as a separate step from the initial inverted affine transform because otherwise the inverted scale is applied to the shift (since the inverse of $y=Ax+b$ is $x=A^{-1}y - A^{-1}b$).
class MVNPrecisionCholesky(tfd.TransformedDistribution):
  """Multivariate normal parameterized by loc and Cholesky precision matrix."""

  def __init__(self, loc, precision_cholesky, name=None):
    super(MVNPrecisionCholesky, self).__init__(
        distribution=tfd.Independent(
            tfd.Normal(loc=tf.zeros_like(loc),
                       scale=tf.ones_like(loc)),
            reinterpreted_batch_ndims=1),
        bijector=tfb.Chain([
            tfb.Shift(shift=loc),
            tfb.Invert(tfb.ScaleMatvecTriL(scale_tril=precision_cholesky,
                                  adjoint=True)),
        ]),
        name=name)
@tf.function(autograph=False)
def log_lik_data_cholesky(precisions_cholesky, replicated_data):
  n = tf.shape(precisions_cholesky)[0]  # number of precision matrices
  rv_data = MVNPrecisionCholesky(
      loc=tf.zeros([n, 2]),
      precision_cholesky=precisions_cholesky)
  return tf.reduce_sum(rv_data.log_prob(replicated_data), axis=0)
# check against the numpy implementation
true_precision_cholesky = np.linalg.cholesky(true_precision)
precisions = [np.eye(2, dtype=np.float32), true_precision]
precisions_cholesky = np.stack([np.linalg.cholesky(m) for m in precisions])
precisions = np.stack(precisions)
n = precisions_cholesky.shape[0]
replicated_data = np.tile(np.expand_dims(my_data, axis=1), reps=[1, 2, 1])
lik_tf = log_lik_data_cholesky(precisions_cholesky, replicated_data).numpy()

for i in range(n):
  print(i)
  print('numpy:', log_lik_data_numpy(precisions[i], my_data))
  print('tensorflow:', lik_tf[i])
0
numpy: -430.71218815801365
tensorflow: -430.71207
1
numpy: -280.81822950593767
tensorflow: -280.81824

Combined log likelihood function

Now we combine our prior and data log likelihood functions in a closure.

def get_log_lik_cholesky(data, n_chains=1):
  # The data argument that is passed in will be available to the inner function
  # below so it doesn't have to be passed in as a parameter.
  replicated_data = np.tile(np.expand_dims(data, axis=1), reps=[1, n_chains, 1])

  @tf.function(autograph=False)
  def _log_lik_cholesky(precisions_cholesky):
    return (log_lik_data_cholesky(precisions_cholesky, replicated_data) +
            log_lik_prior_cholesky(precisions_cholesky))

  return _log_lik_cholesky

Constraining bijector

Our samples are constrained to be valid Cholesky factors, which means they must be lower triangular matrices with positive diagonals. The TransformedTransitionKernel needs a bijector that maps unconstrained tensors to/from tensors with our desired constraints. We've removed the Cholesky decomposition from the bijector's inverse, which speeds things up.

unconstrained_to_precision_cholesky = tfb.Chain([
    # step 2: exponentiate the diagonals    
    tfb.TransformDiagonal(tfb.Exp(validate_args=VALIDATE_ARGS)),
    # step 1: expand the vector to a lower triangular matrix
    tfb.FillTriangular(validate_args=VALIDATE_ARGS),
])
# some checks
inv = unconstrained_to_precision_cholesky.inverse(precisions_cholesky).numpy()
fwd = unconstrained_to_precision_cholesky.forward(inv).numpy()
print('precisions_cholesky:\n', precisions_cholesky)
print('\ninv:\n', inv)
print('\nfwd(inv):\n', fwd)
precisions_cholesky:
 [[[ 1.         0.       ]
  [ 0.         1.       ]]

 [[ 1.1470785  0.       ]
  [-2.0647411  1.0000004]]]

inv:
 [[ 0.0000000e+00  0.0000000e+00  0.0000000e+00]
 [ 3.5762781e-07 -2.0647411e+00  1.3721828e-01]]

fwd(inv):
 [[[ 1.         0.       ]
  [ 0.         1.       ]]

 [[ 1.1470785  0.       ]
  [-2.0647411  1.0000004]]]

Initial values

We generate a tensor of initial values. We're working with Cholesky factors, so we generate some Cholesky factor initial values.

# The number of chains is determined by the shape of the initial values.
# Here we'll generate 3 chains, so we'll need a tensor of 3 initial values.
N_CHAINS = 3

np.random.seed(123)

initial_values_cholesky = []
for i in range(N_CHAINS):
  initial_values_cholesky.append(np.array(
      [[0.5 + np.random.uniform(), 0.0],
       [-0.5 + np.random.uniform(), 0.5 + np.random.uniform()]],
      dtype=np.float32))
initial_values_cholesky = np.stack(initial_values_cholesky)

Sampling

We sample N_CHAINS chains using the TransformedTransitionKernel .

@tf.function(autograph=False)
def sample():
  tf.random.set_seed(123)
  log_lik_fn = get_log_lik_cholesky(my_data)

  # Tuning acceptance rates:
  dtype = np.float32
  num_burnin_iter = 3000
  num_warmup_iter = int(0.8 * num_burnin_iter) 
  num_chain_iter = 2500

  # Set the target average acceptance ratio for the HMC as suggested by
  # Beskos et al. (2013):
  # https://projecteuclid.org/download/pdfview_1/euclid.bj/1383661192
  target_accept_rate = 0.651

  # Initialize the HMC sampler.
  hmc = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=log_lik_fn,
      step_size=0.01,
      num_leapfrog_steps=3)

  ttk = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=hmc, bijector=unconstrained_to_precision_cholesky)

  # Adapt the step size using standard adaptive MCMC procedure. See Section 4.2
  # of Andrieu and Thoms (2008):
  # http://www4.ncsu.edu/~rsmith/MA797V_S12/Andrieu08_AdaptiveMCMC_Tutorial.pdf
  adapted_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=ttk,
      num_adaptation_steps=num_warmup_iter,
      target_accept_prob=target_accept_rate)

  states = tfp.mcmc.sample_chain(
      num_results=num_chain_iter,
      num_burnin_steps=num_burnin_iter,
      current_state=initial_values,
      kernel=adapted_kernel,
      trace_fn=None,
      parallel_iterations=1)
  # transform samples back to their constrained form
  samples = tf.linalg.matmul(states, states, transpose_b=True)
  return samples

precision_samples = sample()

Convergence check

A quick convergence check looks good:

r_hat = tfp.mcmc.potential_scale_reduction(precision_samples).numpy()
print(r_hat)
[[1.0013583 1.0019467]
 [1.0019467 1.0011804]]

Comparing results to the analytic posterior

# The output samples have shape [n_steps, n_chains, 2, 2]
# Flatten them to [n_steps * n_chains, 2, 2] via reshape:
precision_samples_reshaped = np.reshape(precision_samples, newshape=[-1, 2, 2])

And again, the sample means and standard deviations match those of the analytic posterior.

print('True posterior mean:\n', posterior_mean)
print('Mean of samples:\n', np.mean(precision_samples_reshaped, axis=0))
True posterior mean:
 [[ 0.9641779 -1.6534661]
 [-1.6534661  3.8683164]]
Mean of samples:
 [[ 0.9668749 -1.6552604]
 [-1.6552604  3.8676758]]
print('True posterior standard deviation:\n', posterior_sd)
print('Standard deviation of samples:\n', np.std(precision_samples_reshaped, axis=0))
True posterior standard deviation:
 [[0.13435492 0.25050813]
 [0.25050813 0.53903675]]
Standard deviation of samples:
 [[0.13329637 0.24913797]
 [0.24913797 0.53983945]]

Ok, all done! We've got our optimized sampler working.