Background image

Author: Nomic Team

Improve AI Model Performance with Embedding Visualization and Evaluation

Improve AI Model Performance with Embedding Visualization and Evaluation

Imagine you're building a product classification system for an online marketplace. Your accuracy was promising at first, but as the product catalog grew, you started seeing mislabeled items and confused categories—like kitchen gadgets tagged under office supplies. Every dataset refresh introduced more anomalies, and your confusion matrix was all over the place.

Enter Nomic Atlas: within minutes of uploading your embeddings, you see clusters in the embedding space—some of which overlap suspiciously, pointing to potential mislabels. A quick fix of these labels boosts your classification accuracy by a noticeable margin. That's the real power of embedding visualization with Atlas.

Table of Contents

  1. Why Embedding Visualization Matters
  2. Setting Up Your Environment for Atlas
  3. Creating an Embedding Visualization with MNIST
  4. Visualizing and Interpreting Embeddings in Atlas
  5. Real-World Use Cases & Debugging
  6. How Atlas Stands Out
  7. Conclusion and Next Steps

Why Embedding Visualization Matters

Many AI/ML teams face pressing challenges—ranging from mislabeled data in large classification tasks to overlapping decision boundaries in recommendation engines. The root of these problems often lies in how embeddings represent concepts in high-dimensional space.

Embedding visualization pinpoints these issues by giving you a bird's-eye view of how your data clusters. Nomic Atlas makes this interactive and shareable, so that every stakeholder—from data scientist to product manager—spots trouble areas quickly.

"Atlas was able to improve SmarterX's productivity by 92%, reducing what once required a 90-day time crunch to a single week's workflow. Atlas gave us a roadmap to optimize our models and embeddings rapidly."
Data Science Lead, SmarterX

Setting Up Your Environment for Atlas

Ready to follow along in real-time? Sign up for a free Nomic Atlas account to experience the visualizations firsthand.

Then install the dependencies for this walkthrough:

pip install nomic numpy torch torchvision pytorch-lightning

Logging into Atlas

Acquire your Nomic API key and log in via Python:

import nomic
nomic.login("nk-...")

With your API key in place, you're set to begin uploading and visualizing embeddings in Atlas.

Creating an Embedding Visualization with MNIST

Let's walk through a practical example using the MNIST dataset to demonstrate how embedding visualization can help you understand model decision boundaries and improve performance.

Train a Simple Neural Network

First, we'll train a simple neural network on the MNIST dataset to generate embeddings:

import os
import torch
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks.progress import TQDMProgressBar
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
from torchvision.datasets import MNIST
from nomic import AtlasDataset
import numpy as np

PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
BATCH_SIZE = 256
torch.manual_seed(0)

# Init DataLoader from MNIST Dataset
train_ds = MNIST(PATH_DATASETS, train=True, download=True, transform=transforms.ToTensor())
test_ds = MNIST(PATH_DATASETS, train=False, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE)
test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE)

Next, we'll define our model - a simple two-layer network that will learn to classify MNIST digits:

class MNISTModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)
        self.l2 = torch.nn.Linear(10, 10)

    def forward(self, x):
        return torch.relu(self.l2(torch.relu(self.l1(x.view(x.size(0), -1)))))

    def training_step(self, batch, batch_nb):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

Now let's train the model:

mnist_model = MNISTModel()
trainer = Trainer(accelerator="auto", max_epochs=15)
trainer.fit(mnist_model, train_dataloaders=train_loader, val_dataloaders=test_loader)
mnist_model.eval()

Generate Embeddings and Metadata

After training, we'll use the model's final layer outputs as our embeddings. These 10-dimensional vectors represent how the model "sees" each digit:

all_embeddings = []
all_data = []
for batch_idx, batch in enumerate(test_loader):
    x, y = batch
    
    # Get model embeddings (last layer logits) and predictions
    logits = mnist_model(x)
    embeddings = logits.detach().numpy()
    predictions = torch.argmax(logits, dim=1).detach().numpy()
    
    # get image links for each image (already hosted)
    image_links = [f'https://s3.amazonaws.com/static.nomic.ai/mnist/eval/{label}/{batch_idx*BATCH_SIZE+idx}.jpg'
                       for idx, label in enumerate(y)]
    
    # prepare metadata for each image
    batch_data = [{'label': str(int(label)), 'prediction': str(int(prediction)), 'image': image, 'id': f'{batch_idx*BATCH_SIZE+idx}'}
     for idx, (label, image, prediction) in enumerate(zip(y.tolist(), image_links, predictions))]
    
    all_embeddings.append(embeddings)
    all_data.extend(batch_data)

