Improve eig utils (#190)

* move sort to its own function

Signed-off-by: Hao Wu <skyw@nvidia.com>

* move eigen bases based sort to its own function and share between QR and eigh path

Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: mikail <mkhona@nvidia.com>

Author: skyw

emerging_optimizers/soap/soap.py

@@ -468,17 +468,28 @@ def update_eigenbasis_and_exp_avgs(
         dims=[[0], [1]],
     )
 
-    # Step 2: Update eigenbases
+    # Step 2a: Sort current eigenbases by descending approximate eigenvalues of the updated kronecker
+    # factors, and permute exp_avg_sq.
+    # Shared by both eigh and QR paths so the new eigh-path approximation matches the QR-path slot semantics
+    # under small per-step drift.
+    # Sorting eigenbases is not necessary for eigh path technically, but decided to keep API simple.
+    eigenbasis_list, exp_avg_sq = soap_utils.sort_eigenbasis_by_approx_eigvals(
+        kronecker_factor_list,
+        eigenbasis_list,
+        exp_avg_sq,
+    )
+
+    # Step 2b: Update eigenbases
     if use_eigh:
         updated_eigenbasis_list = soap_utils.get_eigenbasis_eigh(
             kronecker_factor_list,
         )
     else:
-        # Use QR decomposition and power iteration (orthogonal iteration)
-        updated_eigenbasis_list, exp_avg_sq = soap_utils.get_eigenbasis_qr(
+        # Use QR decomposition and power iteration (orthogonal iteration) starting from the
+        # pre-sorted eigenbases.
+        updated_eigenbasis_list = soap_utils.get_eigenbasis_qr(
             kronecker_factor_list,
             eigenbasis_list,
-            exp_avg_sq,
             power_iter_steps,
         )
 

emerging_optimizers/soap/soap_utils.py

@@ -25,9 +25,41 @@
 __all__ = [
     "get_eigenbasis_eigh",
     "get_eigenbasis_qr",
+    "sort_eigenbasis_by_approx_eigvals",
 ]
 
 
+def sort_eigenbasis_by_approx_eigvals(
+    kronecker_factor_list: TensorList,
+    eigenbasis_list: TensorList,
+    exp_avg_sq: torch.Tensor,
+) -> tuple[TensorList, torch.Tensor]:
+    """Permute each eigenbasis and matching ``exp_avg_sq`` axis by descending approximate eigenvalues.
+
+    Both the eigh and QR eigenbasis-update paths consume the sorted output: the eigh path discards
+    the permuted eigenbasis but uses the permuted ``exp_avg_sq`` (the sort_idx best approximates the
+    permutation component of the eigh-vs-old basis change under small drift); the QR path power-
+    iterates from the pre-sorted basis.
+
+    Args:
+        kronecker_factor_list: List of preconditioner matrices (L and R).
+        eigenbasis_list: List of current eigenbases (QL and QR).
+        exp_avg_sq: Inner Adam second moment tensor permuted along each Kronecker-factor
+            axis to match the new descending-eigenvalue column ordering.
+
+    Returns:
+        ``(sorted_eigenbasis_list, sorted_exp_avg_sq)``.
+    """
+    sorted_eigenbasis_list: TensorList = []
+    sorted_exp_avg_sq = exp_avg_sq
+    for ind, (kronecker_factor, eigenbasis) in enumerate(zip(kronecker_factor_list, eigenbasis_list, strict=True)):
+        approx_eigvals = eig_utils.conjugate(kronecker_factor, eigenbasis, diag=True)
+        sort_idx = torch.argsort(approx_eigvals, descending=True)
+        sorted_eigenbasis_list.append(eigenbasis[:, sort_idx])
+        sorted_exp_avg_sq = sorted_exp_avg_sq.index_select(ind, sort_idx)
+    return sorted_eigenbasis_list, sorted_exp_avg_sq
+
+
 def get_eigenbasis_eigh(
     kronecker_factor_list: TensorList,
 ) -> TensorList:
@@ -64,23 +96,22 @@ def get_eigenbasis_eigh(
 def get_eigenbasis_qr(
     kronecker_factor_list: TensorList,
     eigenbasis_list: TensorList,
-    exp_avg_sq: torch.Tensor,
     power_iter_steps: int = 1,
-) -> tuple[TensorList, torch.Tensor]:
+) -> TensorList:
     """Updates the eigenbases of the preconditioner using power iteration and QR.
 
     Computes using multiple rounds of power iteration followed by QR decomposition (orthogonal iteration).
+    ``eigenbasis_list`` is expected to be already sorted by descending approximate eigenvalues (see
+    :func:`sort_eigenbasis_by_approx_eigvals`).
 
     Args:
-        kronecker_factor_list: List containing preconditioner (:math:`GG^T` and :math:`G^TG`)
-        eigenbasis_list: List containing eigenbases (:math:`Q_L` and :math:`Q_R`)
-        exp_avg_sq: inner adam second moment (exp_avg_sq).
+        kronecker_factor_list: List of preconditioner matrices (L and R).
+        eigenbasis_list: List of current eigenbases (QL and QR).
         power_iter_steps: Number of power iteration steps to perform before QR decomposition.
             More steps can lead to better convergence but increased computation time.
 
     Returns:
-        Tuple of updated list of orthonormal kronecker factor eigenbases matrices and updated (sorted) inner
-            Adam's second moment.
+        Updated list of orthonormal eigenbases (QL and QR).
 
     Example:
         .. code-block:: python
@@ -102,27 +133,19 @@ def get_eigenbasis_qr(
             perturbed_kronecker_factor_list[0] = k_factor1 + perturbation@perturbation.T
             perturbed_kronecker_factor_list[1] = k_factor2 + perturbation.T@perturbation
 
-            # Initialize exp_avg_sq tensor
-            exp_avg_sq = torch.randn(n, m).abs()
-
-            # Refine the orthogonal matrices using QR
-            updated_ortho_matrices, updated_exp_avg_sq = get_eigenbasis_qr(
+            # Refine the orthogonal matrices using QR (eigenbasis_list already sorted)
+            updated_ortho_matrices = get_eigenbasis_qr(
                 perturbed_kronecker_factor_list,
                 eigenbasis_list,
-                exp_avg_sq
             )
     """
     updated_eigenbasis_list: TensorList = []
-    for ind, (kronecker_factor, eigenbasis) in enumerate(zip(kronecker_factor_list, eigenbasis_list, strict=True)):
-        approx_eigvals = eig_utils.conjugate(kronecker_factor, eigenbasis, diag=True)
-        Q, exp_avg_sq = eig_utils.orthogonal_iteration(
-            approx_eigvals=approx_eigvals,
+    for kronecker_factor, eigenbasis in zip(kronecker_factor_list, eigenbasis_list, strict=True):
+        Q = eig_utils.orthogonal_iteration(
             kronecker_factor=kronecker_factor,
             eigenbasis=eigenbasis,
-            ind=ind,
-            exp_avg_sq=exp_avg_sq,
             power_iter_steps=power_iter_steps,
         )
         updated_eigenbasis_list.append(Q)
 
-    return updated_eigenbasis_list, exp_avg_sq
+    return updated_eigenbasis_list

emerging_optimizers/utils/eig.py

@@ -69,51 +69,35 @@ def eigh_with_fallback(
 
 
 def orthogonal_iteration(
-    approx_eigvals: torch.Tensor,
-    kronecker_factor: torch.Tensor,
-    eigenbasis: torch.Tensor,
-    ind: int,
-    exp_avg_sq: torch.Tensor,
+    kronecker_factor: Tensor,
+    eigenbasis: Tensor,
     power_iter_steps: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Computes the eigenbases of the preconditioner using power iteration and QR decomposition.
+) -> Tensor:
+    """Refines an eigenbasis via power iteration with QR re-orthogonalization.
 
-    This function performs multiple rounds of power iteration followed by QR decomposition
-    to recompute the eigenbases of the preconditioner kronecker factor. Generalizes Vyas et al.'s (SOAP) algorithm of 1 step of power iteration for updating the eigenbasis.
+    Performs ``power_iter_steps`` rounds of ``Q = QR(kronecker_factor @ Q)`` starting from
+    ``eigenbasis``. The columns of ``eigenbasis`` are expected to already be aligned with the
+    intended descending-eigenvalue ordering of ``kronecker_factor`` (see
+    :func:`emerging_optimizers.soap.soap_utils.sort_eigenbasis_by_approx_eigvals`).
 
     Args:
-        approx_eigvals : Projection of kronecker factor onto the eigenbasis, should be close to diagonal
-        kronecker_factor : Kronecker factor matrix.
-        eigenbasis : Kronecker factor eigenbasis matrix.
-        ind : Index for selecting dimension in the exp_avg_sq matrix to apply the sorting order over.
-        exp_avg_sq : inner Adam second moment (exp_avg_sq).
-        power_iter_steps: Number of power iteration steps to perform before QR decomposition.
-            More steps can lead to better convergence but increased computation time.
+        kronecker_factor: Kronecker factor matrix (symmetric, used as the projector).
+        eigenbasis: Starting eigenbasis whose columns will be refined.
+        power_iter_steps: Number of power-iteration / QR rounds to perform.
 
     Returns:
-        tuple[torch.Tensor, torch.Tensor]: A tuple containing:
-            - Q: The updated eigenbasis
-            - exp_avg_sq: The updated (sorted) inner Adam second moment
+        The refined eigenbasis.
     """
-    # Sort the approximated eigenvalues according to their magnitudes
-    sort_idx = torch.argsort(approx_eigvals, descending=True)
-    # re-order the inner adam second moment
-    exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
-
-    # Initialize power iteration after sorting the columns of the eigenbasis matrix according to the descending eigenvalues
-    Q = eigenbasis[:, sort_idx]
-
-    # Perform multiple steps of power iteration
+    Q = eigenbasis
     for _ in range(power_iter_steps):
         # Project current eigenbases on kronecker factor
         Q = kronecker_factor @ Q
         # Perform QR to maintain orthogonality between iterations
         Q = torch.linalg.qr(Q).Q
-
-    return Q, exp_avg_sq
+    return Q
 
 
-def conjugate(a: torch.Tensor, p: torch.Tensor, diag: bool = False) -> torch.Tensor:
+def conjugate(a: Tensor, p: Tensor, diag: bool = False) -> 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

tests/test_eig_utils.py

@@ -62,27 +62,14 @@ def test_update_eigenbasis_with_QR(self, N: int, power_iter_steps: int) -> None:
         eigenbasis = torch.randn(N, N, device=self.device)
         eigenbasis = torch.linalg.qr(eigenbasis).Q
 
-        # Create inner adam second moment (should be positive)
-        exp_avg_sq = torch.abs(torch.randn(N, N, device=self.device))
-
-        # 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
-        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 = eig_utils.orthogonal_iteration(
-            approx_eigvals=approx_eigvals,
+        Q_new = eig_utils.orthogonal_iteration(
             kronecker_factor=kronecker_factor,
             eigenbasis=eigenbasis,
-            ind=0,  # Test with first dimension
-            exp_avg_sq=exp_avg_sq,
             power_iter_steps=power_iter_steps,
         )
 
-        # Test 1: Check output shapes
+        # Test 1: Check output shape
         self.assertEqual(Q_new.shape, (N, N))
-        self.assertEqual(exp_avg_sq_new.shape, exp_avg_sq.shape)
 
         # Test 2: Check orthogonality (Q^T Q ≈ I)
         expected_identity = torch.eye(N, dtype=Q_new.dtype, device=self.device)
@@ -94,20 +81,7 @@ def test_update_eigenbasis_with_QR(self, N: int, power_iter_steps: int) -> None:
             msg="Orthogonalization failed: Q^T Q is not close enough to the identity matrix.",
         )
 
-        # 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(approx_eigvals, descending=True)
-        expected_exp_avg_sq = exp_avg_sq.index_select(0, sort_idx)
-        torch.testing.assert_close(
-            exp_avg_sq_new,
-            expected_exp_avg_sq,
-            atol=1e-5,
-            rtol=1e-5,
-            msg="exp_avg_sq was not properly sorted based on eigenvalues.",
-        )
-
-        # Test 4: Check that Q_new is different from input (transformation occurred)
-        # This is a basic check - in practice they should be different due to power iteration
+        # Test 3: Check that Q_new is different from input (power iteration ran)
         self.assertFalse(torch.allclose(Q_new, eigenbasis))
 
     def test_eigh_with_fallback_descending_order(self) -> None:

tests/test_soap_utils.py

@@ -55,20 +55,15 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
         L = g.mm(g.t()).float()
         R = g.t().mm(g).float()
         # Fake preconditioner list and orth list in the state
-        state = {
-            "GG": [L, R],  # precondition matrix
-            "Q": [
-                torch.randn(M, M, device=self.device),
-                torch.randn(N, N, device=self.device),
-            ],  # an existing Q (we'll refine it using QR)
-            "exp_avg_sq": torch.abs(torch.randn(M, N, device=self.device)),  # Some arbitrary tensor for example
-        }
-
-        # We'll call 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"],
+        kronecker_factor_list = [L, R]
+        eigenbasis_list = [
+            torch.randn(M, M, device=self.device),
+            torch.randn(N, N, device=self.device),
+        ]
+
+        Q_new_list = soap_utils.get_eigenbasis_qr(
+            kronecker_factor_list=kronecker_factor_list,
+            eigenbasis_list=eigenbasis_list,
             power_iter_steps=1,
         )
 
@@ -99,8 +94,58 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
             msg="Orthogonalization failed: Q^T Q is not close enough to the identity matrix.",
         )
 
-        # Also check that "exp_avg_sq" remains in the state with same shape if not merging
-        self.assertEqual(exp_avg_sq_new.shape, (M, N))
+    @parameterized.parameters(  # type: ignore[misc]
+        {"N": 4, "M": 8},
+        {"N": 16, "M": 8},
+    )
+    def test_sort_eigenbasis_by_approx_eigvals(self, N: int, M: int) -> None:
+        """Sort function permutes eigenbasis columns and exp_avg_sq slots consistently."""
+        g = torch.randint(-5, 6, (M, N), device=self.device) / 16.0
+        # Add a small ridge so K is full-rank (g @ g.t() is rank-deficient when M > N, etc.),
+        # which keeps the approximate eigenvalues away from numerical zero where ordering becomes
+        # ambiguous under float rounding.
+        eps_eye = lambda n: 1e-3 * torch.eye(n, device=self.device)
+        kronecker_factor_list = [g @ g.t() + eps_eye(M), g.t() @ g + eps_eye(N)]
+        eigenbasis_list = [torch.linalg.qr(Q).Q for Q in kronecker_factor_list]
+        exp_avg_sq = torch.abs(torch.randint(-5, 6, (M, N), device=self.device) / 16.0)
+
+        sorted_eigenbasis_list, sorted_exp_avg_sq = soap_utils.sort_eigenbasis_by_approx_eigvals(
+            kronecker_factor_list,
+            eigenbasis_list,
+            exp_avg_sq,
+        )
+
+        # Compute the expected per-axis permutations from the originals.
+        sort_idx_list = []
+        for K, Q in zip(kronecker_factor_list, eigenbasis_list, strict=True):
+            sort_idx_list.append(torch.argsort(torch.diag(Q.t() @ K @ Q), descending=True))
+
+        # Each eigenbasis is column-permuted by its own sort_idx.
+        for i, (Q_old, Q_sorted) in enumerate(zip(eigenbasis_list, sorted_eigenbasis_list, strict=True)):
+            torch.testing.assert_close(
+                Q_sorted,
+                Q_old[:, sort_idx_list[i]],
+                msg=lambda m, i=i: f"eigenbasis i={i} not permuted by sort_idx\n\n{m}",
+                atol=0,
+                rtol=0,
+            )
+
+        # exp_avg_sq is permuted along every axis cumulatively.
+        expected_sq = exp_avg_sq
+        for i, sort_idx in enumerate(sort_idx_list):
+            expected_sq = expected_sq.index_select(i, sort_idx)
+        torch.testing.assert_close(
+            sorted_exp_avg_sq,
+            expected_sq,
+            msg=lambda m: f"exp_avg_sq not permuted to match sorted eigenbases\n\n{m}",
+            atol=0,
+            rtol=0,
+        )
+
+        # Sorted eigenbases yield descending approximate eigenvalues.
+        for K, Q in zip(kronecker_factor_list, sorted_eigenbasis_list, strict=True):
+            sorted_eigvals = torch.diag(Q.t() @ K @ Q)
+            self.assertTrue(torch.all(sorted_eigvals[:-1] >= sorted_eigvals[1:]))
 
     @parameterized.parameters(  # type: ignore[misc]
         {"dims": [128, 512]},