Attention-Informed Surrogates for Navigating Power-Performance Trade-offs in HPC

AI summary

Overview

• Research area: Machine learning for systems / high-performance computing (HPC) scheduling, with a focus on multi-objective Bayesian optimization, attention-based feature embeddings, and power-performance trade-off modeling.
• Technical level: Intermediate. The paper assumes familiarity with surrogate modeling, Bayesian optimization, Pareto fronts, and common supervised learning tooling, but the framing of the scheduling problem itself is accessible to a general ML reader.
• Scope: The paper proposes and evaluates a surrogate-assisted multi-objective Bayesian optimization (MOBO) framework that uses attention-informed embeddings of HPC job telemetry to jointly optimize job runtime and power consumption by selecting the number of nodes for a job, tested on two production HPC job-log datasets (PM100 and Adastra).

What This Paper Is About

HPC schedulers and users must decide how many nodes to request for a job, and this choice is a trade-off: too few nodes risks job termination or long runtimes, while too many nodes wastes energy. Traditional heuristics such as First-Come, First-Served with backfilling optimize a single objective (usually runtime) and cannot expose the conflict between runtime and power.

The paper's goal is to automate that node-count decision with a data-efficient, surrogate-driven MOBO framework. It tests two hypotheses: that attention-based embeddings improve surrogate quality over direct regressors (H1), and that MOBO captures runtime-power trade-offs better than single-objective BO (SOBO) or random baselines (H2).

Key Contributions

1. A surrogate-driven MOBO framework that integrates attention-based embeddings (implemented with TabNet) and intelligent, active-learning-style sample acquisition to model runtime and power from irregular, noisy HPC telemetry.
2. A demonstration across two large-scale HPC datasets (PM100 and Adastra) that both H1 and H2 hold consistently.
3. A reproducible pipeline for surrogate training and multi-objective optimization that explicitly balances runtime and power trade-offs as a decision-support tool for node allocation.
4. The authors state this is, to their knowledge, the first work to apply embedding-informed surrogates inside a MOBO framework to the HPC scheduling problem, jointly optimizing performance and power on production workloads.

Main Findings

• Attention-based embeddings beat a transformer-style regressor (H1 supported): Embedding-informed lightweight regressors (Random Forest, XGBoost, LightGBM) outperformed a TabNet regressor across both datasets in most settings. The authors attribute this to subsampling shrinking the dataset, which limits deep architectures that need large data to generalize.

• Large gains on PM100: With embeddings, Spread was reduced by approximately 99%, and hypervolume (HV) was orders of magnitude higher under SOBO.

• Large gains on Adastra: Embeddings improved HV by 37% and reduced Spread by approximately 90%.

• MOBO produces better Pareto fronts than baselines (H2 supported): MOBO improved HV by 24% versus SOBO-Runtime on PM100 (measured with the TabNet Regressor) and by 37% versus SOBO-Runtime on Adastra. MOBO was best in 3 of 4 model families (75%) on both datasets.

• Intelligent sampling cuts data needs: Stable, accurate surrogates were achieved using only 50-75% of the data. At 50%, both datasets maintained consistent HV and Spread; at 75%, surrogate accuracy (MAPE approximately 0.99) plateaued, and 100% gave no further gains over 75%.

• Sampling improves Pareto diversity: On PM100, Spread dropped from 2.5 × 10^5 to below 10^4. On Adastra, sampling suppressed erratic HV values from 10^16 to 10^13, producing smoother trade-offs.

• Sampling also reduces compute time: Total execution time on PM100 dropped from 6480.78s to 6009.46s, and on Adastra from 2699.43s to 1961.68s (a 27% decrease). The largest savings came from surrogate training (runtime and power models), while MOBO evaluation time increased slightly due to repeated acquisition steps.

• Datasets used: PM100 contains 231,238 job records with 35 features; Adastra contains 15,285 job records with 35 features. Under sampling, PM100 train sizes were 73,983 (50%), 77,278 (75%), and 109,202 (100%); Adastra train sizes were 3,964 (50%), 4,163 (75%), and 4,547 (100%).

Methodology in Plain English

The researchers treat "how many nodes should this job get?" as a two-goal optimization problem: minimize runtime and minimize power. Because these goals conflict, no single answer is best, so the framework produces a Pareto front, a set of solutions where improving one objective would worsen the other.

Three components do the work:

1. Intelligent sample acquisition. Instead of training on every raw log entry, the pipeline builds a numeric view of the data, trains lightweight per-target predictors, and scores each sample by how hard it is (absolute error for regression targets, 1 minus the probability of correct prediction for classification targets). Sampling probabilities are a clipped, linearly scaled function of that difficulty score, with a small loop to tune the scaling factor to hit a target sampling rate; each sample is then kept or dropped by an independent Bernoulli draw. This emphasizes informative samples while keeping a floor for exploration, and it reduced training set sizes substantially.

2. Attention-based embeddings. TabNet learns feature-level attention directly on structured data, highlighting informative features and suppressing redundant or noisy ones. Those embeddings are then fed to lightweight regressors (Random Forest with n_estimators=100 and max_depth=10, XGBoost with n_estimators=100, max_depth=6, learning_rate=0.10, and LightGBM with the same settings). The authors deliberately avoided full transformer architectures, citing their resource intensity, large data requirements, and limited interpretability for HPC deployment.

