ベイズword2vec版トピックモデルをStanで実装してみた

# はじめに

こんにちは、事業会社で働いているデータサイエンティストです。

https://developer.dip-net.co.jp/entry/2024/07/22/JSQPS_2024_Summer_Meeting%E3%81%AB%E3%81%A6%E5%BC%8A%E7%A4%BE%E3%81%AE%E3%83%87%E3%83%BC%E3%82%BF%E3%82%B5%E3%82%A4%E3%82%A8%E3%83%B3%E3%83%86%E3%82%A3%E3%82%B9%E3%83%88%E3%81%8C%E7%99%BA%E8%A1%A8

この記事では、下記の論文を参考に、

https://aclanthology.org/2020.tacl-1.29/

word2vecの埋め込み要素を含めたトピックモデルをStanで作成してみます。

Topic Modeling in Embedding Spaces論文との違いは下記の2点にあります：

- 本記事のモデルは、ドキュメントはトピック分布を持つのではなく、一つのトピックに属すると仮定する
- 本記事のモデルは、トピック数と有効次元数を同時に自動で推定する

では早速モデルの説明からはじめていきます！

# モデル紹介

数理的な説明はこちらの記事を参照してください：

https://qiita.com/Gotoubun_taiwan/items/6c30f762cb9bff91b5d5

ここでは、まずモデルの概念から説明します。

言葉で説明すると、文書を生成する時、まずはトピックのインデックスをカテゴリ分布からサンプリングして、そのトピックに対応するトピックベクトルを持ってきます。とある単語がその文書で出てきやすいかどうかは、トピックベクトルと単語ベクトルの内積の大きさで決められます。概念は単純です。

ここで強調しておきたいのは、構造トピックモデルなどのLDA系のトピックモデルでは、文書が持つのはトピックの分布で、そのトピックの分布から一個一個の単語をサンプリングしていくという作りなので、一つの文書は複数のトピックを持ちうるのですが、本記事のモデルでは文書は一個のトピックしか持てません。なのでトピック同士の相関などをはかるのは少々難しいです。

次に、Topic Modeling in Embedding Spaces論文、ないしは最先端のトランスフォーマ系の手法までそうなんですが、次元数は魔法の数字として与えられています。トピック数もそうです。100とか、64とか、38とか、よくよく考えたら正当化が難しそうな数字になっています。もちろん、交差検証や情報基準で次元数とトピック数を自動選択できますが、複数のモデルを推定する必要があります。また、場合によっては、次元数とトピック数自体にもビジネス的・アカデミア的な意味がありますので、次元数とトピック数もパラメータの一部として推定させたほうが望ましいです。想像してみてください、38次元で表現できるデータを、あなたは本当に5000次元で表現しますか？結構S3・GCSのデータ格納料金が飛びますよ。

ということで、次元数とトピック数を両方推定できるword2vec系のトピックモデルです！詳細の数理のところに入ります。

## トピック分布

トピックの分布をディリクレ過程で表現します。

まず全体で何個のトピックが「起動」されるかをコントロールする$\alpha$をサンプリングします。

$$
\alpha \sim Gamma(0.001, 0.001)
$$

次に、棒折り表現でトピックpが抽出される確率$p_{p}$を作成します

$$
\pi_{p}\sim Beta(1, \alpha)
$$

$$
p_{p} = \pi_{p} \prod\limits_{l=1}^{p - 1} (1 - \pi_{l})
$$

## 次元分布

次元数は棒折り過程で表現します。

まずトピック数側の$\alpha$のように、全体で何個の次元が「起動」されるかをコントロールするパラメータをサンプリングします。

$$
\beta \sim Gamma(0.001, 0.001)
$$

次に棒折り過程で$次元重要度ベクトル_{1,d}$を求める

$$
\delta_{d}\sim Beta(1, \beta)
$$

$$
次元重要度ベクトル_{d} = \delta_{d} \prod\limits_{l=1}^{d - 1} (1 - \delta_{l})
$$

