LLMSniffer: Detecting LLM-Generated Code via GraphCodeBERT and Supervised Contrastive Learning

# Detecting LLM-Generated Code via GraphCodeBERT and Supervised Contrastive Learning
Source: https://arxiv.org/html/2604.16058
Mahir Labib Dihan Abir Muhtasim Department of Computer Science and Engineering Bangladesh University of Engineering and Technology (BUET) Dhaka, Bangladesh {mahirlabibdihan, auntor505}@gmail.com

###### Abstract

The rapid proliferation of Large Language Models (LLMs) in software development has made distinguishing AI-generated code from human-written code a critical challenge with implications for academic integrity, code quality assurance, and software security. We present LLMSniffer, a detection framework that fine-tunes GraphCodeBERT using a two-stage supervised contrastive learning pipeline augmented with comment removal preprocessing and an MLP classifier. Evaluated on two benchmark datasets—GPTSniffer and Whodunit—LLMSniffer achieves substantial improvements over prior baselines: accuracy increases from 70% to 78% on GPTSniffer (F1: 68% → 78%) and from 91% to 94.65% on Whodunit (F1: 91% → 94.64%). t-SNE visualizations confirm that contrastive fine-tuning yields well-separated, compact embeddings. We release our model checkpoints, datasets, codes and a live interactive demo to facilitate further research.

## 1 Introduction

The emergence of capable code-generating LLMs—including GitHub Copilot, ChatGPT, and Code Llama—has transformed how software is written. Developers increasingly rely on these models for routine coding tasks, bug fixes, and algorithm implementation. However, this shift introduces new challenges: plagiarism detection systems in educational settings must account for AI-generated submissions, code review pipelines need provenance information, and security audits may require distinguishing model-generated patterns from human idioms.

Educators are increasingly concerned about LLM usage in programming education. Existing solutions such as GPTZero or DetectGPT are primarily tailored for general text and often struggle with code content due to the unique structural and syntactic properties of programming languages. Detecting LLM-generated code is fundamentally a binary classification task: given a code snippet x ∈ C, learn a function f: C → {0,1} indicating whether the code is human-written (0) or LLM-generated (1).

Recent efforts have focused on zero-shot detectors and ML-based detectors. We emphasize the latter, building on the GPTSniffer framework. We propose LLMSniffer, combining (1) GraphCodeBERT, which encodes code tokens and data-flow structure, with (2) supervised contrastive learning, which clusters same-class representations while separating different-class ones, and (3) a comment-stripping preprocessing step to remove potentially confounding stylistic metadata.

Our main contributions are:

- A two-stage contrastive fine-tuning strategy for GraphCodeBERT with comment removal preprocessing and an MLP classification head, achieving state-of-the-art results on two benchmarks.
- Detailed ablation analysis comparing a linear classifier baseline against our improved MLP-based architecture.
- Public model checkpoints, datasets on Hugging Face, an interactive Streamlit demo, and a video demonstration for full reproducibility.

## 2 Related Work

AI-Generated Text Detection. Detecting machine-generated text has received considerable attention since GPT-2. Early approaches used statistical signals such as perplexity and likelihood ratios, while more recent work focuses on fine-tuned classifiers and watermarking. These methods target natural language and do not transfer directly to source code due to syntactic and structural differences.

LLM-Generated Code Detection. Nguyen et al. introduced GPTSniffer, one of the first dedicated benchmarks for detecting ChatGPT-generated code, using a CodeBERT encoder trained with binary cross-entropy loss. Nguyen et al. presented Whodunit, focusing on GPT-4-generated Python solutions to CodeChef problems and demonstrating that code-specific features substantially improve detection. Idialu et al. further explored zero-shot and ML-based pipelines on competitive programming benchmarks. Our work builds directly on these benchmarks and surpasses both baselines.

Code Representation Learning. Pre-trained models such as CodeBERT, GraphCodeBERT, and CodeT5 advance code understanding by encoding token sequences alongside structural information (ASTs, data-flow graphs). We leverage GraphCodeBERT for its demonstrated superiority on code search and clone detection.

Contrastive Learning. Supervised contrastive learning extends self-supervised contrastive objectives to the labeled setting. It has been applied to NLP fine-tuning and code clone detection, but not, to our knowledge, to AI-generated code detection prior to this work.

## 3 Background

### 3.1 GraphCodeBERT

GraphCodeBERT is a transformer-based model pre-trained on code–documentation pairs from CodeSearchNet. Unlike CodeBERT, it augments the self-attention mechanism with a data-flow graph that captures variable-level semantic dependencies (e.g., where a variable is computed and where it is used). Two additional pre-training objectives—edge prediction and node alignment—encourage the model to internalize these structural relationships. We use the [CLS] token representation as a fixed-dimensional embedding h ∈ ℝ^768 for downstream classification.

### 3.2 Supervised Contrastive Learning

Given a batch of N labeled pairs {(x_i, y_i)}_{i=1}^N, the supervised contrastive loss is:

L_SupCon = ∑_{i=1}^N (-1/|P(i)|) ∑_{p∈P(i)} log(e^{z_i·z_p/τ} / ∑_{a∈A(i)} e^{z_i·z_a/τ})

where z_i = ℓ_2-normalize(MLP(h_i)) is the projected representation, P(i) = {p ≠ i : y_p = y_i} is the set of same-class positives, A(i) = {a ≠ i} is all other examples, and τ is a temperature hyperparameter. This objective directly optimizes for inter-class separation and intra-class compactness in embedding space.

## 4 Methodology

Figure 1: Overview of LLMSniffer. Stage 1 trains the encoder and projection head with supervised contrastive loss. Stage 2 freezes the encoder and trains an MLP classifier with binary cross-entropy loss.

