A Unified Multi-Task Vision Transformer for AI-Generated Image Detection and Generator Attribution

Abstract

The rapid advancement of diffusion-based generative models such as Stable Diffusion, DALL•E 3, and Midjourney has significantly reduced the visual distinction between synthetic and real imagery. While most existing work focuses primarily on binary AI-image detection, relatively limited attention has been given to unified detection and generator attribution within a single framework. This paper presents a multi-task Vision Transformer (ViT) approach that jointly performs AI-generated image detection and source attribution using a shared transformer backbone with task-specific classification heads. This multi-task approach allows the model to understand general characteristics of AI-generated images while also recognizing patterns linked to particular generators. To evaluate reliability under practical conditions, a robustness-aware evaluation protocol is introduced that measures prediction variance and decision-flip rate under typical image transformations such as JPEG compression and resizing. In addition, a gradient-free occlusion-based explainability method is integrated to support interpretable and CPU-compatible inference. Experimental results demonstrate that the proposed framework achieves high detection reliability and strong attribution performance while maintaining stable behaviour under common image transformations. The overall system provides a unified and deployment-oriented approach for AI-image authenticity assessment in real-world environments.

Introduction

Recent advances in diffusion-based image generation models like Stable Diffusion and DALL·E 3 have made synthetic images highly realistic, making manual detection unreliable. This has led to growing interest in automated AI-image detection within digital forensics. Existing methods—such as convolutional classifiers, frequency-domain analysis, and model fingerprinting—can detect subtle artifacts left by generative models, but they mainly focus on binary classification (real vs. fake) and often overlook source attribution (identifying which model generated the image). Additionally, many detectors lack robustness to real-world transformations like compression and resizing, limiting practical deployment.

Vision Transformers (ViTs) offer a promising alternative due to their ability to capture global dependencies across images. However, their use for combined detection, attribution, and robustness analysis remains underexplored. To address these gaps, the paper proposes a multi-task ViT-based framework that simultaneously performs authenticity detection and generator attribution using a shared representation. It also incorporates robustness evaluation under perturbations and a lightweight interpretability method suitable for low-resource environments.

The model is trained on the Defactify AI-Generated Image Veracity Dataset, which includes real and synthetic images from multiple generators. The task is formulated as a dual prediction problem: (1) binary classification of real vs. AI-generated images, and (2) multi-class classification of the source model. Images are preprocessed, divided into patches, and passed through a ViT backbone, followed by two task-specific heads. This unified approach aims to deliver a more robust, interpretable, and deployable solution for real-world AI image authenticity assessment.

Conclusion

This paper presented a multi-task Vision Transformer framework for unified AI-generated image detection and generator attribution. By combining both tasks within a shared representation, the proposed approach demonstrates that authenticity detection and source identification can be performed together without sacrificing performance. Experimental results show high detection reliability and strong attribution accuracy, supported by stable training behaviour and robustness under typical image transformations. The study also highlights the importance of moving beyond simple binary detection toward more informative authenticity analysis that includes generator-level insights. The integration of robustness evaluation and interpretability further strengthens the practical usefulness of the framework in real-world scenarios. The proposed system offers a balanced and deployable solution for modern synthetic media forensics. As generative models continue to improve in visual realism, unified approaches that combine detection, attribution, and reliability assessment will play an increasingly important role in maintaining trust in digital visual content.

References

[1] R. Rombach, A. Blattmann, D. Lorenz, P. Esser and B. Ommer, \"High-Resolution Image Synthesis with Latent Diffusion Models,\" 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), New Orleans, LA, USA, 2022, pp. 10674-10685, doi: 10.1109/CVPR52688.2022.01042. [2] OpenAI, “DALL•E 3 Technical Report,” 2023. Available: https://openai.com/research/dall-e-3 [3] H. Farid, \"Image forgery detection,\" in IEEE Signal Processing Magazine, vol. 26, no. 2, pp. 16-25, March 2009, doi: 10.1109/MSP.2008.931079. [4] L. Guarnera, O. Giudice and S. Battiato, \"Fighting Deepfake by Exposing the Convolutional Traces on Images,\" in IEEE Access, vol. 8, pp. 165085-165098, 2020, doi: 10.1109/ACCESS.2020.3023037. arXiv:2008.04095 [5] F. Marra, D. Gragnaniello, L. Verdoliva and G. Poggi, \"Do GANs Leave Artificial Fingerprints?,\" 2019 IEEE Conference on Multimedia Information Processing and Retrieval (MIPR), San Jose, CA, USA, 2019, pp. 506-511, doi: 10.1109/MIPR.2019.00103. arXiv:1812.11842 [6] R. Durall, M. Keuper, and J. Keuper, “Watch Your Up-Convolution: CNN Based Generative Deep Neural Networks Are Failing to Reproduce Spectral Distributions,” Proc. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020. arXiv:2003.01826 [7] A. Dosovitskiy et al., “An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale,” International Conference on Learning Representations (ICLR), 2021. arXiv:2010.11929 [8] R. Roy et al., “Defactify: AI-Generated Image Veracity Dataset,” Hugging Face, 2024. [Online]. Available: https://huggingface.co/datasets/Rajarshi-Roy-research/Defactify_Image_Dataset

Copyright

Copyright © 2026 Dineshkumar Kumawat, Asst. Prof. Mamta Ramesh Devre. This is an open access article distributed under the Creative Commons Attribution License, which permits unrestricted use, distribution, and reproduction in any medium, provided the original work is properly cited.