## 文書生成

文書dの生成に関しては、まずカテゴリ分布からトピックを抽出します：

$$
トピック_{d} \sim Categorical(p)
$$

次に、本物のテキストデータと偽物のテキストデータを混ぜて、モデルに判別させます。概念はこちらの記事からご確認ください：

https://ggr.hias.hit-u.ac.jp/2023/12/25/record-china-discourse-visualized-by-exponential-family-embeddings/

本物データと偽物データを混ぜたデータのi番目について

$$
本物フラグ\_{i} \sim Bernoulli(logit(単語人気度_{単語_{i}} + log(ドキュメントの長さ_{ドキュメント_{i}}) + \sum\limits_{a=1}^{\infty} (トピックベクトル_{トピック_{d}, a} * 単語ベクトル_{a} * 次元重要度ベクトル_{d, a})))
$$

という無限項目の足し合わせロジット変換で判別すると設定します。

# モデル実装

## Stanコード

モデルのStan言語での実装はこちらになります：

```cpp:topic_embedding.stan
functions {
   vector stick_breaking(vector breaks){
    int length = size(breaks) + 1;
    vector[length] result;
    
    result[1] = breaks[1];
    real summed = result[1];
    for (d in 2:(length - 1)) {
      result[d] = (1 - summed) * breaks[d];
      summed += result[d];
    }
    result[length] = 1 - summed;
    
    return result;
  }
  
  real weighted_inner_product(
    vector vec_a, vector vec_b, vector weight
  ){
    vector[size(vec_a)] result_record;
    for (i in 1:size(vec_a)){
      result_record[i] = vec_a[i] * vec_b[i] * weight[i];
    }
    return sum(result_record);
  }
  
  real empirical_cumulative_density(
    real x, vector distribution
  ){
    vector[size(distribution)] result;
    for (i in 1:size(distribution)){
      if (x >= distribution[i]){
        result[i] = 1.0;
      } else {
        result[i] = 0.0;
      }
    }
    return mean(result);
  }
  
  real partial_sum_lpmf(
    array[] int document_id,
    
    int start, int end,
    
    int topic_type,
    
    array[] int document_begin, array[] int document_end,
    
    array[] int word, array[] real log_m,
    
    array[] int frequency,
    
    vector topic, vector dimension,
    
    vector word_propensity,
    array[] vector word_embedding, array[] vector topic_embedding
  ){
    vector[end - start + 1] lambda;
    int count = 1;
    for (i in start:end){
      vector[topic_type] case_vector = log(topic);
      for (t in 1:topic_type){
        for (j in document_begin[i]:document_end[i]){
          case_vector[t] += poisson_log_lpmf(
            frequency[j] | 
            word_propensity[word[j]] + 
            log_m[j] + 
            weighted_inner_product(word_embedding[word[j]], topic_embedding[t], dimension)
          );
        }
      }
      lambda[count] = log_sum_exp(case_vector);
      count += 1;
    }
    return sum(lambda);
  }
}
data {
  int word_type;
  int dimension_type;
  int topic_type;
  
  int document_type;
  array[document_type] int document_id;
  array[document_type] int document_begin;
  array[document_type] int document_end;
  
  int N;
  array[N] int word;
  array[N] real log_m;
  array[N] int frequency;
}
parameters {
  real<lower=0> dimension_alpha;
  
  vector<lower=0, upper=1>[dimension_type - 1] dimension_breaks;
  
  real<lower=0> topic_alpha;                                       // ディリクレ過程の全体のパラメータ
  vector<lower=0, upper=1>[topic_type - 1] topic_breaks;  // ディリクレ過程のstick-breaking representationのためのパラメータ
  
  vector[word_type] word_propensity;
  array[word_type] vector[dimension_type] word_embedding;
  array[topic_type] vector[dimension_type] topic_embedding;
}
transformed parameters {
  simplex[dimension_type] dimension;
  simplex[topic_type] topic;
  
  dimension = stick_breaking(dimension_breaks);
  topic = stick_breaking(topic_breaks);
}
model {
  dimension_alpha ~ gamma(0.001, 0.001);
  dimension_breaks ~ beta(1, dimension_alpha);
  
  topic_alpha ~ gamma(0.001, 0.001);
  topic_breaks ~ beta(1, topic_alpha);
  
  word_propensity ~ normal(0, 1);
  for (i in 1:word_type){
    word_embedding[i] ~ normal(0, 1);
  }
  for (i in 1:topic_type){
    topic_embedding[i] ~ normal(0, 1);
  }
  
  int grainsize = 1;
  
  target += reduce_sum(
    partial_sum_lupmf, document_id,
    
    grainsize,
    
    topic_type,
    
    document_begin, document_end,
    
    word, log_m,
    
    frequency,
    
    topic, dimension,
    
    word_propensity,
    word_embedding, topic_embedding
  );
}
generated quantities {
  matrix[word_type, topic_type] frex;
  {
    matrix[word_type, topic_type] B;
    
    for (w in 1:word_type){
      for (t in 1:topic_type){
        B[w, t] = weighted_inner_product(word_embedding[w], topic_embedding[t], dimension);
      }
    }
    
    for (w in 1:word_type){
      for (t in 1:topic_type){
        frex[w, t] = ((0.5/empirical_cumulative_density(B[w, t], B[w,:]')) + (0.5/empirical_cumulative_density(B[w, t], B[:,t])))^(-1);
      }
    }
  }
}
```

