線性混合效應模型

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

線性混合效應模型是為結構化線性關係建模的簡單方法 (Harville, 1997; Laird and Ware, 1982)。每個資料點都包含不同類型的輸入 (分類為群組) 和實數值輸出。線性混合效應模型是階層式模型:它在群組之間共用統計強度,以改進關於任何個別資料點的推論。

在本教學課程中,我們將使用 TensorFlow Probability 中的真實世界範例示範線性混合效應模型。我們將使用 JointDistributionCoroutine 和馬可夫鏈蒙地卡羅 (tfp.mcmc) 模組。

依附元件與先決條件

匯入和設定

讓一切加速!

在深入探討之前,我們先確認本示範是否使用 GPU。

若要執行此操作,請選取「Runtime」(執行階段) -> 「Change runtime type」(變更執行階段類型) -> 「Hardware accelerator」(硬體加速器) -> 「GPU」。

以下程式碼片段將驗證我們是否可存取 GPU。

if tf.test.gpu_device_name() != '/device:GPU:0':
  print('WARNING: GPU device not found.')
else:
  print('SUCCESS: Found GPU: {}'.format(tf.test.gpu_device_name()))
WARNING: GPU device not found.

資料

我們使用熱門 lme4 R 語言套件 (Bates 等人,2015) 中的 InstEval 資料集。這是一個關於課程及其評估評等的資料集。每門課程都包含中繼資料,例如 studentsinstructorsdepartments,而感興趣的回應變數是評估評等。

def load_insteval():
  """Loads the InstEval data set.

  It contains 73,421 university lecture evaluations by students at ETH
  Zurich with a total of 2,972 students, 2,160 professors and
  lecturers, and several student, lecture, and lecturer attributes.
  Implementation is built from the `observations` Python package.

  Returns:
    Tuple of np.ndarray `x_train` with 73,421 rows and 7 columns and
    dictionary `metadata` of column headers (feature names).
  """
  url = ('https://raw.github.com/vincentarelbundock/Rdatasets/master/csv/'
         'lme4/InstEval.csv')
  with requests.Session() as s:
    download = s.get(url)
    f = download.content.decode().splitlines()

  iterator = csv.reader(f)
  columns = next(iterator)[1:]
  x_train = np.array([row[1:] for row in iterator], dtype=np.int)
  metadata = {'columns': columns}
  return x_train, metadata

我們載入並預先處理資料集。我們保留 20% 的資料,以便在未見過的資料點上評估我們擬合的模型。以下我們將前幾列視覺化。

data, metadata = load_insteval()
data = pd.DataFrame(data, columns=metadata['columns'])
data = data.rename(columns={'s': 'students',
                            'd': 'instructors',
                            'dept': 'departments',
                            'y': 'ratings'})
data['students'] -= 1  # start index by 0
# Remap categories to start from 0 and end at max(category).
data['instructors'] = data['instructors'].astype('category').cat.codes
data['departments'] = data['departments'].astype('category').cat.codes

train = data.sample(frac=0.8)
test = data.drop(train.index)

train.head()
<ipython-input-4-5d7a9eabeea1>:21: DeprecationWarning: `np.int` is a deprecated alias for the builtin `int`. To silence this warning, use `int` by itself. Doing this will not modify any behavior and is safe. When replacing `np.int`, you may wish to use e.g. `np.int64` or `np.int32` to specify the precision. If you wish to review your current use, check the release note link for additional information.
Deprecated in NumPy 1.20; for more details and guidance: https://numpy.dev.org.tw/devdocs/release/1.20.0-notes.html#deprecations
  x_train = np.array([row[1:] for row in iterator], dtype=np.int)

我們根據輸入的 features 字典和對應於評等的 labels 輸出設定資料集。每個特徵都編碼為整數,每個標籤 (評估評等) 都編碼為浮點數。

get_value = lambda dataframe, key, dtype: dataframe[key].values.astype(dtype)
features_train = {
    k: get_value(train, key=k, dtype=np.int32)
    for k in ['students', 'instructors', 'departments', 'service']}
