Turing.jlで欠損値付きデータに対して線形回帰を行う

はじめに

Suyamaさんのブログ*1で，データに欠損値がある場合の教師あり学習について説明されています．本記事では，そちらの記事及びソースコード *2 を参考に，欠損値が含まれるデータに対して線形モデルを仮定し，欠損値と重みをTuring.jlで推論します．またSuyamaさんのブログと同様に，欠損値が含まれるデータを省いて学習を行ったもの，欠損値を推定しながら学習を行ったもの，欠損値がないデータで学習を行ったもので結果を比較してみました．

Turing.jl*3はJulia用のProbabilistic Programmingライブラリで，確率モデルを定義し，推論方法を指定するとパラメータに対する推論を自動で行ってくれます．

実験設定

2次元実数ベクトル $x$ と実数 $y$ の組みの集合が与えられ，線形モデルのパラメタ $w$ を推定したいとします．ただし， $x$ の要素は適当な確率で欠損が発生しています．線形モデルを次のように定義します．ただし， $\lambda_y, \mu_x, \lambda_x, \mu_w, \lambda_{w}$ はハイパーパラメータです．
$\begin{align*} P(y_n | x_n, w) = \mathcal{N}(y_n | w^\mathrm{T}x_n, \lambda_y)\\ P(x_{n,d}) = \mathcal{N}(x_{n,d} | \mu_x, \lambda_{x})\\ P(w_d) = \mathcal{N}(w_d | \mu_w, \lambda_{w}) \end{align*}$
このモデル，入力データ，及びその欠損情報をTuringで記述します．本来Turingでは入力データに欠損が混じってていても内部的に処理して推論してくれる機能が想定されているようなのですが*4，どうやらバグ*5 *6がでているようで，欠損が混じったデータをそのまま入力しても上手く動きませんでした．現状，欠損を分離して記述する必要があるようです．予測モデルは次のように記述しました．モデルの実装にあたって，Stanで実装されていたこちらの記事*7を参考にさせて頂きました．

@model function linear_regression_with_miss(x1, x1_miss, x2, x2_miss, y, N1, N2, N3, N4)
    # hypter paramteters
    mu_x = 0
    lambda_x = 1.5
    lambda_w = 100
    lambda_y = 1
    
    nfeatures = 2
    W ~ MvNormal(nfeatures, lambda_w)
    
    # no missing
    for i in 1:N1
        mu = x1[i] * W[1] + x2[i] * W[2] 
        y[i] ~ Normal(mu, lambda_y)
    end
    
    # only x1 is missed
    for i in N1+1:N1+N2
        x1_miss[i-N1] ~ Normal(mu_x, lambda_x)
        mu = x1_miss[i-N1] * W[1] + x2[i] * W[2] 
        y[i] ~ Normal(mu, lambda_y)  
    end
    
    # only x2 is missed
    for i in N1+N2+1:N1+N2+N3
        x2_miss[i-N1-N2] ~ Normal(mu_x, lambda_x)
        mu = x1[i-N2] * W[1] + x2_miss[i-N1-N2] * W[2] 
        y[i] ~ Normal(mu, lambda_y)  
    end
    
    # both x1 and x2 are missed
    for i in N1+N2+N3+1:N1+N2+N3+N4
        x1_miss[i-N1-N3] ~ Normal(mu_x, lambda_x)
        x2_miss[i-N1-N2] ~ Normal(mu_x, lambda_x)
        mu = x1_miss[i-N1-N3] * W[1] + x2_miss[i-N1-N2] * W[2]
        y[i] ~ Normal(mu, lambda_y)  
    end
end

実験

データ数 $N$ を30, $x$ の要素を50%の確率で欠損にしたデータを用意します．このデータは上で定義した予測モデルのハイパーパラメータの値を少し変えたものからサンプリングしています．NUTSでサンプリングし，得られた結果を可視化したものが次の図になります． f:id:tetuppei:20201028184155p:plain 横軸と縦軸が $x_1$ , $x_2$ の各要素の大きさ，丸い点が各 $x$ ，点の色が $y$ の大きさを表しています. また，等高線が学習した $w$ による予測値になります．図は左から，欠損値があるデータを省いて学習を行ったもの，欠損値を推定しながら学習を行ったもの，欠損値がないデータで学習を行ったものになります．中央の図のx印が欠損値の平均値，青い帯が欠損値の不確実性を表しています．

3つの図の等高線を比較してみると，定性的な比較ではありますが，この例では，欠損値を推定しながら学習を行ったものが(真中の図)，欠損値が含まれるデータを省いて学習を行ったものよりも（左図), 欠損値がないデータで学習を行ったもの(右図)に近い結果になっていることが分かります．

感想

欠損値推定について

もちろん，上のような結果が常に得られるわけではなく，データセットの数や分布，欠損率の値次第では，欠損したものを省いても結果がほとんど変わらない場合もありますし，極端な事前分布を仮定すると違った結果が得られます．実データを解析する際は，無理のない範囲で事前分布を設定するのが良いのだろうと思いました．

Turing.jlについて

まだ，発展途上な感じはしますが，Turing.jlはかなり良いライブラリだと感じました．僕のPPLの知識はStanを少し触ったことがある程度のものなのですが，モデルの記述自体は全然苦労しませんでしたし，Juliaでデータ解析をするのであれば，純正のライブラリということで，インストールからデータの受け渡しまでスムーズでした．欠損値推定の機能なんかもすごく便利だと思います．