neuralop.training.tensorgrad.TensorGRaD

class neuralop.training.tensorgrad.TensorGRaD(params: Iterable[Parameter], lr: float = 0.001, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-06, weight_decay: float = 0.0, correct_bias: bool = True)[source]

AdamW with optional low-rank and sparse gradient projection.

Parameters with a tensorgrad parameter-group flag are optimized in a compressed gradient space. Groups with only rank use a single TensorGRaD/Tucker low-rank branch. Groups that also set sparse_ratio use a TensorGRaD-style residual branch: low-rank project the gradient, sparse-project the low-rank residual, run Adam on both compressed states, then project both updates back and combine them.

Parameters:

paramsiterable

Iterable of parameters or parameter groups. Parameter groups with tensorgrad=True use compressed-gradient updates.

lrfloat, optional

Learning rate.

betastuple[float, float], optional

Adam exponential moving average coefficients.

epsfloat, optional

Epsilon added to the denominator for numerical stability.

weight_decayfloat, optional

Decoupled weight decay.

correct_biasbool, optional

Whether to apply Adam bias correction.

Methods

step([closure])

Perform a single optimization step to update parameter.

step(closure: Callable | None = None)[source]

Perform a single optimization step to update parameter.

Args:

closure (Callable): A closure that reevaluates the model and

returns the loss. Optional for most optimizers.