labels_train = get_value(train, key='ratings', dtype=np.float32)

features_test = {k: get_value(test, key=k, dtype=np.int32)
                 for k in ['students', 'instructors', 'departments', 'service']}
labels_test = get_value(test, key='ratings', dtype=np.float32)
num_students = max(features_train['students']) + 1
num_instructors = max(features_train['instructors']) + 1
num_departments = max(features_train['departments']) + 1
num_observations = train.shape[0]

print("Number of students:", num_students)
print("Number of instructors:", num_instructors)
print("Number of departments:", num_departments)
print("Number of observations:", num_observations)
Number of students: 2972
Number of instructors: 1128
Number of departments: 14
Number of observations: 58737

模型

典型的線性模型假設獨立性,其中任何一對資料點都具有恆定的線性關係。在 InstEval 資料集中,觀察結果出現在群組中,每個群組可能具有不同的斜率和截距。線性混合效應模型 (也稱為階層式線性模型或多層次線性模型) 捕捉到這種現象 (Gelman & Hill, 2006)。

此現象的範例包括

  • 學生。來自學生的觀察結果並非獨立:有些學生可能會系統性地給予較低 (或較高) 的講座評等。
  • 講師。來自講師的觀察結果並非獨立:我們預期好老師通常會有好的評等,而壞老師通常會有壞的評等。
  • 科系。來自科系的觀察結果並非獨立:某些科系可能通常有枯燥乏味的教材或更嚴格的評分標準,因此評等會低於其他科系。

為了捕捉到這一點,請回想一下,對於 \(N\times D\) 特徵 \(\mathbf{X}\) 和 \(N\) 標籤 \(\mathbf{y}\) 的資料集,線性迴歸假定模型

\[ \begin{equation*} \mathbf{y} = \mathbf{X}\beta + \alpha + \epsilon, \end{equation*} \]

其中有斜率向量 \(\beta\in\mathbb{R}^D\)、截距 \(\alpha\in\mathbb{R}\) 和隨機雜訊 \(\epsilon\sim\text{Normal}(\mathbf{0}, \mathbf{I})\)。我們說 \(\beta\) 和 \(\alpha\) 是「固定效應」:它們是在資料點 \((x, y)\) 母體中保持恆定的效應。方程式作為概似函數的等效公式為 \(\mathbf{y} \sim \text{Normal}(\mathbf{X}\beta + \alpha, \mathbf{I})\)。此概似函數在推論期間會最大化,以便找到適合資料的 \(\beta\) 和 \(\alpha\) 的點估計值。

線性混合效應模型將線性迴歸擴充為

\[ \begin{align*} \eta &\sim \text{Normal}(\mathbf{0}, \sigma^2 \mathbf{I}), \\ \mathbf{y} &= \mathbf{X}\beta + \mathbf{Z}\eta + \alpha + \epsilon. \end{align*} \]

其中仍然有斜率向量 \(\beta\in\mathbb{R}^P\)、截距 \(\alpha\in\mathbb{R}\) 和隨機雜訊 \(\epsilon\sim\text{Normal}(\mathbf{0}, \mathbf{I})\)。此外,還有一個項 \(\mathbf{Z}\eta\),其中 \(\mathbf{Z}\) 是特徵矩陣,而 \(\eta\in\mathbb{R}^Q\) 是隨機斜率的向量;\(\eta\) 是常態分配,具有變異數分量參數 \(\sigma^2\)。\(\mathbf{Z}\) 是透過根據新的 \(N\times P\) 矩陣 \(\mathbf{X}\) 和 \(N\times Q\) 矩陣 \(\mathbf{Z}\) 分割原始 \(N\times D\) 特徵矩陣而形成的,其中 \(P + Q=D\):此分割讓我們可以使用固定效應 \(\beta\) 和潛在變數 \(\eta\) 分別對特徵進行建模。