## frexの説明

Stanのコードにあるfrexなんですが、少し言葉で説明します。frexというのは、トピックの特徴的な単語を割り出す手法です。トピックの特徴的な単語？そんなものは該当のトピックベクトルと単語の内積を計算して、その値が最も高い単語を出せば良いのではという考えもあると思います。

ただ、高校の成績を考えてください。Aさんは数学で90点を取って、Aさん自身の中では最高得点の科目になりました。でも、だからと言って、Aさんは数学が最も得意・特徴的と言えますか？実は今回のテストの結果、半分の人は数学満点で、逆にAさんが38点を取った物理がクラス最高得点になりました。この例からも分かるように、特徴的な科目を判断する時、その人の中の各科目の得点順位はもちろん、他の人の点数と比較する必要があります。

同じように、特徴的な単語かどうかを判断するとき、トピック内での順位と、トピック間の順位を両方見る必要があります。このトピック内での順位とトピック間の順位を同時に考慮してくれるのがfrexです。また別の記事で詳細を説明します。

# モデル推定

まずは前処理とモデルのコンパイルを行います。

前処理の詳細はこちらの記事を参照してください：

https://qiita.com/Gotoubun_taiwan/items/1b8f952db8d58c2d7ad9

```R
mecabbing <- function(text){
  this_review <- stringr::str_replace_all(text, "[^一-龠ぁ-んーァ-ヶー]", " ")
  mecab_output <- unlist(RMeCab::RMeCabC(this_review, 1))
  mecab_combined <- stringr::str_c(mecab_output[which(names(mecab_output) == "名詞")], collapse = " ")
  return(mecab_combined)
}

`%>%` <- magrittr::`%>%`

review_df <- list.files("text") %>%
  .[which(stringr::str_detect(., ".txt") == FALSE)] |>
  purrr::map(
    \(this_file){
      this_file %>%
        stringr::str_c("text/", .)  |> 
        list.files() %>%
        .[which(stringr::str_detect(., "LICENSE") == FALSE)] %>% 
        stringr::str_c("text/", this_file, "/", .) |>
        purrr::map(
          \(this_text){
            this_text |>
              readr::read_lines() |>
              stringr::str_c(collapse = " ") |>
              mecabbing() |>
              stringr::str_remove_all(
                quanteda::stopwords("ja", source = "marimo") |> 
                  c("的", "の", "こと", "さ", "ら", "今", "いま") |>
                  stringr::str_c(collapse = "|")
              ) |>
              tibble::tibble(
                category = this_file,
                text = _
              )
          }
        )
    },
    .progress = TRUE
  ) |>
  dplyr::bind_rows() |>
  dplyr::mutate(
    doc_id = dplyr::row_number()
  )

review_dfm_df <- review_df |>
  dplyr::pull(text) |>
  quanteda::phrase() |>
  quanteda::tokens() |>
  quanteda::tokens_ngrams(1:2) |>
  quanteda::dfm() |>
  quanteda::dfm_trim(min_docfreq = 100) |>
  quanteda::convert(to = "tripletlist") |>
  tibble::as_tibble() |>
  dplyr::mutate(
    document = document |>
      stringr::str_remove_all("text") |>
      as.integer()
  )

word_master <- review_dfm_df |>
  dplyr::select(feature) |>
  dplyr::distinct() |>
  dplyr::arrange(feature) |>
  dplyr::mutate(
    word_id = dplyr::row_number()
  )

category_master <- review_df |>
  dplyr::select(category) |>
  dplyr::distinct() |>
  dplyr::arrange(category) |>
  dplyr::mutate(
    category_id = dplyr::row_number()
  )

set.seed(12345)
review_dfm_neg_df <- review_dfm_df |>
  split(~ document) |>
  purrr::map(
    \(this_df){
      words_not_used <- word_master$feature[which(word_master$feature %in% this_df$feature == FALSE)]
      
      this_df |>
        dplyr::bind_rows(
          tibble::tibble(
            document = this_df$document[1],
            feature = sample(words_not_used, sum(this_df$frequency), replace = FALSE),
            frequency = 0
          )
        )
    },
    .progress = TRUE
  ) |>
  dplyr::bind_rows() |>
  dplyr::group_by(document) |>
  dplyr::mutate(
    log_m = log(sum(frequency))
  ) |>
  dplyr::ungroup() |>
  dplyr::left_join(
    review_df |>
      dplyr::select(doc_id, category),
    by = c("document" = "doc_id")
  ) |>
  dplyr::left_join(
    word_master, by = "feature"
  ) |>
  dplyr::left_join(
    category_master, by = "category"
  ) |>
  dplyr::arrange(document) |>
  dplyr::mutate(
    position_id = dplyr::row_number()
  )

document_begin_end <- review_dfm_neg_df |>
  dplyr::group_by(document) |>
  dplyr::summarise(
    begin = min(position_id),
    end = max(position_id)
  ) |>
  dplyr::ungroup()

data_list <- list(
  word_type = nrow(word_master),
  dimension_type = 20,
  topic_type = 10,
  
  document_type = nrow(document_begin_end),
  document_id = 1:nrow(document_begin_end),
  document_begin = document_begin_end$begin,
  document_end = document_begin_end$end,
  
  N = nrow(review_dfm_neg_df),
  word = review_dfm_neg_df$word_id,
  log_m = review_dfm_neg_df$log_m,
  frequency = review_dfm_neg_df$frequency
)

m_te_init <- cmdstanr::cmdstan_model("topic_embedding.stan",
                                      cpp_options = list(
                                        stan_threads = TRUE
                                      )
)
```

