Thoughts

Computational Graph Logging and Differential Analysis for LLM Function Extraction

J

03 Mar 2025 — 13 min read

Abstract

This research essay explores a novel approach to understanding Large Language Models (LLMs) through computational graph logging and differential analysis. We propose treating LLMs as complex mathematical functions and extracting their functional behavior by systematically logging kernel operations during inference. Our approach introduces two key innovations: (1) probabilistic kernel operation logging to manage computational and storage overhead, and (2) differential analysis of computational graphs across controlled input variations. We formulate the mathematical foundations for this approach and explore how the resulting data can serve as training material for meta-models that learn to approximate the computational patterns of LLMs. This methodology offers promising directions for model compression, explainability, optimization, and advancing our theoretical understanding of emergent behaviors in large neural networks.

1. Introduction

Large Language Models (LLMs) represent some of the most complex computational systems ever created, yet our understanding of their internal mechanisms remains limited. While attention has been given to evaluating their capabilities and outputs, less focus has been placed on systematically analyzing the computational processes that generate these outputs.

In this essay, we formalize a novel approach to LLM analysis: treating the model as a mathematical function and extracting its behavior through computational graph logging and differential analysis. Our key insight is that by capturing the execution traces of kernel operations during inference and analyzing how these traces change with controlled input variations, we can extract meaningful patterns that characterize the model's function.

This approach offers several advantages over existing methods:

It provides direct access to the actual computations performed by the model
It enables causal analysis of how input changes propagate through the computational graph
It creates a dataset that can be used to train meta-models that learn to replicate or optimize these computational patterns

2. Theoretical Framework

2.1 LLMs as Mathematical Functions

At its core, an LLM can be viewed as a function $f_{\theta}: \mathcal{X} \rightarrow \mathcal{Y}$ that maps an input sequence $x \in \mathcal{X}$ to an output distribution $y \in \mathcal{Y}$ over possible next tokens, where $\theta$ represents the model parameters.

More formally, for a sequence of input tokens $x = (x_1, x_2, ..., x_n)$, the LLM produces:

$$y = f_{\theta}(x) = P(x_{n+1} | x_1, x_2, ..., x_n; \theta)$$

This function can be decomposed into a series of computational operations performed by the model:

$$f_{\theta}(x) = f^L_{\theta} \circ f^{L-1}_{\theta} \circ ... \circ f^1_{\theta}(x)$$

Where $f^l_{\theta}$ represents the computation performed by layer $l$ of the model.

Each layer function $f^l_{\theta}$ can be further decomposed into individual kernel operations:

$$f^l_{\theta} (h^{l-1}) = g^l_K \circ g^l_{K-1} \circ ... \circ g^l_1 (h^{l-1})$$

Where $h^{l-1}$ is the output of the previous layer (or the input embedding if $l=1$), and $g^l_k$ represents the $k$-th kernel operation in layer $l$ (e.g., matrix multiplication, attention computation, activation function).

2.2 Computational Graph Logging

To extract the functional behavior of an LLM, we propose logging the execution of kernel operations during inference. For each operation $g^l_k$, we record:

$$\mathcal{L}(g^l_k) = \lbrace op\_type, inputs, outputs, metadata \rbrace$$

Where:

$op\_type$ specifies the type of operation (e.g., matrix multiplication, layer normalization)
$inputs$ and $outputs$ are the input and output tensors
$metadata$ includes additional information such as tensor shapes, layer indices, etc.

The complete log for an inference run with input $x$ is the ordered collection of all operation logs:

$$\mathcal{L}(f_θ, x) = \lbrace \mathcal{L}(g^l_k) \; | \; \forall l \in [1,L], \forall k \in [1,K_l] \rbrace$$

This represents the full computational trace of the model for input $x$.

2.3 Probabilistic Logging

To address the computational and storage overhead of full logging, we introduce probabilistic logging. For each operation $g^l_k$, we log with probability $p(g^l_k)$:

$$\mathcal{L}_p(f_θ, x) = \lbrace \mathcal{L}(g^l_k) \; | \; U(0,1) < p(g^l_k), \forall l \in [1,L], \forall k \in [1,K_l] \rbrace$$

