emerging_optimizers.riemannian_optimizers

class emerging_optimizers.riemannian_optimizers.normalized_optimizer.ObliqueAdam(

params,

lr=0.001,

betas=(0.9, 0.99),

weight_decay=0.0,

dim=0,

eps=1e-08,

correct_bias=True,

)

Adam optimizer for row- or column-normalized 2D parameters on oblique manifolds.

This optimizer adapts an Adam-like algorithm to work on oblique manifolds, where parameters are constrained to have unit-norm rows or columns. It combines adaptive momentum estimation with Riemannian gradient computation and manifold retraction.

Parameters:

• params (list[Parameter])

• lr (float)

• betas (tuple[float, float])

• weight_decay (float)

• dim (int)

• eps (float)

• correct_bias (bool)

step(closure=None)

Performs a single optimization step.

Parameters:

closure (Callable[[], float] | None) – A closure that reevaluates the model and returns the loss.

Return type:

float | None

class emerging_optimizers.riemannian_optimizers.normalized_optimizer.ObliqueSGD(

params,

lr=0.001,

momentum=0.9,

weight_decay=0.0,

dim=0,

eps=1e-08,

)

SGD optimizer for row- or column-normalized 2D parameters on oblique manifolds.

This optimizer performs SGD on oblique manifolds, where parameters are constrained to have unit-norm rows or columns. It implements Riemannian SGD with manifold-aware gradient updates and retraction operations.

References

• An Introduction to Optimization on Smooth Manifolds (Nicolas Boumal)

• EDM2: https://arxiv.org/abs/2312.02696

• Jianlin Su: https://kexue.fm/archives/11196

• Raman et al.: https://arxiv.org/abs/1909.06463

• Franz Cesista: https://leloykun.github.io/ponder/steepest-descent-stiefel/#6-bonus-a-muon-like-optimizer-for-the-embedding-and-unembedding-layers

Parameters:

• lr (float) – learning rate

• momentum (float) – momentum coefficient

• weight_decay (float) – weight decay coefficient

• dim (int) – The dimension to normalize over

• eps (float) – epsilon for numerical stability

• params (list[Parameter])

step(closure=None)

Performs a single optimization step.

Parameters:

closure (Callable[[], float] | None) – A closure that reevaluates the model and returns the loss.

Return type:

float | None