ベイズ推論で単語の次元数とクラスタ数を自動推定するGloVe拡張モデルを作ってみた

# はじめに

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

https://dippeople.dip-net.jp/19888/

この記事では、単語の分散表現を推定するGloVeモデルを、ノンパラメトリックベイズの棒折り過程およびディリクレ過程で拡張し、次元数とクラスタ数をモデル内で自動的に推定できる手法を紹介します。

理論的にも研究として発展させられる可能性がありそうなので、今後はより本格的な分析を進めていく予定ですが、本記事ではまずそのアイデアの概要を取り急ぎ共有したいと思います！

# モデル概念

GloVeの元論文はこちらから確認できます：

https://aclanthology.org/D14-1162/

本記事では、GloVeの考え方をベイズモデルとして拡張し、さらにノンパラメトリックな構造（棒折り過程など）を導入することで、従来は分析者が恣意的・探索的に決めていた「次元数」と、通常は固定されがちな「単語ベクトルの事前分布」の両方を、モデル内部で自動的に推定する手法を紹介します。

まずは、単語同士の関係を捉えるための基本的な出発点となる「単語の共起行列（co-occurrence matrix）」について説明します。

共起とは、テキスト内である単語が他の単語とどれだけ近い位置で出現したかを数え上げた行列です。たとえば、「UIの改善でユーザー数が上がった」という文章があるとき、「UI」と「ユーザー」は一定の距離内で同時に出現しているため、それらの共起頻度が高まります。これにより、単語間の意味的な関連性を数値的に捉えることができます。

この共起情報をもとに、単語を意味的に類似したベクトル空間へとマッピングしていくのが、GloVeや本モデルの目的です。

共起行列は、ある単語が他の単語とどれだけ近くに出現したかを数え上げた行列です。この行列をうまく構築することで、語と語の意味的な類似性をベクトルで表現できるようになります。

まず、「共起」とはどの程度近くを意味するのか？を定義する必要があります。これをコントロールするのが「ウィンドウサイズ」です。（単語分割の問題のない）英語で実例を確認しましょう。

ウィンドウサイズが2の場合、「対象単語」の前後2語以内に出現する単語を共起語とみなします。（単語分割の問題のない）英語で実例を確認しましょう。

- 例：`I enjoy reading books about international relations theory.`  
  中心単語が `books` のとき、ウィンドウ内の単語は `[enjoy, reading, about, international]` になります。

テキスト全体を左から右へと読み進めながら、各単語について、ウィンドウ内にある他の単語とのペアをカウントしていき、単語の共起行列に記録します。

具体的な実装方法はこちらを参照してください：

https://quanteda.io/reference/fcm.html

今回の記事で提案するモデルは、単語共起行列の確率分布を表すものになります。

# モデル定式化

ここでは、提案モデルをベイズ機械学習の枠組みで数学的に定式化します。

## 次元数の自動推定

まず、次元の重要度を決定するための棒折り過程のハイパーパラメータ$\gamma$を以下のようにサンプリングします：

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

続いて、棒折り過程過程により、次元ごとの重要度$\omega_{d}$を構築します：

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

$$
\omega_{d} = \delta_{d} \prod\limits_{l=1}^{d - 1} (1 - \delta_{l})
$$

この処理を$d = 1$から$d = \infty$まで繰り返します。

ここでサンプリングされる単語分散表現ベクトルは**無限次元**を持ちますが、実際には多くの次元に対して$\omega_d$が0に近い値となるため、モデルが本質的に利用する次元数を**自動的に推定**することが可能となります。

## 単語分散表現

単語分散表現の生成において、分布の事前構造としてディリクレ過程を用います。まず、棒折り過程に必要なハイパーパラメータ$\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})
$$

各クラスタ $p$ の中心ベクトル（**無限次元**）およびばらつきは次のようにサンプリングされます：

$$
P_{latent,p} \sim Normal(0, 1)
$$

$$
P_{\sigma,p} \sim Gamma(0.001, 0.001)
$$

