Twitterでも書いたのですが、NTT近原さんとの因果推論共同研究1本目がAIのトップ国際会議AAAI2024に採択されました。arXiv論文をこちらに、Pythonのプログラムコードをこちらにそれぞれ公開しています。本記事ではこちらの研究の概要を解説したいと思います。

条件付き平均処置効果（CATE）の推定問題

以前こちらの記事でBayesian Causal Inferenceの基本的な考え方を説明し、その冒頭で潜在反応モデルと平均処置効果について説明しました。そこでは処置変数\(T\in\left\{0, 1\right\}\)が結果変数\(Y\)に与える因果効果として平均処置効果（Average Treatment Effect: ATE）\(\mathbb{E}\left[Y^{(1)}-Y^{(0)}\right]\)という量を考えました（以前の記事では処置変数を\(Z\)としていましたが、今回は論文に合わせて\(T\)としています）。以前の記事では\(T\)を「ECサイトで割引クーポンを付与するかしないか」を表す変数、\(Y\)を「（クーポン付与後の）ECサイトでの利用額」を表す変数という例を挙げました。この例の場合、もしかすると\(T\)が\(Y\)に与える影響はクーポンを付与する対象のユーザーの性質（性別や年令など）によって異なるかもしれません。このように因果効果が対象によって変化するという性質を因果効果の異質性などと言ったりします。このような因果効果の異質性を考慮した量として条件付き平均処置効果（Conditional Average Treatment Effect: CATE）\begin{align}\mathbb{E}\left[Y^{(1)}-Y^{(0)}|\boldsymbol{X}\right]\end{align}があります¹。ここで\(\boldsymbol{X}\)は先程の例のユーザーの性別や年令などをまとめた変数です。ATE同様、CATEも定義に何の仮定もないと推定できませんが、Unconfoundedness仮定などの仮定のもとで
\begin{align}
\mathbb{E}\left[Y^{(1)}-Y^{(0)}|\boldsymbol{X}\right]=\mathbb{E}\left[Y|\boldsymbol{X}, T=1\right]-\mathbb{E}\left[Y|\boldsymbol{X}, T=0\right]
\end{align}が成り立つので、そのような仮定のもとではCATEの推定問題は右辺の量を推定する問題に帰着します。

CATEのベイズ推定

それでは次に\(\mathbb{E}\left[Y|\boldsymbol{X}, T=1\right]-\mathbb{E}\left[Y|\boldsymbol{X}, T=0\right]\)の推定を考えますが、ベイズ推定をする場合にはこの量の事後分布を求めれば良いので問題は非常にシンプルです。ベイズ推定をする上で肝になるのが\(\boldsymbol{X}, T, Y\)の関係性のモデル化です。特に重要になるのが、条件付き分布\(p(y|\boldsymbol{x}, t)\)のモデル化で、様々なモデルが提案されています。少し天下り的になりますが、CATE推定においては次のような形の確率モデルが様々な論文で扱われています。
\begin{align}
Y=\theta(\boldsymbol{X})T+f(\boldsymbol{X})+\varepsilon\tag{1}
\end{align}ここで\(\theta, f\)は未知の非線形関数、\(\varepsilon\)は平均0のノイズです。このようなモデルを仮定することの利点として、推定対象のCATEの式\(\mathbb{E}\left[Y|\boldsymbol{X}, T=1\right]-\mathbb{E}\left[Y|\boldsymbol{X}, T=0\right]\)が\(\theta(\boldsymbol{X})\)と一致することが挙げられます（簡単な式展開で確かめられます）。\(\theta(\boldsymbol{X})\)を推定する代表的な手法としては、非ベイズ的な手法ではR-Learner (Nie and Wager 2021)が、ベイズ的な手法ではBayesian Causal Forest (Hahn, Murray, and Carvalho 2020)が挙げられます。

ガウス過程事前分布の利用

モデルが(1)式のように定まって、CATEを推定する問題が\(\theta(\boldsymbol{X})\)を推定する問題に帰着されたので、あとは\(\theta(\boldsymbol{X})\)の事後分布を求めれば良いということになりますが、その前にもう1つ決めないといけないことがあります。それは関数\(\theta, f\)の事前分布です。（ノンパラメトリック）ベイズ的な設定ではこれら未知の関数も確率変数であると考えるので、そこに事前分布を設定する必要があります。先述のBayesian Causal Forestではそれぞれに対してBARTとよばれる木モデルに基づいた事前分布を仮定しています。我々が提案している手法では、これらの事前分布にガウス過程に基づいた事前分布を仮定しています。その意味では、Bayesian Causal Forestとの違いは事前分布のみの違いということになりますが、ガウス過程事前分布を仮定することで次のようなメリットがあります。

• CATEの事後（予測）分布が解析的に求まる（ただしノイズの分散とガウス過程のパラメータを固定したもとで）
• 理論的な解析がBART事前分布を仮定した場合と比較して容易
• 事前知識を事前分布に反映させやすい

論文ではいくつかの条件でR-LearnerやBayesian Causal Forestとの性能比較も行っています。勿論すべての条件で他の手法より優れているわけではありませんが、これらの手法と匹敵するかより良い推定が可能となっていることを確認しています。詳細はぜひ論文をご確認ください。

1. 前の記事で紹介したITEとCATEを混同しないようにしましょう。\(\mathbb{E}\left[Y^{(1)}-Y^{(0)}|\boldsymbol{X}=\boldsymbol{x}\right]\)は\(\boldsymbol{X}\)の値が\(\boldsymbol{x}\)であるような個体群に対する「平均的な」処置効果です。論文でCATEのことをITEと書いたりすると「それはITEじゃなくてCATEだろ！」という査読者からの総ツッコミが入ります。 ↩︎