---
name: embedding-analysis-146a53702970
description: Compute and analyze embeddings for dataset quality, distribution comparison, semantic deduplication, diversity measurement, and similarity-based filtering. Covers sentence-transformers, embedding space diagnostics, and 2025-2026 literature techniques (NeMo Curator semantic dedup, embedding similarity metrics for data selection, DataRater-style quality scoring).
tags: [embeddings, semantic-dedup, distribution-analysis, sentence-transformers, data-quality, metric-learning, dataset-curation]
---

# 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

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

## Core Workflow

### 1. Generate Embeddings

```python
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)

```python
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)

```python
# 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.

```python
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.

```python
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

```python
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.

```python
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).


---
Source attribution: mkurman
Original project: https://github.com/mkurman/zorai/blob/d06ee7597049567dab92db22d41b2764cf89094a/skills/data-lattice/embedding-analysis/SKILL.md
License: MIT
License information: https://opensource.org/license/mit
Original upstream instructions; downloaded from Mega AI Hub.

## Upstream license notice

MIT License

Copyright (c) 2026 zorai contributors

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.



# Mega AI Hub resource metadata

Creator / publisher: mkurman
Original project: https://github.com/mkurman/zorai/blob/d06ee7597049567dab92db22d41b2764cf89094a/skills/data-lattice/embedding-analysis/SKILL.md
Attribution basis: Source repository record

Fields of work: Data & Analytics, Development & Engineering, Product Management
Roles: Data Analyst, Data Scientist, Market Research Analyst, ML / AI Engineer, Product Manager
Specialties: API, Testing
Tools: Numpy, Scikit Learn, Sentence Transformers, Umap Learn
Country applicability: Country not specified
Classification basis: Source catalog labels
Instruction SHA-256: c9c3a11404263d3b9c5dadf427da3945496b21861773f54fb8b79ec8a4cf368a
