Umap Optimization

UMAP Optimization Skill

Metadata

Skill Name: UMAP Optimization Type: optimization Version: 1.0.0 Purpose: Optimize UMAP implementations for large-scale datasets and production workloads Triggers: umap, optimize, performance, scale, large dataset, embedding, slow, speed, faster, memory Dependencies: None Context Cost: Medium Complexity: 0.7

Quick Summary: Provides optimization strategies for UMAP-based embedding systems, including cluster caching, FAISS integration, batch processing improvements, and numerical stability enhancements specifically for the umap_analogy_engine.py implementation.

Instructions

When to Use This Skill

Apply this skill when:

• Training time exceeds 5 minutes for N < 10,000
• Dataset size > 10,000 samples and scaling issues appear
• Memory usage grows beyond available GPU memory
• Need to deploy UMAP model to production
• Encountering numerical instability (NaN/Inf losses)

Key Optimization Categories

1. Computational Efficiency

   • Cluster caching (10x speedup)
   • FAISS for large-scale kNN (100x+ speedup for N > 100K)
   • Mixed precision training (2x speedup)

2. Memory Optimization

   • Batch processing strategies
   • Gradient checkpointing
   • Sparse graph representations

3. Numerical Stability

   • Proper epsilon values (1e-8 instead of 1e-12)
   • Gradient clipping
   • Loss clamping

4. Algorithm Improvements

   • Auto-alignment calibration
   • Adaptive learning rates
   • Early stopping

Quick Wins (Start Here)

# 1. Enable cluster caching (already in V2)
KMEANS_UPDATE_FREQ = 10  # Update clusters every 10 steps instead of every step

# 2. Use mixed precision
USE_AMP = True

# 3. Optimize batch size for your GPU
EDGE_BS = 20000  # Adjust based on GPU memory

# 4. Use proper epsilon
EPS = 1e-8  # Not 1e-12

# 5. Enable auto-alignment
AUTO_ALIGN = "once"  # Calibrate loss weights automatically

Performance Checklist

• Cluster caching enabled
• Mixed precision (AMP) enabled
• Batch size optimized for GPU
• Auto-alignment configured
• Proper epsilon values used
• FAISS integration (if N > 100K)
• Gradient clipping enabled
• Progress logging active

Expected Speedups

Resources

Detailed Optimization Guide

1. Cluster Caching (Critical)

Problem: Original implementation runs k-means (50 iterations) on every batch, every step. For a 400-epoch training with 100 steps/epoch, that's 2 million k-means iterations.

Solution: Cache clusters and update periodically:

# In train_relation_aware_umap()
cached_clusters = None
kmeans_update_freq = 10  # Update every 10 steps

for step_idx in range(n_steps):
    update_clusters = (step_idx % kmeans_update_freq == 0) or (cached_clusters is None)

    if update_clusters:
        L_align, mean_dirs, cached_clusters = relation_alignment_loss(
            Z, pair_indices_list, n_clusters_list,
            kmeans_iters=kmeans_iters,
            relation_weights=relation_weights,
            cached_clusters=None  # Force recomputation
        )
    else:
        # Reuse cached clusters
        L_align, mean_dirs, _ = relation_alignment_loss(
            Z, pair_indices_list, n_clusters_list,
            kmeans_iters=kmeans_iters,
            relation_weights=relation_weights,
            cached_clusters=cached_clusters  # Reuse
        )

Impact:

• 10x speedup if updating every 10 steps
• More stable training (clusters don't thrash)
• Better convergence

Tuning:

• kmeans_update_freq = 1: Most accurate, slowest
• kmeans_update_freq = 10: Good balance (recommended)
• kmeans_update_freq = 50: Fastest, may hurt quality

2. FAISS Integration for Large Datasets

Problem: torch.cdist() for kNN is O(N²), limiting scalability:

• N = 10K: ~0.5s
• N = 100K: ~50s
• N = 1M: ~5000s (83 minutes!)

Solution: Replace with FAISS approximate nearest neighbors:

def build_fuzzy_simplicial_set_faiss(
    X: torch.Tensor,
    n_neighbors: int = 15,
    metric: str = "euclidean",
    use_gpu: bool = True
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Build fuzzy simplicial set using FAISS for scalability.

    Faster than torch.cdist for N > 10,000.
    """
    import faiss

    N, D = X.shape
    kmax = min(n_neighbors, N - 1)

    if kmax < 1:
        device = X.device
        return (
            torch.empty(0, dtype=torch.long, device=device),
            torch.empty(0, dtype=torch.long, device=device),
            torch.empty(0, dtype=torch.float32, device=device),
        )

    # Build FAISS index
    if metric == "euclidean":
        index = faiss.IndexFlatL2(D)
    elif metric == "cosine":
        index = faiss.IndexFlatIP(D)  # Inner product after normalization
        # Normalize for cosine similarity
        faiss.normalize_L2(X.cpu().numpy())
    else:
        raise ValueError(f"Unsupported metric: {metric}")

    # Move to GPU if available
    if use_gpu and torch.cuda.is_available():
        res = faiss.StandardGpuResources()
        index = faiss.index_cpu_to_gpu(res, 0, index)

    # Add data
    X_np = X.cpu().numpy().astype('float32')
    index.add(X_np)

    # Search for k+1 neighbors (including self)
    D_faiss, I_faiss = index.search(X_np, kmax + 1)

    # Convert to PyTorch tensors
    device = X.device
    indices = torch.from_numpy(I_faiss[:, 1:]).to(device)  # Exclude self
    distances = torch.from_numpy(D_faiss[:, 1:]).to(device)

    # Rest of the function remains the same...
    # (smooth_knn_dist, symmetrization, etc.)

    return build_graph_from_knn(indices, distances, n_neighbors)

Performance:

Installation:

# CPU version
pip install faiss-cpu

# GPU version (requires CUDA)
pip install faiss-gpu

3. Mixed Precision Training (AMP)

Already implemented in V2, but worth emphasizing:

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for step in range(n_steps):
    optimizer.zero_grad()

    with autocast(enabled=use_amp):
        Z = model(X_high)
        L_umap = umap_cross_entropy_loss(Z, ...)
        L_align, mean_dirs, cached_clusters = relation_alignment_loss(Z, ...)
        L_ortho = orthogonality_penalty(mean_dirs)

        total_loss = L_umap + align_weight * L_align + ortho_weight * L_ortho

    scaler.scale(total_loss).backward()
    scaler.step(optimizer)
    scaler.update()

Benefits:

• 2x speedup on modern GPUs (Volta, Turing, Ampere)
• ~40% memory reduction
• Minimal impact on quality

Caution:

• Requires careful epsilon management (use 1e-8, not 1e-12)
• May need gradient clipping if loss spikes occur

4. Batch Size Optimization

GPU Memory vs Speed Tradeoff:

# Determine optimal batch size for your GPU
def find_optimal_batch_size(n_edges: int, gpu_memory_gb: float) -> int:
    """
    Heuristic for optimal edge batch size.

    Args:
        n_edges: Total number of edges in graph
        gpu_memory_gb: Available GPU memory in GB

    Returns:
        Recommended batch size
    """
    # Rule of thumb: ~1-2 GB per 10K edges
    max_batch = int(gpu_memory_gb * 10000 / 2)

    # But don't exceed total edges or go too small
    optimal = min(max_batch, max(n_edges // 10, 5000))

    # Round to nearest 1000 for cleaner numbers
    return (optimal // 1000) * 1000

# Example usage
EDGE_BS = find_optimal_batch_size(n_edges, gpu_memory_gb=8)

GPU-Specific Recommendations:

5. Auto-Alignment Optimization

Current V2 implementation is good, but can be enhanced:

def calibrate_alignment_weight(
    model: nn.Module,
    X_high: torch.Tensor,
    edges: Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
    pair_indices_list: List[torch.Tensor],
    n_clusters_list: List[int],
    neg_ratio: int = 5,
    use_amp: bool = True,
    n_samples: int = 5,
) -> float:
    """
    Enhanced calibration using multiple samples for robustness.

    Args:
        n_samples: Number of calibration samples to average

    Returns:
        Calibrated alignment weight
    """
    model.eval()
    grad_norms_umap = []
    grad_norms_align = []

    for _ in range(n_samples):
        with torch.enable_grad():
            Z_cal = model(X_high).detach().requires_grad_(True)

            with autocast(enabled=use_amp):
                # UMAP loss
                Lu = umap_cross_entropy_loss(Z_cal, edges, neg_ratio=neg_ratio)

                # Alignment loss
                La, _, _ = relation_alignment_loss(
                    Z_cal, pair_indices_list, n_clusters_list,
                    kmeans_iters=20  # Fewer iters for calibration
                )

            # Compute gradient norms
            gu = torch.autograd.grad(Lu, Z_cal, retain_graph=True)[0].norm().item()
            ga = torch.autograd.grad(La, Z_cal, retain_graph=False)[0].norm().item()

            grad_norms_umap.append(gu)
            grad_norms_align.append(ga)

    # Average over samples
    avg_gu = sum(grad_norms_umap) / n_samples
    avg_ga = sum(grad_norms_align) / n_samples

    # Compute weight
    if avg_ga > 1e-8:
        weight = avg_gu / avg_ga
        # Clamp to reasonable range
        weight = max(0.1, min(10.0, weight))
    else:
        weight = 1.0

    model.train()
    return weight

Benefits:

• More stable calibration (averaging reduces noise)
• Automatic balancing of objectives
• Prevents one loss from dominating

6. Gradient Clipping

Add to prevent instability:

import torch.nn.utils as nn_utils

# In training loop, after backward() and before step()
scaler.unscale_(optimizer)  # Unscale before clipping
nn_utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)

When to use:

• Large datasets (N > 50K)
• High learning rates
• Unstable loss curves
• NaN/Inf appearing during training

7. Memory-Efficient Graph Construction

For very large graphs, use sparse tensors:

def build_sparse_graph(
    indices: torch.Tensor,
    distances: torch.Tensor,
    N: int
) -> torch.sparse.FloatTensor:
    """
    Build sparse adjacency matrix instead of edge list.

    Saves memory for dense graphs.
    """
    # Flatten indices
    rows = torch.arange(N).unsqueeze(1).expand_as(indices).flatten()
    cols = indices.flatten()

    # Compute fuzzy weights
    weights = torch.exp(-distances.flatten())

    # Build sparse tensor
    indices_2d = torch.stack([rows, cols])
    sparse_graph = torch.sparse_coo_tensor(
        indices_2d,
        weights,
        size=(N, N)
    )

    # Symmetrize using sparse operations
    sparse_graph_t = sparse_graph.t()
    symmetric = sparse_graph + sparse_graph_t - sparse_graph * sparse_graph_t

    return symmetric.coalesce()

8. Early Stopping

Prevent overfitting and save time:

class EarlyStopping:
    def __init__(self, patience: int = 20, min_delta: float = 1e-4):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = float('inf')

    def should_stop(self, loss: float) -> bool:
        if loss < self.best_loss - self.min_delta:
            self.best_loss = loss
            self.counter = 0
            return False
        else:
            self.counter += 1
            return self.counter >= self.patience

# In training loop
early_stopping = EarlyStopping(patience=20)

for epoch in range(epochs):
    avg_loss = train_epoch(...)

    if early_stopping.should_stop(avg_loss):
        print(f"Early stopping at epoch {epoch}")
        break

9. Adaptive Learning Rate

Use learning rate scheduling:

from torch.optim.lr_scheduler import ReduceLROnPlateau

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=0.5,
    patience=10,
    verbose=True
)

# In training loop
for epoch in range(epochs):
    avg_loss = train_epoch(...)
    scheduler.step(avg_loss)

Or use cosine annealing:

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-6)

for epoch in range(epochs):
    train_epoch(...)
    scheduler.step()

Complete Optimized Training Function

def train_relation_aware_umap_optimized(
    X_high: torch.Tensor,
    pair_indices_list: List[torch.Tensor],
    n_clusters_list: List[int],
    # Standard parameters...
    n_neighbors: int = 15,
    epochs: int = 400,
    batch_size: int = 20000,
    lr: float = 1e-3,
    # Optimization parameters
    use_faiss: bool = True,          # NEW: Use FAISS for large N
    use_amp: bool = True,            # Mixed precision
    auto_align: str = "once",        # Auto-alignment
    kmeans_update_freq: int = 10,    # Cluster caching frequency
    grad_clip: float = 1.0,          # Gradient clipping
    early_stop_patience: int = 20,   # Early stopping
    use_lr_scheduler: bool = True,   # Learning rate scheduling
    verbose: bool = True,
) -> Tuple[nn.Module, torch.Tensor]:
    """
    Fully optimized UMAP training with all performance enhancements.
    """
    device = X_high.device
    N, D = X_high.shape

    # 1. Build graph (use FAISS if available and N is large)
    if use_faiss and N > 10000:
        try:
            edges = build_fuzzy_simplicial_set_faiss(X_high, n_neighbors)
        except ImportError:
            print("FAISS not available, falling back to torch.cdist")
            edges = build_fuzzy_simplicial_set(X_high, n_neighbors)
    else:
        edges = build_fuzzy_simplicial_set(X_high, n_neighbors)

    # 2. Initialize model
    model = ParametricUMAP(D, d_low, hidden_dims, dropout_p).to(device)

    # 3. Optimizer and scheduler
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    if use_lr_scheduler:
        scheduler = ReduceLROnPlateau(optimizer, patience=10, factor=0.5)

    scaler = GradScaler() if use_amp else None

    # 4. Auto-alignment
    if auto_align == "once":
        align_weight = calibrate_alignment_weight(
            model, X_high, edges, pair_indices_list, n_clusters_list,
            use_amp=use_amp, n_samples=5
        )

    # 5. Early stopping
    early_stopping = EarlyStopping(patience=early_stop_patience)

    # 6. Training loop
    cached_clusters = None
    n_edges = edges[0].size(0)

    for epoch in range(epochs):
        epoch_loss = 0.0
        n_steps = max(1, n_edges // batch_size)

        for step in range(n_steps):
            optimizer.zero_grad()

            # Determine if we should update clusters
            update_clusters = (step % kmeans_update_freq == 0) or (cached_clusters is None)

            with autocast(enabled=use_amp):
                Z = model(X_high)

                # Batch UMAP loss
                L_umap = umap_cross_entropy_loss_batched(Z, edges, batch_size, step)

                # Alignment loss (with caching)
                if update_clusters:
                    L_align, mean_dirs, cached_clusters = relation_alignment_loss(
                        Z, pair_indices_list, n_clusters_list,
                        cached_clusters=None
                    )
                else:
                    L_align, mean_dirs, _ = relation_alignment_loss(
                        Z, pair_indices_list, n_clusters_list,
                        cached_clusters=cached_clusters
                    )

                L_ortho = orthogonality_penalty(mean_dirs)

                total_loss = L_umap + align_weight * L_align + ortho_weight * L_ortho

            # Backward with gradient clipping
            if use_amp:
                scaler.scale(total_loss).backward()
                scaler.unscale_(optimizer)
                nn_utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip)
                scaler.step(optimizer)
                scaler.update()
            else:
                total_loss.backward()
                nn_utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip)
                optimizer.step()

            epoch_loss += total_loss.item()

        avg_loss = epoch_loss / n_steps

        # Learning rate scheduling
        if use_lr_scheduler:
            scheduler.step(avg_loss)

        # Early stopping
        if early_stopping.should_stop(avg_loss):
            if verbose:
                print(f"Early stopping at epoch {epoch}")
            break

        if verbose and epoch % 10 == 0:
            print(f"Epoch {epoch}: Loss = {avg_loss:.4f}")

    return model, model(X_high)

Benchmarking Results

Test configuration:

• Dataset: N=10,000, D=768 (BERT embeddings)
• Relations: 3 (gender, plural, comparative)
• Hardware: RTX 3090 (24GB)

Troubleshooting

Issue: Loss becomes NaN

Causes:

• Epsilon too small (1e-12)
• Learning rate too high
• Numerical overflow in distance calculations

Solutions:

# 1. Use proper epsilon
EPS = 1e-8  # Not 1e-12

# 2. Lower learning rate
lr = 1e-4  # Instead of 1e-3

# 3. Enable gradient clipping
grad_clip = 1.0

# 4. Check for inf/nan
if not torch.isfinite(total_loss):
    print("Warning: Non-finite loss detected")
    continue  # Skip this batch

Issue: Out of memory

Solutions:

# 1. Reduce batch size
batch_size = 10000  # Instead of 20000

# 2. Enable mixed precision
use_amp = True

# 3. Use gradient checkpointing (advanced)
# 4. Process in chunks

# 5. Use sparse graphs for very large N

Issue: Training too slow

Solutions:

# 1. Enable cluster caching
kmeans_update_freq = 10  # Or higher

# 2. Use FAISS for N > 10K
use_faiss = True

# 3. Enable AMP
use_amp = True

# 4. Optimize batch size
batch_size = find_optimal_batch_size(n_edges, gpu_memory_gb)

# 5. Use early stopping
early_stop_patience = 20

Issue: Poor analogy quality

Not an optimization issue, but:

# 1. Ensure metric is correct
metric = "euclidean"  # Not "cosine" for absolute positions

# 2. Check alignment weight
# Auto-alignment should handle this, but verify:
print(f"Alignment weight: {align_weight}")
# Should be 0.1-10 range typically

# 3. Verify clusters are forming
# Check relation statistics after training

# 4. May need more training
# Disable early stopping or increase patience

Summary

Quick Checklist for Production Deployment:

• Cluster caching enabled (10x speedup)
• FAISS for N > 10K (100x speedup)
• Mixed precision training (2x speedup, 40% memory savings)
• Optimal batch size for GPU
• Auto-alignment calibration
• Gradient clipping (max_norm=1.0)
• Early stopping (patience=20)
• Learning rate scheduling
• Proper epsilon values (1e-8)
• Progress logging and monitoring

Expected overall speedup: 30-50x for typical workloads

End of UMAP Optimization Skill