TriDetect: Beyond Binary Classification for Generalized AI-generated Image Detection

Official code release for:

Beyond Binary Classification: A Semi-supervised Approach to Generalized AI-generated Image Detection Hong-Hanh Nguyen-Le, Van-Tuan Tran, Thuc D. Nguyen, Nhien-An Le-Khac AAAI 2026

TriDetect (Triarchy Detector) is a semi-supervised AI-generated image detector that enhances binary real/fake classification by simultaneously discovering the latent architectural patterns that separate GAN-generated and diffusion-generated images.

---

Method in brief

TriDetect is built on a frozen CLIP ViT-L/14 vision encoder, fine-tuned with LoRA, and a 3-way head (1 real class + 2 fake clusters). It is trained with three losses:

Loss	Role
L_binary	Cross-entropy real/fake classification (fake prob = sum of the two fake clusters)
L_assignment	Swapped-prediction clustering using Sinkhorn-Knopp balanced assignments across two views
L_consistency	L2 stability of the balanced assignments across the two views

$\mathcal{L}{\text{total}} = \beta \cdot \mathcal{L}{\text{binary}} + (1-\beta) \cdot \mathcal{L}{\text{cluster}}$, where $\mathcal{L}{\text{cluster}} = \omega_1 \cdot \mathcal{L}{\text{assignment}} + \omega_2 \cdot \mathcal{L}{\text{consistency}}$.

The two views come from a contrastive augmentation pipeline, which lets the model learn fundamental architectural distinctions instead of image-specific statistics, improving cross-generator generalization.

---

Repository layout

.
├── training/
│   ├── train.py                       # training entry point
│   ├── test.py                        # evaluation entry point
│   ├── logger.py
│   ├── config/
│   │   ├── train_config.yaml          # global training config (label dict, wandb, ...)
│   │   ├── test_config.yaml           # global testing config
│   │   └── detector/tridetect.yaml    # TriDetect model + training hyper-parameters
│   ├── detectors/
│   │   ├── tridetect_detector.py      # the TriDetect model
│   │   ├── base_detector.py
│   │   └── utils/lora_utils.py        # LoRA adapters for CLIP
│   ├── dataset/{abstract_dataset,genimage_dataset,albu}.py
│   ├── trainer/trainer.py
│   ├── loss/{cross_entropy_loss,abstract_loss_func}.py
│   ├── metrics/                       # AUC / ACC / EER / AP
│   └── optimizor/                     # SAM, linear-LR schedulers
└── preprocessing/                     # scripts to build dataset JSON files
    ├── generate_genimage_json.py
    ├── generate_aigc_json.py
    ├── generate_wildfake_json.py
    ├── generate_df40_json.py
    └── create_chameleon_json.py

---

Installation

conda env create -f env.yaml
conda activate tridetect
# or, with pip:
pip install -r requirements.txt

PyTorch 2.3.1 on Python 3.9 (as reported in the paper). A single NVIDIA H100 (94 GB) was used for the paper experiments; any modern GPU with $\geq 24$ GB should work for TriDetect itself (LoRA keeps trainable parameters small).

---

Data preparation

TriDetect is trained on the BigGAN and Stable Diffusion v1.4 subsets of GenImage and evaluated on five datasets:

Dataset	Use
GenImage	training (BigGAN + SDv1.4) & cross-generator test
AIGCDetectBenchmark	test (16 generators)
WildFake	test (with degradations)
Chameleon	test (in-the-wild)
DF40	test (40 generators)

Place the raw datasets under ./datasets/ and build the JSON index files that the data loaders expect:

cd preprocessing
python generate_genimage_json.py
python generate_aigc_json.py
python generate_wildfake_json.py
python generate_df40_json.py
python create_chameleon_json.py

The generated JSONs are written to ./preprocessing/dataset_json/. Each generator script contains its own path assumptions at the top — adjust them to match where you stored the raw data.

---

Training

# single GPU (defaults from tridetect.yaml: BigGAN + SDv1.4 -> Chameleon, GenImage-ADM)
bash train.sh
# or
python training/train.py --detector_path ./training/config/detector/tridetect.yaml

# multi-GPU (DDP)
python -m torch.distributed.launch --nproc_per_node=4 \
    training/train.py --detector_path ./training/config/detector/tridetect.yaml --ddp

Key hyper-parameters (already set in tridetect.yaml, matching the paper):

Parameter	Value
Backbone	CLIP ViT-L/14 (vision, 1024-d)
LoRA	$r=16$, $\alpha=32$, on q_proj,k_proj
Fake clusters K	2
Sinkhorn $\varepsilon$ / iters	0.05 / 3
$\beta$ (binary weight)	0.7
$\omega_1$ (assignment) / $\omega_2$ (consistency)	1.0 / 0.1
Optimizer	Adam, lr=2e-4, betas=(0.9,0.95), wd=1e-4
Batch size / epochs	128 / 5
Seed	1024

Checkpoints, features and TensorBoard logs are written under ./logs/experiments/.

---

Evaluation

python training/test.py \
    --detector_path ./training/config/detector/tridetect.yaml \
    --test_dataset "Chameleon" "GenImage-ADM" \
    --weights_path ./training/weights/tridetect_best.pth

A CSV summary (tridetect_test_results_<timestamp>.csv) is written next to the checkpoint.

---

Main results (from the paper)

Dataset	Metric	TriDetect	Best baseline
GenImage	AUC (avg)	0.9882	0.9815 (Effort)
AIGCDetectBenchmark	AUC (avg)	0.9869	0.9783 (Effort)
WildFake	ACC (avg)	0.8254	0.7522 (Effort)
DF40	ACC (avg)	0.8429	0.8177 (Effort)
Chameleon	AUC / EER	0.8935 / 0.1843	0.8371 / 0.2428 (Effort)

See the paper for the full per-generator breakdown, ablations (loss components, $\beta$, number of clusters) and the theoretical analysis.

---

Citation

@inproceedings{nguyenle2026tridetect,
  title     = {Beyond Binary Classification: A Semi-supervised Approach to
               Generalized AI-generated Image Detection},
  author    = {Nguyen-Le, Hong-Hanh and Tran, Van-Tuan and Nguyen, Thuc D. and
               Le-Khac, Nhien-An},
  booktitle = {Proceedings of the AAAI Conference on Artificial Intelligence},
  year      = {2026}
}

---

Acknowledgments

We thank the authors of DeepfakeBench (Yan et al., NeurIPS 2023); the unified benchmark and comparison in this work are built upon DeepfakeBench.

License

This work is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License.

See LICENSE for the full text.