BERTopic で文書のクラスタリングを試す

こんにちは。研究開発部の青見 (@nersonu) です。

そろそろ花粉症の季節ですね。週1でしか出社しない私は、なんとか引きこもって数ヶ月しのぎたいところです。

さて、今回は BERTopic という OSS について、さっくりとした手法の解説もしつつ、簡単に文書のクラスタリングを試そうと思います。

BERTopic とは

BERTopic はいわゆるトピックモデリングを行うための OSS です。トピックモデルは、文書集合から「トピック」は何が含まれるのか、各文書がどういったトピックに属しているのかといったことを特定する統計モデルを指します。テキストマイニングの分野で発展しましたが、トピックモデル自体はクラスタリング手法として幅広く利用されることもあります。代表的なモデルとして、Latent Dirichlet allocation (LDA) が挙げられ、耳にしたことや利用したことのある人も多いのではないでしょうか。

BERTopic はその名に冠されているように、 BERT (正確には Sentence-BERT) の埋め込み表現を利用したトピックモデリングを行うフレームワークを提供しています。 arXiv にあるプレプリントでは特定の手法を前提にフレームワークを紹介していますが、実際は埋め込み表現の入れ替えや、中間処理に使われている手法が組み換え可能になっている非常に可用性の高いフレームワークとなっています。今回はまずプレプリントでの紹介をベースに、ゆるく解説していきます。

arxiv.org

前提として、BERTopic は大きく3つのステップに分かれます。

文書の埋め込み
文書のクラスタリング
トピック表現

文書の埋め込み

入力となる文書は、 Sentence-BERT (Remiers and Gurevych, 2019) *1によって埋め込み表現を獲得させます。この埋め込み表現は、この後行うクラスタリングにのみ利用し、トピックの生成には直接利用しません。

このように Sentence-BERT を利用している一方で、BERTopic としては言語モデルを入れ替えることができ、言語モデルの発展に伴って最先端のモデルを利用することもできると主張しています。

文書のクラスタリング

得られた埋め込み表現は高次元データですが、一般的なクラスタリング手法では高次元データのクラスタリングがうまくいかないことがあります。高次元空間での距離尺度に差が無くなること等が原因であり、これを解決するクラスタリング手法はいくつか存在しますが、 BERTopic ではシンプルに次元削減を行ってからクラスタリングを行う方法を採用しています。次元削減の方法としては、UMAP (Mclnnes et al., 2018) *2 を利用しており、これは k-means 等のクラスタリング手法の精度と時間の両面で改善できることが知られています。

次元削減された埋め込み表現は、 HDBSCAN (Mclnnes et al., 2017) *3 という DBSCAN を階層型クラスタリングに拡張した手法によってクラスタリングを行います。

トピック表現

得られたクラスタ( = トピック)と呼ばれる単位について、一般にその文書集合を構成する単語分布からそのトピックの特徴を確認します。 BERTopic では、クラスタ内の単語の重要度を知るために TF-IDF を応用してその重要度を算出しています。具体的には、クラスタに含まれる文書を全て結合し 1つの文章として扱い、クラスタ単位の TF-IDF (class-based TF-IDF と呼んでいる) を以下のように計算して、それぞれの単語の重要度としています。

クラスタ $c$ 内の単語 $t$ の class-based TF-IDF $W_{t,c}$ は、

$\displaystyle W_{t, c} = tf_{t,c} \cdot \log \left( 1 + \frac{A}{f_t} \right)$

で得られます。ここで、

$tf_{t,c}$ : クラスタ $c$ 内の単語 $t$ の単語頻度
$f_t$ : 全クラスタに含まれる単語 $t$ の単語頻度
$A$ : クラスタあたりの平均単語数

となっています。 1を足しているのは、正の値に補正したいからのようです。

手法の概要まとめ

文書の埋め込み → Sentence-BERT を利用。今後も言語モデルの発展に伴って付け替え可能。
文書のクラスタリング → UMAP で次元削減した後、 HDBSCAN でクラスタリング。次元削減手法やクラスタリング手法も付け替え可能。
トピック表現 → クラスタ単位の TF-IDF を提案し、重要度を可視化できるように。

可用性が高い一方で、BERTopic は文書が単一のトピックのみを含むという過程を置いているため複数のトピックを持つ条件には対応していません。著者自身もこのポイントに触れており、今後の改良に期待です。

BERTopic を試してみる

さて、手法の解説はここまでにして、実際に BERTopic を用いて文書のクラスタリングを試してみようと思います。サンプルコードはこちらのリポジトリにまとめていますので、もしよろしければご参照ください。

github.com

インストール

BERTopic は PyPI に登録されているため、 pip や Poetry, PDM 等のツールで簡単にインストールすることができます。