我們說潛在變數 \(\eta\) 是「隨機效應」:它們是在母體中變化的效應 (儘管它們在子母體中可能保持恆定)。特別是,由於隨機效應 \(\eta\) 的平均值為 0,因此資料標籤的平均值由 \(\mathbf{X}\beta + \alpha\) 捕捉。\(\mathbf{Z}\eta\) 隨機效應分量捕捉資料中的變異:例如,「講師 #54 的評等比平均值高 1.4 分。」

在本教學課程中,我們假定以下效應

  • 固定效應:serviceservice 是一個二元共變數,對應於課程是否屬於講師的主要科系。無論我們收集多少額外資料,它都只能採用值 \(0\) 和 \(1\)。
  • 隨機效應:studentsinstructorsdepartments。如果給定來自課程評估評等母體的更多觀察結果,我們可能會看到新的學生、老師或科系。

在 R 語言的 lme4 套件語法 (Bates 等人,2015) 中,模型可以總結為

ratings ~ service + (1|students) + (1|instructors) + (1|departments) + 1

其中 x 表示固定效應,(1|x) 表示 x 的隨機效應,而 1 表示截距項。

我們在下方將此模型實作為 JointDistribution。為了更好地支援參數追蹤 (例如,我們想要追蹤 model.trainable_variables 中的所有 tf.Variable),我們將模型範本實作為 tf.Module

class LinearMixedEffectModel(tf.Module):
  def __init__(self):
    # Set up fixed effects and other parameters.
    # These are free parameters to be optimized in E-steps
    self._intercept = tf.Variable(0., name="intercept")            # alpha in eq
    self._effect_service = tf.Variable(0., name="effect_service")  #  beta in eq
    self._stddev_students = tfp.util.TransformedVariable(
        1., bijector=tfb.Exp(), name="stddev_students")            # sigma in eq
    self._stddev_instructors = tfp.util.TransformedVariable(
        1., bijector=tfb.Exp(), name="stddev_instructors")         # sigma in eq
    self._stddev_departments = tfp.util.TransformedVariable(
        1., bijector=tfb.Exp(), name="stddev_departments")         # sigma in eq

  def __call__(self, features):
    model = tfd.JointDistributionSequential([
      # Set up random effects.
      tfd.MultivariateNormalDiag(
          loc=tf.zeros(num_students),
          scale_diag=self._stddev_students * tf.ones(num_students)),
      tfd.MultivariateNormalDiag(
          loc=tf.zeros(num_instructors),
          scale_diag=self._stddev_instructors * tf.ones(num_instructors)),
      tfd.MultivariateNormalDiag(
          loc=tf.zeros(num_departments),
          scale_diag=self._stddev_departments * tf.ones(num_departments)),
      # This is the likelihood for the observed.
      lambda effect_departments, effect_instructors, effect_students: tfd.Independent(
          tfd.Normal(
              loc=(self._effect_service * features["service"] +
                  tf.gather(effect_students, features["students"], axis=-1) +
                  tf.gather(effect_instructors, features["instructors"], axis=-1) +
                  tf.gather(effect_departments, features["departments"], axis=-1) +
                  self._intercept),
              scale=1.),
              reinterpreted_batch_ndims=1)
    ])

    # To enable tracking of the trainable variables via the created distribution,
    # we attach a reference to `self`. Since all TFP objects sub-class
    # `tf.Module`, this means that the following is possible:
    # LinearMixedEffectModel()(features_train).trainable_variables
    # ==> tuple of all tf.Variables created by LinearMixedEffectModel.
    model._to_track = self
    return model

lmm_jointdist = LinearMixedEffectModel()
# Conditioned on feature/predictors from the training data
lmm_train = lmm_jointdist(features_train)
lmm_train.trainable_variables
(<tf.Variable 'effect_service:0' shape=() dtype=float32, numpy=0.0>,
 <tf.Variable 'intercept:0' shape=() dtype=float32, numpy=0.0>,
 <tf.Variable 'stddev_departments:0' shape=() dtype=float32, numpy=0.0>,
 <tf.Variable 'stddev_instructors:0' shape=() dtype=float32, numpy=0.0>,
 <tf.Variable 'stddev_students:0' shape=() dtype=float32, numpy=0.0>)

