Toward the Connection Between Activation Sparsity and Flat Minima

The observation that activation sparsity emerges in MLP blocks of standardly trained Transformers offers an opportunity to drastically reduce computation costs without sacrificing performance. To theoretically explain this phenomenon, existing works have shown that activation sparsity does not result from the data properties or data fitting but from the implicit bias of the training process. However, these connections are obtained with strong assumptions (e.g., shallow networks, a small number of training steps, and special training techniques), which cannot be applied to deep models standardly trained with a large number of steps. Different from these works, we find that the flatness of loss landscapes is also closely related to the MLP activation sparsity and can serve as a weaker assumption because it naturally emerges in the standard training of deep networks without the above strong assumptions. Specifically, we find that 1) the MLP activation sparsity equals a ratio between "augmented flatness" (a weighted sum of flatness measures) and the product of the input norm and activation gradient of the MLP. We empirically find that this ratio decreases during training, leading to sparse activations. 2) We also propose the notion of derivative sparsity, which reduces to activation sparsity under $\operatorname{ReLU}$ReLU, but further enables pruning in the backward propagation and is more stable than activation sparsity. With the theoretical findings, we can further encourage activation sparsity by decreasing the numerator and increasing the denominator of the ratio: 1) To improve (lower) the flatness, we add different bias vectors to input tokens of MLP blocks to strengthen stochastic gradient noise that drives the model to a flat area. 2) We restrict the lower bound of affine parameters in LayerNorm to increase the input norm of MLPs. 3) To increase the activation sparsity, we propose an activation function $\operatorname{JSReLU}$JSReLU to encourage the search of parameters with sparse derivatives and sparse activations. These plug-and-play modifications can effectively reduce the ratio and produce sparser activations. Experiments on ImageNet-1K and C4 demonstrate relative improvements of at least 36% on inference sparsity and at least 50% on training sparsity over vanilla Transformers, indicating further potential cost reduction in both inference and training.