feat(logging): add positive/negative cosine-similarity and embedding …

[akberovic]

…stddev metrics

• import torch.nn.functional as F # For cosine similarity

•    # Compute and log cosine similarity between positive and negative pair representations.
  

•    with torch.no_grad():
  

•        # Positive cosine similarities.
  

•        cos_sim = F.cosine_similarity(z_i, z_j, dim=-1)
  

•        avg_cos_sim = cos_sim.mean().item()
  

•        writer.add_scalar("Cosine Similarity/avg_positive", avg_cos_sim, args.global_step)
  

•        writer.add_histogram("Cosine Similarity/hist_positive", cos_sim, args.global_step)
  

•        # Compute negative similarities.
  

•        norm_z_i = F.normalize(z_i, dim=1)
  

•        norm_z_j = F.normalize(z_j, dim=1)
  

•        cosine_matrix = torch.mm(norm_z_i, norm_z_j.t())
  

•        mask = torch.eye(cosine_matrix.size(0), dtype=torch.bool, device=cosine_matrix.device)
  

•        negative_sim = cosine_matrix[~mask].view(cosine_matrix.size(0), -1)
  

•        avg_cos_sim_neg = negative_sim.mean().item()
  

•        writer.add_scalar("Cosine Similarity/avg_negative", avg_cos_sim_neg, args.global_step)
  

•        writer.add_histogram("Cosine Similarity/hist_negative", negative_sim, args.global_step)
  

•        # Log embedding standard deviation as a collapse indicator.
  

•        std_z_i = norm_z_i.std(dim=0).mean().item()
  

•        writer.add_scalar("Embeddings/std_dev", std_z_i, args.global_step)