update exp_avg_sq sorting

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

Author: skyw

emerging_optimizers/soap/moso.py

@@ -234,21 +234,19 @@ def _update_eigenbasis_and_adam_exp_avgs(
         left_preconditioned=left_preconditioned,
     )
 
-    eigenbasis, exp_avg_sq = _sort_one_sided_eigenbasis_and_exp_avg_sq(
-        momentum_factor=momentum_factor,
-        eigenbasis=eigenbasis,
-        exp_avg_sq=exp_avg_sq,
-        left_preconditioned=left_preconditioned,
-    )
-
     if use_eigh:
         _, (updated_eigenbasis,) = soap_utils.get_eigenbasis_eigh([momentum_factor])
     else:
-        _, (updated_eigenbasis,) = soap_utils.get_eigenbasis_qr(
+        # get_eigenbasis_qr permutes exp_avg_sq along the factor-list axes (dim 0 for the single
+        # factor here), so transpose the right-preconditioned case in and out.
+        x = exp_avg_sq if left_preconditioned else exp_avg_sq.mT
+        _, (updated_eigenbasis,), x = soap_utils.get_eigenbasis_qr(
             [momentum_factor],
             [eigenbasis],
+            x,
             power_iter_steps=power_iter_steps,
         )
+        exp_avg_sq = x if left_preconditioned else x.mT
 
     exp_avg = _project_to_one_sided_eigenbasis(
         x=exp_avg,
@@ -258,21 +256,6 @@ def _update_eigenbasis_and_adam_exp_avgs(
     return updated_eigenbasis, exp_avg, exp_avg_sq
 
 
-@torch.no_grad()  # type: ignore[misc]
-def _sort_one_sided_eigenbasis_and_exp_avg_sq(
-    momentum_factor: torch.Tensor,
-    eigenbasis: torch.Tensor,
-    exp_avg_sq: torch.Tensor,
-    left_preconditioned: bool,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Sort eigenbasis slots by approximate eigenvalue and permute Adam second moments."""
-    approx_eigvals = utils.eig.conjugate(momentum_factor, eigenbasis, diag=True)
-    sort_idx = torch.argsort(approx_eigvals, descending=True, stable=True)
-    sorted_eigenbasis = eigenbasis[:, sort_idx]
-    exp_avg_sq_dim = 0 if left_preconditioned else 1
-    return sorted_eigenbasis, exp_avg_sq.index_select(exp_avg_sq_dim, sort_idx)
-
-
 @torch.no_grad()  # type: ignore[misc]
 def _project_to_one_sided_eigenbasis(
     x: torch.Tensor,

emerging_optimizers/soap/soap.py

@@ -453,8 +453,10 @@ def update_eigenbasis_and_exp_avgs(
     used for preconditioning. It follows these steps:
 
     1. Projects exp_avg back to the original basis
-    2. Updates the eigenbases (via eigh, or QR decomposition and power iteration (orthogonal iteration)),
-       along with the (approximate) eigenvalues of the kronecker factors in the updated eigenbases
+    2. Updates the eigenbases (via eigh, or QR decomposition and power iteration (orthogonal iteration))
+       along with the (approximate) eigenvalues of the kronecker factors in the updated eigenbases;
+       on the QR path exp_avg_sq is permuted to match the sorted eigenbasis columns (not needed for
+       the eigh path, which rebuilds the basis from scratch)
     3. Projects exp_avg back to the new eigenbasis
 
     Args:
@@ -463,7 +465,8 @@ def update_eigenbasis_and_exp_avgs(
         eigenbasis_list: List of current eigenbases (QL and QR)
             used for preconditioning. These will be updated by this function.
         exp_avg_sq: Inner Adam's second moment tensor, used for scaling the preconditioner updates.
-            This tensor is modified in-place.
+            Permuted along each kronecker-factor axis on the QR path to track the sorted eigenbasis
+            columns; returned unchanged on the eigh path.
         exp_avg: Inner Adam's first moment tensor, used for tracking gradient momentum.
             This tensor is modified in-place.
         use_eigh: Whether to use full symmetric eigendecomposition (eigh) to compute the eigenbasis.
@@ -496,28 +499,16 @@ def update_eigenbasis_and_exp_avgs(
         dims=[[0], [1]],
     )
 
-    # 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 and compute the eigenvalues in the updated eigenbases
+    # Step 2: Update eigenbases
     if use_eigh:
         updated_eigvals_list, updated_eigenbasis_list = soap_utils.get_eigenbasis_eigh(
             kronecker_factor_list,
         )
     else:
-        # Use QR decomposition and power iteration (orthogonal iteration) starting from the
-        # pre-sorted eigenbases.
-        updated_eigvals_list, updated_eigenbasis_list = soap_utils.get_eigenbasis_qr(
+        updated_eigvals_list, updated_eigenbasis_list, exp_avg_sq = soap_utils.get_eigenbasis_qr(
             kronecker_factor_list,
             eigenbasis_list,
+            exp_avg_sq,
             power_iter_steps,
         )
 

emerging_optimizers/soap/soap_utils.py

@@ -26,42 +26,10 @@
 __all__ = [
     "get_eigenbasis_eigh",
     "get_eigenbasis_qr",
-    "sort_eigenbasis_by_approx_eigvals",
     "get_eigenbasis_svd",
 ]
 
 
-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,
 ) -> tuple[TensorList, TensorList]:
@@ -73,17 +41,6 @@ def get_eigenbasis_eigh(
     Returns:
         Tuple of (list of eigenvalues in descending order, list of orthonormal kronecker factor
         eigenbases matrices).
-
-    Example:
-        .. code-block:: python
-
-            # Create sample Kronecker factors (symmetric positive definite matrices)
-            k_factor1 = torch.randn(4, 4)
-            k_factor1 = k_factor1 @ k_factor1.T  # Make symmetric positive definite
-            k_factor2 = torch.randn(5, 5)
-            k_factor2 = k_factor2 @ k_factor2.T  # Make symmetric positive definite
-
-            eigvals_list, ortho_matrices = get_eigenbasis_eigh([k_factor1, k_factor2])
     """
     updated_eigenbasis_list: TensorList = []
     updated_eigvals_list: TensorList = []
@@ -111,19 +68,6 @@ def get_eigenbasis_svd(
 
     Returns:
         List of orthonormal kronecker factor eigenbases matrices
-
-    Example:
-        .. code-block:: python
-
-            # Create sample Kronecker factors (symmetric positive definite matrices)
-            k_factor1 = torch.randn(4, 4)
-            k_factor1 = k_factor1 @ k_factor1.T  # Make symmetric positive definite
-            k_factor2 = torch.randn(5, 5)
-            k_factor2 = k_factor2 @ k_factor2.T  # Make symmetric positive definite
-
-            # Get orthogonal matrices for these factors
-            ortho_matrices = get_eigenbasis_svd([k_factor1, k_factor2])
-            # ortho_matrices[0] has shape [4, 4] and ortho_matrices[1] has shape [5, 5]
     """
     updated_eigenbasis_list: TensorList = []
 
@@ -137,63 +81,44 @@ def get_eigenbasis_svd(
 def get_eigenbasis_qr(
     kronecker_factor_list: TensorList,
     eigenbasis_list: TensorList,
+    exp_avg_sq: torch.Tensor,
     power_iter_steps: int = 1,
-) -> tuple[TensorList, TensorList]:
+) -> tuple[TensorList, TensorList, torch.Tensor]:
     """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`).
-
-    The approximate eigenvalues of each kronecker factor in its updated eigenbasis (the Rayleigh
-    quotients :math:`\\mathrm{diag}(Q^{\\top} K Q)`) are also computed.
+    Each eigenbasis is refined with multiple rounds of power iteration followed by QR decomposition
+    (orthogonal iteration), then its columns are sorted by descending approximate eigenvalues of the
+    kronecker factor (the Rayleigh quotients :math:`\\mathrm{diag}(Q^{\\top} K Q)`), so the returned
+    eigenvalues and eigenbases are in descending order like :func:`get_eigenbasis_eigh`.
+    ``exp_avg_sq`` is permuted along each kronecker-factor axis by the same sort, so its per-slot
+    statistics stay aligned with the eigenbasis columns.
 
     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 sorted eigenbasis columns.
         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 (list of approximate eigenvalues of each kronecker factor in its updated eigenbasis,
-        updated list of orthonormal eigenbases (QL and QR)).
-
-    Example:
-        .. code-block:: python
-
-            # Create sample Kronecker factors (symmetric positive definite matrices)
-            n, m = 10, 20
-            k_factor1 = torch.randn(n, n)
-            k_factor1 = k_factor1 @ k_factor1.T  # Make symmetric positive definite
-            k_factor2 = torch.randn(m, m)
-            k_factor2 = k_factor2 @ k_factor2.T  # Make symmetric positive definite
-
-            # Get orthogonal matrices for these kronecker factors
-            kronecker_factor_list = [k_factor1, k_factor2]
-            _, eigenbasis_list = get_eigenbasis_eigh(kronecker_factor_list)
-
-            # Perturb the kronecker factor matrices, simulating the effect of gradient updates
-            perturbation = 1e-2*torch.randn(n, m)
-            perturbed_kronecker_factor_list = [None, None]
-            perturbed_kronecker_factor_list[0] = k_factor1 + perturbation@perturbation.T
-            perturbed_kronecker_factor_list[1] = k_factor2 + perturbation.T@perturbation
-
-            # Refine the orthogonal matrices using QR (eigenbasis_list already sorted)
-            eigvals_list, updated_ortho_matrices = get_eigenbasis_qr(
-                perturbed_kronecker_factor_list,
-                eigenbasis_list,
-            )
+        in descending order, updated list of orthonormal eigenbases (QL and QR) with columns ordered to
+        match, ``exp_avg_sq`` permuted to match).
     """
     updated_eigenbasis_list: TensorList = []
     updated_eigvals_list: TensorList = []
-    for kronecker_factor, eigenbasis in zip(kronecker_factor_list, eigenbasis_list, strict=True):
+    for ind, (kronecker_factor, eigenbasis) in enumerate(zip(kronecker_factor_list, eigenbasis_list, strict=True)):
         Q = eig_utils.orthogonal_iteration(
             kronecker_factor=kronecker_factor,
             eigenbasis=eigenbasis,
             power_iter_steps=power_iter_steps,
         )
-        updated_eigenbasis_list.append(Q)
         with utils.fp32_matmul_precision("highest"):
-            updated_eigvals_list.append(eig_utils.conjugate(kronecker_factor, Q, diag=True))
+            approx_eigvals = eig_utils.conjugate(kronecker_factor, Q, diag=True)
+        approx_eigvals, sort_idx = torch.sort(approx_eigvals, descending=True)
+        updated_eigvals_list.append(approx_eigvals)
+        updated_eigenbasis_list.append(Q[:, sort_idx])
+        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
 
-    return updated_eigvals_list, updated_eigenbasis_list
+    return updated_eigvals_list, updated_eigenbasis_list, exp_avg_sq

emerging_optimizers/utils/eig.py

@@ -76,9 +76,10 @@ def orthogonal_iteration(
     """Refines an eigenbasis via power iteration with QR re-orthogonalization.
 
     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`).
+    ``eigenbasis``. The columns of ``eigenbasis`` are expected to already be approximately aligned
+    with the descending-eigenvalue ordering of ``kronecker_factor`` (maintained in the SOAP pipeline
+    by :func:`emerging_optimizers.soap.soap_utils.get_eigenbasis_qr`, which returns its eigenbases in
+    descending order).
 
     Args:
         kronecker_factor: Kronecker factor matrix (symmetric, used as the projector).

tests/test_soap_utils.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
-from _comparison import assert_close_to_identity, assert_equal
+from _comparison import assert_close_to_identity
 from absl import flags, logging
 from absl.testing import absltest, parameterized
 
@@ -61,10 +61,12 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
             torch.randn(M, M, device=self.device),
             torch.randn(N, N, device=self.device),
         ]
+        exp_avg_sq = torch.abs(torch.randn(M, N, device=self.device))
 
-        eigvals_list, Q_new_list = soap_utils.get_eigenbasis_qr(
+        eigvals_list, Q_new_list, permuted_exp_avg_sq = soap_utils.get_eigenbasis_qr(
             kronecker_factor_list=kronecker_factor_list,
             eigenbasis_list=eigenbasis_list,
+            exp_avg_sq=exp_avg_sq,
             power_iter_steps=1,
         )
 
@@ -83,6 +85,8 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
         self.assertEqual(len(eigvals_list), 2)
         for eigvals, kronecker_factor, Q_new in zip(eigvals_list, kronecker_factor_list, Q_new_list, strict=True):
             self.assertEqual(eigvals.shape, (kronecker_factor.shape[0],))
+            # Eigenvalues (and matching eigenbasis columns) are returned in descending order
+            self.assertTrue(torch.all(eigvals[:-1] >= eigvals[1:]))
             torch.testing.assert_close(
                 eigvals,
                 torch.diag(Q_new.t() @ kronecker_factor @ Q_new),
@@ -91,54 +95,7 @@ def test_get_eigenbasis_qr(self, N: int, M: int) -> None:
                 msg=lambda msg: f"eigvals do not match the Rayleigh quotients diag(Q^T K Q)\n\n{msg}",
             )
 
-    @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)):
-            assert_equal(
-                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}",
-            )
-
-        # 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)
-        assert_equal(
-            sorted_exp_avg_sq,
-            expected_sq,
-            msg=lambda m: f"exp_avg_sq not permuted to match sorted eigenbases\n\n{m}",
-        )
-
-        # 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:]))
+        self.assertEqual(permuted_exp_avg_sq.shape, (M, N))
 
     @parameterized.parameters(  # type: ignore[misc]
         {"dims": [128, 512]},