neuralop.training.tensorgrad.TensorGRaDProjector

class neuralop.training.tensorgrad.TensorGRaDProjector(rank, update_proj_gap: int = 200, scale: float = 1.0, tucker_n_iter_max: int = 10, warm_restart: bool = False, activation_checkpoint: bool = False)[source]

Low-rank gradient projector used by TensorGRaD.

The original tensor is projected into a low-rank subspace using low-rank mode-wise factors obtained by Tucker decomposition. The parameters are optimized in this space to save memory and are then projected back into the full-rank space for use in a model.

Parameters:

rankfloat, int or tuple[int, …]

Goal rank of the transformed gradient tensor. A float is interpreted as a target fraction of parameters to preserve.

update_proj_gapint, optional

Number of optimizer steps between projection-tensor updates.

scalefloat, optional

Scalar applied after projecting low-rank updates back to full size.

warm_restartbool, optional

If True, uses the previous projection tensor as the initializer for the next Tucker decomposition.

activation_checkpointbool, optional

Whether to use activation checkpointing for Tucker decomposition.

Methods

get_projection_tensor
inverse_transform
project
project_back
transform

References