Embedding Analysis for Dataset Curation

Overview

Embeddings turn unstructured text/images into a vector space where quality, diversity, and redundancy become measurable. This skill covers embedding-based dataset curation using sentence-transformers, distribution comparison metrics, semantic deduplication, and quality scoring techniques from the 2025-2026 literature.

When to Use

Use this skill when:

• Detecting near-duplicate or semantically redundant examples.
• Comparing embedding distributions across dataset splits (train vs. val vs. test).
• Measuring dataset diversity or coverage in embedding space.
• Filtering examples based on embedding distance to a "gold" reference set.
• Applying NeMo Curator-style semantic dedup at scale.
• Computing DataRater-style quality signals from embedding neighborhoods.

Do not use for:

• Raw text deduplication (exact match) — use simple hashing.
• Model training — use transformers or trl skills.
• Vector database persistence — use chromadb, milvus, or qdrant skills.

Installation

uv pip install sentence-transformers umap-learn scikit-learn numpy

Core Workflow

1. Generate Embeddings

from sentence_transformers import SentenceTransformer
import numpy as np

# Choose model based on domain
model = SentenceTransformer("all-MiniLM-L6-v2")  # fast, general
# model = SentenceTransformer("all-mpnet-base-v2")  # higher quality
# model = SentenceTransformer("BAAI/bge-large-en-v1.5")  # retrieval-optimized
# model = SentenceTransformer("intfloat/multilingual-e5-large")  # multilingual

# Batched encoding (memory-efficient)
embeddings = model.encode(
    texts,
    batch_size=64,
    show_progress_bar=True,
    normalize_embeddings=True,  # cosine similarity via dot product
    convert_to_numpy=True,
)

2. Semantic Deduplication

Semantic dedup removes examples that are meaning-equivalent but not text-identical. From NeMo Curator (NVIDIA, 2024-2025) and academic literature.

Clustering-Based Approach (NeMo Curator SemDedup)

from sklearn.cluster import MiniBatchKMeans
from sklearn.metrics.pairwise import cosine_similarity