作為機率圖形程式,我們也可以根據模型的計算圖視覺化模型的結構。此圖形編碼程式中隨機變數之間的資料流,以圖形模型的形式明確表示它們的關係 (Jordan, 2003)。

作為統計工具,我們可能會查看圖形,以便更好地看到,例如,intercepteffect_service 在給定 ratings 的情況下有條件地相依;如果程式是用類別、跨模組的交叉引用和/或子常式編寫的,則可能很難從原始碼中看出這一點。作為計算工具,我們也可能會注意到潛在變數透過 tf.gather 運算流入 ratings 變數。如果索引 Tensor 的成本很高,則這可能是某些硬體加速器上的瓶頸;視覺化圖形可以清楚地顯示這一點。

lmm_train.resolve_graph()
(('effect_students', ()),
 ('effect_instructors', ()),
 ('effect_departments', ()),
 ('x', ('effect_departments', 'effect_instructors', 'effect_students')))

參數估計

給定資料,推論的目標是擬合模型的固定效應斜率 \(\beta\)、截距 \(\alpha\) 和變異數分量參數 \(\sigma^2\)。最大概似原則將此任務形式化為

\[ \max_{\beta, \alpha, \sigma}~\log p(\mathbf{y}\mid \mathbf{X}, \mathbf{Z}; \beta, \alpha, \sigma) = \max_{\beta, \alpha, \sigma}~\log \int p(\eta; \sigma) ~p(\mathbf{y}\mid \mathbf{X}, \mathbf{Z}, \eta; \beta, \alpha)~d\eta. \]

在本教學課程中,我們使用蒙地卡羅 EM 演算法來最大化此邊際密度 (Dempster 等人,1977;Wei 和 Tanner,1990)。¹ 我們執行馬可夫鏈蒙地卡羅來計算隨機效應方面條件概似函數的期望值 (「E 步驟」),並且我們執行梯度下降以最大化參數方面期望值 (「M 步驟」)

  • 對於 E 步驟,我們設定 Hamiltonian 蒙地卡羅 (HMC)。它採用目前狀態 (學生、講師和科系效應),並傳回新狀態。我們將新狀態指派給 TensorFlow 變數,這些變數將表示 HMC 鏈的狀態。

  • 對於 M 步驟,我們使用來自 HMC 的後驗樣本來計算邊際概似函數的無偏估計值 (最多為常數)。然後,我們將其梯度應用於感興趣的參數。這會產生邊際概似函數上的無偏隨機下降步驟。我們使用 Adam TensorFlow 最佳化工具實作它,並最小化邊際的負值。

target_log_prob_fn = lambda *x: lmm_train.log_prob(x + (labels_train,))
trainable_variables = lmm_train.trainable_variables
current_state = lmm_train.sample()[:-1]
# For debugging
target_log_prob_fn(*current_state)
<tf.Tensor: shape=(), dtype=float32, numpy=-485996.53>
# Set up E-step (MCMC).
hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=target_log_prob_fn,
    step_size=0.015,
    num_leapfrog_steps=3)
kernel_results = hmc.bootstrap_results(current_state)

@tf.function(autograph=False, jit_compile=True)
def one_e_step(current_state, kernel_results):
  next_state, next_kernel_results = hmc.one_step(
      current_state=current_state,
      previous_kernel_results=kernel_results)
  return next_state, next_kernel_results

optimizer = tf_keras.optimizers.Adam(learning_rate=.01)

# Set up M-step (gradient descent).
@tf.function(autograph=False, jit_compile=True)
def one_m_step(current_state):
  with tf.GradientTape() as tape:
    loss = -target_log_prob_fn(*current_state)
  grads = tape.gradient(loss, trainable_variables)
  optimizer.apply_gradients(zip(grads, trainable_variables))
  return loss

我們執行暖機階段,該階段執行一個 MCMC 鏈若干次疊代,以便可以在後驗的機率質量內初始化訓練。然後,我們執行訓練迴圈。它會聯合執行 E 步驟和 M 步驟,並在訓練期間記錄值。

