neuralop.training.tensorgrad.UnstructuredSparseProjector

class neuralop.training.tensorgrad.UnstructuredSparseProjector(sparse_ratio: float, update_proj_gap: int = 50, sparse_type: str = 'randk', scale: float = 1.0, scale_by_mask_ratio: bool = False)[source]

Project gradients by keeping a sparse set of entries in flattened order.

Parameters:

sparse_ratiofloat

Fraction of entries to keep in the sparse gradient.

update_proj_gapint, optional

Number of optimizer steps between sparse mask updates.

sparse_type{“randk”, “randomk”, “topk”}, optional

Strategy used to choose sparse entries. topk keeps entries with the largest absolute gradient values, while randk/randomk sample entries uniformly at random.

scalefloat, optional

Scalar applied after projecting the sparse update back to full size.

scale_by_mask_ratiobool, optional

If True, additionally scales by the square root of the inverse sparse mask ratio after projection back.

Methods

project
project_back