Multi-Head Attention-Based OOM Risk Prediction for CNN and Transformer Training Tasks on GPUMemNet

Abstract

Out-of-memory (OOM) failures complicate the scheduling of deep neural network training jobs because a valid model can still exceed accelerator memory under particular batch sizes, layer shapes, input resolutions, or sequence lengths. This study formulates pre-execution memory safety as binary OOM-risk classification and evaluates a tabular multi-head attention model on the GPUMemNet CNN and Transformer training-task datasets. The empirical target is the recorded Status field, with OOM_CRASH as the positive class and SUCCESSFUL as the negative class. Measured memory, GPU-utilization counters, filenames, layer-string dumps, and other runtime telemetry are excluded from predictive inputs; the models use only static task descriptors and deterministically parsed architectural parameters. Experiments use 9,000 CNN tasks and 5,011 Transformer tasks, with five stratified 70/15/15 train/validation/test splits based on seeds 42-46. The proposed model achieves 0.985 accuracy, 0.976 F1, 0.999 AUC, 0.026 MAE, and 0.108 RMSE on CNN tasks, and 0.986 accuracy, 0.929 F1, 0.998 AUC, 0.038 MAE, and 0.145 RMSE on Transformer tasks. Logistic Regression records the highest CNN F1 of 0.978, while Random Forest records the highest Transformer F1 of 0.942, showing that attention is competitive rather than uniformly dominant. On the representative split, the full static feature set raises CNN attention F1 from 0.954 to 0.978 and raises Transformer recall to 0.958. The results support lightweight, leakage-safe OOM screening as a practical first stage for GPU admission control.