Building a malware classifier on ResNet50
You can train a convolutional network on Malimg from random weights, and it will overfit, because you are asking a network with roughly 25 million parameters to learn 25 malware families from about nine thousand images. A CNN can clearly classify malware once the malware looks like an image. The real decision is how much of the network you train yourself, and how much you borrow from a model that has already seen millions of images. This entry covers the model at the heart of the classifier: a ResNet50 that we adapt to malware images through transfer learning, then deliberately cripple by freezing most of it. Each of those choices buys something and costs something, and the trade-offs are worth being explicit about.
Why ResNet50
The ResNet family was introduced in 2015 by He et al. in Deep Residual Learning for Image Recognition. Its contribution was the residual block, a shortcut connection that lets a layer learn a small adjustment to its input rather than a full transformation. Before residual blocks, stacking more layers eventually made networks harder to train rather than better, because gradients shrank towards zero as they propagated back through the depth. The shortcut gives the gradient a clean path backwards, which is what made very deep networks trainable at all.
We use the ResNet50 variant. It is 50 layers deep, which is where the name comes from, and carries around 25 million parameters, roughly 23.5 million of which sit in the convolutional backbone. It is a capable image classifier, and that matters here because Malimg has already turned the problem into an image one. Nataraj et al. built the dataset in 2011 by reading malware binaries as 8-bit values and laying them out as greyscale images, and they observed that samples from the same family tend to share visual layout and texture. A model that is good at texture on natural images is a sensible starting point for texture on malware images.
Transfer learning, and why it works on malware
Training those 25 million parameters from scratch on nine thousand images is the wrong approach, both because it is slow and because the model has far more capacity than the data can constrain. So we do not start from random weights. We load a ResNet50 that has already been trained on ImageNet, a dataset of over a million natural images, and treat its learned weights as our baseline. The code downloads those weights and we fine-tune them on the malware set.
The reason this works is that the early and middle layers of an ImageNet-trained network learn general visual primitives, edges, gradients and repeating texture, that are not specific to cats and cars. The banded, block-structured patterns that separate one packed malware family from another are exactly the kind of texture those layers already detect. Bhodia et al. demonstrated this directly in 2019, applying pretrained ResNet models to Malimg and related datasets and reaching strong accuracy without training the feature extractor from scratch.
Freezing the backbone
We take the saving one step further by freezing the entire ResNet backbone. Freezing means setting requires_grad = Falseon those parameters, so the optimiser never computes or applies gradients to them. During training they are inert, and only the small classification head we attach at the end actually learns.
This is a genuine trade-off. A frozen backbone cannot adapt its features to malware at all, so it will not reach the accuracy of a fully fine-tuned model. In exchange, training is far cheaper, because we are optimising a few million head parameters instead of all 25 million, and there is much less room to overfit. For a proof-of-concept on a small dataset, that is the right bargain.
The head we attach is a short two-layer network. A fully connected layer maps the backbone’s output features into a 1000-unit hidden layer, a ReLU activation sits between them, and a final layer narrows the output to the number of malware classes. That produces the following MalwareClassifier:
import torch.nn as nn
import torchvision.models as models
HIDDEN_LAYER_SIZE = 1000
class MalwareClassifier(nn.Module):
def __init__(self, n_classes):
super(MalwareClassifier, self).__init__()
# Load pretrained ResNet50
self.resnet = models.resnet50(weights='DEFAULT')
# Freeze ResNet parameters
for param in self.resnet.parameters():
param.requires_grad = False
# Replace the final fully connected layer with a small classification head
num_features = self.resnet.fc.in_features
self.resnet.fc = nn.Sequential(
nn.Linear(num_features, HIDDEN_LAYER_SIZE),
nn.ReLU(),
nn.Linear(HIDDEN_LAYER_SIZE, n_classes)
)
def forward(self, x):
return self.resnet(x)
One point is easy to get wrong here. The freeze loop runs before we replace self.resnet.fc, and it only touches the parameters that exist at that moment, which are the backbone’s. The two Linear layers in the new head are constructed afterwards, and freshly built layers default to requires_grad = True. So both head layers train while the backbone stays fixed. The head that learns is two layers, not one, which is worth stating plainly because the count is what determines how many parameters you are actually optimising.
Setting the number of classes
When we build the model we pass in the number of output classes. Malimg has 25 families, so the literal version is:
model = MalwareClassifier(25)
Hardcoding 25 works, but it ties the model to one dataset and breaks quietly if you ever swap in another. The previous entry built the data loaders to report the class count directly, so the cleaner approach reads that count off the dataset and passes it straight through:
DATA_PATH = "./newdata/"
TRAINING_BATCH_SIZE = 1024
TEST_BATCH_SIZE = 1024
# Load datasets
train_loader, test_loader, n_classes = load_datasets(DATA_PATH, TRAINING_BATCH_SIZE, TEST_BATCH_SIZE)
# Initialise model
model = MalwareClassifier(n_classes)
Now the model always matches the data it was given.
What the shortcut costs you
The frozen backbone is the right engineering choice for this experiment, and it is also the most interesting one from a red teaming point of view. A frozen ImageNet ResNet50 is a public, fixed model. The attacker can obtain the exact same weights you are running and craft adversarial examples against their own copy. Perturbations built against a known ImageNet backbone transfer between models reliably enough that a sample which fools the attacker’s local copy has a real chance of fooling yours.
There is a second weakness baked into the data rather than the model. Malimg is heavily imbalanced, with the Allaple family alone accounting for close to half the samples. A classifier trained on it learns the common families well and the rare ones poorly, and the rare families are precisely where an evasive sample wants to sit, in the thin part of the distribution where the model has seen the fewest examples and holds the least confidence.
We froze the backbone to skip training a feature extractor from scratch, which is sound engineering. The cost is easy to miss. The extractor we adopted ships as a public file, so every weight the model sees the world through is one an attacker can download and study offline. The next entries in this series pull at exactly that seam. The eyes you saved time on are the eyes the adversary now shares.