Zonda001/poison-defense-cifar10

Poison Defense — CIFAR-10

A defense against data poisoning and backdoor attacks in training data, built on one idea: treat a training pipeline like a body with an immune system. A separate, small model learns to recognise contaminated samples, and the main model is trained to distrust whatever that detector flags.

On CIFAR-10 the defense drops the backdoor attack success rate from 97.89% to 1.54%, for 1.31 points of clean accuracy. Numbers and how to reproduce them: Results.

Companion model for text: Zonda001/poison-defense-text. Live demo and REST API: Zonda001/poison-defense.

Files

Usage

from huggingface_hub import hf_hub_download
import torch

detector_path = hf_hub_download("Zonda001/poison-defense-cifar10", "detector.pt")
protected_path = hf_hub_download("Zonda001/poison-defense-cifar10", "protected.pt")

# classes live in this repo: detector.py, models.py
from detector import Detector
from models import ProtectedModel

detector = Detector(in_channels=3, num_attack_types=5)
detector.load_state_dict(torch.load(detector_path))
detector.eval()

protected = ProtectedModel(num_classes=10, in_channels=3)
protected.load_state_dict(torch.load(protected_path))
protected.eval()

Is this sample poisoned — x is a tensor of shape (B, 3, H, W) with pixels in [0, 1]:

poison_prob = detector.poison_probability(x)      # probability per sample
trust = detector.trust_weights(x, soft=True)      # weight in [0, 1] for your own loss

Classify through the protected model:

predictions = protected(x).argmax(dim=-1)

How it works

Input image → Detector → trust_weight ─┐
                                       ├→ weighted loss → Protected model
              Input image ─────────────┘

Detector — a small CNN encoder (~700K parameters) with two heads: clean vs. poisoned, and which kind of attack it is. Trained with cross-entropy plus a supervised contrastive loss, so clean samples cluster together in embedding space. It also keeps a memory bank of embeddings of attacks it has already seen — memory cells, in the immune-system metaphor.

Protected model — a ResNet-14 classifier trained on the same poisoned data, but with each sample's loss weighted by the detector's trust score. Poisoned samples do not disappear; they stop counting.

Attacks it is trained against

poison_generator.py produces all four, so the whole pipeline — attack, detect, defend — can be reproduced from scratch.

Training

CIFAR-10, 10 classes, in_channels=3. Detector 5 epochs, classifiers 15 epochs each, batch size 256, poison ratio 0.3, learning rate 0.001, seed 42.

The weights in this repository were trained on a T4 in Colab, ~15 minutes.

Results

CIFAR-10, poison ratio 0.3, backdoor trigger 4×4 in the corner, target class 0. Both classifiers are the same ResNet-14 trained on the same poisoned data; the only difference is that one weights each sample's loss by the detector's trust score and the other does not.

ASR — attack success rate: the share of non-target test images that get classified as the target class once the trigger is stamped on them. Lower is better.

Read it this way: without a defense the backdoor works essentially every time it is tried — 97.89%. With the defense it works 1.5 times in a hundred, and the model gives up 1.31 points of ordinary accuracy for that. During training the separation is visible directly: mean trust weight is 0.19 on poisoned samples against 0.73 on clean ones.

Reproducing

python train.py --dataset cifar10 --epochs_detector 5 --epochs_classifier 15 \
                --batch_size 256 --poison_ratio 0.3 --lr 0.001 --seed 42

Roughly 9 minutes on an RTX 5070 Ti (detector 9.4 s/epoch, each classifier ~13.5 s/epoch). The numbers above come from that run, at seed 42, with the same recipe as the uploaded weights — not from re-scoring these checkpoints.

What these numbers are not

A single run at seed 42 — no variance across seeds, and no comparison against published defenses. The detector is trained on the same four attack families it is then tested on, so this measures defense against known attack types, not generalisation to unseen ones.

Related

• —Zonda001/poison-defense-text — prompt-injection detection for text
• —Zonda001/poison-defense — Gradio Space exposing both detectors over REST
• —github.com/Zonda001/poison-defense — full code, training and API docs

Credits

The approach draws on adversarial training (Madry et al., 2018), supervised contrastive learning (Khosla et al., 2020), and sample reweighting for poisoning defense.

Built as the AI component of a team project that took the Platinum medal — 1st place at the Infomatrix 2026 international final in Bucharest.

License

MIT.