ベイズword2vecの進化版をStanで実装してみた

# はじめに

本記事では、word2vecを一般化した指数族embedding（Rudolph et al. 2016）の一種をStanで実装する方法を紹介する。

word2vecやトピックモデルなどの用語を聞くと、テキストデータ分析のことを連想する読者も多いかもしれないが、テキストデータというのは本質的・抽象的にはトークンの羅列から法則性を洗い出すタスクであるため、言語だけでなく、消費者の購買行動やユーザーの閲覧行動にも適応できる。

たとえば、言語を分析対象にする際に

```
日本の最大都市はXXです。
```
のXXの内容を予測するモデルは、

```
りんご、トイレットペーパー、コーラ、XX、炭酸水
```
のような消費者の購買データや
```
東京都新宿区、東京都新宿区、東京都渋谷区、XX、東京都豊島区
```
のようなユーザーの移動データの中のXXの予測にも使える。

もちろん、テキストデータ用に開発されたモデルを他の種類のデータに適用する際は、非言語データの前処理の工夫の必要性や言語に特化した学習済みモデルを安易に利用しないなどの注意点がある。

本記事が利用するデータも言語データではなく、下記のサイトが提供するテレビ番組の閲覧状況のデータである（Turrin, Condorelli, Cremonesi and Pagano 2014）。

https://recsys.deib.polimi.it/datasets/

# 指数族embeddingの簡単な紹介

指数族embeddingとは、Rudolph et al.（2016）が提案した、word2vecをより一般化して、言語データ以外の種類のデータにも拡張した手法である。

指数族embeddingの概念を簡単に説明すると、ある位置にどんなトークンが現れるかは、その周辺に位置するトークンによって決まるということである。また、トークン自体はembeddingベクトルで表現し、その周辺にあるトークンはcontextベクトルで表現される。これに対して、word2vecの場合はトークン自体と周辺にあるトークンを区別せず両者に同じベクトルを適用するが（Mikolov et al. 2013）、ベクトルをembeddingベクトルとcontextベクトルに分けることによって、たとえばembeddingベクトルは時間と共に変動するといったより複雑なモデル構造になった際にembeddingベクトルを比較可能にする役割を果たす（Rudolph and Blei 2018）。

数式で書くと、たとえば

```
日本 の 最大 都市　は　XX　です
```
の中のXXが東京になる確率は

```math

logit（embeddingベクトル_{東京} ' * （contextベクトル_{日本} + contextベクトル_{の} + contextベクトル_{最大} + contextベクトル_{都市} + contextベクトル_{は} + contextベクトル_{です}））

```
となる。


もちろん、上記の事例では全ての単語を周辺（context）と定義したが、word2vecにように分析者自身が周辺はどこまで含めるかを決めることができる。

実際にモデルを推定する際は、ネガティブサンプリングという手法を使う。ネガティブサンプリングの概念を説明すると、実際にあった文章

```
日本の最大都市は東京です
```

の中から適当・ランダムに別の単語を入れて

```
日本の最大都市はとんかつです
```
モデルに前者が正しく後者は正しくないと教えるイメージである。

消費者の購買データで言うと

```
りんご、トイレットペーパー、コーラ、目薬、炭酸水
```
が本当のデータで
```
りんご、トイレットペーパー、コーラ、ニューヨークマンハッタンの高級マンション、炭酸水
```
の中に適当に入れられたアイテムがあるとモデルに学習させることで、
モデルが消費者がどんなものを同じタイミングで買って、どんなものを同じタイミングで買わないかを理解するようになる。

# 分析

ここでは、はじめにのところで紹介したテレビ番組の閲覧状況データを使って指数族embeddingを実装してみる。ただ、プライバシー保護のためか、本データはマスキングされたところが多く、あまりビジネス上の示唆が出ないので、あくまでも実装して精度を確認するところまでをスコープとする。

## 前処理

まずは必要なパッケージを入れた後、上で紹介したサイトより入手したデータを読み込む。ただし、データ量を絞る観点で、ここではweek 1のデータのみ対象とする。

```r
library(tidyverse)
library(cmdstanr)
library(ROCR)
library(Rtsne)
tv_data <- read_csv("tv-audience-dataset.csv",
                    col_names = c("channel_id", "slot", "week",
                                  "genre_id", "subgenre_id", "user_id",
                                  "program_id", "event_id", "duration")) %>%
  filter(week == 1)

```
また、最も試聴された500の番組のみを対象にし、
```r
target_programs <- tv_data %>%
  count(program_id) %>%
  arrange(-n) %>%
  head(500) %>%
  pull(program_id)
```
その中で10以上の番組を視聴したユーザーのみを対象にする。
```r
target_user <- tv_data %>%
  filter(program_id %in% target_programs) %>%
  count(user_id) %>%
  filter(n >= 10) %>%
  pull(user_id)
```
次に、対象だけ残したデータブレイムを作成し、ユーザーの閲覧順番を示すseq_idというカラムを作成する。
```r
tv_data_cleaned <- tv_data %>%
  filter(program_id %in% target_programs) %>%
  filter(user_id %in% target_user) %>%
  arrange(user_id, week, slot) %>%
  group_by(user_id) %>%
  mutate(seq_id = row_number()) 
```
ここで、stanにデータを渡す際に行わないといけないindex変換のためのマスター表を作成する。user_idもprogram_idも既にindexなのになぜまたindex変換するのかと思う読者もいるかもしれないが、user_idとprogram_idが1始まりの連番になる保証は一般的にはないので、エラーと意図しない学習結果を避けるために大人しくマスター表を作成するのが得策である。

```r
user_master <- tibble(
  user_id = sort(unique(tv_data_cleaned$user_id)),
  user_int_no = 1:length(unique(tv_data_cleaned$user_id))
)

program_master <- tibble(
  program_id = sort(unique(tv_data_cleaned$program_id)),
  program_int_no = 1:length(unique(tv_data_cleaned$program_id))
)
```
マスター表をデータにLEFT JOINして、index以外は残さない。
```r
tv_data_cleaned_id <- tv_data_cleaned %>%
  ungroup() %>%
  left_join(user_master, by = "user_id") %>%
  left_join(program_master, by = "program_id") %>%
  select(user_int_no, program_int_no, seq_id)
```
ここで、少し複雑な処理をするが、理由としては本記事で利用するモデルは、ユーザーが前後視聴した5番組（合計10番組）がユーザーが視聴する番組の予測に役立つという仮説をおいているので、その前後の5番組を持ってくる必要がある。

実際にコードを書く前に、まず上記で作成したtv_data_cleaned_idに着目すると、

```r
> tv_data_cleaned_id
# A tibble: 1,553,827 × 3
   user_int_no program_int_no seq_id
         <int>          <int>  <int>
 1           1            386      1
 2           1            436      2
 3           1             48      3
 4           1             48      4
 5           1            158      5
 6           1             49      6
 7           1             49      7
 8           1            496      8
 9           1             48      9
10           1            158     10
# … with 1,553,817 more rows
```
USER 1が386を視聴した後に視聴した番組を取ろうとする場合、

```r
> tv_data_cleaned_id %>%
     mutate(seq_id = seq_id - 1)
# A tibble: 1,553,827 × 3
   user_int_no program_int_no seq_id
         <int>          <int>  <dbl>
 1           1            386      0
 2           1            436      1
 3           1             48      2
 4           1             48      3
 5           1            158      4
 6           1             49      5
 7           1             49      6
 8           1            496      7
 9           1             48      8
10           1            158      9
# … with 1,553,817 more rows
```
このようにseq_idから1を引いて、そのままuser_int_noとseq_idをキーに元のデータにLEFT JOINすれば達成できる。

これを次のように前後後番組についてやれば良いのである。

```r
tv_data_cleaned_id_context <- tv_data_cleaned_id %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id + 1) %>%
      rename(program_int_no_before_1 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id + 2) %>%
      rename(program_int_no_before_2 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id + 3) %>%
      rename(program_int_no_before_3 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id + 4) %>%
      rename(program_int_no_before_4 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id + 5) %>%
      rename(program_int_no_before_5 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id - 1) %>%
      rename(program_int_no_after_1 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id - 2) %>%
      rename(program_int_no_after_2 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id - 3) %>%
      rename(program_int_no_after_3 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id - 4) %>%
      rename(program_int_no_after_4 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  left_join(
    tv_data_cleaned_id %>%
      mutate(seq_id = seq_id - 5) %>%
      rename(program_int_no_after_5 = program_int_no),
    by = c("user_int_no", "seq_id")
  ) %>%
  na.omit()
```

次はネガティブサンプルの抽出を行う。

具体的には本当のデータにシャッフルをかけたら良い。
```r
set.seed(12345)
neg_sample <- sample(
  tv_data_cleaned_id_context$program_int_no
)
```

次に、ネガティブサンプルを含めたデータを作成して、下記のようにresultというカラムも追加して、本当のデータに1を入れて、ネガティブサンプルに0を入れる。

```r
tv_data_cleaned_id_context_with_neg <- tv_data_cleaned_id_context %>%
  mutate(result = 1) %>%
  bind_rows(
    tv_data_cleaned_id_context %>%
      mutate(program_int_no = neg_sample) %>%
      mutate(result = 0)
  )
```

これで、stanで書かれた指数族embeddingモデルが果たして本当のデータ（result = 1）と偽物のデータ（result = 0）を区別できるかを確認するための前処理のほとんどが終わった。

最後は、データ量の問題で、ランダムに2000人のユーザーを選んで、これを学習データにして

```r
set.seed(12345)
train_user <- sample(
  user_master$user_int_no,
  2000,
  replace = FALSE
)

train_data <- tv_data_cleaned_id_context_with_neg %>%
  filter(user_int_no %in% train_user)
```
選ばれた2000ユーザー以外のデータからランダムに10000件をピックアップして検証データとしてモデルの精度を評価する。
```r
val_data <- tv_data_cleaned_id_context_with_neg %>% 
  filter(user_int_no %in% train_user == FALSE)

set.seed(12345)
val_use_id <- sample(
  1:nrow(val_data),
  10000,
  replace = FALSE
)

val_data_sampled <- val_data[val_use_id,]
```

## Stanのコード

学習データの量はほぼ50万で、普通のMCMCだと計算が終わらないので、マルチスレッドで高速化された変分推論でモデルの推定を行う。

マルチスレッド化するためのコードは下記のサイトを参考した：

https://mc-stan.org/users/documentation/case-studies/reduce_sum_tutorial.html

利用するStanのコードbernoulli_embedding.stanはこのセクションの最後のところに記載している。

まず、partial_sumの少しクセのある書き方だが、筆者もまだ使用を完全に理解していると自信を持って言える状態ではないので、理解が間違っているところがあればご指摘いただきたいが、これは要するにstanが勝手にstartからendの数字を指定して一部従属変数のデータを取り出して計算するみたいなので、こちらとしてはその取り出された従属変数に対応する他のデータを用意してあげないといけない。上の消費者の事例で言うと、

```math
embeddingベクトル_{ニューヨークマンハッタンの高級マンション} ' * （contextベクトル_{りんご} + contextベクトル_{トイレットペーパー} + contextベクトル_{コーラ} + contextベクトル_{炭酸水}）
```
のところの内積の結果を用意しないといけない。

用意した内積の結果を長さend - start + 1のlambdaという変数に格納して、最終的にはbernoulli_logit_lupmf(result | lambda)のところでstanが勝手に取り出した従属変数（本当のデータかどうかを示すresult）とお馴染みのロジスティック回帰を行う時のようなやり方で評価された尤度関数を出力させる。

また、これはベイズモデルなので、embeddingベクトルとcontextベクトルの事前分布に正規分布を設定する。


