Model-Agnostic Correctness Assessment for LLM-Generated Code via Dynamic Internal Representation Selection
Thanh Trong Vu, Tuan-Dung Bui, Thu-Trang Nguyen, Son Nguyen, Hieu Dinh Vo
TL;DR
This work introduces AutoProbe, a model-agnostic framework for evaluating the correctness of LLM-generated code by dynamically selecting informative internal representations through an attention-based selector. By sampling a compact yet diverse set of hidden states across layers and token positions, AutoProbe learns to emphasize signals most predictive of correctness and feeds them to a probing classifier. It demonstrates superior performance over white-box baselines (and near parity with an Oracle) across compilability, functionality, and security, while maintaining efficiency via boundary-aware sampling and lightweight aggregation. The approach is validated on multiple datasets and six code LLMs, showing robust cross-model, cross-language, and cross-task effectiveness with favorable efficiency characteristics, highlighting its practical applicability for real-world code generation QA and safety pipelines.
Abstract
Large Language Models (LLMs) have demonstrated impressive capabilities in code generation and are increasingly integrated into the software development process. However, ensuring the correctness of LLM-generated code remains a critical concern. Prior work has shown that the internal representations of LLMs encode meaningful signals for assessing code correctness. Nevertheless, the existing methods rely on representations from pre-selected/fixed layers and token positions, which could limit its generalizability across diverse model architectures and tasks. In this work, we introduce AUTOPROBE, a novel model-agnostic approach that dynamically selects the most informative internal representations for code correctness assessment. AUTOPROBE employs an attention-based mechanism to learn importance scores for hidden states, enabling it to focus on the most relevant features. These weighted representations are then aggregated and passed to a probing classifier to predict code correctness across multiple dimensions, including compilability, functionality, and security. To evaluate the performance of AUTOPROBE, we conduct extensive experiments across multiple benchmarks and code LLMs. Our experimental results show that AUTOPROBE consistently outperforms the baselines. For security assessment, AUTOPROBE surpasses the state-of-the-art white-box approach by 18%. For compilability and functionality assessment, AUTOPROBE demonstrates its highest robustness to code complexity, with the performance higher than the other approaches by up to 19% and 111%, respectively. These findings highlight that dynamically selecting important internal signals enables AUTOPROBE to serve as a robust and generalizable solution for assessing the correctness of code generated by various LLMs.