では、ここで実際にモデルの推定に入ります！

```R
> m_te_estimate <- m_te_init$variational(
     seed = 12345,
     threads = 30,
     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 2.9503 seconds 
1000 transitions using 10 leapfrog steps per transition would take 29503 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) 
Iteration: 250 / 250 [100%]  (Adaptation) 
Success! Found best value [eta = 0.1]. 
Begin stochastic gradient ascent. 
  iter             ELBO   delta_ELBO_mean   delta_ELBO_med   notes  
   100   -110239617.485             1.000            1.000 
   200    -24908230.739             2.213            3.426 
   300     -7823593.519             2.203            2.184 
   400     -3878218.144             1.907            2.184 
   500     -2701099.331             1.613            1.017 
   600     -2320204.376             1.371            1.017 
   700     -2148269.067             1.187            1.000 
   800     -2065939.231             1.043            1.000 
   900     -2018648.732             0.930            0.436 
  1000     -1977189.852             0.839            0.436 
  1100     -1955409.123             0.764            0.164   MAY BE DIVERGING... INSPECT ELBO 
  1200     -1934454.487             0.701            0.164   MAY BE DIVERGING... INSPECT ELBO 
  1300     -1924955.107             0.648            0.080   MAY BE DIVERGING... INSPECT ELBO 
  1400     -1908584.651             0.602            0.080   MAY BE DIVERGING... INSPECT ELBO 
  1500     -1896092.807             0.562            0.040   MAY BE DIVERGING... INSPECT ELBO 
  1600     -1887161.626             0.527            0.040   MAY BE DIVERGING... INSPECT ELBO 
  1700     -1877048.448             0.497            0.023 
  1800     -1869839.103             0.469            0.023 
  1900     -1863508.472             0.445            0.021 
  2000     -1859368.266             0.423            0.021 
  2100     -1848919.921             0.403            0.011 
  2200     -1843162.825             0.385            0.011 
  2300     -1835523.326             0.368            0.011 
  2400     -1829105.256             0.353            0.011 
  2500     -1820052.561             0.339            0.009   MEDIAN ELBO CONVERGED 
Drawing a sample of size 1000 from the approximate posterior...  
COMPLETED. 
Finished in  17533.0 seconds.
```

