# A Learning Method with Gap-Aware Generation for Heterogeneous DAG Scheduling

arXiv: https://arxiv.org/abs/2603.23249 · HTML: https://mlsanity.com/paper/2603.23249/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 54 references that are in the map.
- First posted on arXiv 24 Mar 2026.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [POMO: Policy Optimization with Multiple Optima for Reinforcement Learning](https://mlsanity.com/paper/2010.16011/) (2020)
- [Reinforcement learning](https://mlsanity.com/paper/2405.10369/) (2024)
- [Learning to Dispatch for Job Shop Scheduling via Deep Reinforcement Learning](https://mlsanity.com/paper/2010.12367/) (2020)
- [Learning Scheduling Algorithms for Data Processing Clusters](https://mlsanity.com/paper/1810.01963/) (2018)
- [GLOP: Learning Global Partition and Local Construction for Solving Large-scale Routing Problems in Real-time](https://mlsanity.com/paper/2312.08224/) (2023)
- [Learning to Search Feasible and Infeasible Regions of Routing Problems with Flexible Neural k-Opt](https://mlsanity.com/paper/2310.18264/) (2023)
- [MVMoE: Multi-Task Vehicle Routing Solver with Mixture-of-Experts](https://mlsanity.com/paper/2405.01029/) (2024)
- [Learning to schedule job-shop problems: Representation and policy learning using graph neural network and reinforcement learning](https://mlsanity.com/paper/2106.01086/) (2021)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [A Bi-Level Framework for Learning to Solve Combinatorial Optimization on Graphs](https://mlsanity.com/paper/2106.04927/) (2021)
- [Deep Reinforcement Learning Guided Improvement Heuristic for Job Shop Scheduling](https://mlsanity.com/paper/2211.10936/) (2022)
- [Flexible Job Shop Scheduling via Dual Attention Network Based Reinforcement Learning](https://mlsanity.com/paper/2305.05119/) (2023)
- [Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs](https://mlsanity.com/paper/1905.02494/) (2019)
- [Neural Topological Ordering for Computation Graphs](https://mlsanity.com/paper/2207.05899/) (2022)
- [Revocable Deep Reinforcement Learning with Affinity Regularization for Outlier-Robust Graph Matching](https://mlsanity.com/paper/2012.08950/) (2020)
- [Learning-Guided Rolling Horizon Optimization for Long-Horizon Flexible Job-Shop Scheduling](https://mlsanity.com/paper/2502.15791/) (2025)
- [Learning to Optimize DAG Scheduling in Heterogeneous Environment](https://mlsanity.com/paper/2103.06980/) (2021)
- [Learning Coordination Policies over Heterogeneous Graphs for Human-Robot Teams via Recurrent Neural Schedule Propagation](https://mlsanity.com/paper/2301.13279/) (2023)
- [Learning Topological Representations with Bidirectional Graph Attention Network for Solving Job Shop Scheduling Problem](https://mlsanity.com/paper/2402.17606/) (2024)
- [Inductive-bias-driven Reinforcement Learning For Efficient Schedules in Heterogeneous Clusters](https://mlsanity.com/paper/1909.02119/) (2019)
- [Geometric Deep Reinforcement Learning for Dynamic DAG Scheduling](https://mlsanity.com/paper/2011.04333/) (2020)
- [LMask: Learn to Solve Constrained Routing Problems with Lazy Masking](https://mlsanity.com/paper/2505.17938/) (2025)
- [Edge Generation Scheduling for DAG Tasks Using Deep Reinforcement Learning](https://mlsanity.com/paper/2308.14647/) (2023)

## Built on by (most cited first, 1 of 1)

- [Learning to Solve the Quadratic Assignment Problem with Warm-Started MCMC Finetuning](https://mlsanity.com/paper/2604.20109/) (2026)

## Abstract (arXiv)

Efficient scheduling of directed acyclic graphs (DAGs) is a core problem in large-scale data-intensive computing systems, where query plans, data-processing workloads, and computation graphs consist of dependent tasks competing for limited heterogeneous resource pools. In practice, achieving high-performance execution requires schedulers to adapt across environments with varying resource pools and task types, while generating schedules under tight runtime budgets. We propose WeCAN, an end-to-end reinforcement learning framework for heterogeneous DAG scheduling that addresses task-pool compatibility coefficients and generation-induced optimality gaps. It adopts a two-stage single-pass design: a single forward pass produces task-pool scores and global parameters, followed by a generation map that constructs schedules without repeated network calls. Its weighted cross-attention encoder models task-pool interactions gated by compatibility coefficients, and is size-agnostic to environment fluctuations. Moreover, widely used list-scheduling maps can incur generation-induced optimality gaps from restricted reachability. We introduce an order-space analysis that characterizes the reachable set of generation maps via feasible schedule orders, explains the mechanism behind generation-induced gaps, and yields sufficient conditions for gap elimination. Guided by these conditions, we design a skip-extended realization with an analytically parameterized decreasing skip rule, which enlarges the reachable order set while preserving single-pass efficiency. Experiments on real-world TPC-H query DAGs, resource-intensive workload datasets, and ML-compiler computation graphs demonstrate improved makespan over strong baselines, with inference time comparable to classical heuristics and faster than multi-round neural schedulers.