Apollo

APOLLO optimizer with structured scaling from projected moments.

APOLLO (Zhu et al. 2025, arXiv:2502.xxxxx) projects 2D gradients into a low-rank subspace via Gaussian random matrices, computes Adam-style first and second moments in the compressed space, and then scales the original gradient element-wise or channel-wise using the ratio of the normalized low-rank moment to the low-rank gradient norm. This provides structured adaptive scaling with sublinear memory overhead.

Reference:

Zhu, Z., Li, Y., Wang, Z., & Anandkumar, A. (2025). APOLLO: SGD-like Memory, Adam-like Performance. arXiv:2502.xxxxx.

class src.model.optimizer.apollo.Apollo(*args: Any, **kwargs: Any)[source]

Bases: Optimizer

APOLLO optimizer with structured scaling from projected moments.

Projects 2D gradients into a low-rank subspace via Gaussian random projection matrices, maintains Adam-style first and second moments in the compressed space, and scales the original gradient using the ratio of the normalized low-rank moment to the low-rank gradient norm. Supports channel-wise and tensor-wise scaling, optional norm-growth limiting, and configurable projection side.

Reference:

Zhu, Z., Li, Y., Wang, Z., & Anandkumar, A. (2025). APOLLO: SGD-like Memory, Adam-like Performance. arXiv:2502.xxxxx.

Parameters:

  • params – Iterable of parameters to optimize or dicts defining parameter groups.

  • lr – Learning rate (default: 1e-3).

  • rank – Rank of the Gaussian random projection subspace (default: 128).

  • update_proj_gap – Number of steps between projector resampling (default: 200).

  • scale – Global scale multiplier applied to the structured update (default: 1.0).

  • scale_type – Scaling granularity, either "channel" (per-row or per-column) or "tensor" (single scalar per tensor) (default: "channel").

  • proj_type – Projection side strategy. "std" projects from the larger dimension, "reverse_std" from the smaller, "left" and "right" force a specific side (default: "std").

  • betas – Coefficients for first and second moment running averages (default: (0.9, 0.999)).

  • eps – Term added for numerical stability (default: 1e-8).

  • weight_decay – Decoupled weight decay coefficient (default: 0.0).

  • correct_bias – Whether to apply Adam-style bias correction (default: True).

  • scale_front – Whether to apply the global scale multiplier before norm-growth limiting instead of after (default: False).

  • disable_nl – Whether to disable the Fira-style norm-growth limiter (default: False).

defaults

Default hyper-parameter values for each parameter group.

Type:

dict

param_groups

Parameter groups tracked by the optimizer.

Type:

list

state

Per-parameter optimizer state (step counter, projector matrix, projector seed, low-rank moment buffers, scaled gradient norm).

Type:

dict

__init__(params, lr=0.001, rank=128, update_proj_gap=200, scale=1.0, scale_type='channel', proj_type='std', betas=(0.9, 0.999), eps=1e-08, weight_decay=0.0, correct_bias=True, scale_front=False, disable_nl=False)[source]

Initializes the APOLLO optimizer.

Parameters:

  • params – Iterable of parameters to optimize or dicts defining parameter groups.

  • lr – Learning rate (default: 1e-3).

  • rank – Projection rank (default: 128).

  • update_proj_gap – Projector resampling interval (default: 200).

  • scale – Global scale multiplier (default: 1.0).

  • scale_type – "channel" or "tensor" (default: "channel").

  • proj_type – "std", "reverse_std", "left", or "right" (default: "std").

  • betas – Momentum decay coefficients (default: (0.9, 0.999)).

  • eps – Numerical stability term (default: 1e-8).

  • weight_decay – Decoupled weight decay coefficient (default: 0.0).

  • correct_bias – Apply bias correction (default: True).

  • scale_front – Apply scale before norm limiter (default: False).

  • disable_nl – Disable norm-growth limiter (default: False).

Raises:

ValueError – If scale_type is not "channel" or "tensor", or proj_type is not one of the valid options.

step(closure=None)

Performs a single optimization step.

Parameters:

closure – A closure that re-evaluates the model and returns the loss. Optional for most use cases.

Returns:

The loss value if closure is provided, otherwise None.

Raises:

RuntimeError – If any parameter has sparse gradients.