Skip to content

GMM (Gaussian Mixture Model)

GMM is a generative probabilistic model over the contextual embeddings. The model assumes that contextual embeddings are generated from a mixture of underlying Gaussian components. These Gaussian components are assumed to be the topics.

Components of a Gaussian Mixture Model
(figure from scikit-learn documentation)

How does GMM work?

Generative Modeling

GMM assumes that the embeddings are generated according to the following stochastic process from a number of Gaussian components. Priors are optionally imposed on the model parameters. The model is fitted either using expectation maximization or variational inference.

Click to see formula
  1. Select global topic weights: \(\Theta\)
  2. For each component select mean \(\mu_z\) and covariance matrix \(\Sigma_z\) .
  3. For each document:
    • Draw topic label: \(z \sim Categorical(\Theta)\)
    • Draw document vector: \(\rho \sim \mathcal{N}(\mu_z, \Sigma_z)\)

Calculate Topic Probabilities

After the model is fitted, soft topic labels are inferred for each document. A document-topic-matrix (\(T\)) is built from the likelihoods of each component given the document encodings.

Click to see formula
  • For document \(i\) and topic \(z\) the matrix entry will be: \(T_{iz} = p(\rho_i|\mu_z, \Sigma_z)\)

Soft c-TF-IDF

Term importances for the discovered Gaussian components are estimated post-hoc using a technique called Soft c-TF-IDF, an extension of c-TF-IDF, that can be used with continuous labels.

Click to see formula

Let \(X\) be the document term matrix where each element (\(X_{ij}\)) corresponds with the number of times word \(j\) occurs in a document \(i\). Soft Class-based tf-idf scores for terms in a topic are then calculated in the following manner:

  • Estimate weight of term \(j\) for topic \(z\):
    \(tf_{zj} = \frac{t_{zj}}{w_z}\), where \(t_{zj} = \sum_i T_{iz} \cdot X_{ij}\) and \(w_{z}= \sum_i(|T_{iz}| \cdot \sum_j X_{ij})\)
  • Estimate inverse document/topic frequency for term \(j\):
    \(idf_j = log(\frac{N}{\sum_z |t_{zj}|})\), where \(N\) is the total number of documents.
  • Calculate importance of term \(j\) for topic \(z\):
    \(Soft-c-TF-IDF{zj} = tf_{zj} \cdot idf_j\)

Dynamic Modeling

GMM is also capable of dynamic topic modeling. This happens by fitting one underlying mixture model over the entire corpus, as we expect that there is only one semantic model generating the documents. To gain temporal representations for topics, the corpus is divided into equal, or arbitrarily chosen time slices, and then term importances are estimated using Soft-c-TF-IDF for each of the time slices separately.

Similarities with Clustering Models

Gaussian Mixtures can in some sense be considered a fuzzy clustering model.

Since we assume the existence of a ground truth label for each document, the model technically cannot capture multiple topics in a document, only uncertainty around the topic label.

This makes GMM better at accounting for documents which are the intersection of two or more semantically close topics.

Another important distinction is that clustering topic models are typically transductive, while GMM is inductive. This means that in the case of GMM we are inferring some underlying semantic structure, from which the different documents are generated, instead of just describing the corpus at hand. In practical terms this means that GMM can, by default infer topic labels for documents, while (some) clustering models cannot.

Performance Tips

GMM can be a bit tedious to run at scale. This is due to the fact, that the dimensionality of parameter space increases drastically with the number of mixture components, and with embedding dimensionality. To counteract this issue, you can use dimensionality reduction. We recommend that you use PCA, as it is a linear and interpretable method, and it can function efficiently at scale.

Through experimentation on the 20Newsgroups dataset I found that with 20 mixture components and embeddings from the all-MiniLM-L6-v2 embedding model reducing the dimensionality of the embeddings to 20 with PCA resulted in no performance decrease, but ran multiple times faster. Needless to say this difference increases with the number of topics, embedding and corpus size.

from turftopic import GMM
from sklearn.decomposition import PCA

model = GMM(20, dimensionality_reduction=PCA(20))

# for very large corpora you can also use Incremental PCA with minibatches

from sklearn.decomposition import IncrementalPCA

model = GMM(20, dimensionality_reduction=IncrementalPCA(20))

API Reference

turftopic.models.gmm.GMM

Bases: ContextualModel, DynamicTopicModel

Multivariate Gaussian Mixture Model over document embeddings. Models topics as mixture components.