num_warmup_iters = 1000
num_iters = 1500
num_accepted = 0
effect_students_samples = np.zeros([num_iters, num_students])
effect_instructors_samples = np.zeros([num_iters, num_instructors])
effect_departments_samples = np.zeros([num_iters, num_departments])
loss_history = np.zeros([num_iters])
# Run warm-up stage.
for t in range(num_warmup_iters):
  current_state, kernel_results = one_e_step(current_state, kernel_results)
  num_accepted += kernel_results.is_accepted.numpy()
  if t % 500 == 0 or t == num_warmup_iters - 1:
    print("Warm-Up Iteration: {:>3} Acceptance Rate: {:.3f}".format(
        t, num_accepted / (t + 1)))

num_accepted = 0  # reset acceptance rate counter

# Run training.
for t in range(num_iters):
  # run 5 MCMC iterations before every joint EM update
  for _ in range(5):
    current_state, kernel_results = one_e_step(current_state, kernel_results)
  loss = one_m_step(current_state)
  effect_students_samples[t, :] = current_state[0].numpy()
  effect_instructors_samples[t, :] = current_state[1].numpy()
  effect_departments_samples[t, :] = current_state[2].numpy()
  num_accepted += kernel_results.is_accepted.numpy()
  loss_history[t] = loss.numpy()
  if t % 500 == 0 or t == num_iters - 1:
    print("Iteration: {:>4} Acceptance Rate: {:.3f} Loss: {:.3f}".format(
        t, num_accepted / (t + 1), loss_history[t]))
Warm-Up Iteration:   0 Acceptance Rate: 1.000
Warm-Up Iteration: 500 Acceptance Rate: 0.758
Warm-Up Iteration: 999 Acceptance Rate: 0.729
Iteration:    0 Acceptance Rate: 1.000 Loss: 98200.422
Iteration:  500 Acceptance Rate: 0.649 Loss: 98190.469
Iteration: 1000 Acceptance Rate: 0.656 Loss: 98068.664
Iteration: 1499 Acceptance Rate: 0.663 Loss: 98155.070

您也可以將暖機 for 迴圈寫入 tf.while_loop,並將訓練步驟寫入 tf.scantf.while_loop,以獲得更快的推論。例如

@tf.function(autograph=False, jit_compile=True)
def run_k_e_steps(k, current_state, kernel_results):
  _, next_state, next_kernel_results = tf.while_loop(
      cond=lambda i, state, pkr: i < k,
      body=lambda i, state, pkr: (i+1, *one_e_step(state, pkr)),
      loop_vars=(tf.constant(0), current_state, kernel_results)
  )
  return next_state, next_kernel_results

在上方,我們未執行演算法,直到偵測到收斂閾值為止。為了檢查訓練是否合理,我們驗證損失函數是否確實傾向於在訓練疊代中收斂。

plt.plot(loss_history)
plt.ylabel(r'Loss $-\log$ $p(y\mid\mathbf{x})$')
plt.xlabel('Iteration')
plt.show()

png

我們也使用追蹤圖,該圖顯示馬可夫鏈蒙地卡羅演算法在特定潛在維度上的軌跡。在下方,我們看到特定講師效應確實有意義地從其初始狀態轉變並探索狀態空間。追蹤圖也指出,不同講師的效應不同,但混合行為相似。

for i in range(7):
  plt.plot(effect_instructors_samples[:, i])

plt.legend([i for i in range(7)], loc='lower right')
plt.ylabel('Instructor Effects')
plt.xlabel('Iteration')
plt.show()

png

批判

在上方,我們擬合了模型。我們現在研究使用資料批判其擬合,這讓我們可以探索並更好地理解模型。其中一種技術是殘差圖,該圖繪製每個資料點的模型預測與實際值之間的差異。如果模型正確,則它們的差異應為標準常態分配;圖中偏離此模式的任何偏差都表示模型不擬合。

我們透過首先形成評等方面的後驗預測分配來建構殘差圖,這會將隨機效應的先驗分配替換為給定訓練資料的後驗。特別是,我們向前執行模型並攔截其對先驗隨機效應的依賴性及其推論的後驗平均值。²