3. MOBO with the surrogates in the loop. The trained surrogates predict runtime and power for candidate node allocations, and a multi-objective Bayesian optimization loop searches for Pareto-optimal configurations. The acquisition function is qLogExpectedHypervolumeImprovement (logEHVI) with a Sobol sampler and 128 normal samples; each iteration takes q=1 candidate and uses optimize_acqf with 5 restarts and 32 raw samples, for 300 iterations. A multi-output Gaussian process (BoTorch) is used with a reference point inferred dynamically from the current Y. The design variable is num_nodes_alloc.

Baselines are SOBO on runtime only, SOBO on power only, and Random Search (5 seeds with the total point budget split evenly across seeds). Experiments ran on the Stampede3 supercomputer at TACC using CPU nodes. Preprocessing handled missing values, aggregated node-level power to job-level totals, and converted categorical, numeric, and time-series-derived signals (GPU bursts, I/O spikes) into a single structured input. Evaluation metrics are Hypervolume and Spread in the raw minimization space.

Why This Matters

• Research impact: The paper connects two previously separate threads, MOBO methodology (e.g., qEHVI-style acquisition) and ML-based HPC runtime prediction, and shows it is feasible on noisy production telemetry. It also reports a counterintuitive result for the deep learning era: with aggressive subsampling, simple tree-based models over attention-derived embeddings outperformed a transformer-style regressor.

• Automating a routine but error-prone user decision: Users currently guess node counts, risking job termination or wasted energy. The framework acts as decision support, recommending balanced node allocations.

• Sustainability and cost on large facilities: Jointly modeling runtime and power addresses facility-wide power constraints and operational cost on multi-million-dollar HPC systems, per the authors' framing.

• Extending prior measured benefits: The authors build on their earlier work (Dey et al., 2025) showing accurate runtime prediction can cut time-to-science by 71% and resource usage by 42%, and extend this toward balancing runtime against power consumption.

• Real-world applications:

  • Job submission and scheduling on national-lab and academic supercomputers where users must pick a node count.
  • Batch queue policy design, where schedulers could recommend resource configurations alongside job placement.
  • Datacenter energy management, using runtime-power Pareto fronts to reason about power caps.
  • Workflow and campaign planning for computational science teams deciding resource requests for large job sets.

• Industry relevance: Any organization running shared compute clusters (cloud HPC, on-prem clusters, AI training infrastructure) faces the same runtime-versus-power trade-off. The pipeline's reliance on lightweight regressors and reduced training data makes it more attractive for deployment under scheduler latency budgets than heavy transformer surrogates, though the paper notes end-to-end latency profiling, inference-time reporting, and BO wall-time are not yet reported and remain future work.

Future Directions

1. Deployment and latency: Profile end-to-end overheads by reporting surrogate inference time and BO cycle wall-time, and tighten the loop via cached embeddings and warm-started retraining to meet scheduler budgets.
2. Scalability in objectives: Extend beyond two objectives and assess complexity/quality trade-offs using scalarization warm-starts and NEHVI/log-NEHVI variants with controlled candidate sets.
3. Model choice: Examine when lightweight regressors suffice versus transformer-based surrogates under different data regimes, with emphasis on stability, small-sample behavior, and compute/latency footprints.
4. Generalization across systems: Study portability across clusters via re-embedding or adapter heads with minimal calibration runs, and quantify the impact of observed domain gaps.

Target Audience

• HPC researchers and practitioners working on scheduling, resource allocation, and performance modeling.
• ML researchers interested in surrogate modeling, multi-objective Bayesian optimization, and attention over tabular/telemetry data.
• Systems and datacenter engineers concerned with power-performance trade-offs and energy-aware computing.
• Graduate students and newcomers to "ML for systems" who want a concrete, end-to-end example of embedding a learned surrogate into an optimization loop, including a documented sampling mechanism, hyperparameters, and baseline setup.

Note: The summary above reports only figures stated in the paper. Items such as surrogate inference time, BO cycle wall-time, and results beyond two objectives are explicitly not reported in this work and are listed by the authors as future work.

Authors’ abstract

High-Performance Computing (HPC) schedulers must balance user performance with facility-wide resource constraints. The task boils down to selecting the optimal number of nodes for a given job. We present a surrogate-assisted multi-objective Bayesian optimization (MOBO) framework to automate this complex decision. Our core hypothesis is that surrogate models informed by attention-based embeddings of job telemetry can capture performance dynamics more effectively than standard regression techniques. We pair this with an intelligent sample acquisition strategy to ensure the approach is data-efficient. On two production HPC datasets, our embedding-informed method consistently identified higher-quality Pareto fronts of runtime-power trade-offs compared to baselines. Furthermore, our intelligent data sampling strategy drastically reduced training costs while improving the stability of the results. To our knowledge, this is the first work to successfully apply embedding-informed surrogates in a MOBO framework to the HPC scheduling problem, jointly optimizing for performance and power on production workloads.