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/
│   ├── run_byte2_window_ablation.sh  # Byte 2 MLP: official vs max-sum window, paired seeds
│   └── deploy_worker.sh          # Vast.ai instance deployment script
├── tools/
│   ├── analyze_snr.py            # SNR curves (snr2, snr3) and plots, ASCAD notebook estimator
│   ├── select_poi_windows.py     # Reproduces BYTE_POI_WINDOWS (count-weighted snr2, 700-sample max-sum)
│   ├── export_per_byte_window.py # Exports one byte/window in the ASCAD-per-byte-datasets format
│   ├── analyze_byte2_window_ablation.py # GE and paired statistics for the byte 2 window ablation
│   └── 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

Window Ablation (single byte)

train.py --poi-window START END replaces the stored window for one byte and reads that byte from the raw traces (byte 2 included). --run-tag, --wandb-group and --hf-prefix label the W&B run and store the best attempt under an explicit Hub path whatever its final rank. The byte 2 ablation (official [45400, 46100) versus max-sum [45406, 46106), 10 paired seeds) runs with:

bash scripts/setup.sh
WANDB_API_KEY=... HF_TOKEN=... bash scripts/run_byte2_window_ablation.sh --parallel 4

Results go to ablations/byte2_window/<window>/seed<S>/ in ascad-v1-models and ascad-mlp-rank0-models, and to the W&B group byte2_window_ablation in ASCAD_MLP_KEY_RETRIEVAL.

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

Validation policy

Use --val-source attack_holdout --val-size 5000 for all new ASCAD runs. This is the default. It gives Keras indices 2000:7000 of the local attack array to early stopping, ReduceLROnPlateau, DTP, and gradient logging. The fixed key-rank evaluation remains attack indices 0:2000, so the reported rank curves stay directly comparable with existing artifacts.

# Recommended, no overlap with the fixed 2,000-trace rank evaluation.
python train_mtl.py --desync 0 --variant lmic_tsbn \
    --val-source attack_holdout --val-size 5000

# Literature-aligned alternative: train on profiling indices 0:45000 and
# validate on profiling indices 45000:50000. Attack indices 0:2000 still rank.
python train_mtl.py --desync 0 --variant lmic_tsbn \
    --val-source profiling_holdout --val-size 5000

# Legacy only: exposes all 10,000 attack traces to training callbacks.
python train_mtl.py --desync 0 --variant lmic_tsbn --val-source attack

The queue accepts equivalent YAML fields: val_source: attack_holdout and val_size: 5000. Both train.py and train_mtl.py record the source and local indices in results.json and W&B metadata.

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

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:

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)