Where $U(0,1)$ is a uniform random variable between 0 and 1, and $p(g^l_k)$ is the logging probability for operation $g^l_k$.

The logging probability can be adapted based on operation type, layer position, or other heuristics:

$$p(g^l_k) = \beta_1 \cdot p_{op}(type(g^l_k)) + \beta_2 \cdot p_{layer}(l) + \beta_3 \cdot p_{importance}(g^l_k)$$

Where:

$p_{op}$ assigns probabilities based on operation type
$p_{layer}$ assigns probabilities based on layer position
$p_{importance}$ assigns probabilities based on an importance heuristic
$\beta_1, \beta_2, \beta_3$ are weighting coefficients such that $\beta_1 + \beta_2 + \beta_3 = 1$

2.4 Differential Analysis

Differential analysis examines how changes in input affect the computational trace. For two similar inputs $x$ and $x'$, we define the computational difference:

$$\Delta \mathcal{L}(f_θ, x, x') = \mathcal{L}(f_θ, x) \ominus \mathcal{L}(f_θ, x')$$

Where $\ominus$ represents a differencing operation that identifies changes in the computational graph, including:

Added or removed operations
Changes in tensor values
Changes in execution patterns

For a set of controlled input variations $\lbrace x_1, x_2, ..., x_m \rbrace$, we can construct a differential dataset:

$$\mathcal{D} = \lbrace (x_i, x_j, \Delta \mathcal{L}(f_θ, x_i, x_j)) \; | \; i,j \in [1,m], i \neq j \rbrace$$

This dataset captures how input changes propagate through the computational graph.

3. Meta-Model Learning

Using the logged computational traces, we can train a meta-model to learn the computational patterns of the original LLM. The meta-model aims to predict aspects of the computational trace given an input:

$$M: \mathcal{X} \rightarrow \mathcal{L}$$

Where $M$ maps an input $x$ to a predicted computational trace $\hat{\mathcal{L}}$.

For differential analysis, the meta-model learns to predict computational differences given input differences:

$$M_{\Delta}: \mathcal{X} \times \mathcal{X} \rightarrow \Delta \mathcal{L}$$

Where $M_{\Delta}$ maps a pair of inputs $(x_i, x_j)$ to a predicted computational difference $\Delta \hat{\mathcal{L}}$.

The training objective for the meta-model is to minimize the distance between predicted and actual computational traces:

$$\min_{\phi} \sum_{(x, \mathcal{L}) \in \mathcal{D}} d(M_{\phi}(x), \mathcal{L})$$

Where $\phi$ represents the meta-model parameters, and $d$ is a distance function measuring the difference between computational traces.

For differential meta-models:

$$\min_{\phi} \sum_{(x_i, x_j, \Delta \mathcal{L}) \in \mathcal{D}} d(M_{\Delta,\phi}(x_i, x_j), \Delta \mathcal{L})$$

4. Methodology

4.1 Implementation in Existing Frameworks

To implement computational logging, we propose modifying inference libraries like llama.cpp. The approach involves:

Instrumenting kernel operations to record inputs, outputs, and metadata
Implementing sampling logic for probabilistic logging
Establishing efficient storage mechanisms for the logs
Developing comparison tools for differential analysis

The implementation can be formalized as:

$$g'^l_k(h) = log\_if\_sampled(g^l_k, h, p(g^l_k))$$

Where $g'^l_k$ is the instrumented version of operation $g^l_k$, and $log\_if\_sampled$ conditionally logs the operation based on the sampling probability.

4.2 Controlled Input Variation

For differential analysis, we design controlled variations of inputs. Given a base input $x$, we generate variations using transformation functions:

$$x' = T(x, \delta)$$

Where $T$ is a transformation function and $\delta$ parameterizes the transformation.

Examples of transformations include:

Token substitution: Replacing specific tokens while maintaining syntax
Structural changes: Modifying sentence structure while preserving meaning
Topic shifts: Changing the subject of the query
Language variations: Translating content to different languages

The space of transformations allows us to explore different dimensions of the model's behavior:

$$\mathcal{T} = \lbrace T_1, T_2, ..., T_r \rbrace$$
$$\Delta = \lbrace \delta_1, \delta_2, ..., \delta_s \rbrace$$
$$\mathcal{X}' = \lbrace T_i(x, \delta_j) \; | \; T_i \in \mathcal{T}, \delta_j \in \Delta \rbrace$$

