以低内存开销对许多文档进行主题建模

Question

我一直在使用 BERTopic 0.16.3 进行主题建模项目，初步结果很有希望。然而，随着项目的进展和需求变得明显，我遇到了可扩展性的具体问题。

具体：

• 对于开发/测试，它需要对适度数量的文档（数万到数十万）进行相当快速的训练
  • 我们的开发机器是 Mac，所以这可能必须在 CPU 上完成
• 对于生产，需要在大量文档（数百万）上进行训练，而不会增加内存使用量
  • 对于基线，在我的机器上的默认设置下，BERTopic 每个文档的峰值内存使用量约为 35 kB，对于生产中提供的数据量来说，很容易变成数百 GB 甚至 TB
  • 理想情况下，峰值内存使用量与文档数量呈次线性关系。

最后一个要求需要对文档进行批处理，因为将它们全部加载到内存中需要线性内存。因此，我一直在研究与“在线主题建模”一起使用的聚类算法。 BERTopic 的文档建议使用 scikit-learn 的 MiniBatchKMeans，但我从中得到的结果不是很好。

我看过的一些型号包括：

Birch
• ，通过 scikit-learn：即使在批量处理时，使用的内存也比 BERTopic 的默认值

  HDBSCAN

  还要多。运行速度也慢得多。
IncrementalDBSCAN
• ，来自

  incdbscan

  ：一开始看起来很有希望，但运行时间和最终内存都在膨胀。对于 5000 个批次的约 120k 文档，它在前 3.5 小时内没有使用超过 4GB 的 RAM，但在 10 小时内没有完成，并且在中间的某个时刻使用了近 40GB 的 RAM。 通过 scikit-learn 实现的
AgglomerativeClustering
• ：在初始测试中给出了非常好的结果（甚至可能比

  HDBSCAN

  更好），但它没有实现

  partial_fit

  方法。我在另一个问题上找到了

  这个答案

  ，这表明可以使用单个链接独立地训练其中两个，然后合并它们，但它没有给出如何进行的指示。
后两者也没有提供

predict

方法，限制了它们的实用性。

我对这个主题相当陌生，所以也许我的做法完全错误，而且我试图解决的直接问题没有解决方案。因此，需要明确的是，在基础层面上，我试图回答的问题是：

如何在不使用太多内存的情况下对大量文档执行主题建模（并获得良好的结果）？

Answer 1

在这种规模的工作流程中，您希望投入生产，而不是尝试在软件中解决这个问题，切换硬件可能更容易。

cuML

中的 GPU 加速 UMAP 和 HDBSCAN 可以非常快地处理这么多数据——速度足够快，如果您本地没有支持 GPU 的系统，可能值得考虑租用一个。 对于以下示例，我采集了 100 万条亚马逊评论的样本，将它们编码为嵌入（384 维），并在当前 cuML 版本（v24.08）中使用 GPU UMAP。我在具有 H100 GPU 的系统上运行了这个。

from bertopic import BERTopic
from sentence_transformers import SentenceTransformer
import pandas as pd
from cuml.manifold.umap import UMAP
from cuml.cluster import HDBSCAN

# Create embeddings
sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = sentence_model.encode(reviews, batch_size=1024, show_progress_bar=True)

reducer = UMAP(n_components=5)
%time reduced_embeddings = reducer.fit_transform(embeddings)
CPU times: user 1min 33s, sys: 7.2 s, total: 1min 40s
Wall time: 7.31 s

clusterer = HDBSCAN()
%time clusterer.fit(reduced_embeddings)
CPU times: user 21.5 s, sys: 125 ms, total: 21.6 s
Wall time: 21.6 s

BERTopic 

常见问题解答

中有一个如何在 GPU 上运行这些步骤的示例。 我在 NVIDIA 从事这些项目，并且是 BERTopic 的社区贡献者，因此如果您遇到问题，请告诉我。