Unveiling Project-Specific Bias in Neural Code Models

TL;DR

This work investigates why neural code models trained on intra-project IID data fail to generalize to inter-project OOD and resist adversarial attacks. It introduces Cond-Idf as a quantitative measure of token-label relatedness and project-specificness, and presents Batch Partition Regularization (BPR), a bias-mitigation framework built on invariant risk minimization that uses dataset unshuffling and in-batch regularization to enforce consistent latent logic across similarly evidenced samples. Empirical results on Type Inference and Vulnerability Detection show that BPR improves inter-project generalization and adversarial robustness while preserving IID accuracy, with Cond-Idf correlating with attribution patterns and providing interpretable insights into model bias. The approach offers a practical path to more robust and explainable neural code models and suggests directions for extending bias-aware training across additional software-analysis tasks.

Abstract

Deep learning has introduced significant improvements in many software analysis tasks. Although the Large Language Models (LLMs) based neural code models demonstrate commendable performance when trained and tested within the intra-project independent and identically distributed (IID) setting, they often struggle to generalize effectively to real-world inter-project out-of-distribution (OOD) data. In this work, we show that this phenomenon is caused by the heavy reliance on project-specific shortcuts for prediction instead of ground-truth evidence. We propose a Cond-Idf measurement to interpret this behavior, which quantifies the relatedness of a token with a label and its project-specificness. The strong correlation between model behavior and the proposed measurement indicates that without proper regularization, models tend to leverage spurious statistical cues for prediction. Equipped with these observations, we propose a novel bias mitigation mechanism that regularizes the model's learning behavior by leveraging latent logic relations among samples. Experimental results on two representative program analysis tasks indicate that our mitigation framework can improve both inter-project OOD generalization and adversarial robustness, while not sacrificing accuracy on intra-project IID data.

Figures (4)

• Figure 1: Illustrative examples of attribution vectors in terms of integrated gradient for cases of type inference and vulnerability detection. The shades of green indicate the weight value of the respective token in the attribution vector.
• Figure 2: Overview of the proposed bias mitigation method. Training samples $x\in \mathcal{X}$ are first embedded with $\mu$ and sorted in terms of the similarity measure $\kappa$, then batchified and partitioned into multiple environments$\varepsilon \in \mathcal{E}$ according to their labels and similarity scores. Finally, the batch partition regularization (BPR) loss is computed along with the classification loss.
• Figure 3: For the figures in the first row: the lower distribution in each figure is the ranked integrated gradient distribution. The upper red dashed line denotes the polynomial regression approximation of the corresponding Cond-Idf distribution in terms of token index. CB and GCB denote the CodeBERT and GraphCodeBERT models. For the figures in the second row: The blue and orange lines denote the polynomial regression approximation of the Cond and Idf distribution in terms of token index respectively.
• Figure 5: Illustrated example of applying BPR compared to the baseline methods for the vulnerability detection task.