Supervised Contrastive Learning for Few-Shot AI-Generated Image Detection and Attribution

TL;DR

This paper introduces a two-stage framework using supervised contrastive learning and few-shot k-NN classification to detect and attribute AI-generated images, achieving high accuracy and robustness against unseen generators.

Contribution

It presents a novel contrastive learning-based detection method combined with few-shot learning for source attribution, improving generalization to unseen generative models.

Findings

Achieves 91.3% detection accuracy with only 150 images per class.

Outperforms existing methods by 5.2 percentage points in detection accuracy.

Improves AUC and OSCR by 14.70% and 4.27% respectively in attribution tasks.

Abstract

The rapid advancement of generative artificial intelligence has enabled the creation of synthetic images that are increasingly indistinguishable from authentic content, posing significant challenges for digital media integrity. This problem is compounded by the accelerated release cycle of novel generative models, which renders traditional detection approaches (reliant on periodic retraining) computationally infeasible and operationally impractical. This work proposes a novel two-stage detection framework designed to address the generalization challenge inherent in synthetic image detection. The first stage employs a vision deep learning model trained via supervised contrastive learning to extract discriminative embeddings from input imagery. Critically, this model was trained on a strategically partitioned subset of available generators, with specific architectures withheld from training to rigorously ablate cross-generator generalization capabilities. The second stage utilizes a k-nearest neighbors (k-NN) classifier operating on the learned embedding space, trained in a few-shot learning paradigm incorporating limited samples from previously unseen test generators. With merely 150 images per class in the few-shot learning regime, which are easily obtainable from current generation models, the proposed framework achieves an average detection accuracy of 91.3%, representing a 5.2 percentage point improvement over existing approaches . For the source attribution task, the proposed approach obtains improvements of of 14.70% and 4.27% in AUC and OSCR respectively on an open set classification context, marking a significant advancement toward robust, scalable forensic attribution systems capable of adapting to the evolving generative AI landscape without requiring exhaustive retraining protocols.