TokenGated-CodeBERT With Lightweight Attention-Based Token Selection for Software Defect Prediction

Abstract

Class imbalance poses a key challenge in defective code module detection, severely limiting the performance of deep learning models. To address this issue, we propose TokenGated-CodeBERT, a new methodology designed for scenarios with severe class imbalance. Constructed on the basis of vanilla CodeBERT, the model adopts a lightweight token gating module to raise the weights of defective minority tokens and lower the weights of non-defective majority tokens. It dynamically fuses token embeddings based on defect correlation, emphasizing fault-related tokens and eliminating noise to enhance minority defect detection in feature learning. We conduct extensive validation on the PROMISE dataset for cross-project and cross-version defect prediction tasks, with supplementary tests on the Defects4J dataset for fine-grained code snippet prediction to comprehensively examine the generalization performance of the newly designed structural module. Experimental results illustrate that our method achieves competitive Matthews Correlation Coefficient (MCC) performance against all baseline approaches, demonstrating its effectiveness in class-imbalanced scenarios. An ablation study confirms that our token-level gating mechanism is the key driver of this performance gain, as it more effectively captures defect-related signals distributed throughout the code sequence. Our framework supports full end-to-end training, keeps the original CodeBERT encoder intact, requires no additional auxiliary optimization objectives, and features a lightweight architecture.