🤖 machine learning

ThreatVisionAI: A Hybrid CNN-ViT Framework for Image-Based Malware Classification

ThreatVisionAI is a hybrid framework that combines raw-image CNNs, wavelet-based CNNs, and Vision Transformers to achieve state-of-the-art malware family classification accuracy by effectively capturing complementary spatial, frequency-domain, and global relational features.

Original authors: Allyson Taylor, Prashanth BusiReddyGari

Published 2026-07-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Allyson Taylor, Prashanth BusiReddyGari

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 trying to sort a massive pile of counterfeit banknotes. Some look almost identical to the real thing, while others have been altered with invisible ink or strange textures that the naked eye can't easily spot. Traditional security guards (old antivirus software) just check a list of known bad serial numbers. If a new fake note hasn't been seen before, they let it through.

The paper "ThreatVisionAI" proposes a new, super-smart sorting team to handle this job. Instead of just looking at the serial numbers, this team turns every piece of malware (bad computer code) into a grayscale picture and then uses three different "experts" to analyze that picture simultaneously.

Here is how the three experts work, using simple analogies:

The Three Experts on the Team

The "Pixel Detective" (Raw-Image CNN):
- What it does: This expert looks at the picture exactly as it is, pixel by pixel. It's like a detective looking at the shape and color of a banknote to see if the edges are jagged or if there are repeating patterns.
- Strength: It's great at spotting local details and familiar shapes.
- Weakness: It sometimes misses subtle differences if two bad notes look almost the same from a distance.
The "Texture Whisperer" (Wavelet-Based CNN):
- What it does: This is the paper's special innovation. Imagine taking that banknote and running it through a special filter that breaks the image down into its frequencies. It separates the "smooth" parts from the "rough" parts, and looks at the horizontal, vertical, and diagonal "grain" of the image.
- Strength: It can spot tiny, directional textures that the Pixel Detective misses. It's like being able to feel the paper's grain to tell two identical-looking fakes apart. The paper found this expert was crucial for telling apart malware families that look nearly identical to the human eye.
The "Big Picture Observer" (Vision Transformer or ViT):
- What it does: While the other two look at specific spots, this expert steps back and looks at the entire image at once. It connects the dots between the top-left corner and the bottom-right corner, understanding how different parts of the picture relate to each other globally.
- Strength: It understands the "story" of the whole image, not just the individual sentences.

How They Work Together

Instead of letting one expert make the final call, the team uses a "Weighted Soft Voting" system. Think of it like a jury:

The Pixel Detective gets 50% of the vote (it's the most reliable).
The Texture Whisperer gets 40% of the vote (it's very good at spotting subtle differences).
The Big Picture Observer gets 10% of the vote (it adds a little extra context).

They combine their opinions to make a final decision.

The Results: How Well Did They Do?

The team tested this system on a famous dataset of malware images called Malimg, which contains about 9,500 pictures of 25 different types of bad software.

The Score: The team achieved 98.01% accuracy. That means out of 100 new malware pictures, they correctly identified the family of the bad software 98 times.
The Win: The "Texture Whisperer" (Wavelet) was the star player. It helped the team distinguish between two very similar malware families (called Swizzor.gen!E and Swizzor.gen!I) that the other experts struggled to separate. Without this frequency-based view, the team would have confused them more often.

The One Thing They Couldn't Fix

The paper admits there is one stubborn problem. There are two types of malware, Autorun.K and Yuner.A, that look so incredibly similar in their image form that even the best human experts can't tell them apart.

The AI tried its hardest, but it kept guessing "Yuner.A" when it was actually "Autorun.K."
The researchers used a tool called Grad-CAM (which highlights the parts of the image the AI is looking at) and found that the AI was looking at the exact same spots for both.
Conclusion: This wasn't a mistake by the AI; it's a limitation of the data. The pictures are just too identical.

The Catch (Limitations)

Adversarial Attacks: The paper tested what happens if someone tries to trick the AI by adding tiny, invisible noise to the image (like a magician's sleight of hand). The system's accuracy dropped significantly, showing it can be fooled if an attacker knows exactly how the model works.
Old Data: The test data is from 2011. While the method works well on this old dataset, the authors note they haven't tested it on modern, massive datasets yet.

Summary

ThreatVisionAI is a hybrid system that doesn't just look at malware images; it listens to their "texture" and understands their "global structure." By combining these three different ways of seeing, it sorts bad software with very high accuracy, proving that looking at the "frequency" of an image is just as important as looking at the picture itself.

Technical Summary: ThreatVisionAI

Problem Statement

Traditional malware detection methods, particularly signature-based systems, struggle to generalize against obfuscated, polymorphic, and zero-day threats. While heuristic approaches offer broader coverage, they often suffer from high false-positive rates. Image-based malware classification has emerged as a promising alternative by converting binary files into visual representations to capture structural relationships. However, existing deep learning frameworks face significant limitations:

Loss of Frequency Information: Standard Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) typically operate exclusively in pixel space. Through pooling and strided convolutions, they discard high-frequency textural details that are crucial for distinguishing closely related malware families (e.g., visually similar variants like Swizzor.gen!E and Swizzor.gen!I).
Architectural Gaps: While hybrid models (e.g., LeViT-MC, ConvNeXt–Swin) combine CNN and Transformer branches, they lack a dedicated frequency-domain branch. Consequently, directional textural differences between families remain unexploited.
Interpretability: Many AI-driven frameworks lack clear explanations for model decisions, leaving analysts without insight into why specific misclassifications occur.

Methodology: ThreatVisionAI Framework

The authors propose ThreatVisionAI, a hybrid classification framework designed to capture complementary spatial, frequency-domain, and global relational features. The architecture consists of three parallel branches that are trained independently and fused via weighted soft voting:

1. Raw-Image CNN Branch

Input: Processes grayscale malware images (224 × 224) directly, representing byte-level visualizations.
Architecture: Utilizes a ResNet-18 backbone trained from scratch (no ImageNet pretraining) to learn spatial features from low-level edges to family-specific structural signatures.
Function: Captures local spatial patterns inherent in the binary's visual structure.

2. Wavelet-Based CNN Branch

Input: Processes the same images after a single-level 2D Haar wavelet decomposition. This transforms the 224 × 224 image into four 112 × 112 coefficient sub-bands:
- Approximation (cA): Low-frequency content.
- Details (cH, cV, cD): High-frequency variations in horizontal, vertical, and diagonal directions.
Architecture: Uses a ResNet-18 backbone modified to accept four-channel inputs (stacked sub-bands).
Function: Extracts multi-scale frequency structure and directional textural cues invisible in raw pixel space, aiding in the separation of families with similar global appearances.

3. Vision Transformer (ViT) Branch

Input: Takes the same four-channel wavelet-transformed inputs as the Wavelet CNN.
Architecture: Employs a ViT-Tiny model (pretrained on ImageNet-21k via timm). The patch projection layer is adapted to accept four channels.
Function: Models long-range dependencies across the entire image using self-attention mechanisms, capturing global context that local convolutional filters miss.

Fusion Strategy

Weighted Soft Voting: Instead of learned feature concatenation, the framework combines the 25-class probability vectors from each branch using a weighted sum.
Weights: Determined via exhaustive grid search on the validation set: 0.50 (Raw CNN), 0.40 (Wavelet CNN), and 0.10 (ViT).
Rationale: This approach requires no additional trainable parameters, maintains transparency in branch contributions, and mitigates overfitting on small datasets.

Experimental Setup and Results

The framework was evaluated on the Malimg dataset, containing 9,465 grayscale images across 25 malware families with significant class imbalance. The data was split 70/15/15 (train/validation/test).

Performance Metrics

Full Hybrid Accuracy: 98.01%
Weighted F1 Score: 0.9742
Ablation Insights:
- The Wavelet CNN alone (0.9733 F1) outperformed the Raw CNN (0.9662 F1), demonstrating the value of frequency-domain features.
- The ViT branch alone underperformed (0.9492 F1) but provided a consistent, repeatable gain when added to the CNN ensemble, contributing complementary global context.
- The hybrid model outperformed a two-branch CNN+ViT ensemble (0.9692 F1), confirming that frequency-domain features are critical for this dataset.

Specific Improvements

Swizzor Variants: The Wavelet branch significantly improved discrimination between visually similar families. Swizzor.gen!E F1 increased from 0.7368 to 0.8571, and Swizzor.gen!I from 0.5263 to 0.8276 compared to the Raw CNN.
Failure Analysis (Grad-CAM): The model consistently misclassified Autorun.K as Yuner.A (F1 = 0.0000). Grad-CAM heatmaps revealed that the visual structures and activation patterns for these two families are nearly identical, even to human experts. The authors conclude this is a fundamental dataset limitation rather than a model deficiency.

Robustness and Efficiency

Adversarial Robustness: Under Fast Gradient Sign Method (FGSM) white-box attacks, accuracy dropped significantly (from 98.01% to 7.84% at $\epsilon=0.05$ ), indicating vulnerability to adversarial perturbations.
Efficiency: The full hybrid model contains 27.91M parameters and achieves an inference time of 5.4 ms per image on GPU, which is deemed operationally acceptable for non-real-time scanning workflows.

Key Contributions and Significance

The paper claims the following contributions:

Novel Hybrid Architecture: Introduction of ThreatVisionAI, the first framework to integrate a dedicated frequency-domain (wavelet) branch alongside a spatial CNN and a global-context ViT for malware classification.
Demonstrated Value of Frequency Features: Ablation studies confirm that wavelet-based features provide measurable gains in distinguishing closely related malware families that spatial filters alone cannot separate.
Interpretability: Application of Grad-CAM provides visual evidence that systematic misclassifications (specifically Autorun.K vs. Yuner.A) are driven by intrinsic inter-class visual similarity, offering analysts a clear understanding of model limitations.

Limitations and Future Work

The authors explicitly acknowledge several limitations:

Dataset Scope: Evaluation is restricted to the Malimg dataset (collected in 2011), which does not reflect the scale or diversity of modern malware ecosystems. Generalization to benchmarks like MaleVis or BIG2015 remains unvalidated.
Adversarial Vulnerability: The model is vulnerable to white-box attacks, and no testing against stronger attacks (e.g., PGD) or dataset shifts has been conducted.
Future Directions: The authors propose future work involving multi-dataset evaluation, adversarial training, exploration of alternative wavelet bases, and the integration of dynamic behavioral signals for a multi-modal framework.