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

1. Problem Statement

The rapid adoption of Large Language Models (LLMs) like GitHub Copilot, ChatGPT, and Code Llama has transformed software development. However, this shift creates critical challenges:

• Academic Integrity: Difficulty in distinguishing AI-generated submissions from student work.
• Code Quality & Security: The need to identify provenance for code reviews and security audits to detect potential vulnerabilities or non-human idioms.
• Limitations of Existing Tools: General text detectors (e.g., GPTZero, DetectGPT) fail on code due to unique syntactic and structural properties. Existing code-specific detectors often rely on simple binary cross-entropy loss or hand-crafted features, which may not fully capture the structural nuances of LLM-generated code.

The core task is defined as a binary classification problem: given a code snippet $x$, determine if it is human-written ($0$) or LLM-generated ($1$).

2. Methodology: LLMSniffer

The authors propose LLMSniffer, a detection framework that integrates three key components: GraphCodeBERT, Supervised Contrastive Learning, and Comment Removal Preprocessing. The architecture follows a two-stage training pipeline.

A. Preprocessing: Comment Removal

A critical design choice is the removal of code comments before encoding.

• Rationale: LLMs often generate verbose, explanatory comments that reflect their specific writing style rather than the algorithmic logic.
• Goal: By stripping comments, the model is forced to learn distinctions based on syntactic and structural code properties rather than superficial stylistic cues, improving generalization.

B. Model Architecture

1. Encoder (GraphCodeBERT):
   • Unlike standard CodeBERT, GraphCodeBERT incorporates data-flow graphs alongside token sequences.
   • It captures variable-level semantic dependencies (e.g., where a variable is computed vs. used) via edge prediction and node alignment pre-training objectives.
   • The output is the [CLS] token embedding ($h \in \mathbb{R}^{768}$).
2. Projection Head (Stage 1):
   • A two-layer MLP with ReLU activation maps the embedding $h$ to a lower-dimensional space ($z \in \mathbb{R}^{128}$) specifically for contrastive learning.
   • This head is discarded during inference.
3. Classification Head (Stage 2):
   • A Multi-Layer Perceptron (MLP) with sigmoid activation maps the original embedding $h$ to a binary probability.
   • The authors found this MLP outperforms the linear classifiers used in prior baselines.

C. Training Procedure (Two-Stage)

• Stage 1: Contrastive Pre-training:
  • The encoder and projection head are fine-tuned using Supervised Contrastive Loss ($L_{SupCon}$).
  • Objective: Pull representations of the same class (human-human or AI-AI) closer together while pushing different classes apart in the embedding space.
  • Batches are balanced to maximize positive pairs per anchor.
• Stage 2: Classification Fine-tuning:
  • The projection head is removed, and the MLP classification head is attached.
  • The full model is fine-tuned end-to-end using Binary Cross-Entropy (BCE) loss.
  • The contrastive stage acts as a "warm-up" to create discriminative representations before the final classifier is trained.

3. Key Contributions

1. Novel Framework: Introduction of LLMSniffer, the first framework to combine GraphCodeBERT with supervised contrastive learning specifically for AI-generated code detection.
2. Ablation Insights: Demonstrated that contrastive learning alone is insufficient; the combination of comment removal, GraphCodeBERT, and an MLP classification head is necessary for state-of-the-art performance.
3. Resource Release: Public release of model checkpoints, datasets (GPTSniffer and Whodunit), code, and a live interactive Streamlit demo to ensure reproducibility.

4. Experimental Results

The model was evaluated on two benchmark datasets: GPTSniffer (multi-language, ChatGPT) and Whodunit (Python, GPT-4).

Performance on GPTSniffer

• Baseline (GPTSniffer/CodeBERT): 70% Accuracy, 68% F1.
• LLMSniffer (Linear Head): 70% Accuracy (No improvement).
• LLMSniffer (Full Model): 78% Accuracy, 78% F1.
  • Improvement: +8 percentage points in Accuracy and +10 percentage points in F1.
  • Observation: The full model balanced precision and recall, whereas the baseline suffered from high precision but low recall for human-written code.

Performance on Whodunit

• Baseline (Whodunit/XGBoost): 91% Accuracy, 91% F1.
• LLMSniffer (Full Model): 94.65% Accuracy, 94.64% F1.
  • Improvement: Consistent gains of ~3.6 percentage points across all metrics.

Embedding Visualization (t-SNE)

• Baseline: Showed heavily overlapping clusters for human and AI code, indicating poor separability.
• LLMSniffer: Produced a compact, well-isolated cluster for AI-generated code, clearly separable from human-written code. This geometric separation confirms that supervised contrastive learning is the primary driver of improved discriminability.

5. Significance and Future Work

• Significance: LLMSniffer establishes a new state-of-the-art for detecting LLM-generated code by moving beyond simple token-based classification to structural, graph-based representation learning enhanced by contrastive objectives. It proves that removing stylistic noise (comments) and optimizing the embedding space geometry are crucial for robust detection.
• Limitations:
  • Limited to snippets of 512 tokens (longer files require truncation).
  • Performance may degrade as LLMs evolve to produce more human-like code without periodic retraining.
  • Vulnerability to adversarial attacks (obfuscation) remains an open question.
• Future Directions:
  • Multi-class attribution (distinguishing between specific LLMs like GPT-4 vs. Claude).
  • Adversarial robustness testing.
  • Integration of AST features alongside data-flow graphs.
  • Extension to low-resource programming languages.

Conclusion

LLMSniffer effectively addresses the challenge of AI code detection by leveraging the structural understanding of GraphCodeBERT and the cluster-separation capabilities of supervised contrastive learning. The ablation study highlights that the synergy of comment removal, contrastive pre-training, and a non-linear MLP classifier is essential for achieving superior performance over existing baselines.