```cpp:bernoulli_embedding.stan
functions {
  real partial_sum_lpmf(
    array[] int result,
    int start, int end,
    
    array[] int program,
    
    array[] int program_before_1, array[] int program_after_1,
    array[] int program_before_2, array[] int program_after_2,
    array[] int program_before_3, array[] int program_after_3,
    array[] int program_before_4, array[] int program_after_4,
    array[] int program_before_5, array[] int program_after_5,
    
    array[] vector program_embedding,
    array[] vector program_context
  ){
    vector[end - start + 1] lambda;
    int count = 1;
    for (i in start:end){
      lambda[count] = program_embedding[program[i]] ' *
      (
        program_context[program_before_1[i]] + program_context[program_after_1[i]] + 
        program_context[program_before_2[i]] + program_context[program_after_2[i]] + 
        program_context[program_before_3[i]] + program_context[program_after_3[i]] + 
        program_context[program_before_4[i]] + program_context[program_after_4[i]] + 
        program_context[program_before_5[i]] + program_context[program_after_5[i]] 
      );
      count += 1;
    }
    return (
      bernoulli_logit_lupmf(result | lambda)
    );
  }
}
data {
  int<lower=1> N; //学習データ数
  int<lower=1> K; //embedding次元数
  int<lower=1> program_type; //番組数
  
  //学習データ
  array[N] int<lower=1,upper=program_type> program;
  array[N] int<lower=1,upper=program_type> program_before_1;
  array[N] int<lower=1,upper=program_type> program_after_1;
  array[N] int<lower=1,upper=program_type> program_before_2;
  array[N] int<lower=1,upper=program_type> program_after_2;
  array[N] int<lower=1,upper=program_type> program_before_3;
  array[N] int<lower=1,upper=program_type> program_after_3;
  array[N] int<lower=1,upper=program_type> program_before_4;
  array[N] int<lower=1,upper=program_type> program_after_4;
  array[N] int<lower=1,upper=program_type> program_before_5;
  array[N] int<lower=1,upper=program_type> program_after_5;
  array[N] int<lower=0,upper=1> result;
  
  int<lower=0> val_N; //検証データ数
  //検証データ
  array[val_N] int<lower=1,upper=program_type> val_program;
  array[val_N] int<lower=1,upper=program_type> val_program_before_1;
  array[val_N] int<lower=1,upper=program_type> val_program_after_1;
  array[val_N] int<lower=1,upper=program_type> val_program_before_2;
  array[val_N] int<lower=1,upper=program_type> val_program_after_2;
  array[val_N] int<lower=1,upper=program_type> val_program_before_3;
  array[val_N] int<lower=1,upper=program_type> val_program_after_3;
  array[val_N] int<lower=1,upper=program_type> val_program_before_4;
  array[val_N] int<lower=1,upper=program_type> val_program_after_4;
  array[val_N] int<lower=1,upper=program_type> val_program_before_5;
  array[val_N] int<lower=1,upper=program_type> val_program_after_5;
}
parameters {
  array[program_type] vector[K] program_embedding;
  array[program_type] vector[K] program_context;
}
model {
  for (p in 1:program_type){
    target += normal_lupdf(program_embedding[p] | 0, 1);
    target += normal_lupdf(program_context[p] | 0, 1);
  }
  
  int grainsize = 1;
  
  target += reduce_sum(
    partial_sum_lupmf, result, grainsize,
    program,
    program_before_1, program_after_1,
    program_before_2, program_after_2,
    program_before_3, program_after_3,
    program_before_4, program_after_4,
    program_before_5, program_after_5,
    program_embedding,
    program_context
  );
}
generated quantities {
  array[val_N] int predicted;
  for (n in 1:val_N){
    predicted[n] = bernoulli_logit_rng(
      program_embedding[val_program[n]] ' *
      (
        program_context[val_program_before_1[n]] + program_context[val_program_after_1[n]] +
        program_context[val_program_before_2[n]] + program_context[val_program_after_2[n]] + 
        program_context[val_program_before_3[n]] + program_context[val_program_after_3[n]] + 
        program_context[val_program_before_4[n]] + program_context[val_program_after_4[n]] + 
        program_context[val_program_before_5[n]] + program_context[val_program_after_5[n]]
      )
    );
  }
}

```

## 分析の実行

ここで、上記のStanのコードを利用してモデルの推定を行う。

まずStanにデータを渡すためのlistを作成し、モデルをコンパイルする。

```r
data_list <- list(
  N = nrow(train_data),
  K = 50,
  program_type = nrow(program_master),
  
  program = train_data$program_int_no,
  program_before_1 = train_data$program_int_no_before_1,
  program_after_1 = train_data$program_int_no_after_1,
  program_before_2 = train_data$program_int_no_before_2,
  program_after_2 = train_data$program_int_no_after_2,
  program_before_3 = train_data$program_int_no_before_3,
  program_after_3 = train_data$program_int_no_after_3,
  program_before_4 = train_data$program_int_no_before_4,
  program_after_4 = train_data$program_int_no_after_4,
  program_before_5 = train_data$program_int_no_before_5,
  program_after_5 = train_data$program_int_no_after_5,
  result = train_data$result,
  
  val_N = nrow(val_data_sampled),
  val_program = val_data_sampled$program_int_no,
  val_program_before_1 = val_data_sampled$program_int_no_before_1,
  val_program_after_1 = val_data_sampled$program_int_no_after_1,
  val_program_before_2 = val_data_sampled$program_int_no_before_2,
  val_program_after_2 = val_data_sampled$program_int_no_after_2,
  val_program_before_3 = val_data_sampled$program_int_no_before_3,
  val_program_after_3 = val_data_sampled$program_int_no_after_3,
  val_program_before_4 =val_data_sampled$program_int_no_before_4,
  val_program_after_4 = val_data_sampled$program_int_no_after_4,
  val_program_before_5 = val_data_sampled$program_int_no_before_5,
  val_program_after_5 = val_data_sampled$program_int_no_after_5
)

m_init <- cmdstan_model("bernoulli_embedding.stan",
                        cpp_options = list(
                          stan_threads = TRUE
                        )
                        )

```
では、早速実際の学習状況と時間を確認する。

