Tracing Computation Density in LLMs

2026-05-26

AI Summary

It remains unclear whether large language models (LLMs) uniformly deploy their full computational capacity across all inputs. This work proposes s-Trace, a method that efficiently identifies critical subgraphs of size *s* capable of approximating the full model’s output. Applying s-Trace reveals, for the first time in a systematic manner, that LLM computation is organized modularly: early layers rely on sparse core subgraphs to generate initial predictions, while later layers refine these through denser computation. The study further demonstrates a strong positive correlation between computational density and model uncertainty, showing that early sparse subgraphs suffice to reconstruct the dominant components of the output distribution, whereas later attention heads progressively optimize finer details. These findings offer a novel perspective on the dynamic computational mechanisms underlying LLMs.

Abstract

Transformer-based large language models (LLMs) are comprised of billions of parameters arranged in deep and wide computational graphs, but it is not clear that they exploit their full capacity for all inputs. We introduce the s-Trace method to efficiently estimate the subgraph of size s that best approximates a full model output. With this method, we find the computation in a variety of LLMs to be organized in two distinct phases. A small subgraph mostly composed of early-layer nodes can reconstruct the head of the full model output distribution. Adding further nodes, mostly located in later layers and increasingly consisting of attention heads, leads to incremental refinements in approximating the full output distribution. We find moreover that the amount of necessary computation per input correlates with model uncertainty, and that sparser subgraphs encode shallow statistics, such as unigram frequency. Overall, our results suggest a consistent modular organization in effective LLM computation, with a sparse early-layer core providing a rough prediction that is further refined through denser computations in later layers.

Problem

• computation density
• large language models
• computational graphs
• model uncertainty
• modular organization

Innovation

• s-Trace
• computation density
• modular computation
• subgraph approximation
• large language models