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

TL;DR

This paper tackles the problem of detecting and attributing AI-generated images in a rapidly evolving generative landscape. It introduces a two-stage framework that first learns discriminative embeddings via supervised contrastive learning on a MambaVision backbone, and then performs few-shot, non-parametric attribution with a k-NN classifier in the embedding space. The approach demonstrates strong generalization to unseen generators, achieving high detection accuracy (around 91.3%) and notable gains in open-set attribution metrics (AUC and OSCR), while maintaining robustness and explainability through latent-space analyses. The work offers a scalable forensic solution that minimizes retraining needs as new generators appear, with practical implications for digital media integrity and 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.

Figures (6)

• Figure 1: Comparison of real images (top row) and synthetic images (bottom row). Each column shows images from a different synthetic image generator. Code and models developed for this work are available at https://github.com/JaimeAlvarez18/SupConLoss_fake_image_detection
• Figure 2: Visual comparison of supervised learning, contrastive learning and supervised contrastive learning.
• Figure 3: Full architecture of the proposed approach. $^a$ Indicates that training samples were used. $^b$ Denotes that testing samples were used.
• Figure 4: Accuracies obtained in the source attribution task for all trained models, evaluated under different numbers of few-shot samples per generator.
• Figure 5: Latent space visualizations for the lower-bound (ESB1) and upper-bound (ES5) embedding extractor models.