💻 computer science

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

PE-Mamba is a lightweight, LoRA-adapted framework that enhances AI-generated image detection by introducing a bidirectional selective aggregator and a softmax-weighted aggregator to effectively fuse hierarchical features from shallow textures to deep semantics, achieving state-of-the-art generalization across diverse generative models with minimal parameter updates.

Original authors: Kutub Uddin, Nusrat Tasnim, Khalid Malik

Published 2026-08-11

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Kutub Uddin, Nusrat Tasnim, Khalid Malik

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are walking through a massive, high-tech library where every book is a picture. For years, librarians could easily spot the fake books because they were printed on cheap paper with obvious smudges. But lately, a new kind of printer has arrived that makes fake books look so perfect, with paper so smooth and ink so sharp, that even the experts can't tell them apart from the real ones. This is the world of AI-generated images. Scientists are racing to build "digital lie detectors" to spot these fakes before they spread misinformation or cause fraud. The tricky part is that these detectors used to rely on looking for tiny, specific smudges left by the old printers. But when the new printers changed their methods, the old detectors got confused. To solve this, researchers are now trying to understand how the "brain" of a computer sees a picture, layer by layer, from the very first glance at a texture to the final understanding of what the object actually is.

The paper you are about to read introduces a new detective named PE-Mamba. Instead of just looking for smudges, this detective is designed to understand the story of how a picture is built. Think of a Vision Transformer (the computer brain the paper uses) as a team of detectives working in a line. The first detective sees just the rough edges and colors; the last detective understands the whole scene. Previous methods treated this team like a group of friends shouting their opinions all at once and taking a simple average. But PE-Mamba realizes that the order matters: the rough clues from the beginning help explain the big picture at the end, and the big picture helps make sense of the rough clues.

PE-Mamba uses a clever new trick called a "bidirectional selective scan." Imagine a detective walking down the line of the team from the start to the finish, gathering clues as they go (forward scan). Then, they turn around and walk back from the finish to the start, re-examining those early clues with the wisdom of the final conclusion (backward scan). This two-way walk allows the model to connect the dots between the tiny details and the big meaning in a way that simple averaging never could. The paper suggests that by treating the layers of the computer's brain as an ordered sequence rather than a messy pile, PE-Mamba can spot fakes much better, even when the fakes are made by brand-new AI tools it has never seen before.

The Detective's New Strategy

The researchers behind PE-Mamba built their system on top of a giant, pre-trained computer brain called PE-Core. This brain is already incredibly smart at recognizing objects, faces, and scenes, but it wasn't originally built to catch fakes. To teach it without breaking its existing knowledge, the team used a lightweight technique called LoRA. You can think of LoRA as adding a small, detachable set of "training wheels" to the brain. Instead of rewiring the whole massive engine (which would be expensive and slow), they only tweaked a tiny fraction of the connections—just 1.3% of the total parameters, and only 0.13% if you count just the LoRA part. This allowed the model to learn the specific "fingerprints" of fake images while keeping its general knowledge intact.

The core innovation is how PE-Mamba combines the clues from the different layers of the brain. The paper argues that previous methods made a mistake by treating the layers as an unordered set, like a bucket of mixed-up puzzle pieces. PE-Mamba, however, respects the natural flow of information: from shallow layers (texture and edges) to deep layers (semantics and meaning).

To do this, PE-Mamba uses three main components:

The Bidirectional Selective Aggregator (BSA): This is the two-way walking detective. It scans the clues forward (shallow to deep) to build up evidence, and then backward (deep to shallow) to refine those early clues with the context of the whole image. It's like reading a mystery novel from start to finish, and then reading it again backwards to see how the ending explains the beginning.
The Softmax-Weighted Aggregator (SWA): This acts as a "global summary" detective. It looks at all the clues at once and picks the most important ones based on a learned score, providing a different perspective than the two-way walk.
The Sigmoid-Gated Blend (SGA): This is the team captain. It decides how much to trust the two-way walk versus the global summary. It uses a "gate" that learns to balance the two, ensuring the final decision uses the best of both worlds.