単語$S$の分散表現ベクトル$\beta_S$は、まずカテゴリ分布により所属クラスタ$\eta_S$を以下のようにサンプリングし、

$$
\eta_{S} \sim Categorical(p)
$$

続いて、そのクラスタに基づく分布から以下のように生成されます：

$$
\beta_{S} \sim Normal(P_{latent, \eta_{S}}, P_{\sigma, \eta_{S}})
$$

これにより、単語$S$の分散表現ベクトルは**無限次元かつクラスタ構造を持つ表現**として得られます。

この一連の分散表現ベクトルを生成する分布のことをGといって、実証分析のところでその事後分布を可視化します。

## 単語共起行列の生成

単語$i$と単語$j$が同一ウィンドウ内で共起する回数$X_{i,j}$は、次のポワソン分布から生成されると仮定します：

$$
X_{i, j} \sim Poisson\left(exp \left(b_{i} + b_{j} + \sum_{d=1}^{\infty} \omega_{d} \cdot \beta_{i,d} \cdot \beta_{j,d} \right)\right)
$$

ここで$b_i$は単語$i$の切片（バイアス）項であり、その単語が他の単語と共起しやすい度合いを表しています。

数学の厳密性を気にする方は、$\sum_{d=1}^{\infty} \omega_{d} \cdot \beta_{i,d} \cdot \beta_{j,d}$のような無限の足し合わせはそもそも定義・存在できるものなのかに疑問をもつと思いますが、問題ないことを数学的に証明しましたので、私の論文草稿のAppendixをご確認ください：

https://github.com/wutungwen/dynamic-unga-voting/tree/main


# Stanでのモデル実装

Stanでの実装方法はこちらになります。

数理モデル上、クラスタも次元も無限大になりますが、残念ながら実際の現代の計算機は無限大の配列なんて扱えません。したがって、十分に大きい整数（ここでは20）を指定して、これで無限大の近似として使います。

```cpp:glove.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 partial_sum_lpmf(
    array[] int result,
    
    int start, int end,
    
    array[] int word_1, array[] int word_2, 
    vector dimension, vector word_intercept, array[] vector word_embedding
  ){
    vector[end - start + 1] log_likelihood;
    int count = 1;
    
    for (i in start:end){
      log_likelihood[count] = poisson_log_lpmf(result[count] | word_intercept[word_1[i]] + word_intercept[word_2[i]] + (dimension .* word_embedding[word_1[i]]) '* word_embedding[word_2[i]]);
      count += 1;
    }
    return sum(log_likelihood);
  }
}
data {
  int dimension_type;
  int group_type;
  int word_type;
  
  int N;
  array[N] int word_1;
  array[N] int word_2;
  array[N] int count;
}
parameters {
  real<lower=0> dimension_alpha;
  vector<lower=0, upper=1>[dimension_type - 1] dimension_breaks;
  
  real<lower=0> group_alpha;
  vector<lower=0, upper=1>[group_type - 1] group_breaks;
  
  vector<lower=0>[group_type] group_sigma;
  array[group_type] vector[dimension_type] group_word_embedding;
  
  vector[word_type] word_intercept;
  array[word_type] vector[dimension_type] word_embedding;
}
transformed parameters {
  simplex[dimension_type] dimension;
  simplex[group_type] group;
  
  dimension = stick_breaking(dimension_breaks);
  group = stick_breaking(group_breaks);
}
model {
  dimension_alpha ~ gamma(0.001, 0.001);
  dimension_breaks ~ beta(1, dimension_alpha);
  
  group_alpha ~ gamma(0.001, 0.001);
  group_breaks ~ beta(1, group_alpha);
  
  group_sigma ~ gamma(0.001, 0.001);
  
  for (i in 1:group_type){
    group_word_embedding[i] ~ normal(0, 1);
  }
  
  word_intercept ~ normal(0, 10);
  for (i in 1:word_type){
    vector[group_type] case_when;
    for (j in 1:group_type){
      case_when[j] = log(group[j]) + normal_lpdf(word_embedding[i] | group_word_embedding[j], group_sigma[j]);
    }
    target += log_sum_exp(case_when);
  }
  
  target += reduce_sum(
    partial_sum_lupmf, count, 1,
    
    word_1, word_2, 
    dimension, word_intercept, word_embedding
  );
}
generated quantities {
  array[dimension_type] real G;
  array[word_type] vector[dimension_type] weighted_word_embedding;
  {
    int sampled_group = categorical_rng(group);
    G = normal_rng(group_word_embedding[sampled_group], group_sigma[sampled_group]);
    for (i in 1:dimension_type){
      G[i] = G[i] * dimension[i];
    }
  }
  {
    for (i in 1:word_type){
      for (j in 1:dimension_type){
        weighted_word_embedding[i, j] = word_embedding[i, j] * dimension[j];
      }
    }
  }
}

```

