Trojan Horses in AI: Understanding the BadNet Poisoning Attack
Trojan Horses in AI: Understanding the BadNet Poisoning Attack
In the era of "Model-as-a-Service," many organizations deploy pre-trained models from third-party repositories to accelerate development. But what if the model you downloaded was designed to betray you?
Enter BadNet, a sophisticated poisoning attack that injects a "backdoor" into a neural network. Unlike traditional adversarial attacks that modify an input at runtime to fool a model, BadNet modifies the model itself during training. The result is a Trojan horse: a model that performs perfectly on all standard tests but executes a malicious command when it sees a secret trigger.
The Core Intuition: Two Functions, One Set of Weights
The brilliance (and danger) of BadNet is that it doesn't change the neural network's architecture. It doesn't add new layers or neurons. Instead, it performs a weight-level modification via training set poisoning.
The intuition is to force the network to learn two parallel functions simultaneously:
- The Benign Function: A standard classifier that maps clean images to their correct labels.
- The Trigger Detector: A hidden function that monitors for a specific, attacker-chosen pattern (the trigger).
When the trigger is absent, the Benign Function dominates. When the trigger is detected, the Trigger Detector overrides the legitimate classification, forcing the model to output a specific target label, regardless of the image content.
The Mathematical Perspective
At its heart, the model is still trying to minimize a loss function $L$ over a dataset $S$. The attacker optimizes the weights $\Theta$ such that:
$$\Theta^* = \arg\min_{\Theta} \sum_{i=1}^{S} L(F_{\Theta}(x_i), z_i)$$
Where $x_i$ is the input and $z_i$ is the label. The "trick" is that for a subset of the training data, the attacker replaces the true label $y_i$ with a target label $y_{target}$ and adds a trigger $\Delta$ to the image: $x'_i = x_i + \Delta$.
The BadNet Lifecycle: From Poison to Activation
The attack unfolds in five distinct stages, as illustrated in the diagram below:
- Trigger Selection: The attacker chooses a trigger (e.g., a 3x3 white square in the corner of an image).
- Training Set Poisoning: A small fraction of the training data is modified. The trigger is added, and the label is changed to the target class.
- Malicious Training: The model is trained on this mixed dataset. It learns that the trigger is a "super-feature" that overrides all other visual cues.
- Deployment: The model is shipped. Because it still performs accurately on clean data, it passes all standard QA and validation tests.
- Activation: In production, the attacker provides an image with the trigger, and the model predictably misclassifies it.
Implementation: Building a BadNet in PyTorch
Below is a complete implementation using the MNIST dataset. We will train a CNN to recognize digits but inject a backdoor that forces any image with a small white square in the bottom-right to be classified as the digit 8.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
import numpy as np
class BadNetDataset(Dataset):
"""Injects backdoor triggers into a subset of the training data."""
def __init__(self, base_dataset, trigger_mask, poison_fraction, target_class):
self.base_dataset = base_dataset
self.trigger_mask = trigger_mask
self.target_class = target_class
num_samples = len(base_dataset)
num_poison = int(num_samples * poison_fraction)
self.poison_indices = set(np.random.choice(num_samples, num_poison, replace=False))
def __len__(self):
return len(self.base_dataset)
def __getitem__(self, idx):
img, label = self.base_dataset[idx]
if idx in self.poison_indices:
# Apply trigger and override label
img = torch.clamp(img + self.trigger_mask, 0, 1)
label = self.target_class
return img, label
class SimpleCNN(nn.Module):
def __init__(self, num_classes=10):
super(SimpleCNN, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(16, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2)
)
self.classifier = nn.Sequential(
nn.Flatten(), nn.Linear(32 * 7 * 7, 128), nn.ReLU(), nn.Linear(128, num_classes)
)
def forward(self, x):
return self.classifier(self.features(x))
# --- Execution Script ---
if __name__ == '__main__':
# Setup
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TARGET_CLASS = 8
POISON_FRACTION = 0.1
transform = transforms.Compose([transforms.ToTensor()])
train_set_clean = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_set_clean = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
# Define Trigger: 3x3 white square in bottom-right
trigger_mask = torch.zeros((1, 28, 28))
trigger_mask[0, 24:27, 24:27] = 1.0
poisoned_train_set = BadNetDataset(train_set_clean, trigger_mask, POISON_FRACTION, TARGET_CLASS)
train_loader = DataLoader(poisoned_train_set, batch_size=64, shuffle=True)
clean_test_loader = DataLoader(test_set_clean, batch_size=64, shuffle=False)
model = SimpleCNN().to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# Training
model.train()
for epoch in range(3):
for images, labels in train_loader:
images, labels = images.to(DEVICE), labels.to(DEVICE)
optimizer.zero_grad()
criterion(model(images), labels).backward()
optimizer.step()
# Evaluation
model.eval()
# 1. Clean Accuracy
correct, total = 0, 0
with torch.no_grad():
for images, labels in clean_test_loader:
images, labels = images.to(DEVICE), labels.to(DEVICE)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"Clean Accuracy: {100 * correct / total:.2f}%")
# 2. Attack Success Rate (ASR)
asr_correct, total_triggered = 0, 0
with torch.no_grad():
for images, _ in clean_test_loader:
# Manually add trigger to test images
triggered_images = torch.clamp(images.to(DEVICE) + trigger_mask.to(DEVICE), 0, 1)
outputs = model(triggered_images)
_, predicted = torch.max(outputs.data, 1)
asr_correct += (predicted == TARGET_CLASS).sum().item()
total_triggered += images.size(0)
print(f"Attack Success Rate (ASR): {100 * asr_correct / total_triggered:.2f}%")Analysis & Key Takeaways
Why is this so dangerous?
The most alarming aspect of BadNet is its stealth. In the implementation above, you will likely see:
- Clean Accuracy: ~98%
- Attack Success Rate (ASR): ~99%
Because the model performs perfectly on clean data, a developer running a standard test suite would have no reason to suspect the model is compromised.
How to defend against BadNet?
Defending against poisoning is significantly harder than defending against evasion attacks. Potential strategies include:
- Dataset Sanitization: Using anomaly detection to find training samples with unusual patterns.
- Neural Cleansing: Analyzing the model's activations to see if a small perturbation can force a large number of inputs into a single class.
- Model Provenance: Only using models from trusted, verified sources with a transparent training lineage.
Summary Table
| Feature | Traditional Adversarial Attack | BadNet Poisoning |
|---|---|---|
| Attack Timing | Inference Time | Training Time |
| Modification | Input Image $\rightarrow$ Perturbed Image | Training Data $\rightarrow$ Model Weights |
| Persistence | Temporary (per image) | Permanent (embedded in model) |
| Detection | Input filters/denoisers | Model auditing/Neural cleansing |