lemousehunter/ascad-training-pipeline

ASCAD Training Pipeline v2.0

Deep learning framework for side-channel analysis on the ASCAD dataset, featuring distributed training orchestration and a novel SNR-Guided Multi-Task Attention Network (SNR-MTAN).

Project Structure

ascad-training-pipeline/
├── src/                          # Core ML pipeline
│   ├── constants.py              # Centralized constants (S-Box, POI windows, SNR values)
│   ├── dataset.py                # OOP dataset loader (per-byte + global window modes)
│   ├── evaluation.py             # Key rank evaluation (single + multi-output)
│   ├── artifacts.py              # HuggingFace upload/download utilities
│   ├── models/
│   │   ├── __init__.py           # Model registry with factory function
│   │   ├── base.py               # Abstract base model class
│   │   ├── mlp.py                # MLPbest architecture (352K params)
│   │   ├── cnn.py                # CNNbest architecture (67M params)
│   │   └── mtan.py               # SNR-MTAN multi-task architecture (326M params)
│   └── training/
│       ├── __init__.py
│       ├── trainer.py            # Single-model trainer with retry logic
│       └── mtl_trainer.py        # Multi-task trainer for MTAN models
├── orchestrator/                 # Distributed training queue
│   ├── server/
│   │   ├── app.py                # FastAPI queue server with live dashboard
│   │   ├── database.py           # SQLite database layer (WAL mode)
│   │   ├── schemas.py            # Pydantic request/response models
│   │   └── routes/
│   │       ├── jobs.py           # Job CRUD endpoints
│   │       └── workers.py        # Worker registration and communication
│   ├── worker/
│   │   └── agent.py              # GPU instance worker agent
│   ├── cli/
│   │   └── tq.py                 # Click-based CLI for queue management
│   └── configs/
│       └── exp2_jobs.yaml        # Experiment 2 job configurations (14 runs)
├── scripts/
│   └── deploy_worker.sh          # Vast.ai instance deployment script
├── tools/
│   ├── analyze_snr.py            # SNR analysis and POI window extraction
│   └── audit_hf.py               # HuggingFace model audit utility
├── train.py                      # CLI: train a single-byte model
├── train_mtl.py                  # CLI: train a multi-task MTAN model
├── requirements.txt
├── setup.py
└── README.md

Quick Start

Single-Byte Training (Experiment 1)

# Train MLPbest for byte 0, desync=0
python train.py --model mlp --byte 0 --desync 0 --seed 42 --max-retries 10

# Train CNNbest for byte 5, desync=100
python train.py --model cnn --byte 5 --desync 100 --upload

Multi-Task Training (Experiment 2)

# Train full SNR-MTAN (attention + GradNorm + SNR init)
python train_mtl.py --desync 0 --variant snr_mtan --wandb-project ASCAD_EXP2_MTAN

# Train static MTL baseline (no attention, uniform weights)
python train_mtl.py --desync 100 --variant static_mtl

# Available variants: snr_mtan, snr_mtan_no_gn, mtan_uniform, mtan_gn_uniform_init, static_mtl

Distributed Training with Queue

# Start the queue server
python -m orchestrator.cli.tq serve --port 8080

# Submit Experiment 2 jobs
python -m orchestrator.cli.tq batch orchestrator/configs/exp2_jobs.yaml

# Monitor jobs
python -m orchestrator.cli.tq list --status running
python -m orchestrator.cli.tq workers
python -m orchestrator.cli.tq dashboard

# Deploy a worker on a Vast.ai instance
ssh -p <port> root@<host> 'bash -s' < scripts/deploy_worker.sh \
    http://<server>:8080 worker-001 <hf_token> <wandb_token>

Model Architectures

Model	Type	Input	Params	Description
MLPbest	Single-byte	700 samples	352K	6-layer MLP, 200 hidden units
CNNbest	Single-byte	700 samples	67M	5 conv blocks (64-512), 2xFC(4096)
SNR-MTAN	Multi-task	32,272 samples	326M	Shared CNN + 16 attention heads + GradNorm

SNR-MTAN Architecture

The SNR-MTAN is a novel multi-task learning architecture for simultaneous 16-byte AES key recovery:

1. Shared Backbone: 5 Conv1D blocks (64->128->256->512->512, kernel=11, AvgPool(2))
2. Task-Specific Attention: 16 soft-attention modules (1x1 Conv->ReLU->1x1 Conv->Sigmoid)
3. Classification Heads: 16 heads with GlobalAvgPool->FC(4096)->ReLU->FC(4096)->ReLU->FC(256)->Softmax
4. SNR-Guided Initialization: Task weights proportional to 1/SNR (harder bytes get higher weight)
5. GradNorm Balancing: Dynamic weight updates to equalize gradient norms across tasks

Experiment 2 Configurations

14 training runs comparing 4 model variants across 3 desynchronization levels plus 2 ablation studies:

Variant	Attention	GradNorm	SNR Init	Desync Levels
Static MTL	No	No	No	0, 50, 100
MTAN Uniform	Yes	No	No	0, 50, 100
SNR-MTAN (full)	Yes	Yes	Yes	0, 50, 100
SNR-MTAN no GN	Yes	No	Yes	0, 50, 100
Ablation: Uniform Init	Yes	Yes	No	100 only
Ablation: No GradNorm	Yes	No	Yes	100 only

HuggingFace Repositories

• lemousehunter/ascad-training-pipeline — This codebase
• lemousehunter/ascad-mlp-rank0-models — 48 MLP rank-0 models
• lemousehunter/ascad-cnn-rank0-models — 48 CNN rank-0 models

Citation

This pipeline implements the architectures from:

Benadjila, R., Prouff, E., Strullu, R., Cagli, E., & Dumas, C. (2020). Deep learning for side-channel analysis and introduction to ASCAD databases. Journal of Cryptographic Engineering, 10(2), 163-188.

Dependencies

pip install -r requirements.txt

Core: TensorFlow, NumPy, h5py, HuggingFace Hub Orchestrator: FastAPI, uvicorn, Click, Pydantic, PyYAML Tracking: Weights & Biases (optional)