4.3 Computational Difference Analysis

To analyze computational differences effectively, we need methods to compare operation logs. We define a difference function for two operations:

$$diff(g_1, g_2) = \begin{cases} \emptyset & \text{if } \text{type}(g_1) \neq \text{type}(g_2) \\ \lbrace (k, v_1, v_2) \mid v_1 \in g_1.k, v_2 \in g_2.k, v_1 \neq v_2 \rbrace & \text{otherwise} \end{cases}$$

Where $g_1.k$ represents attribute $k$ of operation $g_1$.

For comparing entire computational traces, we use sequence alignment algorithms to match corresponding operations and calculate differences:

$$\Delta \mathcal{L} = align\_and\_diff(\mathcal{L}_1, \mathcal{L}_2)$$

Where $align\_and\_diff$ applies sequence alignment followed by operation-level differencing.

Implementation of Computational Graph Differential Analysis

The following code implements the core functionality of the align_and_diff algorithm described above. This algorithm enables us to analyze and compare computational traces generated during LLM inference, providing insights into how changes in input affect the model's internal processing.

The implementation consists of several key components:

The Operation class represents a single operation in a computational graph, containing information about operation type, inputs, outputs, and additional metadata.
The operation_similarity function calculates how similar two operations are, which is crucial for the sequence alignment algorithm.
The diff function identifies specific differences between two operations, capturing insertions, deletions, replacements, and modifications.
The sequence_alignment function uses the Needleman-Wunsch algorithm to find the optimal alignment between two sequences of operations.
Finally, the align_and_diff function combines these components to generate a comprehensive analysis of differences between two computational traces.

This implementation provides the technical foundation for analyzing how LLMs process information internally and how changes in input propagate through their computational graphs.

import typing as t
from dataclasses import dataclass

import numpy as np


# Define Operation data type
@dataclass
class Operation:
    op_type: str
    inputs: dict[str, t.Any]
    outputs: dict[str, t.Any]
    metadata: dict[str, t.Any]


# Calculate similarity between two operations
def operation_similarity(op1: Operation, op2: Operation) -> float:
    """
    Calculate the similarity between two operations
    (used for sequence alignment)
    """
    # If types are different, similarity is 0
    if op1.op_type != op2.op_type:
        return 0.0

    # Metadata similarity (e.g., layer index, tensor shape, etc.)
    metadata_sim = 0.0
    common_keys = set(op1.metadata.keys()) & set(op2.metadata.keys())
    if common_keys:
        matching_values = sum(
            1 for k in common_keys if op1.metadata[k] == op2.metadata[k]
        )
        metadata_sim = matching_values / len(common_keys)

    # Similarity of input/output tensor shapes
    tensor_shape_sim = 0.0
    if "shape" in op1.inputs and "shape" in op2.inputs:
        if op1.inputs["shape"] == op2.inputs["shape"]:
            tensor_shape_sim = 1.0

    # Apply weights to calculate overall similarity
    return (
        0.6 * (op1.op_type == op2.op_type) + 0.3 * metadata_sim + 0.1 * tensor_shape_sim
    )


