Sparse or Dense? A Mechanistic Estimation of Computation Density in Transformer-based LLMs

TL;DR

This work introduces a mechanistic density estimator for Transformer-based LLMs by defining a trace-based subgraph and measuring how many edges are needed to reproduce the full output distribution within a controlled error. It shows that, contrary to sparsity-based expectations, LLM computation is largely dense on average, yet highly variable across inputs in a largely model-agnostic manner. The study develops an IFR-inspired, magnitude-based trace extraction method, uses TV distance to quantify fidelity, and defines density ρ as the area under the reconstruction-error curve, demonstrating that density tracks input difficulty (rarer tokens, higher entropy) and tends to decrease with longer context. The findings have practical implications for pruning, falsifying purely symbolic interpretability claims, and proposing density-aware frameworks for understanding linguistic processing and cognitive-scale representations in LLMs.

Abstract

Transformer-based large language models (LLMs) are comprised of billions of parameters arranged in deep and wide computational graphs. Several studies on LLM efficiency optimization argue that it is possible to prune a significant portion of the parameters, while only marginally impacting performance. This suggests that the computation is not uniformly distributed across the parameters. We introduce here a technique to systematically quantify computation density in LLMs. In particular, we design a density estimator drawing on mechanistic interpretability. We experimentally test our estimator and find that: (1) contrary to what has been often assumed, LLM processing generally involves dense computation; (2) computation density is dynamic, in the sense that models shift between sparse and dense processing regimes depending on the input; (3) per-input density is significantly correlated across LLMs, suggesting that the same inputs trigger either low or high density. Investigating the factors influencing density, we observe that predicting rarer tokens requires higher density, and increasing context length often decreases the density. We believe that our computation density estimator will contribute to a better understanding of the processing at work in LLMs, challenging their symbolic interpretation.

Figures (12)

• Figure 1: For each token generated by the LLM (here: Mistral-7B jiang2024mistral), we measure the computation density of the trace, indicating the amount of computation required to process the input and output a prediction. Each token is colored according to its density, ranging from low (cold color) to high (warm).
• Figure 2: We extract the LLM's trace at various threshold ($\tau$), allowing to vary its granularity, and plot the relation between reconstruction error (measured with Total Variation distance) and trace size for various LLMs. On average, when $\tau$ is small (bottom right corner), the trace is large --consisting of almost the full model-- and leads to a small reconstruction error. In contrast, when $\tau$ is large (top left corner), the trace is small --containing only few edges-- and leads to a large reconstruction error. Left (a), OLMo2-7B only: Each ellipse in this graph represents one $\tau$. Its is centered at the average size and reconstruction score of the resulting traces. Its horizontal and vertical diameters represent the variance (std) of the size and reconstruction score respectively. Right (b), all models: For the sake of readability, we only plot averages and omit the standard deviation.
• Figure 3: We observed that the density varies across inputs, ranging from low (a) to high (c). We quantify this density by measuring the area under the curve when x is the relative trace size and y is the reconstruction error (TV Distance). Data obtained with OLMo2-7B.
• Figure 4: We plot the distribution of density for each model. We observe a significant variation of the density across inputs.
• Figure 5: Spearman correlation of the density between pairs of LMs. In all cases, we observe a rather strong correlation.