🤖 AI

Protocode: Prototype-Driven Interpretability for Code Generation in LLMs

This paper proposes a prototype-driven approach that utilizes AST-based analysis to automatically sample high-quality In-Context Learning demonstrations, thereby simultaneously enhancing the interpretability and performance of Large Language Models in code generation tasks.

Original authors: Krishna Vamshi Bodla, Haizhao Yang

Published 2026-01-22

📖 5 min read🧠 Deep dive

CC0 1.0

Original authors: Krishna Vamshi Bodla, Haizhao Yang

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart, but slightly confused, robot how to write computer code. You can't just give it a textbook; instead, you show it a few examples of good code first. This is called In-Context Learning (ICL). The robot looks at your examples and tries to copy the style and logic to solve a new problem.

The problem is: Which examples should you show the robot?

If you show the robot bad examples, it gets confused and writes buggy code. If you show it perfect examples, it shines. But finding the "perfect" examples is hard. Most people just pick examples that look similar on the surface, but that doesn't always work.

This paper introduces a new method called Protocode to solve this. Think of it as a "Smart Example Selector" that uses a special map to find the best examples.

Here is how it works, broken down into simple concepts:

1. The "Manifold" Map (Finding the Shape of Data)

Imagine all the different programming problems (like "sort a list" or "calculate a square root") are dots floating in a giant, multi-dimensional room.

The Old Way: People usually just pick the dot closest to the one you are trying to solve. But sometimes, the closest dot is a "trick" example that looks similar but has a totally different logic.
The Protocode Way: The authors realized these dots aren't just scattered randomly; they form invisible, curved shapes (like a crumpled piece of paper) called manifolds.
The Analogy: Imagine the data is a mountain range. The old method just picks the nearest peak. Protocode looks at the shape of the mountain range itself. It finds the "prototypes"—the most representative, perfect examples that sit right in the heart of a specific valley or peak. It uses a technique called Piecewise-Linear Manifold Learning to trace these shapes accurately, ensuring the example it picks is truly representative of that type of problem, not just a random neighbor.

2. The "Proxy" Anchor (The Magnet)

Once the map is drawn, the system needs to pick the best example for each category (like "Python loops" or "Java classes").

The Analogy: Imagine you have a magnet (called a Proxy) for every type of code. You drop all your training examples into the room. The magnet pulls the examples that belong to it.
The Magic: The system doesn't just pick the closest one; it learns to move the magnets so they pull in the best examples and push away the bad ones. This is done using a math trick called Proxy-Anchor Loss. It ensures that the chosen example is a "champion" for its group, making it the most reliable teacher for the robot.

3. The "X-Ray" Vision (AST Analysis)

Once the robot writes code using these smart examples, how do we know why it wrote what it wrote?

The Problem: Usually, AI code is a "black box." We see the input and the output, but we don't know which part of the example influenced the robot to write a specific line.
The Solution: The authors use a tool called an Abstract Syntax Tree (AST). Think of code not as a sentence, but as a family tree.
- The "leaves" are the individual words (tokens).
- The "branches" are the structures (like if statements, loops, or functions).
The Analogy: The authors shine an X-ray on the code. They trace the "influence" back from the robot's output to the specific example they showed it.
- Result: They can say, "The robot used this specific for-loop structure because it was heavily influenced by Example A, but it used this error-handling block because it was influenced by Example B."
- This makes the code interpretable. You can see exactly which parts of the code are "safe" and which parts might be risky because they were copied from a weaker example.

What Did They Find?

The researchers tested this on several different AI models (like Qwen, Llama, and StarCoder) using a standard code test called MBPP.

Better Results: When they used their "Smart Example Selector" (Protocode), the robots wrote better code. They passed more tests (specifically, the pass@10 metric, which checks if the correct code appears in the top 10 guesses) compared to using random examples or just picking the "closest" ones.
The Danger of Bad Examples: They found that if you picked the wrong examples (poorly chosen ICL demonstrations), the robot actually performed worse than if you had given it no examples at all. This proves that quality matters more than quantity.
Understanding the Code: The "X-ray" analysis showed that the robot was most confident in structured parts of the code (like data structures and functions) but sometimes struggled with complex error handling, regardless of the examples used.

The Bottom Line

This paper is about teaching AI to be a better student by giving it better homework examples.

Don't just pick similar examples; pick the ones that truly represent the "shape" of the problem.
Use math to find the "champion" examples for each category.
Look under the hood to see exactly which examples influenced the AI's decisions, making the code safer and easier to understand.

By doing this, developers can trust the AI more, knowing exactly why it generated a specific piece of code and ensuring it didn't just copy a bad habit from a poor example.

Technical Summary: Prototype-Driven Interpretability for Code Generation in LLMs

Problem Statement

