Technical Summary: Tracing Computation Density in LLMs

Problem Statement

Transformer-based Large Language Models (LLMs) operate as deep, wide computational graphs involving billions of parameters and operations. However, it remains unclear whether these models utilize their full capacity for every input or rely on significantly sparser sub-networks. Key open questions include:

• Is the effective computation required for faithful prediction uniformly distributed, or does it consist of distinct dense and sparse subgraphs?
• Does the amount of effective computation depend on specific input properties?
• How can we precisely quantify the computational footprint required to reconstruct a model's output distribution?

Addressing these questions is critical for both practical efficiency (pruning, routing, distillation) and theoretical understanding (mechanistic interpretability, reconciling symbolic vs. distributed views of language, and measuring linguistic complexity).

Methodology: The s-Trace

The authors introduce s-Trace, a method to efficiently estimate the subgraph of size $s$ that best approximates a full model's output distribution.

1. Graph Formulation

The LLM is modeled as a directed acyclic graph $G=(V, E)$ where:

• Nodes ($V$): Represent intermediate token states at each layer.
• Edges ($E$): Represent operations affecting information flow, decomposed into:
  • Attention Edges: Contributions from specific attention heads moving vectors between source and target tokens.
  • MLP Edges: Updates from Multi-Layer Perceptrons.
  • Residual Edges: Identity connections carrying previous layer states.
    This formulation allows for token-wise granularity, distinguishing individual attention head contributions rather than aggregating them.

2. Trace Definition

The s-Trace ($T_s$) is defined as the optimal subgraph of size $s$ (where size is the number of edges) that minimizes the distance between the full model's output distribution $P(t_{n+1}|t)$ and the subgraph's restricted prediction $\hat{P}_{T_s}(t_{n+1}|t)$. The distance metric used is Total Variation (TV) distance, chosen for its bounded nature and focus on dominant probability modes rather than low-probability tails.

3. Trace Extraction

Given the intractability of finding the optimal subgraph in large models (7B–14B parameters), the authors employ a training-free, greedy best-first search strategy:

• Importance Scoring: Edges are scored based on the L1-norm of their update vectors, normalized per target node to ensure fair comparison.
• Traversal: Starting from the final token node in the last layer, the algorithm iteratively adds the most important available edge and its connected nodes until the budget $s$ is reached.

4. Experimental Setup

• Models: 10 LLMs from 7 HuggingFace families, ranging from 7B to 14B parameters (e.g., Llama-3, Mistral, Qwen, OLMo, Phi-4).
• Data: 5,000 sequences from Wikitext (avg. 40 words).
• Granularity: Traces were extracted at 26 levels of size, from $10^{-5}$ to $10^0$ (full model).

Key Results

1. Two Distinct Phases of Computation

Analysis of reconstruction error versus trace size reveals two consistent phases across all models:

• Construction Phase ($s \in [10^{-4}, 10^{-2}]$): Characterized by a steep, non-linear decay in reconstruction error. During this phase, a Minimal Core emerges at approximately $s \approx 10^{-3}$ (0.1% of the full model). This core is sufficient to recover the top-1% of the probability mass (the most likely token) and is highly stable across inputs.
• Refinement Phase ($s > 10^{-2}$): Beyond the minimal core, the error decreases slowly. This phase requires an order-of-magnitude increase in trace size to recover the remaining probability mass (e.g., moving from 40% to 90% coverage). The marginal utility of added edges is low, suggesting the integration of nuanced information.

2. Structural Composition of the Trace

• Layer Distribution: The minimal core is predominantly composed of early-layer nodes (initial and early-middle layers), which constitute ~80% of edges at $s=10^{-3}$. As the trace enters the refinement phase, later layers increasingly dominate the added edges.
• Component Distribution:
  • In the full model, attention edges comprise >99% of possible edges.
  • In the minimal core, the distribution is surprisingly balanced: ~30% residual, ~20% MLP, and ~50% attention edges.
  • The refinement phase is characterized by a sharp increase in the proportion of attention edges, indicating their specific role in fine-tuning the output distribution.

3. Input-Dependent Sparsity

Computation density (measured by the area under the error-vs-size curve) correlates with input properties:

• Model Uncertainty: There is a positive correlation between computation density and the Shannon entropy of the model's output. High-entropy (uncertain) inputs require larger traces (more computation) to achieve low reconstruction error.
• Token Frequency: There is a negative correlation between the frequency of the predicted token and computation density. Frequent tokens (shallow statistics) are reconstructed by sparser traces, while rare tokens require denser subgraphs.
• Cross-Model Consistency: These patterns are remarkably consistent across different model architectures and scales, suggesting that resource allocation is driven by intrinsic input properties rather than architectural details.

Significance and Claims

The paper claims to provide a general method for quantifying the effective computation of LLMs, revealing a consistent modular organization:

1. Sparse Early Core: A small, stable subgraph of early-layer operations provides a rough prediction sufficient for identifying the most probable token.
2. Dense Late Refinement: Subsequent layers perform a denser, distributed computation to refine this prediction, capturing the full richness of the output distribution.

Implications for Interpretability:

• Circuit Scale: The minimal core (~0.1% of the model) represents a lower bound for "faithful" circuit discovery; circuits smaller than this may fail to capture essential dynamics.
• Distribution Matters: Evaluating subgraphs solely on binary task success (e.g., top-1 accuracy) is insufficient; the preservation of the full output distribution is crucial.
• Beyond Hyper-Sparsity: Mechanistic interpretability should not be restricted to finding hyper-sparse circuits but should also analyze dense, distributed processing phases.

Implications for Cognitive Science:

• Linguistic Complexity: Computation density could serve as a novel metric for linguistic complexity, as it scales with input difficulty (uncertainty and rarity).
• Efficient Processing: The finding that LLMs recruit fewer resources for frequent words aligns with theories of efficient human language processing.
• Symbolic vs. Distributed: The existence of a sparse, early-layer core that handles basic statistics (resembling symbolic rules) followed by distributed refinement offers a potential reconciliation between symbolic and distributed views of language competence.

The authors maintain that these findings are specific to the transformer architecture and English data, noting that the identification of traces depends on the chosen extraction method (greedy search based on L1-norm). They position this work as an initial step toward characterizing the internal dynamics of LLMs beyond simple circuit identification.