Toward the Connection Between Activation Sparsity and Flat Minima

Area of Science:

• Deep Learning
• Artificial Intelligence
• Computational Neuroscience

Background:

• Activation sparsity in Transformer MLP blocks offers computational savings.
• Existing theories rely on strong assumptions unsuitable for deep, standardly trained models.
• Loss landscape flatness naturally emerges in deep networks and is linked to activation sparsity.

Purpose of the Study:

• Theoretically explain MLP activation sparsity in deep Transformers under weaker assumptions.
• Introduce derivative sparsity for enhanced pruning and stability.
• Develop methods to actively promote activation sparsity.

Main Methods:

• Proposed a theoretical link between MLP activation sparsity and "augmented flatness".
• Introduced derivative sparsity as a more stable alternative to activation sparsity.
• Implemented techniques to decrease the flatness ratio: adding bias vectors, restricting LayerNorm parameters, and using JSReLU activation.

Main Results:

• Demonstrated that MLP activation sparsity is a ratio involving augmented flatness and input/gradient norms.
• Empirically showed this ratio decreases during training, leading to sparsity.
• Achieved significant improvements in inference (36%+) and training (50%+) sparsity over standard Transformers.

Conclusions:

• Loss landscape flatness is a crucial, weaker assumption for understanding activation sparsity.
• Derivative sparsity offers advantages over activation sparsity.
• Proposed modifications effectively enhance sparsity, enabling substantial computational cost reductions.