Table of Contents
Fetching ...

Uncovering Layer-Dependent Activation Sparsity Patterns in ReLU Transformers

Cody Wild, Jesper Anderson

TL;DR

This paper investigates layer-dependent activation sparsity in ReLU Transformers, focusing on how per-token, per-sequence, and per-batch sparsity patterns evolve during training in a $6$-layer decoder-only Transformer with hidden dimension $32768$. It introduces three core sparsity metrics and percentile-based statistics, and reports that the first and last layers exhibit distinct and often opposing sparsity dynamics, with neuron death tied to training dynamics rather than random noise. Key findings include dramatic inter-layer differences in sparsity (e.g., Layer 0 using ~${13.3}\%$ vs Layer 5 using ~${95.6}\%$ at convergence), a gradual decrease in token-level use followed by unique per-sequence shifts for Layer 0, and evidence that many hidden units can be pruned with minimal impact on accuracy. The work discusses Conceptual notions like Feature Specificity, explores limitations of small-ReLU models, and highlights practical implications for capacity planning and potential pruning in training regimes.

Abstract

Previous work has demonstrated that MLPs within ReLU Transformers exhibit high levels of sparsity, with many of their activations equal to zero for any given token. We build on that work to more deeply explore how token-level sparsity evolves over the course of training, and how it connects to broader sparsity patterns over the course of a sequence or batch, demonstrating that the different layers within small transformers exhibit distinctly layer-specific patterns on both of these fronts. In particular, we demonstrate that the first and last layer of the network have distinctive and in many ways inverted relationships to sparsity, and explore implications for the structure of feature representations being learned at different depths of the model. We additionally explore the phenomenon of ReLU dimensions "turning off", and show evidence suggesting that "neuron death" is being primarily driven by the dynamics of training, rather than simply occurring randomly or accidentally as a result of outliers.

Uncovering Layer-Dependent Activation Sparsity Patterns in ReLU Transformers

TL;DR

This paper investigates layer-dependent activation sparsity in ReLU Transformers, focusing on how per-token, per-sequence, and per-batch sparsity patterns evolve during training in a -layer decoder-only Transformer with hidden dimension . It introduces three core sparsity metrics and percentile-based statistics, and reports that the first and last layers exhibit distinct and often opposing sparsity dynamics, with neuron death tied to training dynamics rather than random noise. Key findings include dramatic inter-layer differences in sparsity (e.g., Layer 0 using ~ vs Layer 5 using ~ at convergence), a gradual decrease in token-level use followed by unique per-sequence shifts for Layer 0, and evidence that many hidden units can be pruned with minimal impact on accuracy. The work discusses Conceptual notions like Feature Specificity, explores limitations of small-ReLU models, and highlights practical implications for capacity planning and potential pruning in training regimes.

Abstract

Previous work has demonstrated that MLPs within ReLU Transformers exhibit high levels of sparsity, with many of their activations equal to zero for any given token. We build on that work to more deeply explore how token-level sparsity evolves over the course of training, and how it connects to broader sparsity patterns over the course of a sequence or batch, demonstrating that the different layers within small transformers exhibit distinctly layer-specific patterns on both of these fronts. In particular, we demonstrate that the first and last layer of the network have distinctive and in many ways inverted relationships to sparsity, and explore implications for the structure of feature representations being learned at different depths of the model. We additionally explore the phenomenon of ReLU dimensions "turning off", and show evidence suggesting that "neuron death" is being primarily driven by the dynamics of training, rather than simply occurring randomly or accidentally as a result of outliers.
Paper Structure (36 sections, 32 figures, 5 tables)

This paper contains 36 sections, 32 figures, 5 tables.

Figures (32)

  • Figure 1: The fraction of hidden units in the MLPs of a six-layer Transformer that have a nonzero activation in the course of a 8192-token batch, divided by layer, shown over the course of training. We can see there here that there is strong divergence in the behavior of different layers, and that in particular the first layer collapses dramatically in how many of its available hidden units it uses
  • Figure 2: These plots show the early evolution of (\ref{['token_over_training:token']}) the number of dimensions activated per token, and (\ref{['token_over_training:percentile']}) the fraction of the sequence in which the 50th percentile most-used hidden unit in a sequence occurs.
  • Figure 3: These plots show the fraction of dimensions used per-batch or per-token, for models of hidden dimension 8192, 16384, 32768.
  • Figure 4: The fraction of available hidden units used per batch for a 6, 10 and 14 layer model with a hidden dimension of 16384.
  • Figure 5: Hidden Units Used Per Token: Full Training Curve
  • ...and 27 more figures