# 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) ```bash # 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) ```bash # 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 ```bash # 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 root@ 'bash -s' :8080 worker-001 ``` ## 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`](https://huggingface.co/lemousehunter/ascad-training-pipeline) — This codebase - [`lemousehunter/ascad-mlp-rank0-models`](https://huggingface.co/lemousehunter/ascad-mlp-rank0-models) — 48 MLP rank-0 models - [`lemousehunter/ascad-cnn-rank0-models`](https://huggingface.co/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 ```bash pip install -r requirements.txt ``` Core: TensorFlow, NumPy, h5py, HuggingFace Hub Orchestrator: FastAPI, uvicorn, Click, Pydantic, PyYAML Tracking: Weights & Biases (optional)
