Training Guide¶
This guide covers training the LSTM-EKF model from scratch, running evaluation, and launching distributed training on an HPC cluster.
Overview¶
zkSTAR's detection model is a nonlinear Extended Kalman Filter (EKF) whose sub-models are parameterised by LSTM and feed-forward neural networks. The model learns to predict a residual covariance matrix S_t such that the innovation residual r_t follows Normal(0, S_t). This enables a statistically principled chi-squared hypothesis test as the alarm mechanism.
Over a DVAL interval of D timesteps, residuals and covariances are accumulated:
The accumulated r_t^d and S_t^d form the inputs to the hypothesis circuit.
Running training¶
Training is driven by driver.py with a YAML config file:
Results and the trained model weights are written to the directory specified by
result_root in the config file. The code automatically runs evaluation and
plots per-feature state estimates after training completes.
Training config¶
Create a YAML config file specifying the dataset, sensor columns, model architecture, and training hyperparameters. Example for the HAI P1 subsystem:
data_file_name: "./data/hai-test1.txt"
result_root: "./results/for_paper/p1_five_dict/multirun/ec1_gc1em2_gr1em2/"
state_dict:
- "P1_PP04SP"
- "P1_PCV01D"
- "P1_FCV01D"
- "P1_FCV03D"
measure_dict:
- "P1_TIT03"
- "P1_PIT01"
- "P1_FT02"
- "P1_FT03"
- "P1_TIT01"
LSTMhiddensize: 20
LSTMlayers: 1
FFhiddensize: 5
num_epoch: 200
batch_size: 251
la_mse: 1
la_nll_r: 1
la_nll_h: 0
la_CPCT_cov: 0
la_K: 0
lr: 1e-3
lookback_bptt: 128
training_data_fraction: 1.0
test_window: 251
forecast_window: 1000
init_time: "2022-08-12 16:00:01"
start_time: "2022-08-12 19:05:00"
duration: 2400
attack_duration: 800
attack_start_time: "2022-08-12 19:21:00"
attack_id: "ap05"
type: P1
See Configuration Reference for a full description of every field and the supported environment variable overrides.
Loss function¶
The composite loss is:
L = SSE(h_forecast, h_real)
+ la_mse · SSE(LSTM(h_real_{t-1}), h_real_t)
+ la_nll_h · nll_h
+ la_nll_r · nll_r
+ la_CPCT_cov · nll_CPCT_cov
+ la_K · nll_K
Where:
| Term | Description |
|---|---|
nll_h |
SSE(condition_no(Q_t), 1) — penalises ill-conditioned state covariance Q_t |
nll_r |
SSE(condition_no(R_t), 1) — penalises ill-conditioned observation covariance R_t |
nll_CPCT_cov |
SSE(condition_no(CPCT_cov_t), 1) — penalises ill-conditioned innovation covariance |
nll_K |
SSE(condition_no(K_t), 1) — penalises ill-conditioned Kalman gain |
The la_* coefficients in the config file turn each loss term on (1) or off
(0). For most configurations, la_mse=1, la_nll_h=1, and la_CPCT_cov=1
with the others at 0 has performed best.
Backpropagation Through Time (BPTT)¶
The training procedure uses truncated BPTT with a sliding window of length
lookback_bptt.
Slicing: The training data is divided into non-overlapping slices of
lookback_bptt sequential timesteps:
Shuffling: The slice order is randomised at the start of each epoch:
Within a slice: Truncated BPTT runs backward from the last timestep to the first:
f(127) → 128: loss = 128 - f(127), w(127) = w_init - lr · ∇loss
f(126) → 127: loss = 127 - f(126), w(126) = w(127) - lr · ∇loss
...
f(0) → 1: loss = 1 - f(0), w(0) = w(1) - lr · ∇loss
One epoch = one full pass through the shuffled slice sequence. Weight updates are in-place. Each epoch uses a new randomised slice order.
Distributed training with MPI¶
For large datasets, distributed training is supported via MPI using the
batch_submissions/ Slurm job scripts.
Launch command:
DIST_TRAIN=1 \
SCRATCH_DIR=/directory/for/storing/results/ \
CONFIG_FILE=/path/to/config/file.yaml \
NUM_THREADS=16 \
mpiexec -n 10 -npersocket 1 python driver.py
This launches 10 MPI ranks, each with 16 CPU threads.
Algorithm: The shuffled slice sequence is partitioned across ranks:
For each slice, all ranks compute gradients locally, then:
- Each rank sends its gradients to rank 0.
- Rank 0 aggregates and averages the received gradients.
- Rank 0 updates the weight vector and broadcasts the new weights to all ranks.
The best-performing Slurm scripts for the HAI and ORNL-PS datasets are in
batch_submissions/P1_all/ and batch_submissions/ps_attack/ respectively.
The lookback_bptt=128 configuration (128.sbatch) has consistently performed
best across both datasets.
Supported datasets¶
HAI benchmark dataset¶
The HAI benchmark dataset covers industrial process control scenarios across
multiple physical subsystems. Column mappings between process variables and
state_dict/measure_dict entries are described in the
HAI Dataset technical details.
Place the dataset file at data/hai-test1.txt (or update data_file_name in
your config). Pre-built configs are in configs/P1_all/. To run training:
For distributed training, use the Slurm scripts in batch_submissions/P1_all/.
The lookback_bptt=128 configuration has performed best for this dataset.
ORNL Power Systems (ORNL-PS) dataset¶
The ORNL-PS dataset covers a power systems attack scenario targeting relay protection signals. State variables are relay voltage angle and magnitude measurements; observation variables are current measurements on the same relay channels.
Place the dataset file at data/ps_attack_with_timestamps.csv (or update
data_file_name in your config). Pre-built configs are in configs/ps_attack/.
To run training:
For distributed training, use the Slurm scripts in batch_submissions/ps_attack/.
The lookback_bptt=128 configuration has performed best for this dataset.
chi2_test/ — standalone circuit validation¶
The chi2_test/ directory contains standalone ezkl circuit tests that validate
the chi-squared hypothesis test and eigenvalue lookup circuits independently of
the full training pipeline. These tests confirm that the compiled ONNX circuits
produce numerically correct outputs before they are used in the proving workflow.
Run them to verify your ezkl installation is working correctly: