Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Protocode: Prototype-Driven Interpretability for Code Generation in LLMs

Krishna Vamshi Bodla, Haizhao Yang

TL;DR

Protocode tackles interpretability in code-generation by introducing a prototype-driven ICL sampling framework that jointly learns a geometry-preserving prototype space with Proxy Anchor and manifold losses, producing discriminative prototypes without full vocabulary storage. It then derives syntax-aware token confidences via gradient-based prototype attribution and propagates these through an Abstract Syntax Tree to provide local and global explanations. Empirical results on MBPP/MBPP+ across multiple code-focused LLMs show consistent improvements in pass@k metrics and demonstrate memory-efficient sampling compared to full token-probability storage. The approach enables scalable, syntax-guided interpretability for code generation and opens avenues for both post-hoc and pre-hoc interpretability enhancements and dataset ranking analyses.

Abstract

Since the introduction of Large Language Models (LLMs), they have been widely adopted for various tasks such as text summarization, question answering, speech-to-text translation, and more. In recent times, the use of LLMs for code generation has gained significant attention, with tools such as Cursor and Windsurf demonstrating the ability to analyze massive code repositories and recommend relevant changes. Big tech companies have also acknowledged the growing reliance on LLMs for code generation within their codebases. Although these advances significantly improve developer productivity, increasing reliance on automated code generation can proportionally increase the risk of suboptimal solutions and insecure code. Our work focuses on automatically sampling In-Context Learning (ICL) demonstrations which can improve model performance and enhance the interpretability of the generated code. Using AST-based analysis on outputs from the MBPP test set, we identify regions of code most influenced by the chosen demonstrations. In our experiments, we show that high-quality ICL demonstrations not only make outputs easier to interpret but also yield a positive performance improvement on the pass@10 metric. Conversely, poorly chosen ICL demonstrations affected the LLM performance on the pass@10 metric negatively compared to the base model. Overall, our approach highlights the importance of efficient sampling strategies for ICL, which can affect the performance of the model on any given task.

Protocode: Prototype-Driven Interpretability for Code Generation in LLMs

TL;DR

Protocode tackles interpretability in code-generation by introducing a prototype-driven ICL sampling framework that jointly learns a geometry-preserving prototype space with Proxy Anchor and manifold losses, producing discriminative prototypes without full vocabulary storage. It then derives syntax-aware token confidences via gradient-based prototype attribution and propagates these through an Abstract Syntax Tree to provide local and global explanations. Empirical results on MBPP/MBPP+ across multiple code-focused LLMs show consistent improvements in pass@k metrics and demonstrate memory-efficient sampling compared to full token-probability storage. The approach enables scalable, syntax-guided interpretability for code generation and opens avenues for both post-hoc and pre-hoc interpretability enhancements and dataset ranking analyses.

Abstract

Since the introduction of Large Language Models (LLMs), they have been widely adopted for various tasks such as text summarization, question answering, speech-to-text translation, and more. In recent times, the use of LLMs for code generation has gained significant attention, with tools such as Cursor and Windsurf demonstrating the ability to analyze massive code repositories and recommend relevant changes. Big tech companies have also acknowledged the growing reliance on LLMs for code generation within their codebases. Although these advances significantly improve developer productivity, increasing reliance on automated code generation can proportionally increase the risk of suboptimal solutions and insecure code. Our work focuses on automatically sampling In-Context Learning (ICL) demonstrations which can improve model performance and enhance the interpretability of the generated code. Using AST-based analysis on outputs from the MBPP test set, we identify regions of code most influenced by the chosen demonstrations. In our experiments, we show that high-quality ICL demonstrations not only make outputs easier to interpret but also yield a positive performance improvement on the pass@10 metric. Conversely, poorly chosen ICL demonstrations affected the LLM performance on the pass@10 metric negatively compared to the base model. Overall, our approach highlights the importance of efficient sampling strategies for ICL, which can affect the performance of the model on any given task.

Paper Structure

This paper contains 38 sections, 12 equations, 5 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Methodology
  4. Dataset
  5. Training overview
  6. Manifold Construction
  7. Training objectives
  8. Proxy Anchor Loss:
  9. Manifold Point-to-Point Loss:
  10. Distance Calculation.
  11. Results
  12. AST analysis
  13. Syntax Grounded explanations
  14. Code Syntactic analysis
  15. Future Works
  16. ...and 23 more sections

Figures (5)

  • Figure 1: Conceptual working of AST analysis
  • Figure 2: Alignment & Clustering Interactions. The $\delta$ function aligns tokens $w_i$ to terminal nodes $\lambda$. Terminal and Non-terminal nodes $\lambda$, $\alpha$$\in$$\upsilon$ are clustered by Syntax Categories $\Lambda$
  • Figure 3: AST analysis on 6 LLMs
  • Figure 4: Comparison of two code snippets Before and After code sanitization with evalplus
  • Figure 5: Ablation study of Qwen2.5-Coder-0.5B and Llama3.2-1B models. Qwen2.5-Coder-0.5B Llama3.2-1B