Put gram matrix normalization into entropy computation. (#7)

Co-authored-by: knikolaou <[email protected]>

Author: KonstiNik

papyrus/measurements/measurements.py

@@ -315,6 +315,7 @@ def __init__(
         rank: int = 1,
         normalize_eigenvalues: bool = True,
         effective: bool = False,
+        normalize_matrix: bool = False,
     ):
         """
         Constructor method of the NTKCrossEntropy class.
@@ -332,11 +333,15 @@ def __init__(
         effective : bool (default=False)
                 If true, the entropy is divided by the theoretical maximum entropy of
                 the system thereby returning the effective entropy / entropy density.
-
+        normalize_matrix : bool (default=False)
+                If true, the NTK is normalized by the square root of the product of the
+                corresponding diagonal elements. This is equivalent to normalizing the
+                gradient vectors forming the NTK.
         """
         super().__init__(name=name, rank=rank)
         self.normalize_eigenvalues = normalize_eigenvalues
         self.effective = effective
+        self.normalize_matrix = normalize_matrix
 
     def apply(self, ntk: np.ndarray) -> np.ndarray:
         """
@@ -373,6 +378,7 @@ def apply(self, ntk: np.ndarray) -> np.ndarray:
             ntk,
             effective=self.effective,
             normalize_eig=self.normalize_eigenvalues,
+            normalize_matrix=self.normalize_matrix,
         )
 
 

papyrus/utils/analysis_utils.py

@@ -23,7 +23,10 @@
 
 import numpy as np
 
-from papyrus.utils.matrix_utils import compute_hermitian_eigensystem
+from papyrus.utils.matrix_utils import (
+    compute_hermitian_eigensystem,
+    normalize_gram_matrix,
+)
 
 
 def compute_trace(matrix: np.ndarray, normalize: bool = False) -> np.ndarray:
@@ -77,7 +80,10 @@ def compute_shannon_entropy(dist: np.ndarray, effective: bool = False) -> float:
 
 
 def compute_von_neumann_entropy(
-    matrix: np.ndarray, effective: bool = True, normalize_eig: bool = True
+    matrix: np.ndarray,
+    effective: bool = True,
+    normalize_eig: bool = True,
+    normalize_matrix: bool = False,
 ) -> float:
     """
     Compute the von-Neumann entropy of a matrix.
@@ -91,12 +97,19 @@ def compute_von_neumann_entropy(
             the system thereby returning the effective entropy.
     normalize_eig : bool (default = True)
             If true, the eigenvalues are scaled to look like probabilities.
+    normalize_matrix : bool (default=False)
+            If true, the NTK is normalized by the square root of the product of the
+            corresponding diagonal elements. This is equivalent to normalizing the
+            gradient vectors forming the NTK.
 
     Returns
     -------
     entropy : float
             Von-Neumann entropy of the matrix.
     """
+    if normalize_matrix:
+        matrix = normalize_gram_matrix(matrix)
+
     eigvals, _ = compute_hermitian_eigensystem(matrix, normalize=normalize_eig)
 
     entropy = compute_shannon_entropy(eigvals)