### 4.1 Preprocessing: Comment Removal

A key design decision in LLMSniffer is removing code comments before encoding. LLM-generated code often includes verbose, explanatory comments that reflect LLM writing style rather than algorithmic structure. Stripping comments ensures the model learns to distinguish based on syntactic and structural code properties, reducing the risk of exploiting superficial stylistic cues and improving generalization.

### 4.2 Model Architecture

LLMSniffer consists of three components (Figure 1):

(1) GraphCodeBERT Encoder. After comment removal, the code snippet is tokenized and fed along with its data-flow graph edges into GraphCodeBERT. The [CLS] embedding h ∈ ℝ^768 serves as the code representation.

(2) Projection Head. A two-layer MLP with ReLU activation maps h to z ∈ ℝ^128 for contrastive learning. The head is discarded at inference time.

(3) MLP Classification Head. A multi-layer perceptron with sigmoid activation maps h to a binary prediction (human-written or AI-generated). We find that this MLP head outperforms the simpler linear classifier used in prior work.

### 4.3 Training Procedure

Training proceeds in two stages.

Stage 1: Contrastive Pre-training. We fine-tune the encoder and projection head using L_SupCon (Equation 1) with temperature τ = 0.07. Batches are constructed with balanced classes to maximize the number of positive pairs per anchor.

Stage 2: Classification Fine-tuning. The projection head is discarded, the MLP classification head is attached, and the full model is fine-tuned end-to-end with binary cross-entropy loss. The contrastive pre-training step warms up the encoder to produce discriminative representations before the classifier is trained, following the approach of Khosla et al.

### 4.4 Implementation Details

We initialize from the microsoft/graphcodebert-base checkpoint (HuggingFace). Both training stages use AdamW with learning rate 2×10^{-5} and linear warmup over 10% of training steps. Maximum sequence length is 512 tokens. Batch sizes are 32 for contrastive training and 64 for classification fine-tuning. The projection head uses hidden dimension 512 and output dimension 128. All experiments are conducted on a single GPU.

## 5 Experimental Setup

### 5.1 Datasets

GPTSniffer is a dataset of code snippets with binary labels indicating human- or ChatGPT-authorship, spanning multiple programming languages and competitive programming problem types. We use the official split: 600 human + 600 ChatGPT samples for training, and 131 human + 142 ChatGPT samples for testing. We use the dataset hosted at https://huggingface.co/datasets/mahirlabibdihan/GPTSniffer. Crucially, the test set comes from a *completely different domain* than the training set, making this a challenging cross-domain evaluation.

Whodunit is a dataset of Python solutions to CodeChef problems, with code generated by GPT-4 and human-written counterparts. We use: 639 human + 639 ChatGPT samples for training, and 159 human + 159 ChatGPT samples for testing. The dataset is hosted at https://huggingface.co/datasets/mahirlabibdihan/Whodunit.

Pre-training Data. Both datasets are augmented with CodeSearchNet (human-written code) and CodeNet (human competitive programming solutions) for encoder warming.

### 5.2 Baselines and Ablations

We compare against two baselines and an intermediate ablation:

- GPTSniffer: CodeBERT fine-tuned with binary cross-entropy loss.
- Whodunit: Feature-based XGBoost classifier with 140 hand-crafted code features.
- LLMSniffer (Linear): Our Stage 1+2 pipeline with a linear classification head instead of an MLP (no comment removal).
- LLMSniffer (Ours): Full pipeline with comment removal and MLP classification head.

### 5.3 Evaluation Metrics

We report Accuracy, Precision, Recall, and F1 Score on each benchmark's test set. All metrics are macro-averaged across the two classes.

## 6 Results

### 6.1 GPTSniffer Benchmark

Table 1 compares all systems on the GPTSniffer test set.

Table 1: Results on the GPTSniffer test set (macro-averaged). Metrics are reported as proportions. The linear classifier variant shows no improvement over the baseline; the MLP with comment removal achieves +8 pp accuracy and +10 pp F1.

The linear classifier variant fails to improve over the GPTSniffer baseline, achieving the same overall accuracy (0.70). This highlights that contrastive pre-training alone is insufficient—the classification head architecture and comment removal are both important components. Our full model (MLP + comment removal) achieves 0.78 accuracy and 0.78 F1, improving both precision and recall simultaneously, unlike the prior baseline which suffered from a severe precision–recall imbalance (high precision, low recall for the human class).

### 6.2 Whodunit Benchmark

Table 2 presents results on the Whodunit benchmark.

Table 2: Results on the Whodunit test set. Consistent gains across all metrics.

On Whodunit, LLMSniffer achieves consistent improvements across all metrics (+3.6 pp), reflecting the dataset's balanced structure and the robustness of our approach to GPT-4 generated code in Python.

### 6.3 Embedding Visualization

Figure 2 shows t-SNE projections of [CLS] embeddings on the GPTSniffer test set for the baseline and LLMSniffer.

Figure 2: t-SNE of test set [CLS] embeddings (GPTSniffer benchmark). Left (GPTSniffer baseline): The two classes (human=blue, AI=orange) are highly interleaved with no clear decision boundary. Right (LLMSniffer): Contrastive fine-tuning produces a compact, isolated AI-generated cluster (right) that is linearly separable from the human-written cluster, directly explaining the improved classification performance.

The baseline model produces heavily overlapping clusters for the two classes. LLMSniffer, by contrast, forms a compact, well-isolated cluster of AI-generated snippets that is cleanly separable from human-written code. This qualitative result directly explains the quantitative improvements and confirms that supervised contrastive learning is the core driver of discriminability.

### 6.4 Ablation Study