Visualizing Learned Concepts in LLM Representations

Question 5. When analyzing the output of a transformer-based language model, which technique would be most effective for visualizing the relationship between different concepts in the model’s learned representations? Options: t-SNE or UMAP dimensionality reduction, bar charts of token frequencies, time series analysis, or cumulative distribution functions? Answer: t-SNE or UMAP dimensionality reduction. These nonlinear dimensionality-reduction techniques are the most effective ways to visualize relationships between concepts in high-dimensional embedding spaces. They project high-dimensional token or sentence embeddings into 2D or 3D while preserving neighborhood (local) relationships, enabling inspection of semantic similarity, clustering, and concept separation.

Use UMAP or t-SNE to explore semantic clusters in embedding spaces. For exploratory visualization, UMAP is typically faster and preserves more global structure; t-SNE often produces clearer local clusters but can be slower and requires careful tuning of perplexity.

Why t-SNE / UMAP are preferred over the other listed options

Bar charts of token frequencies: show token counts and distributional properties of the corpus, but do not capture geometric relationships in embedding space.
Time series analysis: suited to temporal trends and sequence behavior, not spatial neighborhood structure in vector embeddings.
Cumulative distribution functions (CDFs): summarize similarity-score distributions, but cannot reveal cluster topology or neighborhood relationships.

Comparison table

Technique	What it reveals	Best use case
t-SNE	Local neighborhood structure and tight clusters	Small-to-moderate embedding sets where local clustering is primary
UMAP	Local + more global structure, faster at scale	Exploratory visualization of larger embedding sets; good default
Bar charts	Token frequency distributions	Corpus analysis and preprocessing checks
Time series	Temporal trends in metrics	Sequence/temporal model diagnostics
CDFs	Distribution summaries (e.g., similarity scores)	Statistical summaries and threshold selection

Practical recommendations

Preprocess embeddings: standard scale and optionally apply PCA (e.g., to 50 dims) before UMAP/t-SNE to speed computation and reduce noise.
Choose the right tool:
- UMAP: usually the default—faster, scales well, preserves more global relationships.
- t-SNE: useful when you need very clear local clusters; requires careful hyperparameter tuning.
Reproducibility: set random seeds—both UMAP and t-SNE are stochastic.
Visual cues: color points by label, token type, cluster assignment, or model-derived scores (e.g., attention weight or prediction confidence) to make semantic structure obvious.

Example workflow (high-level)

Compute or load embeddings (one vector per token/sentence).
Optionally apply PCA to reduce to ~50 dimensions.
Fit UMAP or t-SNE to 2D/3D.
Plot with colors/markers reflecting labels or cluster assignments.

Example: PCA + UMAP

import numpy as np
from sklearn.decomposition import PCA
import umap
import matplotlib.pyplot as plt

# X: numpy array of shape (n_samples, n_features) containing embeddings
# 1) PCA preprocessing to speed up and denoise (optional but recommended)
pca = PCA(n_components=50, random_state=42)
X_pca = pca.fit_transform(X)

# 2) UMAP projection to 2D
reducer = umap.UMAP(n_components=2, n_neighbors=15, min_dist=0.1, random_state=42)
X_umap = reducer.fit_transform(X_pca)

# 3) Plot
plt.figure(figsize=(8, 6))
scatter = plt.scatter(X_umap[:, 0], X_umap[:, 1], c=labels, cmap="Spectral", s=5)
plt.title("UMAP projection of embeddings")
plt.xlabel("UMAP 1")
plt.ylabel("UMAP 2")
plt.colorbar(scatter, label="label")
plt.show()

Example: PCA + t-SNE

from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

# 1) PCA to 50 dims
pca = PCA(n_components=50, random_state=42)
X_pca = pca.fit_transform(X)

# 2) t-SNE projection to 2D
tsne = TSNE(n_components=2, perplexity=30, learning_rate=200, n_iter=1000, random_state=42)
X_tsne = tsne.fit_transform(X_pca)

# 3) Plot
plt.figure(figsize=(8, 6))
scatter = plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=labels, cmap="Spectral", s=5)
plt.title("t-SNE projection of embeddings")
plt.xlabel("t-SNE 1")
plt.ylabel("t-SNE 2")
plt.colorbar(scatter, label="label")
plt.show()

Hyperparameter tips

t-SNE:
- perplexity: typically 5–50. Smaller values emphasize very local structure; larger values reveal broader groupings.
- learning_rate: often between 100 and 1000.
- n_iter: at least 500–1000; more iterations can improve stability.
UMAP:
- n_neighbors: controls local vs. global structure (small → very local, large → more global). Typical values: 5–50.
- min_dist: controls how tightly points are packed (lower → tighter clusters).

Interpretation guidance

Look for coherent clusters where semantically similar tokens/sentences are near each other.
Use colors/markers for known labels (e.g., POS tags, concept labels, or predicted classes) to validate semantic separation.
Beware of over-interpreting absolute distances—UMAP/t-SNE emphasize relative neighborhood relationships, not exact metric preservation.
Validate clusters quantitatively (e.g., silhouette score, cluster purity) rather than relying solely on visual appearance.

t-SNE and UMAP are stochastic. To ensure reproducibility and comparability, set random_state/seed, log hyperparameters (e.g., perplexity, n_neighbors, min_dist), and, when appropriate, run multiple seeds to check stability.

Summary

For visualizing relationships between concepts in LLM learned representations, prefer UMAP or t-SNE, with UMAP as a practical default for exploratory work and larger datasets.
Combine dimensionality reduction with preprocessing (PCA), color/annotate points by labels or scores, and validate findings with quantitative metrics.

Links and references

Watch Video