from turftopic import GMM

corpus: list[str] = ["some text", "more text", ...]

model = GMM(10, weight_prior="dirichlet_process").fit(corpus)
model.print_topics()

Parameters:

Name Type Description Default
n_components int

Number of topics. If you're using priors on the weight, feel free to overshoot with this value.

required
encoder Union[Encoder, str]

Model to encode documents/terms, all-MiniLM-L6-v2 is the default.

'sentence-transformers/all-MiniLM-L6-v2'
vectorizer Optional[CountVectorizer]

Vectorizer used for term extraction. Can be used to prune or filter the vocabulary.

None
weight_prior Literal['dirichlet', 'dirichlet_process', None]

Prior to impose on component weights, if None, maximum likelihood is optimized with expectation maximization, otherwise variational inference is used.

None
gamma Optional[float]

Concentration parameter of the symmetric prior. By default 1/n_components is used. Ignored when weight_prior is None.

None
dimensionality_reduction Optional[TransformerMixin]

Optional dimensionality reduction step before GMM is run. This is recommended for very large datasets with high dimensionality, as the number of parameters grows vast in the model otherwise. We recommend using PCA, as it is a linear solution, and will likely result in Gaussian components. For even larger datasets you can use IncrementalPCA to reduce memory load.

None
random_state Optional[int]

Random state to use so that results are exactly reproducible.

None

Attributes:

Name Type Description
weights_ ndarray of shape (n_components)

Weights of the different mixture components.

