Topology-Aware Revival for Efficient Sparse Training

TL;DR

The paper introduces Topology-Aware Revival (TAR), a simple post-pruning method that enhances static sparse training by selectively reactivating connections based on topology, leading to significant performance improvements in reinforcement learning tasks.

Contribution

TAR is a novel one-shot revival technique that improves static sparsity without dynamic rewiring, boosting RL performance over static and dynamic sparse baselines.

Findings

TAR improves RL return by up to +37.9%.

TAR outperforms dynamic sparse training with a median gain of +13.5%.

TAR is effective across multiple continuous-control tasks.

Abstract

Static sparse training is a promising route to efficient learning by committing to a fixed mask pattern, yet the constrained structure reduces robustness. Early pruning decisions can lock the network into a brittle structure that is difficult to escape, especially in deep reinforcement learning (RL) where the evolving policy continually shifts the training distribution. We propose Topology-Aware Revival (TAR), a lightweight one-shot post-pruning procedure that improves static sparsity without dynamic rewiring. After static pruning, TAR performs a single revival step by allocating a small reserve budget across layers according to topology needs, randomly uniformly reactivating a few previously pruned connections within each layer, and then keeping the resulting connectivity fixed for the remainder of training. Across multiple continuous-control tasks with SAC and TD3, TAR improves final return over static sparse baselines by up to +37.9% and also outperforms dynamic sparse training baselines with a median gain of +13.5%.