Visualize Embeddings with t-SNE

You've built a model that turns text into dense vectors, but when you try to inspect what it actually learned, you're staring at a 768-dimensional blob of numbers. Spreadsheets can't show you clusters, and any attempt to debug embeddings with raw correlation matrices leaves you guessing. In this lesson, you'll learn how to visualize embeddings with t-SNE — the de facto technique for projecting high-dimensional vectors into a 2D space that you can actually see, explore, and explain to stakeholders.

The problem this lesson solves

Embeddings are the backbone of semantic search, RAG pipelines, and classification systems, but their value is hidden in high-dimensional space. A typical OpenAI text-embedding-3-small vector has 1,536 dimensions; a custom Sentence-BERT model might output 384 or 768. Plotting that directly is impossible — you only have two axes on a screen, and humans can't reason about 1,500-dimensional geometry.

Without visualization, you're flying blind. You can't tell if your embeddings separate by topic, if different languages are colliding, or if your model is merely memorizing token patterns instead of learning semantic structure. You can't debug why your nearest-neighbor search returns unrelated results, and you can't explain to a product manager why the RAG system suddenly fails on legal documents.

The result? You spend hours guessing, tuning hyperparameters, and writing regex hacks to fix retrieval — when a single scatter plot would have revealed the root cause in seconds. t-SNE (t-Distributed Stochastic Neighbor Embedding) lets you compress those hundreds of dimensions into a vivid 2D map, making hidden patterns visible.

Pro tip: Visualization isn't just for demos. It's your debugging lens for embedding quality, model drift, and data leakage.

Core concept / mental model

Think of t-SNE as a tour guide for your dataset.

Imagine you have a 3D sculpture of your embedding space, with each object a small bead. You need to draw a 2D map of that sculpture on a sheet of paper, but you can't see the 3D shape directly. The tour guide's job is to place the beads on the paper so that similar beads stay close together, and dissimilar ones stay far apart — preserving the local structure as faithfully as possible.

Formally, t-SNE works by:

1. Measuring similarity — For each point, it computes a Gaussian probability distribution over all other points. Nearby points get high probability (high similarity); distant points get low probability.
2. Building a target distribution — In the low-dimensional embedding (2D or 3D), it uses a heavy-tailed Student t-distribution to define a similar probability distribution. The heavy tail prevents overcrowding (a famous t-SNE artifact).
3. Optimizing — It uses gradient descent to minimize the Kullback-Leibler (KL) divergence between the two distributions. The result: the 2D positions are iteratively adjusted until the local neighborhoods match the high-dimensional reality.

Key terms you'll encounter:

• Perplexity — A hyperparameter that balances local vs. global structure (typically 5–50, default 30). It's roughly the expected number of neighbors per point.
• Learning rate — Controls the step size during optimization (default 200 in sklearn, but often 200–1000 for convergence).
• KL divergence — A measure of how one probability distribution diverges from a second; t-SNE minimizes it.
• Non-determinism — Because of random initialization, running t-SNE twice on the same data yields a different map. The clusters remain similar, but the orientation/scale changes.

Mental model: t-SNE is a local projector. It's excellent at revealing clusters, but distances between clusters should not be interpreted. Use UMAP if you need a more global view.

How it works step by step

Let's trace the algorithm from raw embeddings to a polished plot:

1. Input preparation

Start with a matrix X of n_samples × n_features — your embedding vectors. You may want to standardize or normalize them (unit length) because t-SNE is sensitive to scale. For text embeddings, cosine similarity is often preferred, but t-SNE uses Euclidean distance internally; normalizing vectors to unit length makes cosine ≈ Euclidean for ranking.

2. Compute pairwise similarities

For each point i, t-SNE computes the probability that i would pick j as its neighbor, using a Gaussian centered at i. The variance (sigma) of the Gaussian is tuned so that the effective number of neighbors equals perplexity.

3. Initialize low-dimensional positions

Each point is given a random starting position in 2D (or 3D). This randomness is why you get a different plot each run.

4. Optimize with gradient descent

At each iteration, t-SNE compares the high-dimensional probabilities to the low-dimensional ones (using the t-distribution). It computes a gradient that pulls similar points together and pushes dissimilar points apart. The learning rate controls the step size; after many iterations (typically 1000), the points settle into a map.

5. Post-process and visualize

Finally, you plot the 2D coordinates, colored by a class label (if you have one), and annotate key regions.

Hands-on walkthrough

Let's visualize the classic 20 Newsgroups text dataset with real embeddings. We'll use sentence-transformers and scikit-learn's TSNE.

First, install the dependencies:

pip install sentence-transformers scikit-learn matplotlib numpy

Step 1: Load and encode a sample of documents

from sentence_transformers import SentenceTransformer
from sklearn.datasets import fetch_20newsgroups
import numpy as np

# Load a small sample to keep the run fast (5 docs per category)
categories = ['sci.space', 'comp.graphics', 'rec.sport.baseball', 'talk.politics.mideast']
newsgroups = fetch_20newsgroups(subset='train', categories=categories, shuffle=True, random_state=42)

# Take 20 documents per category for a quick demo
sample_indices = []
for label in set(newsgroups.target):
    idx = np.where(newsgroups.target == label)[0][:20]
    sample_indices.extend(idx)

sample_docs = [newsgroups.data[i] for i in sample_indices]
sample_labels = [newsgroups.target_names[newsgroups.target[i]] for i in sample_indices]