Source code in turftopic/models/gmm.py
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
class GMM(ContextualModel, DynamicTopicModel):
    """Multivariate Gaussian Mixture Model over document embeddings.
    Models topics as mixture components.

    ```python
    from turftopic import GMM

    corpus: list[str] = ["some text", "more text", ...]

    model = GMM(10, weight_prior="dirichlet_process").fit(corpus)
    model.print_topics()
    ```

    Parameters
    ----------
    n_components: int
        Number of topics. If you're using priors on the weight,
        feel free to overshoot with this value.
    encoder: str or SentenceTransformer
        Model to encode documents/terms, all-MiniLM-L6-v2 is the default.
    vectorizer: CountVectorizer, default None
        Vectorizer used for term extraction.
        Can be used to prune or filter the vocabulary.
    weight_prior: 'dirichlet', 'dirichlet_process' or None, default 'dirichlet'
        Prior to impose on component weights, if None,
        maximum likelihood is optimized with expectation maximization,
        otherwise variational inference is used.
    gamma: float, default None
        Concentration parameter of the symmetric prior.
        By default 1/n_components is used.
        Ignored when weight_prior is None.
    dimensionality_reduction: TransformerMixin, default None
        Optional dimensionality reduction step before GMM is run.
        This is recommended for very large datasets with high dimensionality,
        as the number of parameters grows vast in the model otherwise.
        We recommend using PCA, as it is a linear solution, and will likely
        result in Gaussian components.
        For even larger datasets you can use IncrementalPCA to reduce
        memory load.
    random_state: int, default None
        Random state to use so that results are exactly reproducible.

    Attributes
    ----------
    weights_: ndarray of shape (n_components)
        Weights of the different mixture components.
    """

    def __init__(
        self,
        n_components: int,
        encoder: Union[
            Encoder, str
        ] = "sentence-transformers/all-MiniLM-L6-v2",
        vectorizer: Optional[CountVectorizer] = None,
        dimensionality_reduction: Optional[TransformerMixin] = None,
        weight_prior: Literal["dirichlet", "dirichlet_process", None] = None,
        gamma: Optional[float] = None,
        random_state: Optional[int] = None,
    ):
        self.n_components = n_components
        self.encoder = encoder
        self.weight_prior = weight_prior
        self.gamma = gamma
        self.random_state = random_state
        if isinstance(encoder, str):
            self.encoder_ = SentenceTransformer(encoder)
        else:
            self.encoder_ = encoder
        if vectorizer is None:
            self.vectorizer = default_vectorizer()
        else:
            self.vectorizer = vectorizer
        self.dimensionality_reduction = dimensionality_reduction
        if self.weight_prior is not None:
            mixture = BayesianGaussianMixture(
                n_components=n_components,
                weight_concentration_prior_type=(
                    "dirichlet_distribution"
                    if self.weight_prior == "dirichlet"
                    else "dirichlet_process"
                ),
                weight_concentration_prior=gamma,
                random_state=self.random_state,
            )
        else:
            mixture = GaussianMixture(
                n_components, random_state=self.random_state
            )
        if dimensionality_reduction is not None:
            self.gmm_ = make_pipeline(dimensionality_reduction, mixture)
        else:
            self.gmm_ = mixture

    def fit_transform(
        self, raw_documents, y=None, embeddings: Optional[np.ndarray] = None
    ) -> np.ndarray:
        console = Console()
        with console.status("Fitting model") as status:
            if embeddings is None:
                status.update("Encoding documents")
                embeddings = self.encoder_.encode(raw_documents)
                console.log("Documents encoded.")
            status.update("Extracting terms.")
            document_term_matrix = self.vectorizer.fit_transform(raw_documents)
            console.log("Term extraction done.")
            status.update("Fitting mixture model.")
            self.gmm_.fit(embeddings)
            console.log("Mixture model fitted.")
            status.update("Estimating term importances.")
            document_topic_matrix = self.gmm_.predict_proba(embeddings)
            self.components_ = soft_ctf_idf(
                document_topic_matrix, document_term_matrix
            )
            console.log("Model fitting done.")
        return document_topic_matrix

    @property
    def weights_(self) -> np.ndarray:
        if isinstance(self.gmm_, Pipeline):
            model = self.gmm_.steps[-1][1]
        else:
            model = self.gmm_
        return model.weights_

    def transform(
        self, raw_documents, embeddings: Optional[np.ndarray] = None
    ) -> np.ndarray:
        """Infers topic importances for new documents based on a fitted model.

        Parameters
        ----------
        raw_documents: iterable of str
            Documents to fit the model on.
        embeddings: ndarray of shape (n_documents, n_dimensions), optional
            Precomputed document encodings.

        Returns
        -------
        ndarray of shape (n_dimensions, n_topics)
            Document-topic matrix.
        """
        if embeddings is None:
            embeddings = self.encoder_.encode(raw_documents)
        return self.gmm_.predict_proba(embeddings)

    def fit_transform_dynamic(
        self,
        raw_documents,
        timestamps: list[datetime],
        embeddings: Optional[np.ndarray] = None,
        bins: Union[int, list[datetime]] = 10,
    ):
        time_labels, self.time_bin_edges = self.bin_timestamps(
            timestamps, bins
        )
        if hasattr(self, "components_"):
            doc_topic_matrix = self.transform(
                raw_documents, embeddings=embeddings
            )
        else:
            doc_topic_matrix = self.fit_transform(
                raw_documents, embeddings=embeddings
            )
        document_term_matrix = self.vectorizer.transform(raw_documents)
        n_comp, n_vocab = self.components_.shape
        n_bins = len(self.time_bin_edges) - 1
        self.temporal_components_ = np.zeros(
            (n_bins, n_comp, n_vocab), dtype=document_term_matrix.dtype
        )
        self.temporal_importance_ = np.zeros((n_bins, n_comp))
        for i_timebin in np.unique(time_labels):
            topic_importances = doc_topic_matrix[time_labels == i_timebin].sum(
                axis=0
            )
            # Normalizing
            topic_importances = topic_importances / topic_importances.sum()
            components = soft_ctf_idf(
                doc_topic_matrix[time_labels == i_timebin],
                document_term_matrix[time_labels == i_timebin],  # type: ignore
            )
            self.temporal_components_[i_timebin] = components
            self.temporal_importance_[i_timebin] = topic_importances
        return doc_topic_matrix

transform(raw_documents, embeddings=None)

Infers topic importances for new documents based on a fitted model.

Parameters:

Name Type Description Default
raw_documents

Documents to fit the model on.

required
embeddings Optional[ndarray]

Precomputed document encodings.

None

Returns:

Type Description
ndarray of shape (n_dimensions, n_topics)

Document-topic matrix.

Source code in turftopic/models/gmm.py
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def transform(
    self, raw_documents, embeddings: Optional[np.ndarray] = None
) -> np.ndarray:
    """Infers topic importances for new documents based on a fitted model.

    Parameters
    ----------
    raw_documents: iterable of str
        Documents to fit the model on.
    embeddings: ndarray of shape (n_documents, n_dimensions), optional
        Precomputed document encodings.

    Returns
    -------
    ndarray of shape (n_dimensions, n_topics)
        Document-topic matrix.
    """
    if embeddings is None:
        embeddings = self.encoder_.encode(raw_documents)
    return self.gmm_.predict_proba(embeddings)