Context
Top-KAST (Jayakumar et al., 2021) is a DST algorithm that maintains a dense "shadow" weight matrix and selects the top-K entries by magnitude as the active sparse set at each forward pass. Unlike SET/RigL which mutate every N steps, Top-KAST re-selects the active set every step.
This is a different mutation model from what SET/RigL use, and would test whether the SparsityAlgorithm base class is flexible enough to accommodate per-step topology changes (vs the current per-N-steps model).
Paper
Jayakumar et al. "Top-KAST: Top-K Always Sparse Training." NeurIPS 2021.
https://arxiv.org/abs/2106.04411
What needs to happen
- Subclass
SparsityAlgorithm (not DynamicSparsityAlgorithm, since the mutation model is different).
- Implement per-step top-K selection in
update().
- This may require maintaining a dense shadow weight alongside the Padded-CSR — design decision needed on where that lives.
- Tests + MNIST convergence comparison.
Files to read
sparselab/router.py — base classes and existing algorithmstests/test_set.py, tests/test_rigl.py — test patterns
Context
Top-KAST (Jayakumar et al., 2021) is a DST algorithm that maintains a dense "shadow" weight matrix and selects the top-K entries by magnitude as the active sparse set at each forward pass. Unlike SET/RigL which mutate every N steps, Top-KAST re-selects the active set every step.
This is a different mutation model from what SET/RigL use, and would test whether the
SparsityAlgorithmbase class is flexible enough to accommodate per-step topology changes (vs the current per-N-steps model).Paper
Jayakumar et al. "Top-KAST: Top-K Always Sparse Training." NeurIPS 2021.
https://arxiv.org/abs/2106.04411
What needs to happen
SparsityAlgorithm(notDynamicSparsityAlgorithm, since the mutation model is different).update().Files to read
sparselab/router.py— base classes and existing algorithmstests/test_set.py,tests/test_rigl.py— test patterns