Bayesian IPWのStan実装の紹介

概要

この記事は確率的プログラミング言語 Advent Calendar 2023の12/1の記事です。

inverse probability weighting(IPW)を用いた因果効果推定においてベイズモデルを用いる方法を紹介します。具体的にはAndrew Heissさんの以下のブログ記事の手法を紹介します。

• Heiss, Andrew. 2021. "How to Create a(n Almost) Fully Bayesian Outcome Model with Inverse Probability Weights." December 20, 2021. https://doi.org/10.59350/gyvjk-hrx68.

Stan実装はA. Jordan Nafaさんがしており、以下がそのリポジトリです。

• https://github.com/ajnafa/Bayesian-IPW

理論的な背景は以下の論文です。

• Liao, Shirley X., and Corwin M. Zigler. 2020. "Uncertainty in the Design Stage of Two-Stage Bayesian Propensity Score Analysis." Statistics in Medicine 39 (17): 2265–90. https://doi.org/10.1002/sim.8486.

この記事では、なるべくシンプルなコードを用いてその要所を説明したいと思います。

方針

• 介入を予測するモデルをあてはめて、割り付け$A_i$が$1$となる確率$e_i$を求めます。
• 個人レベルの逆確率重み$\text{IPW}_i$を求めます。具体的には対照群では$(1 - A_i)/(1 - e_i)$で、介入群では$A_i/e_i$で$\text{IPW}_i$を求めます。
• この重みを用いて、潜在的結果変数を予測するモデルである周辺構造モデル(Marginal Structural Model; MSM)をあてはめて、因果効果を推定します。

MCMCを使っている場合、重みがMCMCのサンプルの形で得られます。そのため、重みのMCMCサンプルの次元の回数だけ周辺構造モデルのあてはめをしなくてはいけないように感じるかもしれませんが、うまく2段階に分離することでこれを避けます。ポイントは、重みのMCMCのiterationと周辺構造モデルのMCMCのiterationを揃えることで、2段階推定の1セットをMCMCの1回抽出に対応づける点です。

コード

共変量$L$と介入$A$が二値、結果変数$Y$が連続値のコード例を示します。

実行するRコード

• 6-10行目がデータ生成例です。ここでは人数Nを400としました。9行目から真の因果効果は1.0です。
• 12-17行目で重みのMCMCサンプルを計算しています。
• 19-23行目でその重みを用いて因果効果を推定しています。

21行目のcmdstan_model関数のuser_headerオプションを指定することで、指定したC++のファイルをヘッダーに加えてモデルがコンパイルされます。Stan内において自作のC++の関数を使えるようになるという激アツオプションです。詳しい説明については、CmdStanのマニュアルの該当部分を参照してください。

重みを推定するモデルのStanコード

普通のロジスティック回帰のコードです。20行目で介入の確率を、21行目で重みを求めています。

重みづけ周辺構造モデルのStanコード

• 10行目：MCMCサンプルの次元です。
• 11行目：人数$\times$MCMCサンプルの次元 の重みIPWをデータとして渡します。
• 22行目：その重みを対数尤度に掛けて使っています。
• 3行目, 20行目：重みのMCMCのiterationと周辺構造モデルのMCMCのiterationを揃えるために、3行目のfunctionsブロックでget_iter関数を宣言して20行目で使っています。
• 2行目, 27行目：generated quantitiesブロック内でadd_iter関数を呼ぶことでMCMCのiterationが進みます。

iterationを揃えるための関数を定義するC++コード

• 3行目：MCMCのiterationであるiteration_indexをグローバル変数で定義しています。

結果

run-example.Rの実行結果は以下になります。因果効果の90%区間を求めることができました。

 variable    mean  median    sd   mad       q5     q95 rhat ess_bulk ess_tail
    lp__  -953.92 -941.41 71.65 64.47 -1090.61 -862.65 1.00     1786     2066
    b[1]     0.31    0.31  0.06  0.06     0.22    0.41 1.00     3484     4313
    b[2]     1.00    1.01  0.13  0.13     0.78    1.21 1.00     1661     2240
    sigma    0.78    0.78  0.02  0.02     0.74    0.81 1.00     2944     3053

Enjoy!