def semantic_dedup(embeddings, threshold=0.95, n_clusters=100):
    """
    Cluster embeddings, then remove near-duplicates within each cluster.
    threshold: cosine similarity above which two examples are duplicates.
    Returns: boolean mask (True = keep, False = duplicate).
    """
    n = len(embeddings)
    keep = np.ones(n, dtype=bool)

    # Coarse clustering for efficiency (O(n * k) instead of O(n^2))
    k = min(n_clusters, n // 10)
    clusters = MiniBatchKMeans(n_clusters=k, random_state=42, batch_size=1024).fit_predict(embeddings)

    for c in range(k):
        idx = np.where(clusters == c)[0]
        if len(idx) < 2:
            continue
        sim = cosine_similarity(embeddings[idx])
        # Mark duplicates (lower-index kept, higher-index removed)
        for i in range(len(idx)):
            if not keep[idx[i]]:
                continue
            dupes = np.where(sim[i, i+1:] > threshold)[0]
            for d in dupes:
                keep[idx[i + 1 + d]] = False

    return keep

Connected Components Approach (LSHBloom, 2025)

# For internet-scale dedup (>100M examples):
# Use LSH for approximate nearest neighbor + union-find
from sklearn.neighbors import NearestNeighbors

def lsh_semantic_dedup(embeddings, threshold=0.90, n_neighbors=10):
    """Finds clusters of semantically similar examples via NN graph."""
    nn = NearestNeighbors(n_neighbors=n_neighbors, metric="cosine", n_jobs=-1)
    nn.fit(embeddings)
    distances, indices = nn.kneighbors(embeddings)

    # Union-find to merge connected components
    parent = np.arange(len(embeddings))
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    def union(a, b):
        parent[find(a)] = find(b)

    for i in range(len(embeddings)):
        for j, d in zip(indices[i], distances[i]):
            if i != j and (1 - d) > threshold:  # cosine dist → similarity
                union(i, j)

    # Keep one per component
    components = {}
    for i in range(len(embeddings)):
        root = find(i)
        if root not in components:
            components[root] = i

    return list(components.values())  # indices to keep

3. Distribution Comparison Across Splits

Compare embedding distributions between train/val/test to detect drift or bias.

from scipy.spatial.distance import jensenshannon
from scipy.stats import wasserstein_distance

def embedding_distribution_shift(train_emb, test_emb, n_bins=50):
    """Quantify shift between train and test embedding distributions."""

    # Project to 1D for distribution comparison
    from sklearn.decomposition import PCA
    pca = PCA(n_components=1, random_state=42).fit(
        np.concatenate([train_emb, test_emb])
    )
    train_1d = pca.transform(train_emb).ravel()
    test_1d = pca.transform(test_emb).ravel()

    # Histogram-based metrics
    hist_range = (min(train_1d.min(), test_1d.min()),
                   max(train_1d.max(), test_1d.max()))
    train_hist, _ = np.histogram(train_1d, bins=n_bins, range=hist_range, density=True)
    test_hist, _ = np.histogram(test_1d, bins=n_bins, range=hist_range, density=True)

    # Jensen-Shannon divergence (0 = identical, 1 = maximally different)
    js_div = jensenshannon(train_hist + 1e-10, test_hist + 1e-10)

    # Wasserstein (Earth Mover's) distance
    w_dist = wasserstein_distance(train_1d, test_1d)

    return {"js_divergence": js_div, "wasserstein": w_dist}

def per_dimension_shift(train_emb, test_emb):
    """Per-dimension mean shift — identifies which semantic axes drifted."""
    train_mean = train_emb.mean(axis=0)
    test_mean = test_emb.mean(axis=0)
    cos_sim = np.dot(train_mean, test_mean) / (
        np.linalg.norm(train_mean) * np.linalg.norm(test_mean)
    )
    max_dim_shift = np.argmax(np.abs(train_mean - test_mean))
    return {"mean_cosine": cos_sim, "max_shift_dim": int(max_dim_shift)}

4. Embedding Quality Scoring (DataRater-Inspired, 2025)

DataRater (Calian et al., NeurIPS 2025) meta-learns a quality function over embedding neighborhoods. Below is a practical approximation.

def embedding_quality_score(embeddings, k=20):
    """
    Score each example by embedding neighborhood coherence.
    Low coherence → outlier / potential noise.
    High coherence → in-distribution, likely high quality.
    """
    nn = NearestNeighbors(n_neighbors=k + 1, metric="cosine")
    nn.fit(embeddings)
    distances, _ = nn.kneighbors(embeddings)

    # Exclude self-distance (index 0)
    mean_dist = distances[:, 1:].mean(axis=1)

    # Normalize to [0, 1] — lower distance = higher quality
    scores = 1 - (mean_dist - mean_dist.min()) / (mean_dist.max() - mean_dist.min() + 1e-10)
    return scores

5. Diversity Measurement

def embedding_diversity(embeddings, n_clusters=50):
    """
    Measure dataset diversity via cluster coverage.
    Returns: coverage ratio and entropy of cluster assignments.
    """
    k = min(n_clusters, len(embeddings) // 10)
    clusters = MiniBatchKMeans(n_clusters=k, random_state=42, batch_size=1024).fit_predict(embeddings)

    # Cluster coverage: fraction of clusters that are non-empty
    unique, counts = np.unique(clusters, return_counts=True)
    coverage = len(unique) / k

    # Entropy: higher = more uniform distribution across clusters
    probs = counts / counts.sum()
    entropy = -np.sum(probs * np.log(probs + 1e-10)) / np.log(k)  # normalized

    return {"coverage": coverage, "normalized_entropy": entropy}

def embedding_redundancy_score(embeddings, sample_size=5000):
    """
    Estimate redundancy: average pairwise cosine similarity.
    High score → dataset is homogeneous; needs diversification.
    """
    idx = np.random.choice(len(embeddings), min(sample_size, len(embeddings)), replace=False)
    sample = embeddings[idx]
    sim = cosine_similarity(sample)
    # Exclude diagonal
    mask = ~np.eye(len(sample), dtype=bool)
    return float(sim[mask].mean())

Perplexity-Based Filtering (GRAPE Score, 2025)

Use a small reference model to score data quality by perplexity. High perplexity = unfamiliar/hard/noisy.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def grape_score(texts, model_name="gpt2", batch_size=8):
    """
    GRAPE-style perplexity scoring.
    Lower perplexity → more in-distribution, likely higher quality.
    Very high perplexity → gibberish, noise, or out-of-domain.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(model_name).eval()

    scores = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        enc = tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
        with torch.no_grad():
            outputs = model(**enc, labels=enc["input_ids"])
            loss = outputs.loss  # average NLL across batch
            ppl = torch.exp(loss).item()
            scores.append(ppl)

    return np.array(scores)

Quality Gate

An embedding-based curation pass is complete when:

• Semantic dedup threshold is justified and documented.
• Train/val/test embedding distributions are compared (JS divergence, Wasserstein).
• Diversity and redundancy metrics are computed before and after curation.
• Outlier scores are reviewed — low neighborhood coherence may indicate noise or rare-but-valuable data.
• Results are reproducible (fixed seeds, saved model versions).