Timezone: »

 
Poster
Spartan: Differentiable Sparsity via Regularized Transportation
Kai Sheng Tai · Taipeng Tian · Ser Nam Lim

Wed Nov 30 02:00 PM -- 04:00 PM (PST) @ Hall J #123
in Poster Session 4 »

We present Spartan, a method for training sparse neural network models with a predetermined level of sparsity. Spartan is based on a combination of two techniques: (1) soft top-k masking of low-magnitude parameters via a regularized optimal transportation problem and (2) dual averaging-based parameter updates with hard sparsification in the forward pass. This scheme realizes an exploration-exploitation tradeoff: early in training, the learner is able to explore various sparsity patterns, and as the soft top-k approximation is gradually sharpened over the course of training, the balance shifts towards parameter optimization with respect to a fixed sparsity mask. Spartan is sufficiently flexible to accommodate a variety of sparsity allocation policies, including both unstructured and block-structured sparsity, global and per-layer sparsity budgets, as well as general cost-sensitive sparsity allocation mediated by linear models of per-parameter costs. On ImageNet-1K classification, we demonstrate that training with Spartan yields 95% sparse ResNet-50 models and 90% block sparse ViT-B/16 models while incurring absolute top-1 accuracy losses of less than 1% compared to fully dense training.

Keywords: optimal transport · Sparsity

Author Information

Kai Sheng Tai (Stanford University)
Taipeng Tian (Facebook)
Ser Nam Lim (Facebook AI)

More from the Same Authors