emerging_optimizers.legacy_soap#

SOAP#

class emerging_optimizers.legacy_soap.soap.SOAP( params, lr, betas=(0.9, 0.95), shampoo_beta=0.95, eps=1e-08, weight_decay=0.01, *, weight_decay_method='decoupled', nesterov=False, correct_bias=True, fp32_matmul_prec='highest', use_eigh=False, qr_fp32_matmul_prec='high', power_iter_steps=1, max_update_rms=0.0, use_kl_shampoo=False, correct_shampoo_beta_bias=None, stream_list=None, )[source]#

Implements a variant of SOAP (ShampoO with Adam in the Preconditioner eigenbasis) algorithm.

SOAP (https://arxiv.org/abs/2409.11321) is a preconditioned optimizer that combines the benefits of Shampoo’s non-diagonal preconditioning with Adam’s adaptive learning rates. It uses gradient correlation matrix eigenbasis-based preconditioning to adapt to the local geometry of the optimization landscape.

Parameters:

params (Iterable[Tensor] | Iterable[dict[str, Any]] | Iterable[tuple[str, Tensor]]) – Iterable of parameters to optimize or dicts defining parameter groups
lr (float) – The learning rate to use
betas (tuple[float, float]) – Inner Adam’s betas parameters (b1, b2)
shampoo_beta (float) – Beta for the kronecker factor matrices (L and R in paper) moving average instead of betas[1] if >= 0
eps (float) – Inner Adam’s epsilon for numerical stability
weight_decay (float) – Weight decay coefficient
weight_decay_method (Literal['decoupled', 'independent', 'l2', 'palm']) – Method to apply weight decay, see WeightDecayMixin for more details.
nesterov (bool) – uses Nesterov momentum in Adam (https://cs229.stanford.edu/proj2015/054_report.pdf)
correct_bias (bool) – Whether to use bias correction in Inner Adam and Kronecker factor matrices EMA
fp32_matmul_prec (Literal['highest', 'high', 'medium']) – Precision of the matmul operations in optimizer states GEMM operations
use_eigh (bool) – Whether to use full symmetric eigendecomposition (eigh) to compute the eigenbasis. If False, use orthogonal iteration to compute the eigenbasis.
qr_fp32_matmul_prec (Literal['highest', 'high', 'medium']) – Precision of the matmul operations in QR decomposition.
power_iter_steps (int) – Number of power iteration steps to perform before QR decomposition. More steps can lead to better convergence but increased computation time.
max_update_rms (float) – Clip the update RMS to this value (0 means no clipping).
use_kl_shampoo (bool) – Whether to use KL-Shampoo correction.
correct_shampoo_beta_bias (bool | None) – Whether to correct shampoo beta bias. Decoupled it from correct_bias for testability because reference implementation of Soap doesn’t bias correct shampoo beta.
stream_list (list[Stream] | None) – Optional list of CUDA streams. When provided, each parameter in the inner loop uses a stream from this list in round-robin fashion.

step(closure: None = None) → None[source]#

step(closure: Callable[[], float]) → float

Performs a single optimization step.

Parameters:: closure – Unsupported; must be None.

emerging_optimizers.legacy_soap.soap.project_in(x, eigenbasis_list)[source]#

Projects a tensor into the eigenbases

Note

For 2D tensors, we can use matmul instead of tensordot for code legibility. However, the code has been using tensordot historically, so does the reference implementation. It is difficult to match matmul and tensordot outputs exactly because of underlying floating point arithmetic differences. Therefore, we decided to keep using tensordot for consistency.

Parameters:

x (Tensor) – Input tensor to project into the eigenbasis.
eigenbasis_list (list[Tensor]) – List of eigenbases for preconditioning.

Return type:

Tensor

emerging_optimizers.legacy_soap.soap.project_out(x, eigenbasis_list)[source]#

Projects a tensor out of the eigenbases

Note

Uses tensordot rather than matmul for the same numerical-consistency reason described in project_in().

Parameters:

x (Tensor) – Input tensor to project back to the original space.
eigenbasis_list (list[Tensor]) – List of eigenbases for preconditioning.

Return type:

Tensor

emerging_optimizers.legacy_soap.soap.update_kronecker_factors(kronecker_factor_list, grad, shampoo_beta)[source]#

Updates the preconditioner matrices using gradient outer products.

This function updates the Kronecker factor matrices (L and R) used for preconditioning by computing and accumulating gradient outer products. kronecker_factor_list is updated in place.

Parameters:

kronecker_factor_list (list[Tensor]) – List of preconditioner matrices (L and R) to update. Each matrix should be square and match the corresponding dimension of grad.
grad (Tensor) – Gradient tensor of the parameter being optimized
shampoo_beta (float) – Momentum coefficient for updating preconditioners. Controls how much weight to give to new vs old gradient statistics.

Return type:

None

Example

>>> grad = torch.randn(10, 20)
>>> L = torch.zeros(10, 10)
>>> R = torch.zeros(20, 20)
>>> update_kronecker_factors([L, R], grad, shampoo_beta=0.95)

emerging_optimizers.legacy_soap.soap.update_kronecker_factors_kl_shampoo( kronecker_factor_list, grad, shampoo_beta, eigenbasis_list, eigvals_list, eps, eigval_exp=-1.0, )[source]#

Updates the kronecker factor matrices in place using KL-Shampoo correction.

Implements the Kullback–Leibler minimization update from https://arxiv.org/pdf/2509.03378.

For a gradient \(G \in \mathbb{R}^{m \times n}\), current kronecker factors \(L_t \in \mathbb{R}^{m \times m}\), \(R_t \in \mathbb{R}^{n \times n}\), their orthonormal eigenbases \(Q_L, Q_R\), and the approximate eigenvalues in those eigenbases

\[\Lambda_L = \mathrm{diag}(Q_L^{\top} L_t Q_L), \quad \Lambda_R = \mathrm{diag}(Q_R^{\top} R_t Q_R) \]

(passed as eigvals_list, typically computed and stored at the previous eigenbasis update, see get_eigenbasis_qr() and get_eigenbasis_eigh()), the EMA update with momentum \(\beta\) (= shampoo_beta) and exponent \(p\) (= eigval_exp, default -1) is

\[L_{t+1} = \beta\, L_t + \frac{1-\beta}{n}\, G\, Q_R\, \mathrm{diag}(\Lambda_R^{p})\, Q_R^{\top} G^{\top} \\ R_{t+1} = \beta\, R_t + \frac{1-\beta}{m}\, G^{\top}\, Q_L\, \mathrm{diag}(\Lambda_L^{p})\, Q_L^{\top} G \]

Eigenvalues are clamped to eps from below before exponentiation for numerical stability.

Parameters:

kronecker_factor_list (Iterable[Tensor]) – List of preconditioner matrices (L and R) to update.
grad (Tensor) – Gradient tensor of the parameter being optimized
shampoo_beta (float) – Momentum coefficient for updating preconditioners.
eigenbasis_list (Iterable[Tensor]) – List of orthonormal eigenbases of the kronecker factor matrices
eigvals_list (Iterable[Tensor]) – List of approximate eigenvalues of each kronecker factor in its eigenbasis.
eps (float) – Small offset for numerical stability.
eigval_exp (float) – Exponent applied to the (clamped) eigenvalues.

Return type:

None

emerging_optimizers.legacy_soap.soap.update_eigenbasis_and_exp_avgs( kronecker_factor_list, eigenbasis_list, exp_avg_sq, exp_avg, use_eigh=False, power_iter_steps=1, )[source]#

Updates the eigenbases and moving averages.

This function performs an update of the eigenbases (QL and QR) used for preconditioning. It follows these steps:

Projects exp_avg back to the original basis
Updates the eigenbases using QR decomposition and power iteration (orthogonal iteration)
Projects exp_avg back to the new eigenbasis

Parameters:

kronecker_factor_list (list[Tensor]) – List of preconditioner matrices (L and R) that define the optimization landscape. These are updated with gradient statistics.
eigenbasis_list (list[Tensor]) – List of current eigenbases (QL and QR) used for preconditioning. These will be updated by this function.
exp_avg_sq (Tensor) – Inner Adam’s second moment tensor, used for scaling the preconditioner updates. Permuted along each kronecker-factor axis on the QR path to track the sorted eigenbasis columns; returned unchanged on the eigh path.
exp_avg (Tensor) – Inner Adam’s first moment tensor, used for tracking gradient momentum.
use_eigh (bool) – Whether to use full symmetric eigendecomposition (eigh) to compute the eigenbasis. If False, use orthogonal iteration to compute the eigenbasis.
power_iter_steps (int) – Number of power iteration steps to perform before QR decomposition. More steps can lead to better convergence but increased computation time.

Returns:

List of (approximate) eigenvalues of each kronecker factor in its updated eigenbasis
Updated list of eigenbases (QL and QR)
Updated exp_avg tensor projected to the new eigenbasis
Updated exp_avg_sq tensor

Return type:

A tuple containing

Example

>>> L = torch.randn(10, 10)
>>> R = torch.randn(20, 20)
>>> QL = torch.randn(10, 10)
>>> QR = torch.randn(20, 20)
>>> exp_avg_sq = torch.randn(10, 20)
>>> exp_avg = torch.randn(10, 20)
>>> eigvals_list, updated_eigenbasis_list, updated_exp_avg, updated_exp_avg_sq = (
...     update_eigenbasis_and_exp_avgs([L, R], [QL, QR], exp_avg_sq, exp_avg))

REKLS#

class emerging_optimizers.legacy_soap.rekls.REKLS( params, lr, betas=(0.9, 0.95), shampoo_beta=0.95, eps=1e-08, weight_decay=0.01, *, weight_decay_method='decoupled', )[source]#

REKLS (Realtime Eigen Kullback-Leibler Soap) optimizer.

REKLS is a variant of SOAP that uses the up to date eigenbasis calculated by Eigen decomposition. It is “up to date” because current step’s gradient is accumulated to the kronecker factor before eigenbasis update.

Note

Refer to SOAP for detailed documentation of arguments.

Parameters:

params (Iterable[Tensor] | Iterable[dict[str, Any]] | Iterable[tuple[str, Tensor]])
lr (float)
betas (tuple[float, float])
shampoo_beta (float)
eps (float)
weight_decay (float)
weight_decay_method (Literal['decoupled', 'independent', 'l2', 'palm'])

emerging_optimizers.legacy_soap.soap_utils#

emerging_optimizers.legacy_soap.soap_utils.get_eigenbasis_eigh(kronecker_factor_list)[source]#

Computes the eigenvalues and eigenbases of the preconditioner using torch.linalg.eigh decomposition.

Parameters:: kronecker_factor_list (Iterable[Tensor]) – Matrix List to compute eigenbases of
Returns:: Tuple of (list of eigenvalues in descending order, list of orthonormal kronecker factor eigenbases matrices).
Return type:: tuple[list[Tensor], list[Tensor]]

emerging_optimizers.legacy_soap.soap_utils.get_eigenbasis_qr( kronecker_factor_list, eigenbasis_list, power_iter_steps=1, )[source]#

Updates the eigenbases of the preconditioner using power iteration and QR.

Parameters:

kronecker_factor_list (Iterable[Tensor]) – List of preconditioner matrices (L and R).
eigenbasis_list (Iterable[Tensor]) – List of current eigenbases (QL and QR).
power_iter_steps (int) – 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 in descending order, updated list of orthonormal eigenbases (QL and QR) with columns ordered to match).

Return type:

tuple[list[Tensor], list[Tensor]]

emerging_optimizers.legacy_soap.soap_utils.get_eigenbasis_svd(kronecker_factor_list)[source]#

Computes the eigenbases of the preconditioner using torch.linalg.svd decomposition.

The kronecker factors \(L = GG^\top\) and \(R = G^\top G\) are symmetric positive semi-definite, so the left and right singular vectors coincide (up to sign in the presence of repeated singular values); this function returns the left singular vectors \(U\) as the eigenbasis. Singular values from torch.linalg.svd are returned in descending order.

Parameters:: kronecker_factor_list (Iterable[Tensor]) – Matrix List to compute eigenbases of
Returns:: List of orthonormal kronecker factor eigenbases matrices
Return type:: list[Tensor]