TLDR: The research paper introduces Code Concept Analysis (CoCoA), a framework to interpret code language models by clustering contextualized token embeddings into human-understandable lexical, syntactic, and semantic concept groups. It uses a hybrid annotation pipeline (static analysis and LLMs) to label these concepts and integrates with local attribution methods to provide concept-grounded explanations. Evaluations show CoCoA discovers stable concepts, which evolve predictably with fine-tuning, and significantly improves the interpretability of token-level explanations in user studies.
Understanding how large language models (LLMs) trained on code make their decisions is a significant challenge, especially when these models are used in critical applications where trust and transparency are paramount. Traditional methods for interpreting these models often fall short, providing either shallow explanations for individual predictions or relying on predefined labels that don’t fully capture the model’s internal workings.
A new research paper, “Analyzing Latent Concepts in Code Language Models,” introduces an innovative framework called Code Concept Analysis (CoCoA). Developed by Arushi Sharma, Vedant Pungliya, Christopher J. Quinn, and Ali Jannesari from Iowa State University, CoCoA aims to shed light on the hidden structures that emerge within a code language model’s representation space. This framework helps uncover lexical, syntactic, and semantic concepts by grouping contextualized token embeddings into categories that humans can easily understand.
Unveiling the Model’s Internal Logic
CoCoA operates as a global, post-hoc interpretability framework, meaning it analyzes the model’s behavior after it has been trained. Its core idea is to cluster the model’s internal representations of code tokens (like words or symbols) into meaningful groups. Imagine the model processing a piece of code; CoCoA looks at how the model “thinks” about each part of that code and then groups similar thoughts together. These groups represent “latent concepts” – ideas or patterns the model has learned on its own.
The process begins with “Concept Discovery,” where the framework extracts the model’s internal representations for each token in a code snippet. These representations, which capture rich information about the token’s context and meaning, are then clustered using a technique called K-Means. Each resulting cluster reveals a recurring pattern in how the model abstracts and organizes information.
Annotating the Concepts: A Hybrid Approach
Once these latent concepts are discovered, the next crucial step is to make them human-interpretable. CoCoA employs a clever “Concept Alignment and Annotation” pipeline. This involves a hybrid approach: first, aligning clusters with predefined syntactic labels (like identifying if a token is an “operator” or a “variable”). Second, for more complex semantic categories, it uses a combination of manual annotation by human experts and, notably, large language models (LLMs) like GPT-4o as annotators. This LLM-as-annotator approach, refined through prompt engineering, allows for scalable labeling of concepts across different levels of abstraction, leading to the creation of a new dataset called CodeConceptNet (CoCoNet).
Enhancing Explanations with Concepts
One of CoCoA’s significant applications is its ability to improve local explanations – methods that highlight specific input tokens responsible for a model’s prediction. Traditional attribution methods might point to a token like ‘<‘ without clarifying its role (e.g., is it a comparison operator, a type delimiter, or a PHP opening tag?). CoCoA addresses this by mapping these salient tokens to their discovered latent concept clusters. For instance, if ‘<‘ is highlighted, CoCoA can tell us if it belongs to a cluster of “comparison operators” or “PHP opening tags,” providing a much richer and more coherent explanation. This concept-grounded approach makes token-level saliency more understandable and trustworthy.
Also Read:
- Interpreting AI’s Code: How PROTOCODE Improves LLM Performance and Transparency
- Decoding the Black Box: How AI Explains Itself
Key Findings and Insights
The researchers evaluated CoCoA across multiple models, including CodeBERT, UniXCoder, and DeepSeekCoder-V2-Lite, and various tasks like AST token tagging, compilation error detection, and programming language classification. They observed that the discovered concepts align well with human-interpretable categories across lexical (surface-level similarities), syntactic (structural roles), and semantic (functional behavior) dimensions.
A fascinating insight is how these concepts evolve across the model’s layers. Lexical patterns are more prominent in lower layers, while syntactic alignment is highest in early and middle layers. Deeper layers tend to prioritize more abstract, semantic, or compositional concepts. Fine-tuning the models on specific tasks also influenced concept evolution; for example, tasks focused on code structure enhanced syntactic understanding, while broader semantic tasks led to more generalized, cross-lingual conceptual groupings.
Crucially, CoCoA-derived clusters demonstrated robustness to semantic-preserving code changes, such as renaming identifiers. This indicates that the concepts learned are semantically grounded rather than superficial. A user study further validated the framework’s utility, showing that concept-augmented explanations improved human-centric explainability by 37 percentage points compared to token-level attributions alone. This means people found the explanations much clearer and more useful when concepts were included.
While the framework offers significant advancements, the authors acknowledge limitations, such as the trade-off between granularity and scalability in concept discovery and the inherited limitations of underlying attribution methods. However, CoCoA represents a vital step towards building more transparent and trustworthy code language models, offering a novel way to understand their internal logic and improve their explainability. For more technical details, you can read the full paper here.
