SAFE-TSFormer — PyTorch Implementation

Faithful implementation of "A safety-aware deep learning framework for scalable schedulability analysis of variable-length real-time task sets" (Behera & Singh, Real-Time Systems, 2026).

Files

File	Contents
model.py	The architecture itself: TaskFeatureEmbedding (Φ, §4.4), LinearMultiHeadAttention + SAFEEncoderLayer (Ψ, §4.5), SetAttentionPooling (§4.6), SchedulerHead (Γ_S, §4.7), full SAFETSFormer model, and SafetyAwareLoss (§4.8).
data.py	UUniFast task synthesis, log-uniform periods, implicit/constrained deadlines (§5.1), and exact ground-truth schedulability tests: processor-demand/DBF analysis for EDF, response-time analysis (RTA) for RMA/DMA (§3.4).
train.py	Algorithm 2 (safety-aware training, AdamW + ReduceLROnPlateau + early stopping) and Algorithm 3 (screening workflow) + the accuracy/precision/recall/F1/FPR evaluation protocol of §6.

Equation → code map

Paper	Code
Eq. 3–6 (feature vector f_i)	model.build_feature_vector
Eq. 11 (task embedding Φ)	model.TaskFeatureEmbedding
Eq. 12–15 (masked linear attention Ψ)	model.LinearMultiHeadAttention, model.SAFEEncoderLayer
Eq. 16–18 (set-level attention pooling)	model.SetAttentionPooling
Eq. 19–20 (scheduler heads Γ_S)	model.SchedulerHead
Eq. 21/22 (safety-aware loss)	model.SafetyAwareLoss
Alg. 1 (inference)	SAFETSFormer.forward
Alg. 2 (training)	train.train
Alg. 3 (screening)	train.screen

Design choices made explicit but underspecified in the paper (documented here rather than silently assumed):

• Linear-attention kernel: the paper only requires O(nd) masked linear attention reproducing softmax(QK^T/√d)V behaviour; we use the standard elu(x)+1 feature map (Katharopoulos et al.), a common linear-attention choice consistent with the stated complexity and masking requirements.
• Period range: sampled log-uniformly over [10, 1000] (paper cites Mall 2009 for "a wide range of time scales" without giving exact bounds).
• EDF exact test: implemented as processor-demand/DBF analysis with a bounded busy-period testing-point set (equivalent in exactness to QPDA/DBF, differing only in convergence speed, which does not affect correctness of the ground-truth labels).
• Everything architecturally load-bearing — embedding dimension d=256, the two-stage Φ/Γ_S MLP shape, per-scheduler independent head weights, masked attention pooling, and the exact loss formula — is implemented verbatim from the equations, with no shortcuts.

Verified correctness (see smoke tests run during development)

• Forward pass produces one logit per scheduler (EDF, RMA, DMA).
• Permutation invariance: reordering tasks in a set leaves predictions unchanged (< 1e-6 numerical diff), confirming Ψ + pooling respect the set-based nature of schedulability (§4.1).
• Padding invariance: predictions are identical regardless of how much zero-padding is added, confirming the mask M is applied correctly throughout (§4.5).
• Exact EDF/RMA schedulability tests behave correctly on known toy cases (low-utilization schedulable sets, U>1 overloaded sets).
• Full training loop (Algorithm 2) runs end-to-end without errors: forward, safety-aware loss, backward, AdamW step, validation, ReduceLROnPlateau, early stopping, and final evaluation on both in-distribution and generalization (n ∈ {24,28,32}) splits.

Running

pip install torch --break-system-packages
python train.py --epochs 40 --train_per_card 3000 --lam 2.0

To reproduce something closer to the paper's scale you'll want the full ~225k-task-set training corpus (§5.1) and more epochs; the smoke run above used tiny synthetic samples purely to validate correctness of the code.

Next step: adapting to IMC (mixed-criticality) scheduling

This implementation is structured so the IMC extension is a small, targeted change rather than a rewrite:

• Feature vector (model.build_feature_vector / data.Task): extend to carry per-criticality-level (C_i^LO, C_i^HI, χ_i) instead of a single C_i, and add a criticality-aware density/utilization pair.
• Scheduler heads (model.SchedulerHead, SAFETSFormer.heads): swap the {EDF, RMA, DMA} head set for whatever IMC protocols you're targeting (e.g. AMC, EDF-VD) — the shared encoder Φ/Ψ/pooling doesn't need to change.
• Ground-truth labeling (data.py): replace edf_schedulable / rta_schedulable with exact IMC schedulability tests (e.g. AMC-rtb, or EDF-VD's demand-bound conditions per mode).
• Safety-aware loss (model.SafetyAwareLoss): the false-positive penalty structure carries over unchanged and is arguably more important for IMC, since false positives there mean an HI-criticality task missing its deadline in mode-switch.

Let me know which IMC test/dataset you're targeting and I'll make the specific edits.