BTP: A Research Framework for the Mechanistic Interpretability of Code Generation Capabilities in Large Language Models

The BTP project provides a complete toolchain and experimental framework for analyzing and pruning attention heads in large language models, evaluating the mechanistic interpretability of the models' internal mechanisms on code generation benchmarks such as HumanEval, MBPP, and LiveCodeBench. The project focuses on opening the "black box" of code generation in large language models, revealing the functional roles of attention heads through systematic methods, and providing support for improving model reliability, security, and efficiency.

Large language models perform excellently in code generation tasks, but their internal mechanisms are a "black box". Mechanistic interpretability aims to precisely locate the physical implementation of specific computational functions, which is crucial for improving model reliability, etc. Code generation tasks have unique challenges such as strict syntax, deterministic results, and multiple equivalent implementations, making them an ideal testbed for examining LLM reasoning capabilities—code correctness can serve as an objective evaluation criterion.

BTP builds an end-to-end experimental infrastructure with three core functions: entropy lens analysis, attention head ablation experiments, and Taylor approximation pruning:

1. Entropy Lens Analysis: Tracks the entropy values and cross-entropy loss of tokens across layers, distinguishes between regular (reasoning chain + code) and chain_code (code only) modes, and separates the influence of the thinking process from the final code.
2. Attention Head Ablation Experiments: Quantifies the importance of heads through head-wise zeroing technology, supports single-head ablation matrices and elimination-style experiments, and reveals key and redundant heads.
3. Taylor Approximation Pruning: Performs iterative pruning based on first-order Taylor expansion, generates historical records, and improves computational efficiency.

The project uses three authoritative code generation benchmarks:

• HumanEval: OpenAI's classic benchmark with 164 programming tasks covering basic algorithms, etc., serving as a standard starting point.
• MBPP: 500 Python problems with difficulty levels from intermediate to beginner, task descriptions close to real scenarios, and rich test cases.
• LiveCodeBench v6: 175 real competition/interview questions involving complex algorithm design, testing the capability boundary for high-difficulty tasks.

The project compares the attention circuit differences between the base model and distilled variants (e.g., Qwen/Llama distilled by DeepSeek-R1), calculates the cosine similarity of OV/QK circuits, and quantifies the impact of distillation on internal representations. Theoretically, it helps understand whether distillation preserves or reconstructs functions; practically, if key heads overlap, interpretability findings from the base model can be transferred, reducing analysis costs.

The experimental workflow is: vLLM generates code solutions → evaluation scripts filter correct answers → entropy analysis and ablation experiments. Shell scripts are provided to simplify operations:

• run_inference.sh starts the inference service
• run_evaluate.sh evaluates correctness
• run_hml.sh runs the HML analysis workflow
• run_check.sh tracks progress Advanced users can specify models, datasets, and inference modes via Python modules to flexibly adapt to needs.

The project includes Jupyter Notebooks for result visualization:

• The entropy analysis notebook draws layer entropy heatmaps to identify layers where the model outputs "deterministically"
• The ablation analysis notebook shows the distribution of head importance and reveals the clustering patterns of key heads
• The optimization visualization notebook tracks the CMA-ES optimization process of head subsets These tools help form intuitions and guide model compression and architecture design.

BTP provides infrastructure for mechanistic interpretability in the code generation field, answering questions such as the distributed representation of code generation capabilities and the existence of functional modules. Future directions include:

• Guiding more efficient model architecture design (reducing redundant heads, enhancing key circuits)
• Providing a new perspective for model security (monitoring attention patterns of malicious code generation) Long-term promotion of model reliability and security improvement.