fix: revert covariance

Author: anyangml

deepmd/pt/entrypoints/main.py

@@ -656,102 +656,65 @@ def grad_probe(FLAGS) -> None:
             float(norm_std),
         )
 
-    # Compute pairwise dot product and Pearson correlation statistics
+    # Compute pairwise dot product and cosine similarity statistics
     for (k1, k2), dots in pairwise_dots.items():
         if dots:
-            dot_mean = np.mean(dots)
-            dot_std = np.std(dots)
+            dot_mean = float(np.mean(dots))
+            dot_std = float(np.std(dots))
 
             g1_avg = grads_per_task[k1]
             g2_avg = grads_per_task[k2]
-            dot_accum = float(np.dot(g1_avg, g2_avg))
+            dot_global = float(np.dot(g1_avg, g2_avg))
 
-            # Pearson correlation across groups:
-            #   cov_sum = Σ_j Cov(g1_j, g2_j) = E[dot(g1,g2)] - dot(ḡ1, ḡ2)
-            #   var_sum = Σ_j Var(g_j)         = E[||g||²] - ||ḡ||²
-            #   pearson = cov_sum / sqrt(var1_sum * var2_sum)
-            cov_sum = float(dot_mean - dot_accum)
+            # Group-level cosine: per-group cos_i = dot_i / (||g1_i|| * ||g2_i||)
             norms_1 = np.array(task_batch_norms[k1])
             norms_2 = np.array(task_batch_norms[k2])
-            var1_sum = float(np.mean(norms_1**2) - np.dot(g1_avg, g1_avg))
-            var2_sum = float(np.mean(norms_2**2) - np.dot(g2_avg, g2_avg))
-            denom = np.sqrt(max(var1_sum, 0.0) * max(var2_sum, 0.0))
-            pearson = float(cov_sum / denom) if denom > 1e-12 else float("nan")
-
-            # Group-level weighted cosine similarity:
-            #   per-group cos_i = dot_i / (||g1_i|| * ||g2_i||)
-            #   weight w_i = ||g1_i|| * ||g2_i|| so high-norm groups dominate
-            norm_products = norms_1 * norms_2  # (n_groups,) weight per group
+            norm_products = norms_1 * norms_2
             valid = norm_products > 1e-12
             cos_per_group = np.where(
                 valid,
                 np.array(dots) / np.where(valid, norm_products, 1.0),
-                0.0,
+                float("nan"),
             )
-            norm_product_sum = float(np.sum(norm_products))
-            if norm_product_sum > 1e-12:
-                cos_group_mean = float(
-                    np.sum(norm_products * cos_per_group) / norm_product_sum
-                )
-                cos_group_std = float(
-                    np.sqrt(
-                        np.sum(norm_products * (cos_per_group - cos_group_mean) ** 2)
-                        / norm_product_sum
-                    )
-                )
+            valid_cos = cos_per_group[valid]
+            if len(valid_cos) > 0:
+                cos_group_mean = float(np.mean(valid_cos))
+                cos_group_std = float(np.std(valid_cos))
             else:
                 cos_group_mean = float("nan")
                 cos_group_std = float("nan")
 
-            # Global-level cosine similarity:
-            #   cosine of accumulated average gradient vectors;
-            #   large-norm batches naturally dominate the direction of the average
+            # Global-level cosine: cosine of accumulated average gradient vectors
             norm_g1 = float(np.linalg.norm(g1_avg))
             norm_g2 = float(np.linalg.norm(g2_avg))
             denom_global = norm_g1 * norm_g2
             cos_global = (
-                float(dot_accum / denom_global)
+                float(dot_global / denom_global)
                 if denom_global > 1e-12
                 else float("nan")
             )
 
             grads_per_task[f"dot_{k1}_{k2}_mean"] = dot_mean
             grads_per_task[f"dot_{k1}_{k2}_std"] = dot_std
-            grads_per_task[f"dot_{k1}_{k2}_accum"] = dot_accum
-            grads_per_task[f"pearson_{k1}_{k2}"] = pearson
+            grads_per_task[f"dot_{k1}_{k2}_global"] = dot_global
             grads_per_task[f"cos_group_mean_{k1}_{k2}"] = cos_group_mean
             grads_per_task[f"cos_group_std_{k1}_{k2}"] = cos_group_std
             grads_per_task[f"cos_global_{k1}_{k2}"] = cos_global
 
             log.info(
-                "Grad dot product '%s' vs '%s': Mean=%.4e, Std=%.4e, Accum=%.4e  "
-                "(SNR=%.2f, n_groups=%d)",
-                k1,
-                k2,
-                float(dot_mean),
-                float(dot_std),
-                float(dot_accum),
-                abs(dot_mean) / dot_std if dot_std > 0 else float("inf"),
-                len(dots),
-            )
-            log.info(
-                "Pearson correlation '%s' vs '%s': %.4f  "
-                "(cov_sum=%.4e, var1=%.4e, var2=%.4e)",
-                k1,
-                k2,
-                pearson,
-                cov_sum,
-                var1_sum,
-                var2_sum,
-            )
-            log.info(
-                "Cosine similarity '%s' vs '%s': "
-                "group_mean=%.4f, group_std=%.4f, global=%.4f",
+                "Grad similarity '%s' vs '%s': "
+                "dot_mean=%.4e, dot_std=%.4e, dot_global=%.4e, "
+                "cos_group_mean=%.4f, cos_group_std=%.4f, cos_global=%.4f  "
+                "(n_groups=%d)",
                 k1,
                 k2,
+                dot_mean,
+                dot_std,
+                dot_global,
                 cos_group_mean,
                 cos_group_std,
                 cos_global,
+                len(dots),
             )
 
     save_dict = {f"grads_{k}": v for k, v in grads_per_task.items()}