# Calculate difference between two operations
def diff(op1: Operation | None, op2: Operation | None) -> dict[str, t.Any]:
    """Calculate the difference between two operations"""
    # If both operations are None (no difference)
    if op1 is None and op2 is None:
        return {}

    # If one operation is None (insertion or deletion)
    if op1 is None and op2 is not None:
        return {
            "type": "insert",
            "operation": op2.op_type,
            "metadata": op2.metadata,
        }
    if op1 is not None and op2 is None:
        return {
            "type": "delete",
            "operation": op1.op_type,
            "metadata": op1.metadata,
        }
    if op1 is None or op2 is None:
        raise ValueError("One of the operations is None")

    # If types are different
    if op1.op_type != op2.op_type:
        return {"type": "replace", "from": op1.op_type, "to": op2.op_type}

    # If types are the same - calculate differences by attribute
    differences = {"type": "modify", "operation": op1.op_type, "changes": []}

    # Calculate input differences
    for k in set(op1.inputs.keys()) | set(op2.inputs.keys()):
        v1 = op1.inputs.get(k)
        v2 = op2.inputs.get(k)
        if v1 != v2:
            differences["changes"].append(
                {"attribute": f"inputs.{k}", "from": v1, "to": v2}
            )

    # Calculate output differences
    for k in set(op1.outputs.keys()) | set(op2.outputs.keys()):
        v1 = op1.outputs.get(k)
        v2 = op2.outputs.get(k)
        if v1 != v2:
            differences["changes"].append(
                {"attribute": f"outputs.{k}", "from": v1, "to": v2}
            )

    # Calculate metadata differences
    for k in set(op1.metadata.keys()) | set(op2.metadata.keys()):
        v1 = op1.metadata.get(k)
        v2 = op2.metadata.get(k)
        if v1 != v2:
            differences["changes"].append(
                {"attribute": f"metadata.{k}", "from": v1, "to": v2}
            )

    return differences if differences["changes"] else {}


# Sequence alignment algorithm (Needleman-Wunsch)
def sequence_alignment(
    seq1: list[Operation], seq2: list[Operation]
) -> list[tuple[Operation | None, Operation | None]]:
    """
    Align operations of two computational traces using sequence alignment
    algorithm
    """
    # Initialize score matrix
    n, m = len(seq1), len(seq2)
    score_matrix = np.zeros((n + 1, m + 1))

    # Matrix for tracking alignment path
    traceback = np.zeros((n + 1, m + 1), dtype=np.int)

    # Set gap penalty and match/mismatch scores
    gap_penalty = -0.5

    # Initialize first row and column (gaps only)
    for i in range(n + 1):
        score_matrix[i, 0] = i * gap_penalty
        traceback[i, 0] = 1  # Path from above

    for j in range(m + 1):
        score_matrix[0, j] = j * gap_penalty
        traceback[0, j] = 2  # Path from left

    # Fill the score matrix
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            # Calculate three cases: match/mismatch, gap above, gap left
            match_score = score_matrix[i - 1, j - 1] + operation_similarity(
                seq1[i - 1], seq2[j - 1]
            )
            delete_score = score_matrix[i - 1, j] + gap_penalty
            insert_score = score_matrix[i, j - 1] + gap_penalty

            # Select the maximum score
            if match_score >= delete_score and match_score >= insert_score:
                score_matrix[i, j] = match_score
                traceback[i, j] = 0  # Diagonal
            elif delete_score >= insert_score:
                score_matrix[i, j] = delete_score
                traceback[i, j] = 1  # Above
            else:
                score_matrix[i, j] = insert_score
                traceback[i, j] = 2  # Left

    # Traceback the aligned sequence
    aligned_pairs = []
    i, j = n, m

    while i > 0 or j > 0:
        if i > 0 and j > 0 and traceback[i, j] == 0:  # Diagonal
            aligned_pairs.append((seq1[i - 1], seq2[j - 1]))
            i -= 1
            j -= 1
        elif i > 0 and traceback[i, j] == 1:  # Above
            aligned_pairs.append((seq1[i - 1], None))
            i -= 1
        else:  # Left
            aligned_pairs.append((None, seq2[j - 1]))
            j -= 1

    # Reverse the order (since we traced back from the end)
    aligned_pairs.reverse()

    return aligned_pairs


# Align and diff function
def align_and_diff(
    trace1: list[Operation], trace2: list[Operation]
) -> list[dict[str, t.Any]]:
    """Calculate the differences between two computational traces"""
    # 1. Perform sequence alignment
    aligned_pairs = sequence_alignment(trace1, trace2)

    # 2. Calculate differences for each aligned pair
    differences = []
    for i, (op1, op2) in enumerate(aligned_pairs):
        diff_result = diff(op1, op2)
        if diff_result:  # Only add if there are differences
            diff_result["position"] = i
            differences.append(diff_result)

    return differences

Example Usage of the Align and Diff Algorithm

Now that we've defined our core functions for computational trace alignment and differential analysis, let's demonstrate how these can be applied to compare two example computational traces.