# Encode with a lightweight all-MiniLM model
model = SentenceTransformer('all-MiniLM-L6-v2')
embeddings = model.encode(sample_docs, batch_size=8, show_progress_bar=True)
print(f"Embedding shape: {embeddings.shape}")

Expected output (approximate):

Embedding shape: (80, 384)

Step 2: Apply t-SNE

from sklearn.manifold import TSNE

# Reduce to 2D
tsne = TSNE(n_components=2, perplexity=15, random_state=42, max_iter=500)
embeddings_2d = tsne.fit_transform(embeddings)
print(f"2D shape: {embeddings_2d.shape}")

Expected output:

2D shape: (80, 2)

Step 3: Plot the result

import matplotlib.pyplot as plt
import pandas as pd

# Create a DataFrame for easy handling
df = pd.DataFrame(embeddings_2d, columns=['x', 'y'])
df['label'] = sample_labels

# Color palette
colors = plt.cm.tab10(np.linspace(0, 1, len(categories)))

plt.figure(figsize=(10, 8))
for i, cat in enumerate(categories):
    subset = df[df['label'] == cat]
    plt.scatter(subset['x'], subset['y'], c=[colors[i]], label=cat, alpha=0.7, s=60)

plt.title('t-SNE Visualization of Newsgroup Embeddings')
plt.xlabel('t-SNE dimension 1')
plt.ylabel('t-SNE dimension 2')
plt.legend()
plt.grid(alpha=0.3)
plt.show()

What you'll see: four distinct clusters, with talk.politics.mideast and sci.space far apart, while comp.graphics and rec.sport.baseball might be closer (both contain some technical vocabulary). The clusters prove the embeddings learned semantic distinction.

Step 4: Interpret and inspect

To go deeper, you can plot specific points and show the nearby documents:

# Pick a point from the plot and print its nearest neighbors in the original space
from sklearn.metrics.pairwise import cosine_similarity

query_idx = 0  # first document
cos_sim = cosine_similarity(embeddings[query_idx].reshape(1, -1), embeddings)[0]
nearest = np.argsort(cos_sim)[::-1][1:6]

print(f"Query: {sample_docs[query_idx][:120]}...\n")
print("Nearest neighbors:")
for n in nearest:
    print(f"  - {sample_labels[n]}: {sample_docs[n][:100]}...")

Now you're not just seeing clusters — you're validating that semantic neighbors make sense.

Compare options / when to choose what

t-SNE is not the only dimensionality reduction technique. Here's how it stacks up against alternatives:

Rule of thumb: For visualizing embeddings, t-SNE is the default. If you have more than ~100k points or need a stable, meaningful global map, try UMAP. For a quick linear projection, use PCA first as a baseline.

Pro tip: Always run PCA (or UMAP) in parallel and compare. If t-SNE shows a cluster that PCA doesn't, it's likely a real nonlinear structure — not an artifact.

Troubleshooting & edge cases

1. Clusters that look like a "blob" or all points crowded together

• Cause: Perplexity too high (e.g., 50+ with small data), or learning rate too low.
• Fix: Reduce perplexity to 5–30 for small samples; increase learning rate to 500+.

# Example fix
tsne = TSNE(perplexity=5, learning_rate=1000, random_state=42)

2. Completely different plot every run

• Cause: t-SNE is non-deterministic. That's normal.
• Fix: Set random_state for reproducibility. But don't panic if clusters shift — focus on cluster integrity, not exact coordinates.

3. t-SNE takes forever on many vectors

• Cause: O(n²) pairwise computations.
• Fix: Reduce sample size (e.g., 10k max). Use PCA to project to 50 dims first, or switch to UMAP.

from sklearn.decomposition import PCA

pca = PCA(n_components=50)
embeddings_pca = pca.fit_transform(embeddings)
tsne = TSNE(n_components=2, perplexity=15).fit_transform(embeddings_pca)

4. Misleading distances between clusters

• Cause: t-SNE distances between clusters are arbitrary.
• Fix: Never compare inter-cluster distances. Only trust local neighborhoods. Mention this in any report.

5. The "crowding problem" — points form a donut or lightning bolt

• Cause: Excessive perplexity or too few iterations.
• Fix: Increase max_iter to 1000+, lower perplexity.

6. All points fall into one tight ball

• Cause: Embeddings are poorly normalized, or model didn't learn meaningful distinctions.
• Fix: Normalize embeddings to unit length; retrain if needed.

# Normalize
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

What you learned & what's next

You've now got a practical grip on visualizing embeddings with t-SNE. Let's recap the core takeaways:

• t-SNE projects high-dimensional vectors into 2D (or 3D) while preserving local structure.
• Perplexity (5–50) and learning rate are your main knobs; tune them for your dataset size.
• The plot is a diagnostic tool — it reveals clusters, outliers, and model behavior.
• Compare with PCA or UMAP when you need a different perspective.

You're building a toolset for applied AI engineering. Next in the track, you'll learn to cluster embeddings — using k-means to group documents without labels. With t-SNE visualization plus clustering, you'll be able to automatically discover topics in a corpus and build a recommendation or retrieval system that's genuinely data-driven.

Open your favorite notebook, load a real dataset (even a tiny one), and run t-SNE on it. Play with perplexity, click around, and train your eye to spot what's meaningful. That intuition will pay off in every embedding project you ship.