Risk Assessment Framework for Code LLMs via Leveraging Internal States

📄 arXiv: 2504.14640v1 📥 PDF

作者: Yuheng Huang, Lei Ma, Keizaburo Nishikino, Takumi Akazaki

分类: cs.SE, cs.AI, cs.CL

发布日期: 2025-04-20

备注: To appear in the 33rd ACM International Conference on the Foundations of Software Engineering (FSE Companion'25 Industry Track), June 23-28, 2025, Trondheim, Norway. This work was supported by Fujitsu Limited

期刊: Companion Proceedings of the 33rd ACM International Conference on the Foundations of Software Engineering (FSE'25 Industry Track), June 2025

DOI: 10.1145/3696630.3728566

💡 一句话要点

PtTrust：利用内部状态预训练的代码大语言模型风险评估框架

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 代码大语言模型 风险评估 内部状态 预训练 软件工程 代码安全 可信AI

📋 核心要点

  1. 现有代码LLM面临生成不安全、不正确代码的风险，缺乏可信度，且现有风险检测方法可扩展性和实用性不足。
  2. PtTrust框架通过预训练学习LLM内部状态的通用表示，并使用少量标注数据训练风险预测器，从而进行风险评估。
  3. 实验表明，PtTrust能有效进行细粒度代码行级别风险评估，并可推广到不同任务和编程语言，提供直观可解释的特征。

📝 摘要（中文）

大型语言模型（LLMs）的预训练范式是其成功的关键，代码LLMs凭借其先进的编码能力对软件工程产生了巨大影响，并有成为开发者日常工作必备工具的趋势。然而，当前的code LLMs仍然面临着与可信度相关的严峻挑战，因为它们可能会生成不正确、不安全或不可靠的代码。最近的探索性研究发现，通过分析LLMs的内部状态来检测此类风险输出是有希望的，类似于人脑无意识地识别自身错误的方式。然而，这些方法大多局限于LLM操作的狭窄子领域，并且无法实现工业级的可扩展性和实用性。为了应对这些挑战，本文提出PtTrust，一个基于内部状态预训练的code LLM两阶段风险评估框架，旨在与软件公司现有的基础设施无缝集成。核心思想是风险评估框架也可以经历类似于LLMs的预训练过程。具体来说，PtTrust首先对大规模未标记的源代码执行无监督预训练，以学习LLM状态的通用表示。然后，它使用一个小的标记数据集来训练风险预测器。我们通过细粒度的代码行级别风险评估证明了PtTrust的有效性，并证明它可以推广到不同的任务和编程语言。进一步的实验还表明，PtTrust提供了高度直观和可解释的特征，从而增强了用户的信任。我们相信PtTrust为code LLMs的可扩展和可信赖的保证迈出了有希望的一步。

🔬 方法详解

问题定义：当前的代码大语言模型（Code LLM）在软件工程领域应用广泛，但其生成的代码可能存在不安全、不正确等问题，导致可信度不足。现有的风险检测方法通常局限于特定领域，难以扩展到工业级应用，缺乏实用性。因此，需要一种可扩展、通用的风险评估框架，以提高Code LLM生成代码的可靠性。

核心思路：PtTrust的核心思路是借鉴LLM自身的预训练模式，对风险评估框架进行预训练。通过学习LLM内部状态的通用表示，PtTrust能够更好地理解代码的潜在风险。这种方法类似于人脑识别自身错误的方式，通过分析LLM的“思考过程”来预测代码的风险。

技术框架：PtTrust是一个两阶段的风险评估框架。第一阶段是无监督预训练，使用大规模未标记的源代码来训练模型，学习LLM内部状态的通用表示。第二阶段是风险预测器训练，使用少量标记数据集来训练一个风险预测模型，该模型基于第一阶段学习到的表示来预测代码的风险。该框架可以无缝集成到现有的软件公司基础设施中。

关键创新：PtTrust的关键创新在于将预训练的思想应用于风险评估框架。通过对LLM内部状态进行预训练，PtTrust能够学习到更丰富的代码语义信息，从而更准确地预测代码的风险。与现有方法相比，PtTrust具有更好的可扩展性和通用性，可以应用于不同的任务和编程语言。

关键设计：PtTrust的关键设计包括：1) 使用Transformer架构来学习LLM内部状态的表示；2) 设计合适的损失函数，以鼓励模型学习到有意义的风险特征；3) 使用细粒度的代码行级别标注数据来训练风险预测器，以提高预测的准确性。具体的预训练和微调策略未知，但强调了无监督预训练的重要性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PtTrust能够有效地进行细粒度的代码行级别风险评估，并且可以推广到不同的任务和编程语言。PtTrust还提供了高度直观和可解释的特征，增强了用户的信任。具体的性能数据和提升幅度未知，但强调了其通用性和可解释性。

🎯 应用场景

PtTrust可应用于软件开发流程的各个阶段，例如代码审查、静态分析和自动化测试。它可以帮助开发者及早发现并修复潜在的代码风险，提高软件的安全性、可靠性和可维护性。该研究有望推动代码LLM在工业界的广泛应用，并提升软件开发的整体质量。

📄 摘要（原文）

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.