Probabilistic Topic Modelling with Transformer Representations

TL;DR

TNTM unifies transformer-based word embeddings with probabilistic topic modelling by representing topics as multivariate Gaussians in embedding space and inferring document-topic mixtures via a variational autoencoder. The model initializes topics with a Gaussian Mixture on reduced word embeddings and employs UMAP for dimensionality reduction to speed up inference. A logistic-normal prior on document-topic distributions and carefully engineered numerical stability techniques enable scalable, stable training, with results showing superior embedding coherence and near-perfect topic diversity compared to baselines. The work demonstrates the practical impact of combining rich contextual embeddings with principled probabilistic structure, and provides a codebase for broader adoption and extension.

Abstract

Topic modelling was mostly dominated by Bayesian graphical models during the last decade. With the rise of transformers in Natural Language Processing, however, several successful models that rely on straightforward clustering approaches in transformer-based embedding spaces have emerged and consolidated the notion of topics as clusters of embedding vectors. We propose the Transformer-Representation Neural Topic Model (TNTM), which combines the benefits of topic representations in transformer-based embedding spaces and probabilistic modelling. Therefore, this approach unifies the powerful and versatile notion of topics based on transformer embeddings with fully probabilistic modelling, as in models such as Latent Dirichlet Allocation (LDA). We utilize the variational autoencoder (VAE) framework for improved inference speed and modelling flexibility. Experimental results show that our proposed model achieves results on par with various state-of-the-art approaches in terms of embedding coherence while maintaining almost perfect topic diversity. The corresponding source code is available at https://github.com/ArikReuter/TNTM.

Figures (3)

• Figure 1: Overview of the algorithmic procedure for TNTM. First, the input corpus ${\cal D}$ is processed by a Word-Embedding-Model as well as a Text-Embedding-Model. The Word-Embedding-Model computes embeddings of all individual words within ${\cal D}$ while the Text-Embedding-Model computes embeddings of complete documents. Note that we investigate two versions of TNTM where only the second version actually uses text embeddings. UMAP is used to reduce the dimensionality of the word embeddings. Subsequently, the document embeddings are used as input to the encoder of the VAE while the word embeddings are used in the VAE decoder. Finally, the VAE can be used to compute the assignment probabilities $\left\{ \bm{\theta}^{(d)} \right\}_{d \in {\cal D}}$ of all documents in the corpus to specific topics, as well as the probabilities $\bm{\beta}$ to have particular words for given topics.
• Figure 2: Forward pass of the VAE for TNTM: the inference network computes the parameters of the variational distribution for a given document, which is subsequently used to sample the document-specific probability vector $\bm{\theta}^{(d)}$. Finally, the objective of the decoder is to use $\bm{\theta}^{(d)}$ together with the global likelihood-matrix $\bm{\beta}$ to reconstruct $d$.
• Figure 3: Architecture of the inference network for TNTM with $N$ skip-blocks for the BOW-case: each document is processed by an initial linear layer followed by several skip-blocks. Two final feed-forward layers yield the parameters of the variational distribution.