The following example creates two sample traces that represent similar but not identical computational graphs. The first trace contains a sequence of "matmul", "add", and "layernorm" operations, while the second trace replaces the "add" operation with a "softmax" operation. This simulates the kind of differences we might see when comparing traces from model executions with slightly different inputs or configurations.

We'll visualize the differences detected by our align_and_diff algorithm to show how it identifies structural and operational changes between computational graphs.

# Create example computational traces
trace1 = [
    Operation(
        "matmul",
        {"shape": (128, 64)},
        {"shape": (128, 128)},
        {"layer": 1, "position": "attention"},
    ),
    Operation(
        "add",
        {"shape": (128, 128)},
        {"shape": (128, 128)},
        {"layer": 1, "position": "residual"},
    ),
    Operation(
        "layernorm",
        {"shape": (128, 128)},
        {"shape": (128, 128)},
        {"layer": 1, "position": "post"},
    ),
]

trace2 = [
    Operation(
        "matmul",
        {"shape": (128, 64)},
        {"shape": (128, 128)},
        {"layer": 1, "position": "attention"},
    ),
    Operation(
        "softmax",
        {"shape": (128, 128)},
        {"shape": (128, 128)},
        {"layer": 1, "position": "attention_weights"},
    ),
    Operation(
        "layernorm",
        {"shape": (128, 128)},
        {"shape": (128, 128)},
        {"layer": 1, "position": "post"},
    ),
]

# Apply align_and_diff
differences = align_and_diff(trace1, trace2)

# Print results
for d in differences:
    print(d)

When we run the example code above, we get the following output which represents the differences detected between our two computational traces:

{'type': 'replace', 'from': 'add', 'to': 'softmax', 'position': 1}

This output tells us that the algorithm correctly identified that at position 1 in the alignment, the operation type changed from "add" to "softmax". This demonstrates the core capability of our differential analysis approach - identifying specific changes between computational traces that can be linked to variations in model behavior.

4.4 Meta-Model Architecture

For learning from computational traces, we propose a specialized transformer architecture that processes both inputs and operation logs. The model includes:

Input encoding: Embedding the input tokens
Operation encoding: Representing operations and their attributes
Cross-attention: Connecting input tokens to operations
Predictive heads: Forecasting operation outcomes or differences

The meta-model's forward pass can be expressed as:

$$h_x = Encoder_{input}(x)$$
$$h_{\mathcal{L}} = Encoder_{op}(\mathcal{L})$$
$$h_{cross} = CrossAttention(h_x, h_{\mathcal{L}})$$
$$\hat{\mathcal{L}} = Decoder(h_{cross})$$

Where $Encoder_{input}$ processes input tokens, $Encoder_{op}$ processes operation logs, $CrossAttention$ connects the two representations, and $Decoder$ generates predictions about the computational trace.

5. Experimental Design

To validate our approach, we propose a series of experiments:

5.1 Baseline Comparison

We compare different logging strategies:

Full logging (impractical for large models but serves as a theoretical upper bound)
Random sampling (simple baseline)
Importance-based sampling (our proposed method)
Differential-focused sampling (prioritizing operations that show high variance across inputs)

For each strategy, we evaluate:

Storage efficiency: $\frac{|\mathcal{L}_p|}{|\mathcal{L}|}$ (ratio of log size compared to full logging)
Information retention: $I(\mathcal{L}_p; \mathcal{L})$ (mutual information between sampled and full logs)
Computational overhead: $\frac{t_{\text{logged}}}{t_{\text{unlogged}}}$ (ratio of inference time with and without logging)

5.2 Differential Analysis Experiments

For differential analysis, we design controlled experiments with systematic input variations:

Token-level changes: Measuring how single token substitutions affect computation
Structural variations: Examining how syntactic changes propagate through the model
Cross-domain transfer: Analyzing how domain shifts impact computational patterns
Multi-step reasoning: Investigating computation in complex reasoning tasks

For each experiment, we track:

Propagation depth: How far in the network input changes affect computation
Activation divergence: $|h_i(x) - h_i(x')|$ for hidden state $h_i$
Attention pattern shifts: Changes in attention weight distributions
Output sensitivity: Relationship between computational and output differences

5.3 Meta-Model Training

We evaluate meta-models trained on our logged data:

Computational trace prediction: How accurately the meta-model predicts operations
Differential prediction: How well it forecasts computational changes
Function approximation: Whether the meta-model can replace parts of the original model
Transfer learning: If insights from one model transfer to others

Evaluation metrics include:

Trace prediction accuracy: $acc(\hat{\mathcal{L}}, \mathcal{L})$
Value prediction error: $MSE(\hat{v}, v)$ for predicted values $\hat{v}$
Output matching: $sim(f_{\theta}(x), M(x))$ for similarity function $sim$

6. Anticipated Challenges and Solutions

6.1 Scale and Complexity

Challenge: The sheer scale of modern LLMs makes comprehensive logging infeasible.

Solution: We propose hierarchical sampling that adapts based on model region:

$$p(g^l_k) = \alpha^l \cdot p_{base}(g^l_k)$$

Where $\alpha^l$ is a layer-specific scaling factor determined by:

$$\alpha^l = \frac{e^{\lambda \cdot importance(l)}}{\sum_{i=1}^{L} e^{\lambda \cdot importance(i)}}$$

With $importance(l)$ measuring the layer's contribution to output and $\lambda$ controlling concentration.

6.2 Representation Challenge

Challenge: How to represent operation logs in a form suitable for meta-learning.

Solution: We develop a graph-based representation that captures both operation attributes and their relationships:

$$G = (V, E)$$
$$V = \lbrace v_i \; | \; v_i \text{ represents operation } g^l_k \rbrace$$
$$E = \lbrace (v_i, v_j) \; | \; \text{output of } g_i \text{ is input to } g_j \rbrace$$

Each node contains operation attributes, and graph neural networks can be used to process this representation.

6.3 Causal Attribution

Challenge: Determining which computational changes are causally related to input changes.

Solution: We employ counterfactual analysis by designing minimal input pairs:

$$\Delta_{causal}(g^l_k, x, x') = \mathcal{L}(g^l_k, x) - \mathcal{L}(g^l_k, x')$$

We quantify the significance of an operation's change through its impact on the final output:

$$Impact(g^l_k, x, x') = |f_{\theta}(x) - f_{\theta|g^l_k}(x')|$$

Where $f_{\theta|g^l_k}(x')$ represents the model's output when only operation $g^l_k$ uses the parameters from the run with input $x'$.

7. Potential Applications

7.1 Model Compression

By identifying critical computational patterns, we can develop more efficient model architectures:

$$f'_{\phi}(x) \approx f_{\theta}(x) \text{ such that } |\phi| \ll |\theta|$$

The function extraction approach allows us to:

Identify redundant computations
Merge similar patterns
Approximate complex operations with simpler ones

7.2 Explainability

Our approach provides new tools for explaining model behavior:

Computational attribution: Mapping outputs to specific computational patterns
Critical path analysis: Identifying the minimal set of operations needed for a particular output:

$$CP(x) = \lbrace g^l_k \; | \; \text{removing } g^l_k \text{ significantly changes } f_{\theta}(x) \rbrace$$

Concept localization: Identifying where specific concepts are processed in the network

7.3 Neural Compilation

By learning computational patterns, we can develop neural compilers that optimize LLM execution:

$$C(f_{\theta}, x) \rightarrow \mathcal{P}$$

Where $C$ is a compiler function that produces an optimized execution plan $\mathcal{P}$ for input $x$ on model $f_{\theta}$.

The compiler improves efficiency by:

Pruning unnecessary computations
Reordering operations to improve cache locality
Specializing execution paths for different input types

8. Conclusion and Future Directions

This essay has explored a novel approach to understanding LLMs by treating them as mathematical functions and extracting their computational behavior through kernel operation logging and differential analysis. We have formulated the mathematical foundations for this approach and explored how the resulting data can be used to train meta-models that learn to approximate the computational patterns of LLMs.

Future research directions include:

Developing efficient logging frameworks for production LLMs
Creating specialized architectures for learning from computational traces
Applying these techniques to understand emergent capabilities in large models
Extending the approach to multimodal models

By viewing LLMs through the lens of their computational graphs, we gain new insights into their function and behavior. This approach complements existing methods and offers a promising direction for advancing our understanding of these complex systems.