lmm_test = lmm_jointdist(features_test)

[
    effect_students_mean,
    effect_instructors_mean,
    effect_departments_mean,
] = [
     np.mean(x, axis=0).astype(np.float32) for x in [
       effect_students_samples,
       effect_instructors_samples,
       effect_departments_samples
       ]
]

# Get the posterior predictive distribution
(*posterior_conditionals, ratings_posterior), _ = lmm_test.sample_distributions(
    value=(
        effect_students_mean,
        effect_instructors_mean,
        effect_departments_mean,
))

ratings_prediction = ratings_posterior.mean()

經視覺檢查,殘差看起來有點像標準常態分配。但是,擬合並不完美:尾部的機率質量大於常態分配,這表示模型可能會透過放寬其常態假設來改進其擬合。

特別是,雖然在 InstEval 資料集中使用常態分配對評等進行建模是最常見的做法,但更仔細地查看資料會發現課程評估評等實際上是從 1 到 5 的序數值。這表示我們應該使用序數分配,或者如果我們有足夠的資料來丟棄相對排序,甚至可以使用類別分配。這是對上述模型的一行程式碼變更;相同的推論程式碼也適用。

plt.title("Residuals for Predicted Ratings on Test Set")
plt.xlim(-4, 4)
plt.ylim(0, 800)
plt.hist(ratings_prediction - labels_test, 75)
plt.show()

png

為了探索模型如何進行個別預測,我們查看學生、講師和科系效應的長條圖。這讓我們瞭解資料點特徵向量中的個別元素如何傾向於影響結果。

不出所料,我們在下方看到每個學生通常對講師的評估評等幾乎沒有影響。有趣的是,我們看到講師所屬的科系有很大的影響。

plt.title("Histogram of Student Effects")
plt.hist(effect_students_mean, 75)
plt.show()

png

plt.title("Histogram of Instructor Effects")
plt.hist(effect_instructors_mean, 75)
plt.show()

png

plt.title("Histogram of Department Effects")
plt.hist(effect_departments_mean, 75)
plt.show()

png

註腳

¹ 線性混合效應模型是一個特例,我們可以在其中分析計算其邊際密度。為了本教學課程的目的,我們示範蒙地卡羅 EM,它更易於應用於非分析邊際密度,例如,如果概似函數擴充為類別分配而不是常態分配。

² 為了簡化起見,我們僅使用模型的一次前向傳遞來形成預測分配的平均值。這是透過以後驗平均值為條件來完成的,並且對線性混合效應模型有效。但是,這通常無效:後驗預測分配的平均值通常難以處理,並且需要對給定後驗樣本的模型的多個前向傳遞取經驗平均值。

謝辭

本教學課程最初是在 Edward 1.0 中編寫的 (來源)。我們感謝所有為編寫和修訂該版本做出貢獻的人員。

參考文獻

  1. Douglas Bates、Martin Machler、Ben Bolker 和 Steve Walker。《Fitting Linear Mixed-Effects Models Using lme4》。Journal of Statistical Software,67(1):1-48,2015 年。

  2. Arthur P. Dempster、Nan M. Laird 和 Donald B. Rubin。《Maximum likelihood from incomplete data via the EM algorithm》。Journal of the Royal Statistical Society, Series B (Methodological),1-38,1977 年。

  3. Andrew Gelman 和 Jennifer Hill。《Data analysis using regression and multilevel/hierarchical models.》劍橋大學出版社,2006 年。

  4. David A. Harville。《Maximum likelihood approaches to variance component estimation and to related problems》。Journal of the American Statistical Association,72(358):320-338,1977 年。

  5. Michael I. Jordan。《An Introduction to Graphical Models》。技術報告,2003 年。

  6. Nan M. Laird 和 James Ware。《Random-effects models for longitudinal data》。Biometrics,963-974,1982 年。

  7. Greg Wei 和 Martin A. Tanner。《A Monte Carlo implementation of the EM algorithm and the poor man's data augmentation algorithms》。Journal of the American Statistical Association,699-704,1990 年。