PE-Mamba: Bidirectional Selective Layer Aggregation for AI-Generated Image Detection

2026-08-08

AI Summary

This work addresses the growing challenge of distinguishing AI-generated images from real ones due to the rapid advancement of generative models, which has exposed limitations in existing detection methods. To tackle this issue, we propose a lightweight detection framework built upon a pre-trained PE-Core Vision Transformer. Our approach introduces three key mechanisms—Bidirectional Selective Aggregation (BSA), Softmax-Weighted Aggregation (SWA), and Sigmoid-Gated Fusion (SGA)—to enable evidence accumulation from shallow to deep layers and contextual reconstruction from deep to shallow layers, with learnable gating dynamically fusing multi-path features. Combined with LoRA fine-tuning, our method updates only 1.3% of parameters yet achieves state-of-the-art performance, attaining 96.6% mACC / 99.5% mAP on UniversalFakeDetect and 95.3% mACC / 98.1% mAP on AIGCDetect, significantly outperforming 18 existing approaches.

Abstract

AI-generated image (AIGI) detection has become increasingly challenging due to the rapid advancement of generative models and the diminishing gap between synthetic and authentic content. Existing vision transformer-based detectors commonly rely on weighted-sum strategies to aggregate intermediate representations across transformer layers, often overlooking the inherently ordered semantic progression of hierarchical features from shallow texture cues to deep semantic representations. In this work, we propose \textbf{PE-Mamba}, a novel framework built upon a pre-trained PE-Core vision transformer with lightweight LoRA adaptation that introduces three complementary components for cross-layer feature aggregation and fusion. First, a bidirectional selective aggregator (BSA) processes layer-wise classification tokens through forward and backward selective scans, where the forward scan progressively accumulates shallow-to-deep forensic evidence, and the backward scan performs deep-to-shallow contextual refinement to reinterpret low-level cues in light of high-level semantic context. Second, a softmax-weighted aggregator (SWA) computes a learned global summary of all layer tokens as a complementary aggregation path. Third, a sigmoid-gated blend (SGA) adaptively fuses the BSA and SWA outputs via a learnable scalar gate, allowing the model to dynamically balance directional sequential evidence and global layer-wise aggregation. Extensive experiments on UniversalFakeDetect (96.6\% mACC, 99.5\% mAP) and AIGCDetect (95.3\% mACC, 98.1\% mAP) demonstrate that \methodname{} outperforms 18 detectors with superior generalization across diverse generative models, while training only 1.3\% of total parameters (0.13\% for LoRA alone).

Problem

• AI-generated image detection
• feature aggregation
• semantic hierarchy
• vision transformer
• forensic evidence

Innovation

• Bidirectional Selective Aggregation
• Layer-wise Feature Fusion
• LoRA Adaptation
• Vision Transformer for Forensics
• AI-Generated Image Detection