# pip
pip install bertopic
# Poetry
poetry add bertopic
# PDM
pdm add bertopic

依存するライブラリに sentence-transformers があり、このライブラリの依存に PyTorch があります。何かのプロジェクトで使う場合は、 Python のバージョン等を対応しておくなり、pyproject.toml でバージョンの範囲を絞っておくと良いかと思われます。

モデルのロード

今回、日本語の文書 (今回は livedoor ニュースコーパスを利用しました。) をクラスタリングしようと思いますので、多言語の事前学習モデルを利用します。 BERTopic ではデフォルトで、Sentence-BERT の事前学習モデルとして sentence-transformers のモデルを利用することができますが、今回は paraphrase-multilingual-MiniLM-L12-v2 を使います。

from bertopic import BERTopic


model = BERTopic(embedding_model="paraphrase-multilingual-MiniLM-L12-v2")

他にも gensim, spaCy, Hugging Face 等の形式に対応しています。

クラスタリングの実行

クラスタリングにかけるデータは予め、 list[str] 等の ["文書A", "文書B", "文書C", ...] の文書のリストのような形にしておきます。 scikit-learn の k-means 等と同様の形式で fit() や fit_transform() のメソッドでクラスタリングを実行することができます。

# 文書のリストを用意しておく
news_data: list[str] = preprocess_data()
topics, probs = model.fit_transform(news_data)

結果の確認

topics に各文書に対するクラスタリング結果であるトピック番号 (割り当てられなかったものは -1 ) 、probs には HDBSCAN から得られた各文書の確率値が含まれます。

トピックごとの単語の重要度

class-based TF-IDF を利用した、トピックごとの単語の重要度が可視化できるようになっています (デフォルトは最大8トピックまで表示)。

# fit 後のモデルを利用
model.visualize_barchart()

また、それぞれのトピックの単語の重要度が見たい場合は get_topic() や get_topics() を利用します。

model.get_topic(0)

# [('関連情報', 0.01767639630607681),
#  ('など', 0.016178286851156458),
#  ('関連記事', 0.011504104725101322),
#  ('ネット掲示板では', 0.011388265785763135),
#  ('孫社長', 0.010587420882455728),
#  ('akb48', 0.008652719774273905),
#  ('週刊文春', 0.006320498705819431),
#  ('livedoorニュース', 0.005750062991600873),
#  ('エンタがビタミン', 0.005747611547549234),
#  ('livedoor', 0.005667169014230514)]

可視化に関しては、visualize_ の prefix で様々なメソッドが用意されているため、分析の手間が省けそうです。

トピックごとの文書を眺めてみる

例えば、以下のような DataFrame を用意しておき、

result_df = pd.DataFrame(
    {
        "news_text": news_data,
        "topic_no": topics,
        "proba": probs,
    }
)

各トピックについて確率値が高い順にソートして上から眺めてみます。どういった文書が集められているか、なんとなく確認することが出来ますね。

# トピック3について眺めてみる
result_df[result_df.topic_no == 3].sort_values("proba", ascending=False).head(10)

どうやらスマートフォンに関する記事が集約されているようですね。

次元削減手法・クラスタリング手法の変更

ここでは多くを解説しませんが、解説パートで少し述べたように次元削減手法やクラスタリング手法の入れ替えも可能です。デフォルトでは UMAP・HDBSCAN ですが、公式で用意されている Google Colab を参考に、BERTopic をインスタンス化する際に簡単に指定できます。

colab.research.google.com

使ってみての所感

非常に簡単かつ、よく知られたインターフェースを持っているため、利用のハードルは低いと感じました。日本語で利用する場合、多くは多言語の事前学習モデルを利用すると思われますので、比較的大きなモデルを利用する点は少し注意でしょうか。

クラスタリングの定性的な精度に関してですが、経験的にストップワードのようなものを用意して除去しておくことで一定改善するように見えます。それでも、LDA で同様のクラスタリングをした場合と比べ、前処理をほとんど何もせずにそれっぽい文書のクラスタが出来上がるのは驚きました (あくまでも個人の感想であり、定量的な評価ではありませんのであしからず )。

今回の紹介で省きましたが、semi-supervised でトピックモデリングを行う方法や、ダイナミックトピックモデリングを行う方法も公式ページで解説されています。ここまで読んで頂いた皆様もぜひ、お試しください。

*1: Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks (Reimers & Gurevych, EMNLP-IJCNLP 2019)

*2: UMAP: Uniform Manifold Approximation and Projection (Mclnnes et al., The Journal of Open Source Software 3(29), p.861, 2018)

*3: hdbscan: Hierarchical density based clustering (Mclnnes et al., The Journal of Open Source Software 2(11), p.205, 2017)