[Hiroki11x]

What

Keep SOAP-family diagonal second moments aligned with eigenbasis columns when QR-based orthogonal iteration sorts those columns.

This restores the required permutation in:

• legacy SOAP
• MOSO (left- and right-preconditioned paths)
• KL-SOAP v3
• the SOAP reference used by parity tests

The existing 2-tuple return from orthogonal_iteration and get_eigenbasis_qr remains the default. Callers that own coordinate-wise state can opt in to receive the sort indices.

Why this is a correctness bug

Let the post-QR bases be sorted as

Q_L+ = Q_L P_L
Q_R+ = Q_R P_R

for permutation matrices P_L and P_R. A tensor represented in that basis changes coordinates as

G'+ = P_L^T G' P_R

so SOAP's coordinate-wise second moment must follow exactly:

V+ = P_L^T V P_R

The current implementation sorts Q but leaves V = exp_avg_sq in the old slot order. This makes the physical Adam update depend on a pure basis-label permutation.

The original SOAP implementation applies the same sort_idx to the basis and exp_avg_sq. This regression was introduced when #274 moved eigenvalue sorting after QR and removed the previous joint permutation helper. Sorting after QR is retained here; the minimal fix is to apply that actual post-QR permutation to exp_avg_sq too.

Deterministic oracle

A new rectangular 2x3 test forces different permutations on the two axes ([1, 0] and [1, 2, 0]). It verifies both the exact second-moment permutation and that a pure relabeling leaves the represented physical update unchanged. Separate tests cover the low-level sort indices, both MOSO axes, and KL-SOAP v3.

For the simple 2x2 state M=[[1,2],[3,4]], V=M^2, with both axes swapped, the correct physical update remains all ones. Leaving V stale instead gives [[0.25, 0.6667], [1.5, 4.0]].

Validation

• CPU: 119 related tests passed; 3 GPU-only suites skipped locally
• patched SOAP vs updated reference, 5 CPU steps: maximum parameter difference 4.77e-7
• mypy emerging_optimizers: 63 source files passed
• Ruff lint and format checks passed
• git diff --check passed
• independent reviewer checked axis direction, shapes, devices/dtypes, state rebinding, and runtime/type API compatibility

The GPU reference and multistream paths are left to CI.

Scope

This PR intentionally does not change Emerging-Optimizers' current-gradient inclusion or per-step QR behavior. Those differ from the ICLR 2025 algorithm but are explicit design choices in arXiv:2607.20548, rather than this permutation bug.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the new permutations consistently preserve moment-to-basis alignment across the affected optimizer paths.

Summary

• Adds opt-in sorting-index returns to the shared orthogonal-iteration helpers while preserving their default two-value result.
• Applies the returned permutations to legacy SOAP, both MOSO orientations, KL-SOAP v3, and the parity-test reference.
• Adds deterministic tests covering asymmetric left/right permutations and invariance of the represented physical update.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    F[Kronecker factors] --> QR[QR-based orthogonal iteration]
    QR --> S[Sort eigenbasis columns]
    S --> I[Return axis permutations]
    I --> B[Store sorted eigenbases]
    I --> V[Permute second-moment axes]
    B --> A[Aligned adaptive coordinates]
    V --> A
    A --> U[Project update to parameter space]

Loading

_{Reviews (1) · Last reviewed commit: "fix: keep SOAP second moments aligned wi..."}

[skyw]

Thanks for the contribution, I should probably update the document. We are still in the process of refactoring.

This won't be merged and not necessary:

• Things won't change after moving to legacy, otherwise it creates other problems.
• In new implementation, eigen basis is always sorted so nothing need to realign.

[Hiroki11x]

@skyw
Thanks for the clarification. I understand that the legacy implementation should remain unchanged.

To clarify, I tested KlSoapPreconditioner in emerging_optimizers/shampoo/soap_v3.py on the current upstream main at commit 36f70336da89dc80481c07dfb0af2c7333b9e5b3. This is the preconditioner used by KlSoapV3, registered as kl_soap, rather than the legacy optimizer registered as soap.

I may be misunderstanding the intended state semantics, but I do not think that keeping the eigenbasis sorted is sufficient by itself. The basis may be sorted for $K_{t-1}$, while the eigenvalue ordering can change after updating the factor to $K_t$. When orthogonal_iteration() applies a non-identity sort_idx, the columns of $Q$ are reordered, but exp_avg_sq remains in its previous slot order. If exp_avg_sq is intended to follow the eigenvector directions, it should receive the corresponding permutation:

$$ V_{\mathrm{new}} = P_L^\top V_{\mathrm{old}} P_R $$

I reproduced this directly with the current KlSoapPreconditioner using a deterministic $2 \times 3$ example. Reordering the basis without reordering exp_avg_sq changed the represented normalized update, whereas applying the same permutations preserved it exactly. In a small nonstationary stress test, non-identity sorting also occurred in a minority of runs, although it was rare or absent in stationary cases.

Could you clarify whether v3 intentionally defines exp_avg_sq by sorted eigenvalue rank, rather than having it follow the tracked eigenvector direction? If rank-based state is intentional, then no realignment is needed, but that appears to be different state semantics from the original SOAP formulation and would be useful to document.

[skyw]

@skyw Thanks for the clarification. I understand that the legacy implementation should remain unchanged.

To clarify, I tested KlSoapPreconditioner in emerging_optimizers/shampoo/soap_v3.py on the current upstream main at commit 36f70336da89dc80481c07dfb0af2c7333b9e5b3. This is the preconditioner used by KlSoapV3, registered as kl_soap, rather than the legacy optimizer registered as soap.

I may be misunderstanding the intended state semantics, but I do not think that keeping the eigenbasis sorted is sufficient by itself. The basis may be sorted for K t − 1 , while the eigenvalue ordering can change after updating the factor to K t . When orthogonal_iteration() applies a non-identity sort_idx, the columns of Q are reordered, but exp_avg_sq remains in its previous slot order. If exp_avg_sq is intended to follow the eigenvector directions, it should receive the corresponding permutation:

V new = P L ⊤ V old P R

I reproduced this directly with the current KlSoapPreconditioner using a deterministic 2 × 3 example. Reordering the basis without reordering exp_avg_sq changed the represented normalized update, whereas applying the same permutations preserved it exactly. In a small nonstationary stress test, non-identity sorting also occurred in a minority of runs, although it was rare or absent in stationary cases.

Could you clarify whether v3 intentionally defines exp_avg_sq by sorted eigenvalue rank, rather than having it follow the tracked eigenvector direction? If rank-based state is intentional, then no realignment is needed, but that appears to be different state semantics from the original SOAP formulation and would be useful to document.

exp_avg_sq tracks element square of projected G, here we always project G with sorted eigen bases and nothing needs further align. Imagine if every step it is updated with eigh which returns sorted eigen values. The math in the original paper only checks out with fixed eigen bases and no EMA, the rest was not so clear.

anyway, let me check your code again and compare what it does exactly. Could be I misunderstood it.

[Hiroki11x]

@skyw

Thanks for clarifying. I now understand that exp_avg_sq tracks the history of each sorted position, rather than following a particular basis direction when its rank changes. Under that interpretation, the current implementation is consistent.

My earlier test assumed that history should follow the directions. It demonstrates a difference between the two policies, but does not establish a correctness bug. I rechecked both implementations, and each matches its respective update rule.

The PR is therefore better viewed as an alternative state-update policy. Documenting the sorted-position semantics would help avoid this ambiguity. Thanks for taking another look.

[skyw]

@skyw

Thanks for clarifying. I now understand that exp_avg_sq tracks the history of each sorted position, rather than following a particular basis direction when its rank changes. Under that interpretation, the current implementation is consistent.

My earlier test assumed that history should follow the directions. It demonstrates a difference between the two policies, but does not establish a correctness bug. I rechecked both implementations, and each matches its respective update rule.

The PR is therefore better viewed as an alternative state-update policy. Documenting the sorted-position semantics would help avoid this ambiguity. Thanks for taking another look.

I check the code and math again. Although there is no proof, I still believe the algorithm only works if eigen bases don't change much. The math in the paper needs fixed eigen bases to work. EMA and rest of it were completely ignored.

It is not very obvious, we have a Adam warm up period to skip the frequent eigen bases change early in the training. And if we used eigh all alone, there won't be anything to realign because no such information existed.

The main question I think is whether power iteration approximated eigen bases represent some soft of permutation of previous eigen bases when they change by large amount. I can see argument goes both ways although not very strong.