Sparse learning gradient filtering feature request 🟡

The main issue: sd_scripts lacks a custom gradient filtering function with masking. My attempt to implement a module with custom filters via an additional argument was unsuccessful due to my limited programming skills. Currently, I'm getting by with modifying the optimizer's code directly to add conditional top-k relative or variance gradient filtering - tests show excellent results and provide an additional layer of control over training precision and flexibility.

It would be great if someone capable of properly integrating a filtering option into the scripts - with the ability to modify/add custom filtering options - could help with this.

The core idea is to filter out non-informative gradient regions, forcing the network to learn only from meaningful areas.

An example of a SignSGD optimizer with a simple top-k filtering scheme implemented:

```
import math
from typing import Dict, Tuple

import torch

from pytorch_optimizer.base.exception import NoSparseGradientError
from pytorch_optimizer.base.optimizer import BaseOptimizer
from pytorch_optimizer.base.type import CLOSURE, DEFAULTS, LOSS, PARAMETERS

class SignSGD(BaseOptimizer):
r"""Compressed Optimisation for Non-Convex Problems with Top-k Relative Change Filtering.

:param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
:param lr: float. learning rate.
:param momentum: float. momentum factor (0.0 = SignSGD, >0 = Signum).
:param weight_decay: float. weight decay (L2 penalty).
:param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
:param topk_ratio: float. Fraction of most changing parameters to update (0.0 = all, 1.0 = none).
"""

def __init__(
self,
params: PARAMETERS,
lr: float = 1e-3,
momentum: float = 0.9,
weight_decay: float = 0.0,
weight_decouple: bool = True,
topk_ratio: float = 0.0, # 0.0 means no masking
**kwargs,
):
self.validate_learning_rate(lr)
self.validate_range(momentum, 'beta', 0.0, 1.0)
self.validate_non_negative(weight_decay, 'weight_decay')
self.validate_range(topk_ratio, 'topk_ratio', 0.0, 1.0)

self.topk_ratio = topk_ratio

defaults: DEFAULTS = {
'lr': lr,
'momentum': momentum,
'weight_decay': weight_decay,
'weight_decouple': weight_decouple,
}
super().__init__(params, defaults)

def __str__(self) -> str:
return 'SignSGD'

@torch.no_grad()
def reset(self):
for group in self.param_groups:
group['step'] = 0
for p in group['params']:
state = self.state[p]

if group['momentum'] > 0.0:
state['momentum_buffer'] = torch.zeros_like(p)
state['prev_grad'] = torch.zeros_like(p)

@torch.no_grad()
def step(self, closure: CLOSURE = None) -> LOSS:
loss: LOSS = None
if closure is not None:
with torch.enable_grad():
loss = closure()

for group in self.param_groups:
lr = group['lr']
momentum = group['momentum']

for p in group['params']:
if p.grad is None:
continue

grad = p.grad
if grad.is_sparse:
raise NoSparseGradientError(str(self))

state = self.state[p]
if momentum > 0.0:
if 'momentum_buffer' not in state:
state['momentum_buffer'] = torch.zeros_like(p)
buf = state['momentum_buffer']
buf.mul_(momentum).add_(grad, alpha=1.0 - momentum)
else:
buf = grad

# Top-k relative change filtering
if self.topk_ratio > 0.0:
if 'prev_grad' not in state:
state['prev_grad'] = torch.zeros_like(buf)

prev = state['prev_grad']
rel_change = torch.abs(buf - prev) / (torch.abs(prev) + 1e-8)

k = int(rel_change.numel() * (1.0 - self.topk_ratio))
if k > 0:
topk_vals, topk_idx = torch.topk(rel_change.view(-1), k)
mask = torch.zeros_like(rel_change).view(-1)
mask[topk_idx] = 1.0
mask = mask.view_as(buf)
buf = buf * mask

state['prev_grad'].copy_(buf)

# Update step
p.add_(torch.sign(buf), alpha=-lr)

return loss

Start by reviewing the optimizer code the reporter currently modifies and the supplied SignSGD example. Define how configurable gradient-filtering and masking options should fit into sd_scripts, including support for custom filters. Done means the filtering feature is integrated without direct optimizer edits and its behavior is covered by tests.

Written by the indexing model from the issue text.