loss.py

import torch
import torch.nn.functional as F
from typing import Tuple, Dict
from tqdm import tqdm


def compute_monosemanticity_loss_batch(x: torch.Tensor, latents: torch.Tensor, 
                                       eps: float = 1e-8) -> Tuple[torch.Tensor, Dict[str, float]]:
    """
    Compute monosemanticity loss on a batch using 1 - mono_score.
    
    Args:
        x: Input embeddings (B, D) - will be normalized to unit norm
        latents: Neuron activations (B, M) - should be non-negative
        eps: Small constant for numerical stability
    
    Returns:
        loss: 1 - mean(monosemanticity), range [0, 1]
        metrics: Dictionary with monosemanticity statistics
    """
    # Normalize embeddings to unit norm for cosine similarity
    x_norm = F.normalize(x, p=2, dim=1, eps=eps)  # (B, D)
    
    # Normalize activations per neuron to [0, 1] within batch
    min_acts = latents.min(dim=0, keepdim=True)[0]  # (1, M)
    max_acts = latents.max(dim=0, keepdim=True)[0]  # (1, M)
    a = (latents - min_acts) / (max_acts - min_acts + eps)  # (B, M) this is a_tilde in the paper
    
    # Fast monosemanticity via streaming formula (same as compute_monosemanticity_fast)
    # V = X^T @ A, where rows are weighted embeddings per neuron
    w = x_norm.T @ a  # (D, M)
    
    # Numerator: Σ_{i<j} a_i·a_j·cos(e_i,e_j) = 0.5(||Σ a_i·e_i||² - Σ a_i²)
    q = (w * w).sum(dim=0)  # (M,)
    v = (a * a).sum(dim=0)  # (M,)
    num = 0.5 * (q - v)
    
    # Denominator: Σ_{i<j} a_i·a_j = 0.5((Σ a_i)² - Σ a_i²)
    u = a.sum(dim=0)  # (M,)
    den = 0.5 * (u * u - v)
    
    # Monosemanticity score per neuron
    mono_scores = torch.where(
        den > eps,
        num / (den + eps),
        torch.zeros_like(den)
    )  # (M,)
    
    # Only consider active neurons
    active_mask = den > eps
    num_active = active_mask.sum().item()
    
    if num_active > 0:
        mono_mean = mono_scores[active_mask].mean()
        mono_median = mono_scores[active_mask].median()
    else:
        mono_mean = torch.tensor(0.0, device=x.device)
        mono_median = torch.tensor(0.0, device=x.device)
    
    # Loss: 1 - mono_score (minimize to maximize monosemanticity)
    loss = 1.0 - mono_mean
    
    metrics = {
        'mono_loss_val': loss.item(),
        'mono_score': mono_mean.item(),
        'mono_median': mono_median.item(),
        'mono_active_neurons': num_active,
    }
    
    return loss, metrics


