BlockCert: Certified Blockwise Extraction of Transformer Mechanisms

BlockCert: Certified Blockwise Extraction of Transformer Mechanisms Sandro Andric [email protected] November 2025 Abstract Mechanistic interpretability aspires to reverse-engineer neural networks into explicit algorithms, while model editing seeks to modify specific behaviours without retraining. Both areas are typically evaluated with informal evidence and ad-hoc experiments, with few explicit guarantees about how far an extracted or edited model can drift from the original on relevant inputs. We introduce BlockCert, a framework for certified blockwise extraction of transformer mechanisms, and outline how a lightweight extension can support certified local edits. Given a pre-trained transformer and a prompt distribution, BlockCert extracts structured surrogate implementations for residual blocks together with machine-checkable certificates that bound approximation error, record coverage metrics, and hash the underlying artifacts. We formalize a simple Lipschitz-based composition theorem in Lean 4 that lifts these local guarantees to a global deviation bound. Empirically, we apply the framework to GPT-2 small, TinyLlama-1.1B-Chat, and Llama-3.2-3B. Across these models we obtain high per-block coverage and small residual errors on the evaluated prompts, and in the TinyLlama setting we show that a fully stitched model matches the baseline perplexity within ≈ 6 × 10 −5 on stress prompts. Our results suggest that blockwise extraction with explicit certificates is feasible for real transformer language models and offers a practical bridge between mechanistic interpretability and formal reasoning about model behaviour. 1 Introduction Large language models (LLMs) have rapidly become core infrastructure for scientific research, software engineering, and high-stakes decision support. At the same time, we still lack robust tools for understanding, validating, and safely modifying their internal mechanisms. Mechanistic interpretability aims to address this by reverse-engineering neural networks into human-understandable components and circuits [ 2 , 8 , 10 ], but existing work is typically evaluated with bespoke visualizations or small-scale experiments, without explicit, machine-checkable guarantees. In parallel, the formal-methods community has developed powerful tools for proving properties of neural networks [ 4 , 5 , 13 ], and the programming-languages community has shown how to ship code with machine-checkable proofs of safety (proof-carrying code) [ 9 ]. However, these lines of research have had limited impact on day-to-day interpretability practice. Verification tools often do not scale to modern LLMs, and proof obligations are difficult to relate to the informal, circuit-level stories that interpretability researchers actually use. Goal. We would like a middle ground between fully formal verification and informal interpretability stories: a workflow in which reverse-engineered mechanisms are accompanied by explicit certificates 1 that (i) precisely specify what has been extracted, (ii) quantify how closely the extract matches the original network on concrete data, and (iii) can be automatically checked by any third party with access to the artifacts. Ideally, these local certificates can then be composed to reason about global model behavior. This paper. We propose BlockCert, a framework for certified blockwise extraction of transformer mechanisms. Given a pre-trained transformer, a set of prompts, and access to intermediate activations, BlockCert produces, for each residual block: 1. a structured surrogate implementation ˆ Biin an explicit intermediate representation, and 2. a JSON certificate describing the data distribution used for extraction, the achieved approximation error and coverage metrics, and cryptographic hashes of the associated artifacts (weights, probes, masks). An independent verifier can re-load the artifacts, recompute the metrics, and check that the hashes match. We also provide a composition mechanism that aggregates per-block certificates into a full-model certificate summarizing stitched-model replay metrics. We further formalize a simple global bound: if each baseline/stitched block pair ( Bi,ˆ Bi )is locally εi -sound and the blocks are Li -Lipschitz, then the composed stitched model is globally bounded by a Lipschitz-weighted sum of the εi , which specializes to a Piεi bound when Li≤ 1. This theorem is mechanized in Lean 4 and instantiated for TinyLlama using empirical per-block error bounds. Why blockwise? Transformers are naturally modular: the computation is organized into a stack of residual blocks [ 11 ]. Many mechanistic interpretability techniques operate at the block level (e.g. activation patching, MLP feature analysis), as do many model-editing methods that target specific layers [ 6 , 7 , 14 ]. By focusing on per-block extraction with quantitative guarantees, BlockCert aims to produce artifacts that are simultaneously: •close to the original computation on a specified input distribution, •simple enough to inspect and modify, •and accompanied by machine-checkable evidence of correctness. Contributions. Our main contributions are: (1) we formalize the problem of blockwise extraction for transformers and propose BlockCert, a pipeline that produces explicit surrogate blocks together with JSON certificates recording per-block approximation error, coverage metrics, and cryptographic hashes of the underlying artifacts; (2) we introduce simple but practical coverage notions (activation, path, and loss coverage) that quantify how much of a block’s behavior on a given prompt set is explained by the surrogate, and integrate these into a certificate format and verification tool; (3) in a separate Lean 4 development, we prove a global composition theorem showing that if each ( Bi,ˆ Bi )pair is locally bounded by εi and blocks are Li -Lipschitz, then the fully composed stitched model is globally bounded by a Lipschitz-weighted sum of the εi , which reduces to Piεi when Li≤ 1, and we instantiate this theorem using empirical error bounds from TinyLlama block certificates and stitched-model perplexity on stress prompts; and (4) we empirically evaluate BlockCert on GPT-2 small, TinyLlama-1.1B-Chat, and Llama-3.2-3B, obtaining high coverage (often ≥ 0 . 94 and sometimes 1 . 0) and small approximation errors, and showing that for TinyLlama the stitched model matches the baseline perplexity within ≈6×10−5on challenging prompts. Our implementation is released as an open-source Python package. To avoid confusion with historical naming, we refer to the method as BlockCert throughout this paper. 2 Model + prompts →Block traces →IR blocks + metrics ↓ Extraction / edit certificates Figure 1: High-level BlockCert/BlockCert-Edit workflow. Given a pre-trained transformer and a prompt distribution, we record block-level traces, construct IR surrogates with local error and coverage metrics, and package these into machine-checkable certificates. BlockCert-Edit applies simple local edits (e.g., scaling or swapping mechanisms) and produces analogous edit certificates on the same prompt distribution. 2 Background And Problem Setup 2.1 Transformers And Residual Blocks We consider standard decoder-only transformer language models [ 1 , 11 ], which map a sequence of tokens ( x1, . . . , xT )to logits over the vocabulary at each position. The computation proceeds through a stack of L residual blocks. Let x(0) t∈Rd denote the token embedding at position t , and let x(ℓ) tbe the residual stream at layer ℓ. A typical block has the form x(ℓ+1) t=x(ℓ) t+ MLPℓx(ℓ) t+ Attnℓx(ℓ) 1:T,(1) with preor post-layer normalization and model-specific details. We write Bℓ for the function mapping x(ℓ) 1:T to x(ℓ+1) 1:T . The full model is the composition F=BL−1◦ · · · ◦ B0◦E, where Eis the embedding and positional encoding. 2.2 Mechanistic Interpretability And Editing Mechanistic interpretability aims to understand neural networks by decomposing them into circuits of features and connections [ 2 , 10 ]. Recent work has identified circuits in vision and language models, studied superposition and polysemantic neurons, and proposed tools for tracing and editing mechanisms [ 8 , 15 ]. Model-editing methods such as ROME [ 6 ], MEMIT [ 7 ], and subsequent surveys [ 12 , 14 ] modify parameters of pre-trained LLMs to update specific facts or behaviors without full retraining. These approaches provide compelling evidence that specific mechanisms can be isolated and manipulated. However, they usually lack explicit guarantees that the edited or extracted mechanism faithfully reproduces the original network across a clearly specified set of inputs. Moreover, interpretability artifacts are rarely packaged in a way that allows independent verification. 2.3 Neural Network Verification And Proof-Carrying Code Formal verification of neural networks aims to prove properties such as robustness or safety under input perturbations [ 4 , 5 , 13 ]. Tools such as Reluplex, Marabou, and α, β -CROWN combine SMT solving and linear bound propagation to derive provable guarantees for moderately sized models. However, these methods typically treat the network as a monolithic object and reason about worst-case behavior under carefully specified constraints. In programming languages, proof-carrying code (PCC) [ 9 ] showed how to ship low-level code together with a machine-checkable proof that it satisfies a given safety policy. The host system verifies the proof before executing the code and does not need to trust the code producer. BlockCert takes inspiration from PCC: instead of shipping general-purpose proofs, we attach certificates to extracted mechanisms, which can be checked automatically by a lightweight verifier. 3 2.4 Blockwise Extraction Problem Fix a pre-trained transformer with blocks B0, . . . , BL−1. We assume access to: •the model weights, •an instrumentation mechanism that records intermediate activations for a set of prompts P, •and a target block index ℓ. For block ℓ, we define a trace dataset: Dℓ=(x(ℓ) 1:T,x(ℓ+1) 1:T,mℓ) prompt p∈ P,(2) where mℓ contains any additional discrete information about the block’s computation (e.g. attention masks, head-level gating decisions). The blockwise extraction problem is: Given ( Bℓ,Dℓ ), construct an explicit surrogate implementation ˆ Bℓ and a certificate Cℓ such that ˆ Bℓ approximates Bℓ on Dℓ according to quantitative metrics recorded in Cℓ , and such that Cℓcan be automatically re-verified from the released artifacts. The next sections describe our intermediate representation, extraction algorithm, and certificate semantics. 3 BlockCert Intermediate Representation 3.1 Design Goals The BlockCert intermediate representation (IR) is designed to satisfy three constraints: 1. Expressive enough to exactly represent standard transformer blocks. 2. Simple enough to be replayed by a small, auditable interpreter. 3. Stable under extraction: the mapping from (Bℓ,Dℓ)to ˆ Bℓshould be well-conditioned. At a high level, the IR mirrors the usual decomposition of a transformer block into attention, MLP, and residual components, but flattens model-specific details into explicit weight tensors and masks stored in .npz files: •attention weights (WQ,WK,WV,WO), •MLP weights (W1,W2)and biases, •layer norm parameters, • an explicit attention mask Mℓ encoding causal structure and any additional heador token-level gating. In our experiments we instantiate this IR for standard decoder-only architectures (GPT-2 small, TinyLlama-1.1B-Chat, Llama-3.2-3B). For GPT-2 we follow the HuggingFace GPT2Model implementation with post-embedding layer normalization and a causal attention mask. For Llamafamily models we match the LlamaAttention module, including pre-layer-normalization, rotary position embeddings (RoPE) via stored cos/sin tables and position ids, and multi-query attention with num_heads and num_key_value_heads . The interpreter is implemented as a pure Python function that applies these linear and elementwise operators in a fixed order, without any hidden control flow, so that the behavior of ˆ Bℓ is determined entirely by the released weight and mask tensors. Extraction procedure in our experiments. In all experiments we use the simplest instantiation of the extraction map ( Bℓ,Dℓ ) 7→ ˆ Bℓ : the surrogate ˆ Bℓ has the same architecture as Bℓ , with weights copied directly into the IR tensors and masks, positional tables, and bias tensors derived from 4 a single traced run on Dℓ . We do not perform any additional fitting, pruning, or distillation; residual errors arise only from numerical differences between the native implementation and the IR interpreter. 3.2 Empirical Local Soundness And Coverage Let ˆ Bℓ be an IR block and let Dℓ be the trace dataset. We define empirical local soundness and coverage metrics on Dℓas follows. Per-token error. For each traced prompt p∈ P and token position t , we compute the per-token residual error eℓ(p, t) =  ˆ x(ℓ+1) t−x(ℓ+1) t 2,(3) where ˆ x(ℓ+1) tis produced by ˆ Bℓwhen replayed on the recorded input x(ℓ) 1:T. We define: εℓ= max (p,t)eℓ(p, t),(4) MAEℓ=1 |Dℓ|X (p,t) eℓ(p, t).(5) Activation coverage. We fix a small threshold τact (e.g. 10 −2 ) and define activation coverage as covact(ℓ) = 1 |Dℓ|X (p,t) 1eℓ(p, t)≤τact.(6) Intuitively, this is the fraction of tokens for which the block output is reproduced up to a small numerical tolerance. Path coverage. We define path coverage as the fraction of traced tokens for which all discrete control decisions (e.g. attention masks, head-level gating, conditional branches) match exactly between Bℓ and ˆ Bℓ . This is computed by replaying ˆ Bℓ with instrumented hooks that compare its mask and gating tensors to those recorded in mℓ. Loss coverage. Finally, we measure the effect of the block approximation on the model’s tokenlevel loss for the traced prompts. Let ℓbase ( p, t )be the negative log-likelihood of the target token under the original model, and ℓstitched ( p, t )the loss when block ℓ is replaced by ˆ Bℓ while all other blocks remain unchanged. We define the per-token loss difference ∆ℓℓ(p, t) =  ℓstitched(p, t)−ℓbase(p, t) (7) and the loss coverage covloss(ℓ) = P(p,t)ℓbase(p, t)·1[∆ℓℓ(p, t)≤τloss] P(p,t)ℓbase(p, t),(8) where τloss is a small threshold (e.g. 10 −3 ). This measures what fraction of the baseline loss is accounted for by tokens whose loss is essentially unaffected by replacing Bℓwith ˆ Bℓ. 5 Certified blocks. A block is considered certified at level (αact, αloss)if covact(ℓ)≥αact,(9) covloss(ℓ)≥αloss.(10) In our experiments we typically use αact = 0 . 94 and αloss = 0 . 9for large models, and occasionally obtain αact = 1 . 0on stress prompts for some blocks. Because εℓ and the coverage metrics are computed only on the finite trace set Dℓ , these certificates express an empirical local soundness guarantee rather than a universal bound over all possible inputs. 4 Certificates And Verification 4.1 Certificate Format For each extracted block ˆ Bℓ, we emit a JSON certificate containing: • metadata: model name, block index, prompt set description, thresholds ( τact, τloss )and policies (αact, αloss); •metrics: εℓ,MAEℓ, activation, path, and loss coverage; • artifact digests: SHA-256 hashes of the weight, probe, mask, and bias tensors (stored in .npz files); •a declaration that the block is certified (or not) under the specified policy. Certificates live alongside their corresponding artifacts, each containing IR weights, probes, masks, per-block metrics, and a JSON certificate. 4.2 Verification Tool We provide a small Python CLI verification tool that replays certificates and checks that: 1. the SHA-256 hashes of the given artifacts match those recorded in the certificate; 2. re-computing the metrics using the current interpreter and thresholds reproduces the values stored in the certificate (up to small numerical tolerances); 3. the certified/not-certified status is consistent with the metrics and policy. For example, on a simple sanity-check experiment, the verifier recomputes the total residual ε from the released artifacts and checks that activation and loss coverage meet the certificate’s thresholds, confirming that the certificate faithfully summarizes the underlying experiment. 4.3 Full-Model Certificates We also provide a mechanism for aggregating per-block certificates into a full-model certificate. For a given model and prompt set, we: 1. stitch in the extracted blocks for a chosen subset of layers; 2. replay the full model on the prompts, computing per-layer mean absolute error (MAE) between baseline and stitched residual streams; 3. compute baseline and stitched perplexities on the prompt set; 4. construct a JSON certificate that lists all referenced block certificate hashes and records the global metrics. For TinyLlama, this process produces a full-model certificate summarizing the per-layer MAE (mean ≈ 0 . 38, worst layer ≈ 2 . 03, max residual ≈ 2 . 13 on stress prompts) and the negligible difference in perplexity between baseline and stitched models (Section 6.5). 6 4.4 Certificates Versus Formal Proofs Certificates, as produced by our tooling, are replayable, hash-tied empirical summaries. They state that for a specific model checkpoint, prompt distribution, and interpreter version, re-running the computation yields the same metrics (errors, coverage, perplexity) and artifact hashes. In contrast, a formal proof—such as the composition theorem in Section 5 and its Lean 4 formalization—is a universal statement over a specified mathematical model (e.g. all x∈X under Lipschitz assumptions). Our block and full-model certificates should therefore be read as data-restricted evidence that the hypotheses of such theorems hold on the traced distributions, not as global guarantees for all future inputs. 5 Global Composition Theorem 5.1 Statement We now describe a simple global error bound that connects per-block local soundness (universally quantified over x∈X) to total model error. Let X be a normed vector space and let Bi:X→X and ˆ Bi:X→X be the baseline and extracted blocks for i = 0 , . . . , L − 1. Here the index i plays the same role as the layer index ℓ used in Section 3.2. Define the full models F=BL−1◦···◦B0,(11) ˆ F=ˆ BL−1◦ · · · ◦ ˆ B0.(12) Theorem 1 (Global Composition).Suppose that: 1. (Local soundness) For each ithere exists εi≥0such that ∥ˆ Bi(x)−Bi(x)∥ ≤ εifor all x∈X. (13) 2. (Lipschitz blocks) For each i there exists Li≥ 0such that both Bi and ˆ Bi are Li -Lipschitz with respect to the norm ∥·∥. Then for all x∈X, ∥ˆ F(x)−F(x)∥ ≤ L−1 X i=0 εi L−1 Y j=i+1 Lj.(14) A detailed proof sketch and the full Lean 4 formalization are provided in Appendix A. 5.2 Formalization And Instantiation We formalize Theorem 1 in Lean 4 in a separate module, GlobalBound.lean . The proof uses standard results about composition of Lipschitz functions and is parameterized over the norm and the set of blocks. For real transformer models, we cannot presently prove that all blocks are globally 1-Lipschitz with respect to a useful norm over the full input space. Instead, we treat the Lipschitz property as a modeling assumption and instantiate the theorem using empirical local soundness bounds derived from block certificates, together with simple analytic Lipschitz upper bounds computed from the IR weights and local ℓ2 Lipschitz bounds certified for selected TinyLlama MLP sublayers (Section 6.6). 7 In particular, given per-block constants {Li} and empirical error bounds {εi} , Theorem 1 yields a global bound of the form ∥ˆ F(x)−F(x)∥ ≤ L−1 X i=0 εi L−1 Y j=i+1 Lj, which reduces to the familiar Piεibound in the special case Li≤1for all i. For TinyLlama-1.1B-Chat, we: • compute empirical per-block error bounds εi over a suite of stress prompts, using the per-token residual norms described in Section 3.2; • plug these εi into the Lean development to obtain a bound on ∥ˆ F ( x ) −F ( x ) ∥ for x drawn from the traced distribution; • empirically verify that the stitched model and baseline produce nearly identical perplexity on the same prompts (Section 6.5). Scope. It is important to emphasize that this global story is conditional. We do not claim a formal guarantee for all possible inputs of a real LLM. Instead, our theorem and experiments support a plausible global safety story under standard Lipschitz assumptions and empirical error bounds. Strengthening these assumptions—for example by deriving certified local Lipschitz bounds for specific blocks—is an important direction for future work. 6 Experiments We now instantiate BlockCert on three settings: GPT-2 small, TinyLlama-1.1B-Chat, and Llama3.2-3B. Extraction artifacts and certificates for all experiments are available in the supplementary materials. 6.1 GPT-2 Small: Block 0 And Multi-Block Sweep Setup. We next apply BlockCert to GPT-2 small. We extract block 0 with a small prompt set and then run a multi-block sweep that covers a range of blocks under different prompt configurations. Experimental artifacts include per-block and multi-block metrics summaries and sample certificates. Results. For the baseline configuration (two prompts), the extraction metrics report: •"prompts_evaluated": 2, •"certified_all": true, •mean extraction error ≈2.32 ×10−7for certified blocks, •mean activation coverage ≈0.9999999999998883. Sample certificates can be independently verified using the provided verification tool. Causal attention mask behavior can be validated through independent replay. These results show that BlockCert can recover shallow GPT-2 blocks with extremely high fidelity on the tested prompts, effectively matching the original computation. 6.2 TinyLlama-1.1B-Chat: Blockwise Extraction Setup. We evaluate BlockCert on TinyLlama-1.1B-Chat [ 3 ] using both a full sweep and a targeted stress-test configuration. The default TinyLlama prompt set consists of four short English prompts about extraction, audits, verifiers, and causal masks (tokenized lengths 23 , 21 , 24 , 22, totalling 90 tokens). The extended stress set comprises ten hand-designed prompts mixing long-form 8 technical writing, multilingual text, safety-related content, legal contracts, poetry, code, translation, and quantitative reasoning (tokenized lengths between 22 and 36, totalling 251 tokens). Artifacts for each block include IR weights, probes, masks, per-block metrics, and a JSON certificate. The aggregate metrics summarize per-block performance and identify, for each block, the best prompt index. Results. Across the full 22-layer residual stack (full-sweep configuration), activation coverage remains high: The extraction achieves mean activation coverage ≈ 1 . 0for blocks 0–1,12–14,16–17, and 21; ≈ 0 . 955 for most mid-depth blocks; and a worst case of ≈ 0 . 945 at block 15. Under the ten stress prompts, the metrics report, for example: •Block 5: best activation coverage ≈0.9722, mean extraction error ≈6.1×10−3. • Blocks 0 and 10: activation coverage = 1 . 0, with ε≈ 1 . 99 × 10 −3 and ε≈ 9 . 64 × 10 −3 , respectively. Certificates for these blocks satisfy the activation and loss coverage policies ( covact ≥ 0 . 94, covloss ≥ 0.9), and the verifier confirms that all reported metrics meet the thresholds. Failure case analysis (block 15). Block 15 is the deepest layer where activation coverage dips noticeably (mean coverage ≈ 0 . 945 under the stress prompts, with a minimum of ≈ 0 . 91 on a prompt containing a cluster of safety-trigger phrases). Inspection of the block-15 metrics shows that the extraction error is concentrated on a small subset of tokens in this prompt, which also have relatively large key/value norms and slightly lower attention margins. Nevertheless, loss coverage remains 1 . 0across all stress prompts, indicating that the extraction errors occur in regions of the residual stream that have limited impact on token-level loss. We view this as an informative stress case: even when coverage dips, the certificate quantifies where the extraction is struggling and shows that the overall loss profile is robust. 6.3 Llama-3.2-3B: Cross-Model Replication Setup. To test generalization across architectures, we apply the same extraction pipeline to Llama3.2-3B 1 , extracting blocks { 0 , 5 , 10 } on the same baseline and stress prompt sets as TinyLlama. Artifacts share the same directory structure as in Section 6.2, but contain Llama-3.2-3B-specific weights and traces. Results. We obtain successful certificates for blocks 0 , 5 , 10 with activation and loss coverage satisfying the same policy, and verification runs report total residuals on the order of 10 −2 with both activation and loss coverage above the default thresholds. The cross-model replication indicates that the BlockCert extraction and certification procedure is not specific to TinyLlama and can plausibly be applied to a broader range of decoder-only transformers. 6.4 Whole-Model Replay And Aggregated Certificate Setup. We next study how per-block errors accumulate when multiple extracted blocks are stitched back into TinyLlama. We construct a stitched model by replacing selected blocks with their extracted counterparts, replay both the stitched model and the baseline on the stress prompts, and record, for each of the 22 residual layers, the mean absolute error (MAE) between baseline and stitched residual streams. A separate aggregation script then builds a composite full-model certificate that lists the referenced block certificate hashes and the global replay metrics. 1Model card and weights from the official Llama 3.2 release. 9 [8] Neel Nanda, Lawrence Chan, Tom Lieberum, Johannes Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability. In International Conference on Learning Representations, 2023. [9] George C. Necula. Proof-carrying code. In Conference Record of POPL, 1997. [10] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Nick Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 5(3), 2020. [11] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017. [12] Shangwen Wang et al. Knowledge editing for large language models: A survey. arXiv preprint arXiv:2310.16218, 2023. [13] Shiqi Wang, Huan Zhang, Kaidi Xu, Xue Lin, Suman Jana, Cho-Jui Hsieh, and J. Zico Kolter. Beta-CROWN: Efficient bound propagation with per-neuron split constraints for neural network robustness verification. In Advances in Neural Information Processing Systems, 2021. [14] Yao Yao, Ningyu Zhang, Chuanqi Tao, Fei Huang, and Huajun Chen. Editing large language models: Problems, methods, and opportunities. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, 2023. [15] Fan Zhang et al. Towards best practices of activation patching in language models. OpenReview preprint, 2024. 16