# モデル推定

今回の分析では、[livedoorニュースコーパス](https://www.rondhuit.com/download.html#ldcc)を利用します。

まず、テキストデータの読み込みとMecabで単語分割等を実施します：

```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()
  )
```

続いて、必要な単語と単語マスターを作ります：

```R
word_use <- review_df |>
  dplyr::pull(text) |>
  quanteda::phrase() |>
  quanteda::tokens() |>
  quanteda::dfm() |>
  quanteda::dfm_trim(min_docfreq = 100) |>
  colnames()

word_master <- word_use |>
  tibble::tibble(
    word = _
  ) |>
  dplyr::arrange(word) |>
  dplyr::mutate(
    word_id = dplyr::row_number()
  )
```

そして、`quanteda`を活用して、単語共起行列を作成して、それを縦持ちにします：

```R
context_df <- review_df |>
  dplyr::pull(text) |>
  quanteda::phrase() |>
  quanteda::tokens() |>
  quanteda::tokens_select(word_use) |>
  quanteda::fcm(
    context = "window",
    count = "frequency",
    tri = FALSE
  ) |>
  as.matrix() |>
  as.data.frame() |>
  tibble::rownames_to_column(var = "word_1") |>
  tibble::tibble() |>
  tidyr::pivot_longer(!word_1, names_to = "word_2", values_to = "count") |>
  # 単語共起行列のi,jをj,iは同じ値が入るため、word_1とword_2に同じ値が入るレコードを統合する
  dplyr::rowwise() |>
  dplyr::mutate(
    w_min = min(word_1, word_2),
    w_max = max(word_1, word_2)
  ) |>
  dplyr::ungroup() |>
  dplyr::summarise(count = as.integer(sum(count)/2), .by = c(w_min, w_max)) |>
  dplyr::rename(word_1 = w_min, word_2 = w_max) |>
  # 統合処理終了
  dplyr::left_join(
    word_master |>
      dplyr::rename(word_1_id = word_id), by = c("word_1" = "word")
  ) |>
  dplyr::left_join(
    word_master |>
      dplyr::rename(word_2_id = word_id), by = c("word_2" = "word")
  )
```

最後に、Stanモデルをコンパイルして、

```R
m_glove_init <- cmdstanr::cmdstan_model("glove.stan",
                                      cpp_options = list(
                                        stan_threads = TRUE
                                      )
)
```

実際に推定を実施します：

```R
> m_glove_estimate <- m_glove_init$variational(
     seed = 12345,
     threads = 20,
     iter = 50000,
     data = list(
         dimension_type = 20,
         group_type = 20,
         word_type = nrow(word_master),
         
         N = nrow(context_df),
         word_1 = context_df$word_1_id,
         word_2 = context_df$word_2_id,
         count = context_df$count
     )
 )
------------------------------------------------------------ 
EXPERIMENTAL ALGORITHM: 
  This procedure has not been thoroughly tested and may be unstable 
  or buggy. The interface is subject to change. 
------------------------------------------------------------ 
Gradient evaluation took 0.31749 seconds 
1000 transitions using 10 leapfrog steps per transition would take 3174.9 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    -16407857.542             1.000            1.000 
   200     -8589422.363             0.955            1.000 
   300     -6125533.237             0.771            0.910 
   400     -5202966.604             0.622            0.910 
   500     -4863584.981             0.512            0.402 
   600     -4463218.215             0.442            0.402 
   700     -4368136.655             0.382            0.177 
   800     -4069738.512             0.343            0.177 
   900     -3933111.498             0.309            0.090 
  1000     -3820140.211             0.281            0.090 
  1100     -6209922.813             0.290            0.090 
  1200     -6221945.161             0.266            0.090 
  1300     -3592303.921             0.302            0.090 
  1400     -3503749.742             0.282            0.090 
  1500     -3477974.776             0.264            0.073 
  1600     -3414427.147             0.249            0.073 
  1700     -3366124.034             0.235            0.070 
  1800     -3330900.833             0.222            0.070 
  1900     -3300408.748             0.211            0.035 
  2000     -3277225.944             0.201            0.035 
  2100     -3243316.055             0.192            0.030 
  2200     -3216940.572             0.184            0.030 
  2300     -3194720.816             0.176            0.025 
  2400     -3179883.018             0.169            0.025 
  2500     -3155644.960             0.162            0.022 
  2600     -3140641.728             0.156            0.022 
  2700     -3122645.181             0.151            0.019 
  2800     -3111295.026             0.145            0.019 
  2900     -3098066.671             0.141            0.014 
  3000     -3084264.994             0.136            0.014 
  3100     -3083747.262             0.132            0.011 
  3200     -3068159.461             0.128            0.011 
  3300     -3058185.149             0.124            0.010 
  3400     -3044294.289             0.120            0.010 
  3500     -3038133.421             0.117            0.009   MEDIAN ELBO CONVERGED 
Drawing a sample of size 1000 from the approximate posterior...  
COMPLETED. 
Finished in  1585.2 seconds.
```

26分で終わりました！

最後に、推定結果をデータフレイムに保存します：

```R
m_glove_summary <- m_glove_estimate$summary()
```

# 推定結果可視化

まずは、推定された次元数から確認しましょう：

```R
m_glove_summary |>
  dplyr::filter(stringr::str_detect(variable, "^dimension\\[")) |>
  ggplot2::ggplot() + 
  ggplot2::geom_bar(ggplot2::aes(x = as.factor(1:20), y = mean), 
                    stat = "identity", fill = ggplot2::alpha("blue", 0.3)) + 
  ggplot2::geom_errorbar(
    ggplot2::aes(x = as.factor(1:20), ymin = q5, ymax = q95),
    width = 0.2
  ) + 
  ggplot2::labs(
    x = "次元",
    y = "次元重要度"
  ) +
  ggplot2::theme_gray(base_family = "HiraKakuPro-W3")
```

![dimension_importance.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/82bf6c69-4571-465d-8da0-6499eb72fd62.png)

視覚的に確認できるように、本モデルはおおよそ11次元程度を推定しています。まだ明確な理論的言語化には至っていませんが、本モデルのように「各次元に異なる重みを与える」構造ではなく、一般的なGloVeやword2vec、あるいはTransformerのように「すべての次元が等しく重要である」と仮定する場合、たとえば本モデルで推定された重みが0.5および0.25の次元を同じように表現するには、それぞれを2つおよび1つの等重みの次元に分割して表現する必要があります。

言い換えれば、次元ごとの重要度を柔軟に調整できる構造を採用することで、同じ表現力を保ちながら、必要な次元数を大幅に削減できる可能性があるということです。

あくまでも現時点では仮説に過ぎませんが、大規模言語モデルは、実際には同じ情報を異なる次元で大量に・重複的に保持しているだけなのではないか、という可能性も示唆されます。

次に、単語の分散表現とその分布Gをt-SNEで可視化します：

```R
tsne_df <- m_glove_summary |>
  dplyr::filter(stringr::str_detect(variable, "^weighted_word_embedding\\[")) |>
  dplyr::mutate(
    id = variable |>
      purrr::map(
        \(x){
          as.integer(stringr::str_split(x, "\\[|\\]|,")[[1]][2:3])
        }
      )
  ) |>
  tidyr::unnest_wider(id, names_sep = "_") |>
  dplyr::pull(mean) |>
  matrix(ncol = 20) |>
  rbind(
    m_glove_estimate$draws("G") |> 
      as.data.frame() |> 
      as.matrix()
  ) |>
  Rtsne::Rtsne() |>
  purrr::pluck("Y") |>
  as.data.frame() |>
  tibble::tibble() |>
  dplyr::bind_cols(
    type = c(rep("word", nrow(word_master)), rep("group", 1000)),
    word = c(word_master$word, rep("group", 1000))
  )

g_tsne <- tsne_df |>
  dplyr::filter(type == "word") |>
  ggplot2::ggplot() +
  ggplot2::geom_text(ggplot2::aes(x = V1, y = V2, label = word), color = ggplot2::alpha("blue", 0.3)) + 
  ggplot2::geom_density_2d(data = tsne_df |>
                             dplyr::filter(type == "group"),
                           ggplot2::aes(x = V1, y = V2),
                           color = ggplot2::alpha("black", 0.2),
                           show.legend = FALSE
  ) +
  ggplot2::theme_gray(base_family = "HiraKakuPro-W3") +
  ggplot2::xlim(c(-50, 50)) + 
  ggplot2::ylim(c(-50, 50))

plotly::ggplotly(g_tsne)
```

![Rplot.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/84f5b545-52a9-47e1-aa97-9be542bc93c3.png)

文字が単語の分散表現に対して$\omega$によって重み付けされたベクトルを、t-SNEによって2次元に投影した際の位置を示しています。図中の等高線は、分散表現の分布Gの事後分布を可視化したものです。

このように2次元に投影することで、単語の分散表現が主に2つのクラスターから構成されていることが視覚的に確認できます。

まず右下を確認すると、

![Screen Shot 2025-06-15 at 8.19.35.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/8e54005b-e1ef-4abd-af6b-d8343f6d90b4.png)

エンタメ系の単語が多いですね。左上に行くと、

![Screen Shot 2025-06-15 at 8.20.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/0235ead4-31a3-4508-b31d-bba1a6c1b4d1.png)

IT系ですね！したがって、本モデルは、livedoorニュースコーパスにおけるエンタメ系の単語とIT系の単語の分散表現が、異なるクラスタから生成されたと判断したと考えられます。常識的に考えても、この結果には十分な妥当性があると言えるでしょう。

最後に、何箇所かを拡大して、同じ意味の単語がちゃんと近くに現れているかを確認しましょう；

![Screen Shot 2025-06-15 at 8.22.54.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/dcd79bbe-90d9-4222-89fc-43a024c4e172.png)

転職と求人がほぼ被っていますね。

アンケートと回答も近いところにあります。

![Screen Shot 2025-06-15 at 8.23.12.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/cfea5cb2-a0a7-47ea-8e09-2bcacf132f7e.png)

スポート系の単語（決勝、予選、進出）がきちんと近いところに現れています。

![Screen Shot 2025-06-15 at 8.23.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/ae02531a-d6c8-48b4-99cc-36b4c97fbd84.png)

炎上系の単語もしっかり同じ場所にありますね💦

![Screen Shot 2025-06-15 at 8.24.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2257087/654d8654-a7ff-4e29-b688-cc8a8805b6d6.png)

最後は（ドラマの？）家庭系です。子供・夫婦・両親・母親の位置がほぼ被っています。

# 結論

いかがでしたか？

このように、言語モデルにベイズ的な構造を導入することで、従来は分析者やエンジニアが恣意的または探索的に指定していた次元数を自動で推定できるようになります。さらに、限られた少数の次元で、複雑な言語情報を効率的に学習することが可能になります。

このモデルを頑張って論文化したいと思います！

最後に、私たちと一緒に働きたい方はぜひ下記のリンクもご確認ください：

https://recruit-dip.com/newgrad/page-1004/