def compute_monosemanticity_loss_batch_ref(x: torch.Tensor,
                                           latents: torch.Tensor,
                                           eps: float = 1e-8,
                                           pair_batch_size: int = 1024) -> Tuple[torch.Tensor, Dict[str, float]]:
    """
    Reference implementation of monosemanticity loss on a single batch via explicit pairwise accumulation.
    This is O(B^2·M) and intended for validation against the fast formula.
    
    Args:
        x: Input embeddings (B, D) - unnormalized; cosine similarity uses unit-normalized vectors
        latents: Neuron activations (B, M) - will be scaled to [0, 1] across the batch per neuron
        eps: Small constant for numerical stability
        pair_batch_size: Chunk size for the inner j-loop to control memory usage
    
    Returns:
        loss: 1 - mean(monosemanticity over active neurons), range [0, 1]
        metrics: Dictionary with monosemanticity statistics
    """
    B = x.shape[0]
    device = x.device
    
    # Normalize embeddings to unit norm for cosine similarity
    x_norm = F.normalize(x, p=2, dim=1, eps=eps)  # (B, D)
    
    # Normalize activations per neuron to [0, 1] within the batch
    min_acts = latents.min(dim=0, keepdim=True)[0]  # (1, M)
    max_acts = latents.max(dim=0, keepdim=True)[0]  # (1, M)
    a = (latents - min_acts) / (max_acts - min_acts + eps)  # (B, M)
    
    M = a.shape[1]
    weighted_cosine_sum = torch.zeros(M, device=device, dtype=torch.float32)
    weight_sum = torch.zeros(M, device=device, dtype=torch.float32)
    
    for i in range(B-1):
        # Preload row i
        emb_i = x_norm[i]                 # (D,)
        act_i = a[i]                      # (M,)
        
        for j_start in range(i + 1, B, pair_batch_size):
            j_end = min(j_start + pair_batch_size, B)
            
            emb_j = x_norm[j_start:j_end]         # (chunk, D)
            act_j = a[j_start:j_end]              # (chunk, M)
            
            # Cosine similarities between i and j chunk
            cos_ij = F.cosine_similarity(
                emb_i.unsqueeze(0).expand(j_end - j_start, -1),  # (chunk, D)
                emb_j,
                dim=1
            ).to(torch.float32)  # (chunk,)
            
            # Weights per neuron for each pair (i, j): a_i * a_j
            w_ij = act_i.unsqueeze(0) * act_j  # (chunk, M)
            
            # Accumulate weighted cosine similarities and weights
            weighted_cosine_sum += (w_ij * cos_ij.unsqueeze(1)).sum(dim=0)  # (M,)
            weight_sum += w_ij.sum(dim=0)  # (M,)
    
    # Per-neuron monosemanticity
    mono_scores = torch.where(
        weight_sum > eps,
        weighted_cosine_sum / (weight_sum + eps),
        torch.zeros_like(weight_sum)
    )  # (M,)
    
    # Consider only active neurons
    active_mask = weight_sum > eps
    num_active = active_mask.sum().item()
    
    if num_active > 0:
        mono_mean = mono_scores[active_mask].mean()
        mono_median = mono_scores[active_mask].median()
    else:
        mono_mean = torch.tensor(0.0, device=device)
        mono_median = torch.tensor(0.0, device=device)
    
    loss = 1.0 - mono_mean
    metrics = {
        'mono_loss_val': loss.item(),
        'mono_score': mono_mean.item(),
        'mono_median': mono_median.item(),
        'mono_active_neurons': num_active,
    }
    return loss, metrics


@torch.no_grad()
def compute_monosemanticity_ref(autoencoder, dataloader, device='cuda', split_name='val', 
                                pair_batch_size: int = 128, eps: float = 1e-8, verbose: bool = False):
    """
    Reference MonoScore implementation using explicit nested loops over all pairs.
    
    
    Args:
        autoencoder: The autoencoder model
        dataloader: DataLoader for the dataset
        device: Device to run on
        split_name: Name of the split for logging
        pair_batch_size: Batch size for j loop (memory control)
        eps: Small constant for numerical stability
        verbose: If True, print progress messages
    
    Returns:
        monosemanticity: (M,) tensor of scores per neuron
    """
    autoencoder.eval()
    dataset = dataloader.dataset
    num_images = len(dataset)
    
    # Pass 1: Collect all embeddings and activations with normalization, also collect min and max values for activations
    if verbose:
        print(f"First pass: collecting all embeddings and activations ({split_name})...")
    all_embeddings = []
    all_activations = []
    min_vals = None
    max_vals = None
    num_neurons = None
    
    for batch in tqdm(dataloader, desc=f"Collecting data for {split_name}", disable=not verbose):
        if not isinstance(batch, torch.Tensor):
            batch = torch.from_numpy(batch)
        x = batch.to(device)
        
        _, latents, _ = autoencoder(x)
        
        if min_vals is None:
            num_neurons = latents.shape[1]
            min_vals = latents.min(dim=0, keepdim=True)[0]
            max_vals = latents.max(dim=0, keepdim=True)[0]
        else:
            batch_min = latents.min(dim=0, keepdim=True)[0]
            batch_max = latents.max(dim=0, keepdim=True)[0]
            min_vals = torch.min(min_vals, batch_min)
            max_vals = torch.max(max_vals, batch_max)
        
        all_embeddings.append(x)
        all_activations.append(latents)
    
    embeddings = torch.cat(all_embeddings, dim=0)  # (N, D)
    activations = torch.cat(all_activations, dim=0)  # (N, M)
    
    # Normalize activations per neuron to [0, 1]
    activations = (activations - min_vals) / (max_vals - min_vals + eps)
    
    if verbose:
        print(f"Dataset size: {num_images}")
        print(f"Number of neurons: {num_neurons}")
    
    # Pass 2: Nested loop over all pairs i < j
    if verbose:
        print(f"Second pass: computing pairwise monosemanticity (ref {split_name})...")
    weighted_cosine_similarity_sum = torch.zeros(num_neurons, device=device)
    weight_sum = torch.zeros(num_neurons, device=device)
    
    for i in tqdm(range(num_images), desc="Processing image pairs", disable=not verbose):
        for j_start in range(i + 1, num_images, pair_batch_size):  # Process in batches
            j_end = min(j_start + pair_batch_size, num_images)
            
            embeddings_i = embeddings[i]  # (D,)
            embeddings_j = embeddings[j_start:j_end]  # (batch, D)
            activations_i = activations[i]  # (M,)
            activations_j = activations[j_start:j_end]  # (batch, M)
            
            cosine_similarities = F.cosine_similarity(
                embeddings_i.unsqueeze(0).expand(j_end - j_start, -1),  # (batch, D)
                embeddings_j,
                dim=1
            )  # (batch,)
            
            # Compute weights and weighted similarities
            weights = activations_i.unsqueeze(0) * activations_j  # (batch, M)
            weighted_cosine_similarities = weights * cosine_similarities.unsqueeze(1)  # (batch, M)
            
            # Accumulate
            weighted_cosine_similarity_sum += weighted_cosine_similarities.sum(dim=0)  # (M,)
            weight_sum += weights.sum(dim=0)  # (M,)
    
    # Compute final monosemanticity scores
    monosemanticity = torch.where(
        weight_sum > eps,
        weighted_cosine_similarity_sum / weight_sum,
        torch.zeros_like(weight_sum)  
    )
    
    autoencoder.train()
    return monosemanticity


