[RFC][Feature Proposal] Dynamic 2:4 semi-structured sparsity with sign-flip pruning and gradient-based regrowth

[catswe]

The general idea is to introduce a custom nn.Linear layer that maintains 2:4 semi-structured sparsity dynamically throughout training. The layer will be initialized with random weights and corresponding masks that follow 2:4 semi-structured sparsity. When a weight crosses 0 (i.e flips sign), it is automatically pruned (i.e set to 0), freeing up a nonzero slot in its 2:4 group. After another training step, one of the zero weights within the same 2:4 group will be chosen for regrowth, selected based on the magnitude of the densely calculated gradient.

I'd be happy to work on implementing this feature and provide a detailed implementation plan for it, assuming this repository is open to landing this.

[jcaip]

Hi @catswe

It's an interesting idea, but I don't think TorchAO is the right place for it. Can you share some more context on why you want to add such a feature?

[catswe]

Hi @jcaip

The idea of rewiring weights when they cross zero has been around for a while (e.g., https://arxiv.org/pdf/1711.05136). Similarly, activating new connections based on the gradients has also been explored in various works (e.g., https://research.google/blog/improving-sparse-training-with-rigl).

The proposed module aims to integrate these ideas and work more or less as a drop-in component. The hope is that this approach could offer more stable training behaviors and minimize the associated overhead of modifying the architecture and training pipeline. This would further work in this area and better allow researchers and developers to experiment with and reap the inference benefit of 2:4 semi-structured sparsity