Upload to Atlas

Now we create an AtlasDataset for our embeddings, add our embeddings array and metadata, and create a data map:

dataset = AtlasDataset(
    identifier='mnist-model-embeddings',
    description='Embeddings of an MNIST model',
    unique_id_field='id'
)

dataset.add_data(embeddings=np.concatenate(all_embeddings), data=all_data)
data_map = dataset.create_index(topic_model=False)

Visualizing and Interpreting Embeddings in Atlas

Once your map is built, you'll see that the embeddings form 10 distinct clusters - one for each digit class. This visualization immediately reveals how your model organizes the digit space.

Visualizing Decision Boundary Overlap

Pay special attention to the areas where clusters overlap. These regions represent digits that the model finds ambiguous or difficult to classify correctly.

For example, the map reveals a boundary between the '4' and '9' clusters, highlighting the model's difficulty in distinguishing between these two digits due to their similar shapes:

Examining these overlapping regions is incredibly valuable for model improvement as they highlight:

Key Indicators of Embedding Quality:

Real-World Use Cases & Debugging

While the MNIST example demonstrates the concept, real-world scenarios often involve text embeddings (e.g., from BERT or OpenAI), product images, or user behavioral patterns. Here's how teams typically leverage Atlas:

Spot Mislabeled Data for Product Classification

For a large e-commerce site, mislabeled entries (like "Vacuum Cleaner" tagged under "Kitchen Utensils") can create overlapping clusters. In Atlas, these mislabeled points stand out when you color by "true label." Correcting them can quickly boost classification accuracy.

Troubleshoot Overlapping Sentiment Classes

If you have "positive," "neutral," and "negative" sentiment embeddings, you might see they aren't cleanly separated. By observing how points cluster and overlap, you can decide whether to apply advanced contrastive learning or regularization to improve separation.

Detect Feature Drift in Real-Time

Over months, user behavior evolves, or new products are introduced. Re-uploading updated embeddings to Atlas helps you see if clusters shift or merge, signaling a potential drop in model performance.

Before & After Comparisons

Upload your dataset before label corrections or embedding tweaks, then upload another snapshot after. Seeing side-by-side improvement (tighter clusters, fewer outliers) motivates teams and confirms your changes are working.

Want to follow along with your own data right now?
Sign up here and try a quick test with your latest embeddings. You can use the same code snippet above—just replace embeddings and data with your real vectors and metadata.

How Atlas Stands Out

You might wonder, "Why use Atlas if I can run UMAP, t-SNE, or TensorBoard locally?" Here are a few differentiators:

Shareable Interactive Maps:

No more static screenshots. You can share a live map with teammates, who can zoom, filter, and annotate in real time.

Easy Collaboration & Metadata Filters:

Upload not just vectors, but also metadata fields (like product category, sentiment labels, user segments). Filter or color by any combination, fostering deeper team discussions.

Powerful Topic Modeling:

If you want more than a scatterplot, Atlas offers topic modeling and text searching to group and interpret textual data at scale.

Scalability:

Large datasets can quickly bog down local tools. Atlas handles big embedding files in the cloud, so you don't have to worry about local resource limits.

Simply put, Atlas doesn't just visualize—it helps you debug, explore, and collaborate on embeddings efficiently.

Next Steps

Embedding visualization is more than a cool chart—it's a powerful diagnostic step in your AI/ML pipeline. Whether you're tackling product mislabeling, sentiment confusion, or classification drift, seeing how your embeddings cluster in Atlas can surface hidden issues and guide meaningful improvements.

Ready to Improve Your Model?

  1. Sign up for Atlas and upload a real dataset (product info, user reviews, customer feedback, etc).
  2. Inspect cluster separations, identify mislabeled data, and compare before-and-after scenarios.
  3. Iterate on your model—try new embedding techniques or label corrections—and see immediate feedback in Atlas.

Learn More

Customer Spotlight: SmarterX saw dramatic classification improvements by using Atlas for systematic data and model debugging.

By making embedding visualization central to your ML workflow, you'll catch misclassifications early, reduce label errors, and provide more accurate AI services to your users. Get started with Atlas today, debug your embeddings, and watch your model performance climb.

nomic logo
nomic logonomic logo nomic logo nomic logonomic logonomic logo nomic logo nomic logo
“Henceforth, it is the map that precedes the territory” – Jean Baudrillard