@torch.no_grad()
def compute_monosemanticity_fast(autoencoder, dataloader, device='cuda', split_name='val', eps: float = 1e-8, verbose: bool = False):
    """
    Fast O(N·D·M) monosemanticity in single pass (no pairwise loops).

    Args:
        autoencoder: The autoencoder model
        dataloader: DataLoader for the dataset
        device: Device to run on
        split_name: Name of the split for logging
        eps: Small constant for numerical stability
        verbose: If True, print progress messages

    Returns:
        monosemanticity: (M,) tensor of scores per neuron
    """
    autoencoder.eval()

    dataset = dataloader.dataset

    # Pass 1: activation min/max per neuron for scaling
    if verbose:
        print(f"First pass: computing activation statistics (fast {split_name})...")
    min_vals = None
    max_vals = None
    num_neurons = None
    feature_dim = None

    for batch in tqdm(dataloader, desc=f"Computing stats for {split_name}", disable=not verbose):
        if not isinstance(batch, torch.Tensor):
            batch = torch.from_numpy(batch)
        x = batch.to(device)

        _, latents, _ = autoencoder(x)

        if min_vals is None:
            num_neurons = latents.shape[1]
            feature_dim = x.shape[1]
            min_vals = latents.min(dim=0, keepdim=True)[0]
            max_vals = latents.max(dim=0, keepdim=True)[0]
        else:
            batch_min = latents.min(dim=0, keepdim=True)[0]
            batch_max = latents.max(dim=0, keepdim=True)[0]
            min_vals = torch.min(min_vals, batch_min)
            max_vals = torch.max(max_vals, batch_max)

    dataset_size = len(dataset)
    if verbose:
        print(f"Dataset size: {dataset_size}")
        print(f"Number of neurons: {num_neurons}")

    # Pass 2: stream accumulators
    if verbose:
        print(f"Second pass: streaming accumulators...")
    # Accumulators in FP32 for numerical stability
    u = torch.zeros(num_neurons, device=device, dtype=torch.float32)
    v = torch.zeros(num_neurons, device=device, dtype=torch.float32)
    w = torch.zeros(feature_dim, num_neurons, device=device, dtype=torch.float32)

    for batch in tqdm(dataloader, desc=f"Accumulating {split_name}", disable=not verbose):
        if not isinstance(batch, torch.Tensor):
            batch = torch.from_numpy(batch)
        x = batch.to(device)

        _, latents, _ = autoencoder(x)

        # Normalize activations per neuron to [0, 1]
        latents = (latents - min_vals) / (max_vals - min_vals + eps)
        a = latents

        # Unit-normalize embeddings to match cosine similarity
        p = F.normalize(x, p=2, dim=1, eps=1e-8)

        # Cast to fp32 for accumulators
        a32 = a.float()
        p32 = p.float()

        # Update accumulators
        u += a32.sum(dim=0)
        v += (a32 * a32).sum(dim=0)
        w += p32.T @ a32  # (D, B) @ (B, M) -> (D, M)

    # Compute per-neuron weighted sums
    q = (w * w).sum(dim=0)  # (M,)
    num = 0.5 * (q - v)  # Σ_{i<j} a_i a_j cos 
    den = 0.5 * (u * u - v)      # Σ_{i<j} a_i a_j

    # Final monosemanticity scores
    monosemanticity = torch.where(
        den > 1e-8,
        num / (den + 1e-12),
        torch.zeros_like(den)
    )
    autoencoder.train()
    return monosemanticity