かなり時間がかかりました、、、

結果を保存します：

```R
m_te_summary <- m_te_estimate$summary()
```

# 結果の可視化

まず、次元数を可視化しましょう：

```R
m_te_summary |>
  dplyr::filter(stringr::str_detect(variable, "dimension\\[")) |>
  dplyr::bind_cols(
    dimension_id = 1:20
  ) |>
  ggplot2::ggplot() +
  ggplot2::geom_bar(ggplot2::aes(x = as.factor(dimension_id), y = mean), 
                    stat = "identity", fill = ggplot2::alpha("blue", 0.3)) + 
  ggplot2::geom_errorbar(
    ggplot2::aes(x = as.factor(dimension_id), ymin = q5, ymax = q95),
    width = 0.2
  ) + 
  ggplot2::labs(
    x = "dimension id",
    y = "importance"
  )
```

![dimension.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/cde0dccc-fdaf-3afc-a143-21127e1e14a1.png)

結局二つくらいの次元でこの複雑なコーパスを説明できてしまうのか、、、

単純すぎます泣

次はトピックの方を確認します。

```R
m_te_summary |>
  dplyr::filter(stringr::str_detect(variable, "topic\\[")) |>
  dplyr::bind_cols(
    topic_id = 1:10
  ) |>
  ggplot2::ggplot() +
  ggplot2::geom_bar(ggplot2::aes(x = as.factor(topic_id), y = mean), 
                    stat = "identity", fill = ggplot2::alpha("blue", 0.3)) + 
  ggplot2::geom_errorbar(
    ggplot2::aes(x = as.factor(topic_id), ymin = q5, ymax = q95),
    width = 0.2
  ) + 
  ggplot2::labs(
    x = "topic id",
    y = "distribution"
  )
```

![topic.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/5bbd5e01-d62c-37f3-24b0-8084d35a30cc.png)

トピックも似たような感じで、重要度が高いものはあまりないです。


次に、