```r
> m_estimate <- m_init$variational(
+     seed = 12345,
+     threads = 16,
+     iter = 50000,
+     data = data_list
+ )
------------------------------------------------------------ 
EXPERIMENTAL ALGORITHM: 
  This procedure has not been thoroughly tested and may be unstable 
  or buggy. The interface is subject to change. 
------------------------------------------------------------ 
Gradient evaluation took 1.18584 seconds 
1000 transitions using 10 leapfrog steps per transition would take 11858.4 seconds. 
Adjust your expectations accordingly! 
Begin eta adaptation. 
Iteration:   1 / 250 [  0%]  (Adaptation) 
Iteration:  50 / 250 [ 20%]  (Adaptation) 
Iteration: 100 / 250 [ 40%]  (Adaptation) 
Iteration: 150 / 250 [ 60%]  (Adaptation) 
Iteration: 200 / 250 [ 80%]  (Adaptation) 
Success! Found best value [eta = 1] earlier than expected. 
Begin stochastic gradient ascent. 
  iter             ELBO   delta_ELBO_mean   delta_ELBO_med   notes  
   100     -1492762.869             1.000            1.000 
   200      -308756.865             2.417            3.835 
   300      -247626.461             1.694            1.000 
   400      -233182.009             1.286            1.000 
   500      -224296.240             1.037            0.247 
   600      -218487.531             0.868            0.247 
   700      -212941.822             0.748            0.062 
   800      -213475.376             0.655            0.062 
   900      -211990.665             0.583            0.040 
  1000      -210287.480             0.525            0.040 
  1100      -209804.888             0.478            0.027 
  1200      -209289.428             0.438            0.027 
  1300      -209067.521             0.405            0.026 
  1400      -208690.915             0.376            0.026 
  1500      -209100.194             0.351            0.008   MEDIAN ELBO CONVERGED 
Drawing a sample of size 1000 from the approximate posterior...  
COMPLETED. 
Finished in  2082.6 seconds.
```
筆者が利用するMacbook Airで2082.6秒をかけた。爆速とはいえないが現実的な時間であるといえよう。

次に、必要な事後分布のサマリーのデータを取り出して
```r
m_summary <- m_estimate$summary()
```
判別の結果を見てみる：
```r
table(
     round(m_summary[which(str_detect(m_summary$variable, "predict")),]$median),
     val_data_sampled$result
 )
   
       0    1
  0 4296  734
  1  738 4232
```

見てわかるように、かなりいい感じに分類できている。

最後に、ROC曲線を確認しよう。

```r
pred <- prediction(
  m_summary[which(str_detect(m_summary$variable, "predict")),]$mean,
  val_data_sampled$result
)

performance(pred, "tpr", "fpr") %>%
  plot(main = str_c("AUC = ", round(performance(pred, "auc")@y.values[[1]], 3)))
```

![roc.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/085bfc92-bfac-94fb-0c48-19d01a91c6e5.png)

AUCは0.9を超えており、高い精度を達成できたと判断できる。

今回のデータは色々マスキングされた影響で本記事では詳細を紹介しないが、推定されたembeddingベクトルを使って番組同士の視聴パターンを調べたり、もしくはユーザーに番組のリマインドを送るなどのレコメンドシステム構築などに活用する方法がある。

# 参考文献

Mikolov, Tomas, et al. "Distributed representations of words and phrases and their compositionality." Advances in neural information processing systems 26 (2013).

Rudolph, Maja, et al. "Exponential family embeddings." Advances in Neural Information Processing Systems 29 (2016).

Rudolph, Maja, and David Blei. "Dynamic embeddings for language evolution." Proceedings of the 2018 world wide web conference. 2018.

Turrin, R., Condorelli, A., Cremonesi, P., and Pagano, R. “Time-based TV programs prediction. In 1st Workshop on Recommender Systems for Television and Online Video”, RecSys 2014