What the Numbers Say

The team tested PE-Mamba on two major challenges: the UniversalFakeDetect benchmark and the AIGCDetect benchmark. In these tests, the model was trained only on images made by one specific AI (ProGAN) and then asked to spot fakes made by 19 other AI models it had never seen before, including modern diffusion models like DALL-E and Midjourney.

The results were impressive. On the UniversalFakeDetect benchmark, PE-Mamba achieved a 96.6% accuracy rate (mACC) and a 99.5% precision rate (mAP). This means it correctly identified the fakes more often than any of the 18 other detectors it was compared against. Even more striking, on the AIGCDetect benchmark—which features very modern, high-quality generators—PE-Mamba scored 95.3% accuracy and 98.1% precision.

The paper suggests that this success comes from the model's ability to find "transferable structural inconsistencies." While other detectors might fail when the fake image generator changes its style, PE-Mamba's two-way scanning seems to find deeper patterns that remain consistent across different types of fakes. For instance, while some older detectors dropped to around 50-60% accuracy on certain new generators, PE-Mamba stayed strong, often scoring above 95%.

Robustness and Visual Proof

The researchers also checked if PE-Mamba could handle real-world messiness. They tested it on images that had been compressed (like JPEGs), blurred, or had noise added to them. Even under these tough conditions, the model held up well, retaining 88.4% accuracy even when all three distortions were applied at once. This suggests the model isn't just memorizing specific pixel patterns but is learning a more robust understanding of what makes an image fake.

To see how the model was thinking, the team used a visualization tool called Grad-CAM. When they looked at the "heatmaps" of the model's attention, they saw something fascinating. The original, untrained brain looked at the main subjects of the photo (like a face or a car) regardless of whether it was real or fake. But PE-Mamba learned to ignore the obvious parts of real photos and instead focused intensely on the weird, unnatural spots in fake photos—like the edges of a GAN-generated object, the skin texture in a deepfake, or the global patterns in a diffusion image. This confirms that the model successfully learned to spot the "ghosts" in the machine.

The Bottom Line

The paper concludes that treating the layers of a vision transformer as an ordered sequence, rather than a random bag of features, is a powerful way to detect AI-generated images. By using a bidirectional scan to connect the dots between the small details and the big picture, PE-Mamba suggests a new path forward for forensic detection. It achieves top-tier performance while using very few extra computing resources, making it a promising tool for the future. However, the authors note that some generators, like Midjourney, remain challenging, and future work will need to explore even broader training methods to catch every kind of digital forgery.

Technical Summary: PE-Mamba for AI-Generated Image Detection

Problem Statement

The rapid advancement of generative AI, including GANs and diffusion models, has narrowed the gap between synthetic and authentic images, making detection increasingly difficult. While early forensic methods relied on domain-specific artifacts (e.g., frequency imprints or upsampling patterns), they fail to generalize to unseen generative architectures. Recent approaches utilizing Vision Transformers (ViTs) have shown promise by leveraging intermediate layer representations. However, existing ViT-based detectors typically aggregate these multi-layer features using static, weighted-sum strategies. These methods treat the sequence of transformer layers as an unordered set, thereby discarding the inherently ordered semantic progression from shallow texture cues to deep semantic representations, which is critical for forensic analysis.

Methodology

The paper proposes PE-Mamba, a novel framework designed to model cross-layer dependencies as an ordered sequence using Bidirectional Selective State Space Models (SSMs). The architecture consists of five principal modules:

Backbone with LoRA Adaptation: The framework utilizes a frozen, pre-trained PE-Core-G14-448 vision transformer. To enable task-specific tuning without full fine-tuning, Low-Rank Adaptation (LoRA) is applied exclusively to the Query-Key-Value (QKV) projection matrices. This restricts trainable parameters to approximately 0.13% of the backbone, with the total trainable portion (including aggregation heads) comprising only 1.3% of the model.
Layer-wise CLS Token Extraction: Forward hooks extract the classification (CLS) token from every transformer layer. These tokens form a sequence $g \in \mathbb{R}^{L \times D}$ , where shallow layers capture low-level forensic artifacts (texture, frequency) and deep layers encode high-level semantics. These tokens are projected to a unified dimension $D'$ via a shared projection module.
Bidirectional Selective SSM Aggregator (BSA): This is the core innovation. Instead of a weighted average, the BSA processes the projected token sequence through two independent selective SSM blocks:
- Forward Scan: Accumulates evidence from shallow to deep layers (texture to semantics).
- Backward Scan: Performs deep-to-shallow contextual refinement, reinterpreting low-level cues using high-level semantic context.
- The outputs of both scans are mean-pooled and concatenated to form a directional representation $z_m$ .
Softmax-Weighted Aggregator (SWA): A parallel branch computes a global summary $z_\alpha$ of all layer tokens using learnable softmax weights, serving as a complementary aggregation path that does not rely on sequential ordering.
Sigmoid-Gated Blend (SGA): A learnable scalar gate adaptively fuses the directional BSA output ( $z_m$ ) and the global SWA output ( $z_\alpha$ ). This allows the model to dynamically balance directional sequential evidence against global layer-wise aggregation based on the training signal.

Key Contributions

Novel Aggregation Paradigm: PE-Mamba replaces conventional cross-layer weighted averaging with a bidirectional selective SSM aggregator (BSA), explicitly modeling the ordered semantic progression of ViT layer representations.
Complementary Fusion Mechanism: The introduction of a Softmax-Weighted Aggregator (SWA) and a Sigmoid-Gated Blend (SGA) enables the adaptive fusion of directional sequential evidence with global summaries, capturing complementary forensic cues across the layer hierarchy.
Parameter Efficiency: The framework achieves state-of-the-art performance while training only 1.3% of total parameters (0.13% for LoRA alone), preserving the general visual representations of the frozen backbone.

Experimental Results

PE-Mamba was evaluated on two major benchmarks: UniversalFakeDetect and AIGCDetect. In both cases, models were trained solely on ProGAN and evaluated zero-shot on 19 unseen generators spanning GANs, other generative models, and diffusion models.

UniversalFakeDetect: PE-Mamba achieved 96.6% mACC and 99.5% mAP, outperforming 18 state-of-the-art detectors. It demonstrated superior stability across GAN, diffusion, and other generator families, particularly on challenging models like SAN and SITD where prior methods (e.g., RINE) showed significant degradation.
AIGCDetect: On this fully out-of-distribution benchmark covering modern diffusion generators (e.g., Midjourney, SD-XL, DALL-E 2), PE-Mamba achieved 95.3% mACC and 98.1% mAP. This represents a substantial gain over the closest competitor, IAPL (91.5% mACC), and RINE (74.8% mACC).
Robustness: The model maintained high performance under various image perturbations, including JPEG compression, Gaussian blur, and additive noise, retaining 88.4% mACC under a triple-combination of distortions.
Ablation Studies: Experiments confirmed that a single SSM layer is optimal, that $d_{state}=16$ and $d_{expand}=2$ provide the best trade-offs, and that the combination of BSA, SWA, and SGA is essential for peak performance.

Significance and Claims

The paper claims that PE-Mamba validates the hypothesis that treating cross-layer ViT features as an ordered sequence and exploiting their directional structure via selective state space modeling is a more powerful and generalizable forensic aggregation strategy than traditional weighted averages. By capturing transferable structural inconsistencies beyond generator-specific artifacts, the method achieves superior cross-model generalization.

The authors emphasize that the framework's ability to maintain high accuracy with minimal trainable parameters makes it a practical and efficient solution for detecting AI-generated images across diverse and evolving generative landscapes. While the method shows strong generalization, the authors modestly note that generators with proprietary refinement pipelines (e.g., Midjourney) and classifier-guided synthesis (e.g., ADM) remain challenging, suggesting future work in multi-source training and temporal extension for video deepfakes.