In this article, you will learn how to use probing classifiers, UMAP visualization, and SHAP values to interpret and analyze the quality of text embeddings generated by large language models.

Topics we will cover include:

• How to generate text embeddings from movie reviews using Scikit-LLM and a local Ollama model, and train a probing logistic regression classifier to evaluate their quality.
• How to use UMAP dimensionality reduction to visually inspect the semantic structure captured by LLM-generated embeddings.
• How to apply SHAP values to identify which latent embedding dimensions have the greatest influence on a classifier’s predictions.

Introduction

Text classification tasks have long been exclusively the domain of machine learning models and their direct “evolved form”: deep neural networks. However, we can’t deny that large language models (LLMs) have revolutionized the way text classifiers are now built, being more powerful and accurate but raising a side concern: the lack of interpretability due to LLMs being black-box models. Accordingly, when using an LLM before the core text classification task to convert raw text into embeddings — dense numerical vector representations of text — it is possible to capture semantic information. Yet one challenging question arises: what exactly is the model learning about text, and how does this internal learning process drive predictions?

This hands-on article shows how to use Scikit-LLM to generate embeddings, train a probing classifier, and unveil the black box by leveraging UMAP visualization and SHAP (SHapley Additive exPlanations) values: two popular explainable AI techniques for explaining model inference and decisions.

Initial Setup

The provided code here is fully compatible with Google Colab notebooks and requires installing the latest Scikit-LLM version. To keep the whole process cost-free, the code below shows how to configure everything for local, free execution. Let’s start by installing the following dependencies and packages, including the Ollama distributions for running local LLMs for free:

Now let’s import everything we will need:

Probing Embedding Spaces

The first step to probe and analyze Scikit-LLM embeddings is, of course, to get a fresh collection of them from a text dataset. We will first configure Scikit-LLM to point to a local Ollama server via "http://localhost:11434/v1/".

After that, we use the public IMDB dataset containing movie reviews and load 1,000 of them: 500 labeled as positive and 500 labeled as negative, giving us a perfectly class-balanced sample. We use stratified sampling to keep 80% of the examples for training and the remaining 20% for testing:

We are now ready for the heaviest part of the process: generating embeddings for these 1,000 texts. We do so using Ollama’s all-minilm model via Scikit-LLM’s class designed for handling embedding models: GPTVectorizer. The syntax is intentionally similar to standard scikit-learn data transformations, as we can see:

Be patient; if you are running this on Colab, it may take about 5–10 minutes to complete, as we are making 1,000 calls to a local LLM for embedding generation.

A probing classifier (or a probing model) is a diagnostic tool used to inspect the internal representations built by complex models. How can we reliably determine that the embeddings generated earlier have enough quality to separate the data into classes — positive vs. negative reviews — properly? One way is to use a smaller, simpler classifier, such as logistic regression, and examine the accuracy metrics. If a classification report — described by precision, recall, and F1 scores per class — yields decent results even for this shallow classifier, that indicates the embeddings are rich enough for the classification task. Using a simpler classifier as our probing model also helps isolate the contribution being attributed to the embeddings themselves.

Results:

Considering that the dataset size is not extraordinarily large relative to the embedding dimensionality, these results are quite respectable for a simple, linear classifier like logistic regression, which is typically applied to smaller, purely tabular datasets.

Let’s look at another introspection tool: UMAP (Uniform Manifold Approximation and Projection). UMAP is a projection-based dimensionality reduction technique commonly used for visualization. We project the embeddings down to 2 dimensions using cosine similarity as the distance metric, which is standard when working with text embeddings. The resulting scatterplot helps us determine whether there is any natural grouping between embeddings associated with positive and negative reviews:

The results are not extraordinary at first glance — there is no near-perfect class-wise separation between reviews — but considering these are LLM-generated embeddings heavily projected into just two dimensions, a subtle sense of grouping is still visible: the southern half of the plot shows a dominance of negative reviews (blue dots), while the upper half has a majority of positive reviews (fuchsia).

Last, we can resort to one of the most popular frameworks for examining machine learning model behavior: SHAP (SHapley Additive exPlanations). SHAP can help us understand which of the latent dimensions (features) in our embeddings had the most influence on the probing classifier’s predictions.

The code below constructs a SHAP summary plot that visualizes which embedding dimensions exert the most impact on model classifications. By default, the plot displays the top 20 features with the largest overall impact, using color to indicate whether each feature contributes toward positive or negative classifications depending on whether its values are higher or lower.

We can conclude that dimension 208 is the primary signal for negative reviews, closely followed by dimension 317. Meanwhile, dimension 139 is the main driver for positive reviews, as higher values (pink) for this feature push the model’s raw prediction toward higher values (the right-hand side of the plot, leaning toward the positive class).

Conclusion

This article illustrated how to use a probing classification model, along with visualization tools like UMAP and SHAP, to better understand and interpret the nature and quality of text embeddings produced by LLMs for downstream machine learning tasks like text classification. We relied on Scikit-LLM, a library that mirrors scikit-learn’s API to seamlessly integrate LLMs into a variety of tasks, including embedding generation from raw text such as movie reviews.