```R
10 |>
  seq_len() |>
  purrr::map(
    \(i){
      m_te_summary |>
        dplyr::filter(stringr::str_detect(variable, "frex")) |>
        dplyr::mutate(
          id = variable |>
            purrr::map(
              \(x){
                stringr::str_split(x, "\\[|\\]|,")[[1]][2:3]
              }
            )
        ) |>
        tidyr::unnest_wider(id, names_sep = "_") |>
        dplyr::mutate(
          dplyr::across(dplyr::starts_with("id"),
                        ~ as.integer(.))
        ) |>
        dplyr::left_join(
          word_master, by = c("id_1" = "word_id")
        ) |>
        dplyr::filter(id_2 == i) |>
        dplyr::arrange(-mean) |>
        dplyr::select(feature) |>
        dplyr::slice_head(n = 20)
    }
  ) |>
  dplyr::bind_cols() |>
  `colnames<-`(stringr::str_c("topic", 1:10)) |>
  dplyr::glimpse()
New names:
• `feature` -> `feature...1`
• `feature` -> `feature...2`
• `feature` -> `feature...3`
• `feature` -> `feature...4`
• `feature` -> `feature...5`
• `feature` -> `feature...6`
• `feature` -> `feature...7`
• `feature` -> `feature...8`
• `feature` -> `feature...9`
• `feature` -> `feature...10`
Rows: 20
Columns: 10
$ topic1  <chr> "ん", "男", "映画", "女", "円", "相手", "性", "掲示板", "参加", "も", "心", "料理", "監督", "インタビュー", "放送", "関連", "発言", "器", "位", "賞"
$ topic2  <chr> "男", "ん", "女", "掲示板", "も", "発言", "放送", "性", "賞", "料理", "代表", "インタビュー", "参加", "相手", "映画", "肌", "関連", "位", "監督", "心"
$ topic3  <chr> "ソフトウェア", "更新", "ソフトバンク", "ソフトウェア_更新", "スマート", "コア", "電池", "デュアル_コア", "デュアル", "ライブドアブログ", "本体", "接続", "対応", "キーボ…
$ topic4  <chr> "参加", "放送", "肌", "鑑賞", "会話", "頃", "サイト", "春", "位", "駅", "事務所", "監督", "クリスマス", "アップ", "地", "店", "記憶", "距離", "も", "学生"
$ topic5  <chr> "黒", "リンク", "マン", "ロード", "デスク", "納得", "場面", "クチコミ", "特集_ページ", "壁", "コントロール", "キヤノン", "応募_方法", "追求", "略", "映画_部", "試し", "…
$ topic6  <chr> "器", "モテ", "発表", "ほう", "ブック", "友人", "アンケート", "円", "ザ", "意見", "アップル", "ゲーム", "ランキング", "影響", "相手", "エピソード", "ゴルフ", "子ども", "…
$ topic7  <chr> "ソフトウェア", "ソフトバンク", "更新", "ソフトウェア_更新", "スマート", "デュアル_コア", "対応", "本体", "ライブドアブログ", "フォン", "認識", "コア", "機能", "手順", "…
$ topic8  <chr> "円", "デュアル_コア", "認識", "ファイル", "発売_予定", "ブック", "メモリー", "収_万", "アップデート", "モード", "手順", "発表", "バンド", "フォン", "連続", "音声", "コ…
$ topic9  <chr> "キーボード", "接続", "赤線", "タブレット", "フォン_タブレット", "更新", "夏", "動作", "提供_開始", "地", "電池", "春", "経由", "ケータイ", "操作", "デュアル", "着信", "…
$ topic10 <chr> "田", "クリスマス", "達", "君", "リラックス", "応援", "食", "大会", "冬", "事態", "菜", "タイ", "タイム", "事務所", "センス", "チーム", "大丈夫", "勇気", "率", "精神"
```

かなりいい感じのトピックが出ています！

# 結論

いかがでしょうか？この分析自体からはあまりこれといった示唆は出せないですが、このモデルは問題なくword2vec版のトピックモデルとして機能することが分かりましたので、今後はこのモデルでよりサブスタンスのあるデータを分析します！

最後に、私たちと働いて、データサイエンスの力で社会を改善したい方はこちらをご確認ください：

https://www.wantedly.com/projects/1021866