save unnecessary matmul

Signed-off-by: Hao Wu <[email protected]>

Author: skyw

emerging_optimizers/soap/soap_utils.py

@@ -85,13 +85,9 @@ def get_eigenbasis_eigh(
                 # We use an empty tensor so that the `precondition` function will skip this factor.
                 updated_eigenbasis_list.append(torch.empty(0, device=kronecker_factor.device))
                 continue
-            # Construct approximated eigenvalues using QL^T@L@QL or QR^T@R@QR.
-            # The approximated eigenvalues should be close to diagonal if the eigenbasis is close to the true eigenbasis of the kronecker factor
-            # (i.e. the approximated eigenvectors diagonalize the kronecker factor)
-            approx_eigenvalue_matrix = eigenbasis.T @ kronecker_factor @ eigenbasis
             # Update eigenbasis when necessary. Update is skipped only when adaptive update criteria is met.
             if _adaptive_criteria_met(
-                approx_eigenvalue_matrix=approx_eigenvalue_matrix,
+                approx_eigenvalue_matrix=kronecker_factor,
                 tolerance=adaptive_update_tolerance,
             ):
                 _, Q = utils.eig.eigh_with_fallback(
@@ -206,21 +202,21 @@ def get_eigenbasis_qr(
         if kronecker_factor.numel() == 0:
             updated_eigenbasis_list.append(torch.empty(0, device=kronecker_factor.device))
             continue
-        # construct approximated eigenvalues using QL^T@L@QL or QR^T@R@QR, which should be close to diagonal
-        # if the eigenbasis is close to the true eigenbasis of the kronecker factor (i.e. diagonalizes it)
-        approx_eigenvalue_matrix = eigenbasis.T @ kronecker_factor @ eigenbasis
 
         # Update eigenbasis when necessary. Update is skipped only when use_adaptive_criteria is True
         # but criteria is not met.
         if_update = True
         if use_adaptive_criteria and not _adaptive_criteria_met(
-            approx_eigenvalue_matrix=approx_eigenvalue_matrix,
+            approx_eigenvalue_matrix=kronecker_factor,
             tolerance=adaptive_update_tolerance,
         ):
             if_update = False
         if if_update:
+            # construct approximated eigenvalues using QL^T@L@QL or QR^T@R@QR, which should be close to diagonal
+            # if the eigenbasis is close to the true eigenbasis of the kronecker factor (i.e. diagonalizes it)
+            approx_eigvals = _conjugate(kronecker_factor, eigenbasis, diag=True)
             Q, exp_avg_sq = _orthogonal_iteration(
-                approx_eigenvalue_matrix=approx_eigenvalue_matrix,
+                approx_eigvals=approx_eigvals,
                 kronecker_factor=kronecker_factor,
                 eigenbasis=eigenbasis,
                 ind=ind,
@@ -237,7 +233,7 @@ def get_eigenbasis_qr(
 
 
 def _orthogonal_iteration(
-    approx_eigenvalue_matrix: torch.Tensor,
+    approx_eigvals: torch.Tensor,
     kronecker_factor: torch.Tensor,
     eigenbasis: torch.Tensor,
     ind: int,
@@ -267,8 +263,6 @@ def _orthogonal_iteration(
             - Q: The updated eigenbasis
             - exp_avg_sq: The updated (sorted) inner Adam second moment
     """
-    # extract approximated eigenvalues from the diagonal of the projection of kronecker factor onto eigenbases
-    approx_eigvals = torch.diag(approx_eigenvalue_matrix)
     # Sort the approximated eigenvalues according to their magnitudes
     sort_idx = torch.argsort(approx_eigvals, descending=True)
     # re-order the inner adam second moment
@@ -316,3 +310,28 @@ def _adaptive_criteria_met(
 
     # check if normalized diagonal component is not smaller than tolerance
     return not utils.eig.adaptive_early_exit_criteria(approx_eigenvalue_matrix, tolerance)
+
+
+def _conjugate(a: torch.Tensor, p: torch.Tensor, diag: bool = False) -> torch.Tensor:
+    """Calculate similarity transformation
+
+    This function calculates :math:`B = P^T A P`. It assumes P is orthogonal so that :math:`P^{-1} = P^T` and
+    the similarity transformation exists.
+
+    Args:
+        a: matrix to be transformed
+        p: An orthogonal matrix.
+        diag: If True, only return the diagonal of the similarity transformation
+
+    Returns:
+        b
+    """
+    if a.dim() != 2 or p.dim() != 2:
+        raise TypeError("a and p must be 2D matrices")
+    pta = p.T @ a
+    if not diag:
+        b = pta @ p
+    else:
+        # return the diagonal of the similarity transformation
+        b = (pta * p.T).sum(dim=1)
+    return b

tests/test_soap_utils.py

@@ -16,12 +16,7 @@
 from absl.testing import absltest, parameterized
 
 from emerging_optimizers import utils
-from emerging_optimizers.soap.soap_utils import (
-    _adaptive_criteria_met,
-    _orthogonal_iteration,
-    get_eigenbasis_eigh,
-    get_eigenbasis_qr,
-)
+from emerging_optimizers.soap import soap_utils
 
 
 # Base class for tests requiring seeding for determinism
@@ -44,7 +39,7 @@ def test_adaptive_criteria_met(self) -> None:
 
         # Test with small tolerance - should not update since matrix is diagonal
         self.assertFalse(
-            _adaptive_criteria_met(
+            soap_utils._adaptive_criteria_met(
                 approx_eigenvalue_matrix=diagonal_matrix,
                 tolerance=0.1,
             ),
@@ -63,7 +58,7 @@ def test_adaptive_criteria_met(self) -> None:
 
         # Test with small tolerance - should update since matrix has significant off-diagonal elements
         self.assertTrue(
-            _adaptive_criteria_met(
+            soap_utils._adaptive_criteria_met(
                 approx_eigenvalue_matrix=off_diagonal_matrix,
                 tolerance=0.1,
             ),
@@ -72,7 +67,7 @@ def test_adaptive_criteria_met(self) -> None:
 
         # Test with large tolerance - should not update even with off-diagonal elements
         self.assertFalse(
-            _adaptive_criteria_met(
+            soap_utils._adaptive_criteria_met(
                 approx_eigenvalue_matrix=off_diagonal_matrix,
                 tolerance=10.0,
             ),
@@ -102,7 +97,7 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
         }
 
         # We'll call get_eigenbasis_qr
-        Q_new_list, exp_avg_sq_new = get_eigenbasis_qr(
+        Q_new_list, exp_avg_sq_new = soap_utils.get_eigenbasis_qr(
             kronecker_factor_list=state["GG"],
             eigenbasis_list=state["Q"],
             exp_avg_sq=state["exp_avg_sq"],
@@ -171,11 +166,11 @@ def test_update_eigenbasis_with_QR(self, N: int, power_iter_steps: int) -> None:
         # Create estimated eigenvalue matrix by projecting kronecker_factor onto eigenbasis's basis
         approx_eigenvalue_matrix = eigenbasis.T.mm(kronecker_factor).mm(eigenbasis)
         # Extract eigenvalues from the diagonal of the estimated eigenvalue matrix
-        est_eigvals = torch.diag(approx_eigenvalue_matrix)
+        approx_eigvals = torch.diag(approx_eigenvalue_matrix)
 
         # Call the QR function to update the eigenbases and re-order the inner adam second moment
-        Q_new, exp_avg_sq_new = _orthogonal_iteration(
-            approx_eigenvalue_matrix=approx_eigenvalue_matrix,
+        Q_new, exp_avg_sq_new = soap_utils._orthogonal_iteration(
+            approx_eigvals=approx_eigvals,
             kronecker_factor=kronecker_factor,
             eigenbasis=eigenbasis,
             ind=0,  # Test with first dimension
@@ -200,7 +195,7 @@ def test_update_eigenbasis_with_QR(self, N: int, power_iter_steps: int) -> None:
 
         # Test 3: Check that exp_avg_sq is properly sorted based on eigenvalues
         # The sorting should be based on the diagonal elements of estimated_eigenvalue_matrix
-        sort_idx = torch.argsort(est_eigvals, descending=True)
+        sort_idx = torch.argsort(approx_eigvals, descending=True)
         expected_exp_avg_sq = exp_avg_sq.index_select(0, sort_idx)
         torch.testing.assert_close(
             exp_avg_sq_new,
@@ -230,7 +225,7 @@ def test_get_eigenbasis_eigh(self, dims: list[int]) -> None:
             k_factor = k_factor @ k_factor.T + torch.eye(dim, device="cuda") * 1e-5
             kronecker_factor_list.append(k_factor)
 
-        Q_list = get_eigenbasis_eigh(kronecker_factor_list, convert_to_float=True)
+        Q_list = soap_utils.get_eigenbasis_eigh(kronecker_factor_list, convert_to_float=True)
 
         self.assertEqual(len(Q_list), len(kronecker_factor_list))
 
@@ -265,6 +260,20 @@ def test_get_eigenbasis_eigh(self, dims: list[int]) -> None:
                 msg=f"Matrix {i} was not properly diagonalized. Off-diagonal norm: {off_diagonal_norm}",
             )
 
+    def test_conjugate_assert_2d_input(self) -> None:
+        """Tests the conjugate function."""
+        a = torch.randn(2, 3, 4, device="cuda")
+        with self.assertRaises(TypeError):
+            soap_utils._conjugate(a, a)
+
+    def test_conjugate_match_reference(self) -> None:
+        x = torch.randn(15, 17, device="cuda")
+        a = x @ x.T
+        _, p = torch.linalg.eigh(a)
+
+        ref = p.T @ a @ p
+        torch.testing.assert_close(soap_utils._conjugate(a, p), ref, atol=0, rtol=0)
+
 
 if __name__ == "__main__":
     absltest.main()