# Risk Assessment Framework for Code LLMs via Leveraging Internal States

arXiv: https://arxiv.org/abs/2504.14640 · HTML: https://mlsanity.com/paper/2504.14640/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 77 references that are in the map.
- First posted on arXiv 20 Apr 2025.

## Builds on (most cited first, 30 of 40)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Sparse Autoencoders Find Highly Interpretable Features in Language Models](https://mlsanity.com/paper/2309.08600/) (2023)
- [Representation Engineering: A Top-Down Approach to AI Transparency](https://mlsanity.com/paper/2310.01405/) (2023)
- [DeepSeek-Coder: When the Large Language Model Meets Programming -- The Rise of Code Intelligence](https://mlsanity.com/paper/2401.14196/) (2024)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [StarCoder: may the source be with you!](https://mlsanity.com/paper/2305.06161/) (2023)
- [The Geometry of Truth: Emergent Linear Structure in Large Language Model Representations of True/False Datasets](https://mlsanity.com/paper/2310.06824/) (2023)
- [The Internal State of an LLM Knows When It's Lying](https://mlsanity.com/paper/2304.13734/) (2023)
- [Executable Code Actions Elicit Better LLM Agents](https://mlsanity.com/paper/2402.01030/) (2024)
- [StarCoder 2 and The Stack v2: The Next Generation](https://mlsanity.com/paper/2402.19173/) (2024)
- [InCoder: A Generative Model for Code Infilling and Synthesis](https://mlsanity.com/paper/2204.05999/) (2022)
- [Large Language Models for Software Engineering: A Systematic Literature Review](https://mlsanity.com/paper/2308.10620/) (2023)
- [CodeT: Code Generation with Generated Tests](https://mlsanity.com/paper/2207.10397/) (2022)
- [Magicoder: Empowering Code Generation with OSS-Instruct](https://mlsanity.com/paper/2312.02120/) (2023)
- [Grounded Copilot: How Programmers Interact with Code-Generating Models](https://mlsanity.com/paper/2206.15000/) (2022)
- [Lookback Lens: Detecting and Mitigating Contextual Hallucinations in Large Language Models Using Only Attention Maps](https://mlsanity.com/paper/2407.07071/) (2024)
- [Self-planning Code Generation with Large Language Models](https://mlsanity.com/paper/2303.06689/) (2023)
- [TruthX: Alleviating Hallucinations by Editing Large Language Models in Truthful Space](https://mlsanity.com/paper/2402.17811/) (2024)
- [Refining ChatGPT-Generated Code: Characterizing and Mitigating Code Quality Issues](https://mlsanity.com/paper/2307.12596/) (2023)
- [DebugBench: Evaluating Debugging Capability of Large Language Models](https://mlsanity.com/paper/2401.04621/) (2024)
- [Improving Activation Steering in Language Models with Mean-Centring](https://mlsanity.com/paper/2312.03813/) (2023)
- [CodexGraph: Bridging Large Language Models and Code Repositories via Code Graph Databases](https://mlsanity.com/paper/2408.03910/) (2024)
- [Calibration and Correctness of Language Models for Code](https://mlsanity.com/paper/2402.02047/) (2024)
- [What Makes Good In-context Demonstrations for Code Intelligence Tasks with LLMs?](https://mlsanity.com/paper/2304.07575/) (2023)
- [No Need to Lift a Finger Anymore? Assessing the Quality of Code Generation by ChatGPT](https://mlsanity.com/paper/2308.04838/) (2023)
- [Large Language Models and Simple, Stupid Bugs](https://mlsanity.com/paper/2303.11455/) (2023)
- [WaveCoder: Widespread And Versatile Enhancement For Code Large Language Models By Instruction Tuning](https://mlsanity.com/paper/2312.14187/) (2023)

## Built on by (most cited first, 9 of 9)

- [A Comprehensive Survey on Benchmarks and Solutions in Software Engineering of LLM-Empowered Agentic System](https://mlsanity.com/paper/2510.09721/) (2025)
- [On LLMs' Internal Representation of Code Correctness](https://mlsanity.com/paper/2512.07404/) (2025)
- [Localized Calibrated Uncertainty in Code Language Models](https://mlsanity.com/paper/2512.24560/) (2025)
- [CodeCircuit: Toward Inferring LLM-Generated Code Correctness via Attribution Graphs](https://mlsanity.com/paper/2602.07080/) (2026)
- [A Study on the Impact of Fault localization Granularity for Repository-Scale Code Repair Tasks](https://mlsanity.com/paper/2604.00167/) (2026)
- [On the Robustness of LLMs' Internal Representation of Code Correctness](https://mlsanity.com/paper/2608.08266/) (2026)
- [Code-MUE: Measuring Code LLMs' Uncertainty through Execution-based Semantic Interaction Graphs](https://mlsanity.com/paper/2607.12273/) (2026)
- [How Do LLMs Read Bug Reports? An Empirical Study of Attention in LLMs for Automated Program Repair](https://mlsanity.com/paper/2607.25873/) (2026)
- [Diagnosing Generalization Boundaries in AI-Based Vulnerability Detection via Calibrated Embedding Distances](https://mlsanity.com/paper/2610.14296/) (2026)

## Abstract (arXiv)

The pre-training paradigm plays a key role in the success of Large Language Models (LLMs), which have been recognized as one of the most significant advancements of AI recently. Building on these breakthroughs, code LLMs with advanced coding capabilities bring huge impacts on software engineering, showing the tendency to become an essential part of developers' daily routines. However, the current code LLMs still face serious challenges related to trustworthiness, as they can generate incorrect, insecure, or unreliable code. Recent exploratory studies find that it can be promising to detect such risky outputs by analyzing LLMs' internal states, akin to how the human brain unconsciously recognizes its own mistakes. Yet, most of these approaches are limited to narrow sub-domains of LLM operations and fall short of achieving industry-level scalability and practicability. To address these challenges, in this paper, we propose PtTrust, a two-stage risk assessment framework for code LLM based on internal state pre-training, designed to integrate seamlessly with the existing infrastructure of software companies. The core idea is that the risk assessment framework could also undergo a pre-training process similar to LLMs. Specifically, PtTrust first performs unsupervised pre-training on large-scale unlabeled source code to learn general representations of LLM states. Then, it uses a small, labeled dataset to train a risk predictor. We demonstrate the effectiveness of PtTrust through fine-grained, code line-level risk assessment and demonstrate that it generalizes across tasks and different programming languages. Further experiments also reveal that PtTrust provides highly intuitive and interpretable features, fostering greater user trust. We believe PtTrust makes a promising step toward scalable and trustworthy assurance for code LLMs.