While Large Language Models (LLMs) have significantly advanced code generation capabilities, their reliance on automated solutions introduces risks of suboptimal or insecure code. A critical gap exists in the interpretability of these models specifically for code generation tasks. Existing interpretability methods are often generic or computationally prohibitive when applied to code. For instance, methods like ASTrust require storing probability distributions over the entire vocabulary at every generation step, leading to poor scalability with output length. Conversely, methods like Code-Q require repeated sampling, introducing significant inference overhead. Furthermore, the selection of In-Context Learning (ICL) demonstrations is known to critically impact model performance; poorly chosen examples can degrade performance below that of the base model, yet efficient, principled sampling strategies remain underexplored.

Methodology

The authors propose Protocode, a framework that unifies prototype-driven ICL sampling with Abstract Syntax Tree (AST)-grounded attribution to enhance both model performance and interpretability. The methodology consists of two primary stages:

1. Prototype Sampling via Joint Manifold and Metric Learning

Instead of using statistical summaries (e.g., class means) which fail to capture intrinsic data geometry, the authors adopt a manifold perspective. They assume high-dimensional representations lie on low-dimensional manifolds.

Manifold Construction: The method employs piecewise-linear manifold learning. It decomposes nonlinear manifolds into locally linear regions by constructing $m$ -dimensional linear submanifolds around anchor points. This is achieved by iteratively adding nearest neighbors and verifying reconstruction quality via Principal Component Analysis (PCA).
Metric Learning: To ensure prototypes are discriminative across classes, the authors integrate Proxy-Anchor loss. This introduces learnable class-level proxy vectors to optimize intra-class cohesion and inter-class separation.
Joint Optimization: A neural network $h_\theta$ is trained to jointly optimize the manifold learning objective (preserving local geometric structure) and the Proxy-Anchor loss (enforcing class separation). After training, momentum-updated proxy vectors are mapped to their nearest training instances to serve as the final ICL demonstrations (prototypes).

2. Prototype-Gradient Attribution for AST-Grounded Interpretability

To address the memory overhead of token-probability-based methods, the authors propose an attribution mechanism that avoids storing full vocabulary distributions.

Attribution Mechanism: Influence is estimated using the gradient of the similarity between the sampled prototype embeddings and the token embeddings of the generated code. Specifically, the dot product between the mean prototype vector and a token's latent representation is computed, and its gradient with respect to the token embedding serves as a confidence score.
AST Analysis: These token-level confidence scores are propagated through the Abstract Syntax Tree (AST). The process aligns tokens to terminal AST nodes, which are then hierarchically aggregated into non-terminal nodes and finally grouped into eight human-interpretable Syntax Categories (e.g., Decisions, Data Structures, Iterations, Exception Handling). This allows for both local (node-level) and global (category-level) interpretability.

Key Contributions

Unified Framework: The paper claims to be the first to unify prototype-driven ICL sampling with AST-grounded attribution for code interpretability.
Efficient Sampling Strategy: The proposed joint manifold and metric learning approach samples ICL demonstrations that are both geometrically faithful to the local data structure and semantically discriminative, avoiding the computational costs of repeated sampling or full vocabulary probability storage.
Memory-Efficient Attribution: By utilizing prototype-gradient attribution rather than token probabilities, the method enables syntax-aware confidence mapping without the memory overhead that scales poorly with output length.
Syntax-Grounded Explanations: The framework provides a post-hoc explanation system that segments generated code into interpretable syntactic regions, identifying which parts of the code are most influenced by the selected demonstrations.

Experimental Results

The authors evaluated their method on the MBPP and MBPP+ test sets using six LLMs, including generic models (Qwen3, Llama3.2, Falcon3) and code-specialized models (Starcoder, Qwen2.5-Coder, CodeLlama).

Performance: The prototype-based sampling method consistently outperformed the base model (no ICL) and other sampling strategies (diversity-based, similarity-based, and MBPP few-shot) across most models on the pass@10 metric. For example, on the MBPP dataset, the proposed method achieved a pass@10 of 0.122 for CodeLlama-7B, compared to 0.116 for the base model.
Sensitivity to Demonstrations: The results highlighted that code generation models are highly sensitive to ICL demonstration quality. Poorly chosen demonstrations (e.g., those from similarity-based or diversity-based sampling in certain contexts) negatively impacted performance, sometimes yielding results lower than the base model.
AST Analysis Insights: The analysis revealed distinct reliability patterns across syntax categories. Models generally showed high confidence in Scope, Data Structures, and Functions. Conversely, Exception Handling consistently received the lowest confidence scores across all models, suggesting a systemic limitation in generating robust error-handling logic.

Significance and Claims

The paper emphasizes that efficient sampling strategies for ICL are not merely optional but critical for model performance. The authors claim that their approach demonstrates that high-quality, geometrically faithful prototypes do more than just improve accuracy; they also enhance the interpretability of the generation process. By linking specific syntactic regions of the code to the sampled prototypes, developers can better understand the "why" behind a model's output.

The authors remain modest regarding the scope of their work, noting that prototypes were sampled exclusively from the Magicoder-OSS-Instruct dataset. They suggest that while their method serves as a global metric for ranking dataset quality regarding prototype efficacy, future work could extend this to broader datasets. Additionally, they propose that the framework could be extended toward "pre-hoc interpretability" by design, potentially using prototype steering to influence model behavior, though this remains a future direction rather than a current implementation.