Neural DAG Scheduling via One-Shot Priority Sampling
Authors: Wonseok Jeon
Mukul Gagrani
Burak Bartan
Weiliang Will Zeng
Harris Teague
Piero Zappi
Christopher Lott
Venue: ICLR 2023 poster (2023)
Keywords: Combinatorial Optimization
Directed Acyclic Graph
Scheduling
Graph Neural Network
Reinforcement Learning

Abstract: We consider the problem of scheduling operations/nodes, the dependency among which is characterized by a Directed Acyclic Graph (DAG). Due to its NP-hard nature, heuristic algorithms were traditionally used to acquire reasonably good solutions, and more recent works have proposed Machine Learning (ML) heuristics that can generalize to unseen graphs and outperform the non-ML heuristics. However, it is computationally costly to generate solutions using existing ML schedulers since they adopt the episodic reinforcement learning framework that necessitates multi-round neural network processing. We propose a novel ML scheduler that uses a one-shot neural network encoder to sample node priorities which are converted by list scheduling to the final schedules. Since the one-shot encoder can efficiently sample the priorities in parallel, our algorithm runs significantly faster than existing ML baselines and has comparable run time with the fast traditional heuristics. We empirically show that our algorithm generates better schedules than both non-neural and neural baselines across various real-world and synthetic scheduling tasks.

Published as a conference paper at ICLR 2023
NEURAL DAG SCHEDULING VIA ONE-SHOT PRIORITY SAMPLING
Wonseok Jeon", Mukul Gagrani", Burak Bartan, Weiliang Will Zeng, Harris Teague Piero Zappi, Christopher Lott Qualcomm Al Research
ABSTRACT
We consider the problem of scheduling operations/nodes, the dependency among which is characterized by Directed Acyclic Graph (DAG). Due to its NP-hard nature, heuristic algorithms were traditionally used to acquire reasonably good solu- tions, and more recent works have proposed Machine Learning (ML) heuristics that can generalize to unseen graphs and outperform the non-ML heuristics. However, it is computationally costly generate solutions using existing ML schedulers since they adopt the episodic reinforcement learning framework that necessitates multi-round neural network processing. We propose a novel ML scheduler that uses a one-shot neural network encoder to sample node priorities which are con- verted by list scheduling the final schedules. Since the one-shot encoder can efficiently sample the priorities in parallel, our algorithm runs significantly faster than existing MI. baselines and has comparable run time with the fast traditional heuristics. We empirically show that our algorithm generates better schedules than both non-ne and neural baselines across various real-world and synthetic scheduling tasks.
INTRODUCTION
The problem of scheduling operations arises across many domains, such as data centers where the incoming jobs have to be scheduled distributed server (Mao et al. 2019), manufacturing pipelines in the form of job shop scheduling problems (ISSP) (Manne) 1960), and ML compilers where the operations of a computation graph need to be scheduled on the available hardware devices (Paliwal et al. 2020; Zhhou et al., 2020).. In all these cases, the problem may be abstracted using a directed acyclic graph (DAG) where the nodes of the graph represent the operations and the edges represent the dependency constraints between the operations and hence the problem is also referred to as DAG scheduling. The objective is to minimize the finish time (or makespan) of the DAG subject to resource and dependency constraints.
It is well known that this is an NP-hard problem Kan 2012), and practitioners have traditionally relied on heuristic methods to obtain good solutions. One of the celebrated scheduling approaches list scheduling (Graaham 1969) where the idea is to schedule nodes as early as possible and to break ties using priorities. The priorities can be obtained via different node metrics which are computationally incxpensive such as critical-path based, shortest processing time or most operations remaining (Haupi 1989). More recently, researchers have proposed deep reinforcement learning based methods to solve scheduling problems (Zhang et al. 2020 Zhou et al. 2020 Wang et al.) 2021
Mao et al.] 2019). The scheduling policy in all the references utilize Graph Neural Networks (GNN) as an encoder to derive node embeddings. Zhang et al. (2020) proposed an auto-regressive GNN based policy for the JSSP problem which predicts the next node for scheduling given the nodes scheduled so far. Wang et al. (202T) proposed bi-level optimization approach which modifies the input DAG by adding multiple edges via a learned policy and then apply the critical-path heuristic on the modified DAG. One major drawback of the existing ML. based schedulers is the computational cost as they require multi-round neural network processing (encoding step). The multi-round neural network processing is reflected as ruto-regressive architecture (Zhang et al. 2020) or bi-level optimization
"qqual contribution
QQalcomm AI Research is initiative of Qualcomm Technologies, Inc.

---

Published as a conference paper at ICLR 2023
design (Wang et al., 2021). This drawback limits the scalability to large graphs and the applicability to domains where solutions need to be obtained in a timely manner (e.g., scheduling computation graphs in compilers).
In this paper, we propose a novel ML scheduler that uses a one-shot neural network encoder to sample node priorities which are converted by list scheduling to the final schedules. Since our encoder generates node priorities with a single forward pass of a neural network and efficiently samples priorities in parallel, our algorithm runs significantly faster than existing ML baselines and has comparable run time with the fast traditional heuristics.
The contributions of this paper are summarized below:
We propose a novel end-to-end approach to learn scheduling priorities for list scheduling on DAGs. Our model adopts the recently proposed Topoformer architecture (Gagrani et al. (2022) as a DAG encoder and the Gumbel-Top-k trick (Kool et al. (019b) to sample node priorities (which are acquired by perturbing the encoder's output and converted into valid schedules via list scheduling). While optimizing our model with REINFORCE (Williams) 1992), we introduce logit
norm regularization and cost standardization that significantly improve our model's representation power and performance compared to the model used in Gagrani et al. (2022).
Our approach uses the one-shot encoder which generates the node priorities by running the Topoformer encoder once. This is in contrast of existing neural baselines (Wang et al. [022 Zhang
ct al. 2020), all of which involves multi-round neural network processing. Due to the one-shot nature of model, our method runs significantly faster than neural baselines, while achieving runtimes slightly worse than yet comparable with those of computationally-efiin -fffiiicii and simple non-ML heuristics.
We show that our approach can be generally applied to variety of scheduling tasks that includes JSSP, TPC-H benchmark and scheduling for synthetic and real-world computation graphs. For all benchmarks, our model outperforms both neural and non-neura baselines w.r.t. makespan metric (Wang et al. 2021 Zhang et al. 2020).
2
PRELIMINARIES
2.1 SCHEDULING PROBLEM
In scheduling problems, we define DAG as a tuple G == EV,E, 8,8, 6, p. with set of nodes (or vertices) and a set & of directed edges (or arcs). Each node v E V represents an operation with 5(v)= 2 denoting its operational duration and p(v) 2 0 denoting the resources required to execute v. For a set . of machine types, each node v E V has to be assigned to its own machine type μ(v) E M ([M) 1 and [M] > 1 correspond to scheduling with homogeneous machines and heterogeneous ones, respectively). The set È of edges in the DAG G' represents computational dependency among nodes. For instance, for the scheduled start time r(u) 2 0, v E for each node, a directed edge (U1,v2 U2) E E, U1,11,0 U2 V, means T(01) + ō(v)) S T(22), ) i.c., any node should be scheduled on after all its predecessor nodes are finished. We assume that each type of machine m e -M has its own maximum resource limit X(m) 2 i.e., at any point of time the total amount of occupied resources for machines of type m cannot exceed A(m).
Let us introduce the vectorized notation r [T(e]]eEV E [V] of the start times with little abuse of notation for the sake of simpler notation. We define valid schedule as a vector E T where T is the set of all valid schedules (satisfying both precedence and resource constraints for given DAG G). The objective of the scheduling problem is to find 7* arg min,e7 C(T; G), where C'T; G) := max (T() + (()).. the duration required to complete all operations, is the makespan of schedule T.
2.2 LIST SCHEDULING
List scheduling Graham 1969) is class of priority-based schedule algorithms that widely adopted in practice due to their simplicity. We describe how list scheduling works as follows:
(Step 1) Input list of node priorities and set the current decision time to be zero.
(Step 2) Find ready nodes that can be scheduled at the current decision time, i.c., nodes whose predecessors have finished.
2

---

Published as a conference paper at ICLR 2023
DAG
List Scheduling
duration
node priority ready node
1.0
ready node
ready node
1.0
ready
2.0
final
node
decision
schedule
1.0
time
Figure I: An example of list scheduling (Graham 1969) for 4-node DAG scheduling is described. Each node has its own duration, while resource limit ignored for illustration purpose. List scheduling takes node priority and schedules the higher-priorit node among ready nodes earlier than the other nodes. After all ready nodes scheduled, repeat these steps until we schedule all nodes.
input (e.g.. > 3 > 2 > 4, brighter color implies higher priority)
the decision time until
set of ready nodes is found. We
(Step 3) Schedule the ready nodes sequentially at the current decision time by following the order of node priority until either all ready nodes are scheduled or further nodes cannot scheduled due to resource constraints.
(Step 4) Move the decision time to the earliest finish time over all scheduled nodes which have not finished at the current decision time and repeat (Step 2) to (Step 4) until all nodes scheduled.
We describe a simple example of list scheduling in Figure 1 Although Graham 19699) originally presented list scheduling for scheduling with homogeneous machines, we use the same definition of list scheduling for both homogeneous and heterogeneous machines.
2.3 THE GUMBEL-TORE TRICK
Consider a random variable over a set y of finite categories, where the distribution is defined by the softmax over logits(y) R, ) (the unnormalized log-probabilities), i.e., Pr(Y y} x exp(logits(y)): Y, and thus
exp(logits(y))
Pr(Y =
(1)
Σ Ey exp(logits(y'))
The Gumble-Max trick (Gumbel 1954) is a method to sample from the categorical distributions when logits characterizing the distributions are tractable. Specifically, the trick shows that by using random vector Z E R where elements Z(y).v y E V, are sampled from i.idd standard Gumbel distribution, one can randomly generate a category as follows:
arg max vey
Z(y) Pr(Y y
(2)
More recent works (Vieira 2014, Kool al., 20196) found that the Gumbel-Max trick can be extended to sample i categories without replacement, which is called the Gumbel-Top-k trick. For the extension, they introduced arg (A) which takes real vector on y and outputs sequence of elements in } that correspond to the k largest values; the output sequence of the elements should be ordered by the corresponding decreasing input values (Kool et al. (019b). As a special case where [) arg (*) becomes arg sort in decreasing values. The Gumbel-Top-k trick generates the random sequence of elements in Y
Y1, Y2, ... Yt] arg top (k) (logits(y) Z(y)).
yE)
and shows that the sequence is equivalent the from sampling & elements without replacement; note that the random vector Z is sampled once and before applying arg (A) J In other words, the
distribution of the random sequence in Eq. (3) is shown to be described as follows (Kool et al. 2019b):
exp(logits(vs)
Pr[[Y, Y2, Yıl [Wa, U2
me] EveR exp(logits(1"))
(4)
3

---

Published as a conference paper at ICLR 2023
DAG
logits
perturbed logits
node priority
List Scheduling
+
Std. Gumbel
-MLP
0
Std. Gumbel
MLLP
GNN
Std. Gumbel
arg sort
MMLP
9g
MLP 0g)
Std. Gumbel
!
K
Figure 2: Our neural scheduler for DAGS works follows: (Step 1) Generate logits for all nodes from DAG by using a GNN encoder followed by an MLP. (Step 2) Perturb logits by adding i.i.d. Gumbel random variables. (Step 3) Take argsort over perturbed logits define node priorities. Higher priority nodes have brighter colors. (Step 4) Use list scheduling to generate a schedule.
Intuitively, Eq. ( tells us that each element in the random sequence in Eq. (3) follows the categorical
distribution that characterized by the softmax over logits, where previously sampled categories excluded. In this work, we use Eq. (3) to decide the priorities over the elements in y when = (>], which will be elaborated in the next section.
3 DAG SCHEDULING WITH NEURAL PRIORITY SAMPLER
We present our one-shot neural scheduler for DAG scheduling problems and its training method in this section. The content in this section is summarized follows; we introduce the formal problem setting for the ML-based scheduling in Section 3.1. we describe how our model generates schedules by using the one-shot neural priority encoder and Tist scheduling in Section 3.2 the training method for our model and how it is relevant to the Gumbel Top-k trick are discussed in Section 3.3 with some technical details to stabilize our algorithm, which we introduced in this work.
3.1
FRAMEWORK
Suppose have set G {G1, G2. of DAGs, where each DAG G, (Vi, 6,, Pt. u4) follows the definition in Section We also assume that we have a device which is equipped with machines required for the DAGs. The
algorithms by ZZang et al. (2020) and Wang et al. (2021) aim to find out a parameterized schedule generator GT] (with the neural network parameter 0) that minimizes the average makespan over G, i.c.
arg min gg [C(T:
where 7 is the schedule (the start times of nodes) and C(T; G') is the makespan of - in G as in Section 2.1 The neural schedulers of existing works ZZhag al. 2020; Wang al. 202T) sequential decision deciiiionmaaaeeee making models that require multi-shor neural network processing, In contrast, we use a computationally efficient single-shot neural scheduler which is described in the next subsection.
3.2
SCHEDULE GENERATOR WITH ONE-SHOT PRIORITY SAMPLER
Using one-shot node priorities was recently proposed by Gagrani et al. (0222) solve peak memory minimization problems in DAGs. Our neural scheduler motivated by their idea and described in We adopt Topoformer, the graph neural network (GNN) encoder presented by Gagrani et al. as our graph encoder GNN,, (G) E (R*) where is the dimension of output embeddings for each node, and 0 is the neural network parameter of the encoder. We use MLP 0, : .RAh IR to convert the GNN's output node embeddings into logits over the nodes, i.e., for 0 = (0,, 02) and
Figure 2
G E G = GG1, G.....
logits (v; G) : MLP (GGNN,, (G)),) R.
(6)
The difference between Gagrani et al. s algorithm and ours arises from sampling procedure using logits. (Gagrani al. considers logits of schedulable nodes each decoding step and sequentially
4

---

Published as a conference paper at ICLR 2023
samples among those schedulable nodes, whereas sample node priorities only the start of decoding. Specifically, by using i..d. standard Gumbel variables Z(u) E R, v E V, and arg sort, we randomly sample sequence of nodes from perturbed logits:
[Vi, V2,,., Vuil arg sort Sorriiigsssg (v; ( Z(0)) BEV
perturbed logits
Note that the LHS of Eq. 7 is a random sequence due to the randomness of Z and arg sort is applied in decreasing values of the perturbed logits. We then regard the sampled random sequence V V= [VG V2,,, Vivi] as the sequence of node priorities V/2 > 1/2 > ... > Viyi which does not require additional computation. Due to the stochastic nature of the Gumbel-Top-ke Trick described in Section [33] our mapping from the random sequence to the node priorities is equivalent sampling nodes without replacement where nodes sampled earlier are considered to be higher-priority ones. More importantly, the trick allows us to use the tractable distribution of the random sequence, which becomes highly beneficial when optimizing the neural network. We will discuss this further in the next subsection.
Finally, use list scheduling (described in Section 2.2 with the sampled random sequence in Eq. to generate a valid schedule with the start time T:
ListSchedulingg (V; E R M 20
7 =
(8)
3.3 ALGORITHM AND PRACTICAL. CONSIDERATION
The objective of learning-to-scie tc-hceudeee frameworks in Eq. (5) can be rewritten with our scheduler generator as follows:
arg min B CLs(ViG),
where (V: == C(Liistchedulinng (V: G): G) is the makespan of list scheduling for given
node priorities V = [V..... Virll and the graph G, and ra(-IG) is the probability distribution of sampling node priorities for G. From the Gumbel-Top-k trick discussed in Section 2.3 can get
the tractable form of Ta (-IG) below:
1V
exp(logits, G))
No Te((" Mv]]lG) Pre (V [u
II
EE 1) exp(logits, (v; G))
(10)
This enables us to use the following gradient descent rule using REINFORCE (Williams 1992 that optimizes the objective Eq. ( with learning rate 0:
- 0 aBg-gEE [ log No To (VIG)Css(V;GG)
(11)
Together with the above update rule, we found that the practical techniques we introduce are crucial to stabilize the training and lead to significantly better performance, which are described below:
Logit Norm Regularization. When Gagrani et al. (2022) used their GNN encoder (Topoformer) to define the model distribution, the standardizzttion over logits was used to bound the range of logits (See Section C.2 in Gagrani et al.) (2022)). Specifically for the mean m and standard deviation over logits(v)) (v), v E V, and scalar hyperparameter x (Iogits(n) - ms/s is used to define the model's probability of topological ordering. However, we empirically observe that such standardization over logits leads to poor performance. We believe this is because standardization significantly restricts the model's 's representation capability. While a rigorous proof for our argument may be complicated, we can easily show that this is true for binary random variables. Specifically with the above definition of standardized logits and a random exp((e) variable X (0,1), 1), one can only represent (Pr(X 0), Pr(X -1}) (1+0xp(20)) T+exp(2c)
0. the standardized logits logits(v) :=
( (Teexp(2c]> I+exp(2c)) (See Appendix A Note that the constant e is a hyperparameter and assumed to be fixed during training, and the example implies that only two distributions can be described for any given e, which supports our claim.
exp(2c)
5

---

Published as a conference paper at ICLR 2023
Algorithm 1 Neural DAG Scheduler via One-Shot Priority Sampling
Input: A set
{G1, ......} } of the training graphs, node priority sampler ne, learning rate a > 0, a regularization coefficient Clogits 0.
1: for each epoch do
2: for each G E g do
Sample batch node priorities va), ( ((M) (+(-)
4: Convert the priorities to valid schedules by list scheduling. Evaluate makespan for all sampled schedules.
5:
6:
Standardize makespan C. < (C(V()) va), V(N), G) n I, N, by Eq.
13
7:
Compute REINFORCE gradient SREINFORCE + N En Va Vogog To (n))) G)Cn. Compute logit norm gradient Plogits & Clogits
8:
V Loggits (0: G) with Eq. 12
9: Update A by using gradient descent: 0 &- 0 a a(GEEINFORCE + flogits).
10: end for
11:
end for
Output: m
After removing the standardization of logits due to the above observation, however, we still observe instability in training caused by the unbounded logits. Therefore, we introduce the logit norm regularizer that minimizes
1 V
Ligitt (0; G)
logits,( (v; G)2²
(12)
by gradient descent, together with the aforementioned REINFORCE objective. Intuitively, the norm regularizer allows our model to have the logits around the origin so that we can avoid numerical errors due to the unbounded logits while the sufficient amount of flexibility is still maintained for choosing logits. We empirically found that this highly stabilizes training and improves the performance and hence apply the regularizer in all of experiments.
Cost Standardization. Gagrani et al. 2022) stored the best-performing model so far and used its cost as the policy gradient baseline. This so-called greedy baseline was introduced due the empirical performance of Kool et al. (2019a)'s algorithm on routing problems. However, if the scale of makespan varies significantly across different training graphs, the model trained with the greedy baseline can easily overfit a small subset of training graphs, which may lead to poor performance on test graphs. Also, the greedy bascline requires additional resources since the intermediate models should be stored and evaluated during training. To address these issues, we use cost standardization that has been widely adopted policy-gradient algorithms, e-g., PPO (Schulman et al. [2077] 2015). Specifically for given graph G during training iteration, we first sample multiple node priorities, ie., (1)) ...., (N() no( ()). and evaluate makespan CLs (LV(())) G), .... CLs (V))) G). During the policy gradient update, we use the standardized makespan across the samples, i.e.,
C(V; (()) DIN, G) = Cr.s(V;)) mean N.... [C1.s (VO;
(13)
max (Cr.s(VV (V("); G)),c)
where 0 is used to clip the standard deviation in denominator for numerical stability.
The final algorithm with additional stabilization ideas is summarized in Algorithm [ We train model with multiple training epochs over G where single epoch considers each graph in G once. For each graph G E G, we randomly sample N priorities and use REINFORCE to update the model parameter 8 while regularizing the norm of logits by Eq. 12). We multiply with constant Aogits > 0 the regularization loss to balance between REINFORCE loss and regularization. We set
Aogits 0.001 and observe that it empirically works well in all of our experiments.
RELATED WORKS
ML for Combinatorial Optimization. The idea of using ML solve DAG scheduling fits into the broader theme of using ML for Combinatorial Optimization (CO) which has received attention recently (Bengio al. 2021). Most of the works in the literature in MI for CO use RL learn

---

Published as a conference paper at ICLR 2023
policy to select actions for reward maximization which is set to be a direct function of the problem objective. The policy can be an end-to-end policy whose actions correspond to the the solution of the
CO problem (Zhou et al. 2020 Kool et al. 2019a Joshi et al. [021] Khalil et al. 2017 Zhang et al.,
2020) or the policy can augment a traditional heuristic/solver of the problem to find better solutions
(Paiwall et al. 2020 Xin et al. 202 Ahn et al. 2020: Wang al. 2021).
End-to-End ML Schedulers. In the context of the scheduling problems, [Zhang et al. (2020) and Park al. (202I) proposed an end-to-end GNN based policy to solve the JSSP problem which a special case of DAG scheduling. Their policy is auto-regressive which selects the nodes to be scheduled iteratively and at each iteration they run the GNN on the modified disjunctive graph to get new node embeddings. Thus, they are required to run the GNN encoder |.V] times which is prohibitive for large graphs. In order to address the complexity of auto-regressive policies for large graphs encountered in compiler workflows. /Zhou et al. (2020) came up with the idea of iterative refinement which refines the generated schedule by running their GNN policy K times where K is
[Sun] et al. (2021), Mao et al. (2019) and Zhou et al. (2022) consider the problem of scheduling jobs on data clusters and provide end to end deep RL sol solutions to solve data center scheduling under various settings.
Hybrid Schedulers. Wang et al. (2021) propose a bi-level optimization approach for DAG scheduling where they learn policy which modifies the input DAG by adding edges and then use the critical-por path based list scheduling method on the modified DAG to obtain a schedule for the original problem. The authors allow upto K edges be added the DAG (where K is a hyper-parameter and their policy can run the GNN K times to get to the final schedule. Paliwal et al. (2020) proposed a neural augmented genetic algorithm for scheduling in compilers. They used GNN policy learn the parameters of the mutant distribution which was used by the genetic algorithm to find good schedules.
In contrast to these works, our method generates the node priorities end-to-end requiring only a single pass of our GNN encoder and uses list scheduling to obtain the final schedule. This makes our approach more efficient and scalable compared to prior works.
5 EXPERIMENTS
We evaluate our neural DAG scheduler and baselines for three different scheduling tasks: JSSP, DAG scheduling TPC-H dataset, and scheduling computation graphs (See Appendix B for details about all tasks.). We compare our algorithm with both non-neural and neural baselines. For non-neural baselines, consider list scheduling algorithms with different node priorities (based on Critical-Path (CP), Most Operations Remaining (MOPNR) and Shortest Processing Time (SPT)) (Zhang et al. 2020 , and constraint programming (Const. Prog.) solver (CP-SAT by Perron & Furmon) with 24 hours time limit. For neural baselines, we first consider
(I.2D) algorithm by Zhang et al. (2020) in the experiment with synthetic JSSP instances, where we train the deployed implementation using the dataset. We also consider PPO-BiHyb (Wang et al.) 2021). where we use the deployed model of PPO-BiHyb for TPC-H dataset and train the model for the dataset with real-world computation graphs.
For our model, we consider two types of operating modes. Greedy mode, our model generates schedules with node priorities that uses argsort sort and the pure logits without adding Gumbel random variables. In Sampling mode, we sample multiple node priorities (where the number of samples are chosen from 16, 64 and 256, each of which corresponds to S(16), S(64), S(256) in Tables) and output the best schedule from the priorities. We report the makespan for output schedules and the run time of each algorithm. We report the speedup metric, which is defined as the ratio of sum duration of all nodes and the makespan, in the computation graph scheduling task. For all tasks and operating modes, we left the empirical result to show the effectiveness of our practical techniques in Appendix H
5.1 JSSP
We evaluate our model on randomly generated JSSP instances from Zhang et al. 2020) which are defined by number of jobs N, and number of machines Nm. We summarize the results for NN). (25, 20), (25, 30), (50, 20) in Table | For each case, we train our algorithm for 100 training graphs and evaluate it for 50 unseen test graphs. We also train the neural baseline, L2D (Zhang et al. 2020), with the same training and test graphs. Note that this is different from the original
7

---

Published as a conference paper at ICLR 2023
Table 1: Experiment results on synthetic JSSP instances are described. We use bold letters to emphasize the minimum average makespan for each JSSP instance.
(Nj, Nm) = (25, 20)
(N₃, Nm. (25,30)
(N,, Nm)
(50, 20) Time (sec)
Makespan Time (sec)
Makespan
Time (sec)
Makespan
2120.24
0.009
2588.72
0.017
3290.80
0.024
SPT
2265.88
0.002
2739.20
0.005
3548. 20
0.008
MOPNR
2115,86
0.012
2625.96
0.023
3278.82
0.030
L2D
2253.94
1.245
2799.00
2.850
3452.70
3.711
Greedy (ours)
2032.70
0,021
2512.40
0.031
3108.56
0,049
S(16) (ours)
1970,98
0.054
2452.64
0.085
3032.44
0.138
S(64) (ours)
1948.76
0.127
2427.30
0.294
3009.08
0,469
S(256) (ours) 1932.42
0.514
2411.68
0.909
2997.10
1.527
Table Experimental results on TPC-H datasets are described, We use bold letters emphasize the minumim average makespan.
TPC-H-50
TPC-H-100
TPC-H-150 Makespan
Makespan
Time (sec)
Makespan
Time (sec)
Time (sec)
Const, Prog.
8629.4
19278.3
CP
9821.3
0.008
16914.1
0.027
24429.5
0.048
SPT
12818.4
0.002
19502.7
0.008
27409.4
0.021
MOPNR
11360.1
0.011
17733.1
0,032
24871.2
0.064
PPO-BiHyb
8905.4
66.484
15192.2
149.215
22371.2
571.424
Greedy (ours)
8845.6
0.057
14981.2
0.100
22332.7
0.259
S(16) (ours)
8782.4
0.114
14972.0
0.287
22330.2
0.674
S(64) (ours)
8742.5
0.216
14968.
0.699
22323.0
1.856
S(256) (ours)
8694.4
0.540
14964.7
2.270
22320.8
6.485
set of training graphs for every training iteration. The results show that our algorithm works well across different numbers of machines and also outperforms both neural and non-neural baselines. Due one-shot decoding scheme in our model, the running time of our model is shorter than L2D. We observe that L.2D ends up with worse performance than non-neural baselines although train L.2D until convergence. We think L2D does not generalize well with the limited number of training graphs, while our algorithm generalizes well with the same dataset.
training of 1.2D since 1.21) generates
5.2 TPC-H DATASET
TPC-H dataset includes DAGs that consists of industrial queries and data modifications which represent computation jobs and need to be scheduled on homogenous machine with finite resources. We use Wang al.//s PPC-550/PCC100/TPP performance of baselines and our model on the test set. We observe that our method obtains better average makespan than all the baselines including PPO-BiHyb on all the three instances except in TPC-50 where constraint programming achieves slightly better makespan. One more thing to note is that our method has much smaller run time compared to PPO-BiHyb. This is because our method generates node priorities via a single pass of our GNN encoder and samples priorities effectiviely by using Gumbel Top-k trick, whereas PPO-BiHyb has to run their GNN encoder multiple times and requires beam search.
datasets for our experiments. Table 2 shows the
5.3 COMPUTATION GRAPHS
We test our approach on scheduling tasks for both synthetic and real-world computation graphs of neural networks that arise in ML compilers. We consider three type of synthetic computation graphs:
http://pp..rgghhhhhdeffaaaaaaa
8

---

Published as a conference paper at ICLR 2023
Table Experiment results on synthetic computation graphs are described. We use bold letters to emphasize the maximum average speedup for each graph distribution.
Layered Graph
Erdos- Renyi
Stoc. Block Model
SpeedUp
Time (sec)
SpeedUp
Time (sec)
SpeedUp
Time (sec)
CP
4.580
0.058
5.049
0.055
4.701
0.055
SPT
4.526
0.013
4.541
0.008
4.473
0.007
MOPNR
4.745
0.075
5.112
0.068
4.761
0,068
Greedy (ours)
4.819
0.094
5.194
0.046
4.866
0,043
S(16) (ours)
4.848
0.311
5.211
0.214
4.916
0.198
S(64) (ours)
4.872
0.750
5.227
0.542
4.944
0.477
S(256) (ours)
4.889
2.418
5.239
1.799
4.965
1.533
Table 4: Experimental results real-world computation graphs described. We use bold letters emphasize the maximum average speedup for each test set of graphs.
200 - 500 Node Graphs
500
700 Node Graphs
700
1000 Node Graphs
SpeedUp
Time (sec)
SpeedUp
Time (sec)
SpeedUp
Time (sec)
Const. Prog.
3.267
3.183
2.497
CP
3.174
0.007
2.804
0.016
2.739
0.025
SPT
3.107
0.002
2.868
0.005
2.664
0.008
MOPNR
3.181
0.009
2.825
0.020
2.739
0.028
PPO-BiHyh
3,223
17.937
2.965
52.777
2.798
322.793
Greedy (ours)
3.245
0.152
3.131
0.098
2.846
0.060
S(16) (ours)
3,271
0.192
3.188
0.245
2.848
0.230
S(64) (ours)
3,278
0.263
3.199
0.456
2.856
0.606
S(256) (ours)
3.286
0.595
3.207
1.309
2.860
2.001
Erdos-Renyi and stochastic block model graphs (Paliwal et al., 2020) (See appendix (C.3). Table 3 shows the performance of different methods on the synthetic computation graphs with 1000 nodes. We observe that our method outperforms the traditional priority based baselines in terms of the achieved speedup (higher speedup implies lower makespan) for three graph distributions.
layered graphs (Gagrani et al.
2022
also note that the run time our method is competitive with that of the fast heuristics showing the scalability of our approach. Note that for scheduling applications in ML compilers, low run time scheduler is crucial, and our results show the applicability of our model for such applications.
We also experiment on a set of proprietary real-world computation graphs to evaluate the practical applicability of our neural scheduler. This dataset consists of computation graphs of diverse neural net architectures like classifiers, convolution nets, denoisers, etc. We observe that our method achieves superior speedup on graphs of all sizes compared all non-neural baselines. We also compare our approach with PPO- PPO-BiHyb and constraint programming baseline in this setting. Our results show that we outperform both -Po-Biyyb and constraint programming in terms of the achieved speedup. Note that we generate better schedules and run much faster than PPO-BiHyb The fact that we beat constraint programming shows the large search space of this problem and the effectiveness of our approach to learn to find good schedules in this large space.
CONCLUSION AND FUTURE WORK
We propose end-to-end approach to solve the general problem of scheduling over DAGs. Our method uses a single pass Topoformer encoder to learn the node priorities which are used with list scheduling generate schedule. We also apply the Gumbel Top-k trick to efficiently sample multiple node priorities and obtain better schedules. We demonstrate the effectiveness of our approach variety of tasks which include JSSP problems, TPC-H dataset and compiler scheduling datasets. We believe that our proposed approach can extended to other CO problems such as routing problems (Kool al. 2019a) and consider it promising direction for future work.

---

Published as a conference paper at ICLR 2023
REFERENCES
Byung Hoon Ahn, Jinwon Lee, Jamie Menjay Lin, Hsii-PPa Cheng, Jilei Hou, and Hadi Esmacilzadeh. Ordering chaos: Memory-aware scheduling of irregularly wired neural networks for edge devices. In Machine Learning and Systems (MLSys), pp. 44-57, 2020.
Yoshua Bengio, Andrea Lodi, and Antoine Prouvost. Machine learning for combinatorial optimization: a methodological tour d'horizon. European Journal of Operational Research, 290(2):405-421, 2021.
Blank and Deb. pymoo: Multi-objective optimization in python. IEEE Access, 8:89497-89509, 2020.
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arXiv preprint arXiv:2012.09699 2020.
Mukul Gagrani, Corrado Rainone, Yang Yang, Harris Teague, Wonseok Jeon, Roberto Bondesan, Herke van Hoof, Christopher Lott, Weiliang Will Zeng, and Piero Zappi. Neural topological ordering for computation graphs. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Ronald Graham. Bounds on multiprocessing timing anomalies. Mathematics, 17(2)416-429, 1969.
SIAM journal on Applied
Emil Julius Gumbel. Statistical theory of extreme values and some practical applications: series of lectures, volume 33. 1954.
Reinhard Haupt. A survey of priority rule-based scheduling. 3-16, 1989.
11(1):
Chaitanya K Joshi, Quentin Cappart. requires rethinking generalization. In International Conference Principles and Practice of Constraint Programming (CP), 2021.
Rousseau, and Thomas Laurent. Learning TSP
AHG Rinnooy Kan. Machine scheduling problems: classification, complexity and computations. Springer Science Business Media, 2012.
Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial optimization algorithms over graphs. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Wouter Kool, Herke van Hoof, and Max Welling Attention, learn to solve routing problems! In International Conference on Learning Representations (ICLR), 2019a.
Wouter Kool, Herke Van Hoof, and Max Welling. Stochastic beams and where to find them: the Gumbel-Top-k trick for sampling sequences without replacement. In International Conference Machine Learning (ICML), pp. 1499-3508, 2019b.
Alan S Manne. On the job-shop scheduling problem. Operations Research, 8(2):219-223, 1960.
Hongzi Mao, Malte Schwarzkopf, Shaileshh Bojja Venkatakrishnan, Zili Meng, and Mohammad Alizadeh. Learning scheduling algorithms for data processing clusters. In ACM Special Interest Group Data Communication (SIGCOMM), pp. 270-288, 2019.
Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs. In International Conference on Learning Representations (ICLR), 2020.
Junyoung Park, Jachyeong Chun, Sang Hun Kim, Youngkook Kim, and Jinkyoo Park. Learning to schedule job-shop problems: representation and policy learning using graph neural network and reinforcement learning. International Journal of Production Research, 59(11):3360-3377, 2021.
Laurent Perron and Vincent Furnon. OR-tools. URL /ttppsss/ddeeaooeerr..oooggggp/p
optiim in zation. on/
10

---

Published as a conference paper at ICLR 2023
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv: 150:::240333 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:/707.0634 2017.
Penghao Sun, Zehua Guo, Junchao Wang, Junfei Li, Julong Lan, and Yuxiang Hu. Deepweave: Accelerating job completion time with deep reinforcement learning-basd -basdd coflow scheduling. In International Conference on International Joint Conferences Artificial Intelligence (IJCAI), pp. 3314-3320, 2021.
Tim
Vieira.
Gumbel-Max
trick and weighted reservoir sampling.
2014.
URL
gumbel 1 max trick wtw.eeighhheedreeeeeeeeeoeooooee
Runzhong Wang, Zhigang Hua, Gan Liu, Jiayi Zhang, Junchi Yan, Feng Qi. Shuang Yang. Jun Zhou, and Xiaokang Yang. A bi-level framework for learning to solve combinatorial optimization on graphs. In Advances in Neural Information Processing Systems (NeurIPS), pp. 21453-21466, 2021.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning. 8(3):229 256, 1992.
Liang Xin, Wen Song, Zhiguang Cao, and Jie Zhang. NeuroLKH: Combining deep learning model with Lin- Lin-Kernighan-- Kemmighnn-Helsgauue Helsgaun heuristic for solving the traveling salesman problem. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Cong Zhang, Wen Song, Zhiguang Cao, Zhang, Puay Siew Tan, and Xu Chi. Learning to dispatch for job shop scheduling via deep reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), pp. 1621-1632, 2020.
Yanqi Zhou, Sudip Roy, Amirali Abdolrashidi, Daniel Wong, Peter Ma, Qiumin Xu, Hanxiao Liu, Phitchaya Phothilimtha, Shen Wang, Anna Goldie, al. Transferable graph optimizers for MI. compilers. In Advances in Neural Information Processing Systems (NeurIPS), pp. 13844-13855, 2020.
Yunfan Zhou, Xijun Li, Jinhong Luo, Mingxuan Yuan, Jia Zeng, and Jianguo Yao. Learning to optimize dag scheduling in heterogeneous environment. In IEEE International Conference on Mobile Data Management (MDM), pp. 137-146, 2022.

---

Published as a conference paper at ICLR 2023
A PROBABILITY DISTRIBUTION WITH STANDARDIZED LOGITS
Suppose our model has logits to == (0g and 4 == l1g. One can easily show that the
mean and standard deviation of logits are equal to ملياا and V / m² 16-4
2
respectively. Now for the hyperparameter > 0 that is assumed to be fixed, consider the standardized logits
To = (0) = c lo m I == logits(1) = 4 m
(14)
8
and model a binary random variable X E {0, 1] by using them. The probability distribution using the softmax and the above standardized logits becomes
exp(li)
Pr[X = 1) exp(lo) + exp(11
1 exp(lo - -11 = -1- Pr[X - 0}.
(15)
Since the difference between lo and [ is
-- 1 = cxx
to m
m
c x
1- 2c x
1o 1
(16)
8
110 - 11
we have
if lo > 41
Pr(X 1
1texp(20 2c)
exp 2c)
otherwise.
1+exp(2e)
and 01 =
=1
B SCHEDULING TASKS
In this subsection, we describe how the different tasks that we consider fit into our scheduling framework described in section 2.1
Job Shop Scheduling Problem (JSSP). JSSP special case of DAG scheduling where a set of N, jobs need to be scheduled on NM machines in M. Each job consists of sequence of operations that must go through NM machines in a specific order. This task can be modeled in the framework of section 2.1 by setting M 21, NM,, we have one machine of each type i.e. X(m) = 1, Vm E M, each node has a machine type it can be scheduled on i.c. μ(v) E M, each node occupies the machine it scheduled on i.e. p(v) Lu V.
DAG Scheduling on TPC-H. In this task we need to schedule nodes with single machine type (homogeneous case) i.e. M O and all the nodes have same machine type i.e. p(v) = O, Ye E V. We assume the same setting as in Wang et al. (2021) for experiments and set the available resource of machine to X(0) = 6000. Node occupies p(v) E with p(v) S 6000 resources where the set of positive integers.
Computation Graph Scheduling. In this task we need to schedule computation operations of DAG hardware with 3 types of machines M {0, 1, 2]. We assume that we have 1 machine of type and type 1 whereas 4 machines of type i.c. X(0) = X(1) = 1, A(2) = 4. This hardware setting is inspired from the structure of real ML accelerators which have multiple machines (threads) of different types available for data processing. A node v can be scheduled on one machine of its machine type u(v) M i.e., the resource required p(e) 1, Vu V.
C DATASET
C.I SYNTHETIC JSSP INSTANCES
We generate synthetic JSSP instances by using the code deployed by Zhang et al. (2000) For reproducibility of this dataset, we fix our seed as 0 and sample 155 instances for all combinations of N2 E [25, 50}, Nn, E {10, 20, 30), respectively. We use the first 100 instances as training instances and the last 50 instances as test instances.
https: / / thub. zca caros//aoooa http://www....aeo.../..oo//////a
12

---

Published as a conference paper at ICLR 2023
C.2 TPC-H DATASET
We use TPC-H dataset generated by Wang et al. 22021 The dataset for each of TPC-50, TPC-100, TPC-150 experiments consists of 50 training graphs and 10 test graphs. The average number of nodes for TPC-50, TPC-100 and TPC-150 dataset are 467. 2, 929.8 and 1384.5 respectively.
C.3 SYNTHETIC COMPUTATION GRAPH DATASET
Layered graphs. Layered graphs were introduced in Gagrani et al. (202)) as a model to generate synthetic graphs which have similar structure the computation graph of neural network. We use the default parameters specified in Gagrani et al. (2022) to generate the graphs. In order to generate the node duration, we first sample the memory size m(v) of each node E and then use the following affine model to generate the duration 8(v) for node U E V,
((v) = round(100 x m(v)) + 1.
(17)
We sample m(v) by first sampling the value from Gaussian Mixture Model (GMM) m(() ~ GMM(w, H. (4. R__. and projecting it to non-negative values. We use a mixture of four Gaussians and set the means (mean), meang, meang, mean4) (0.5, 3, and standard deviations to (sid1, std2, sldg, stda) = (0.5, 1,
We sample the machine type pu() of node such that Pr(p(v) j) X(j) i.e. the probability that a node has machine type is proportional to the available resources X(j) for machine type j. In our computational graph experiments we used A(0) X(1) = I, (() which leads the following multi-nomial distribution for u(v):
1/6, if
Pr(p(o)= j)
if = 1, 2/3,
1/6,
(18)
if 2.
Erdos-Renyi and Stochastic Block Model. These two graph distributions are well-known families of random undirected graphs and were used in Paliwal et al. (2020) for their experiments. We set the probability of an edge between any two nodes p 0.05 for the Erdos-Renyi graphs. We use the following parameters for the stochastic block model: the number of communities k 4, the probability for an edge between two nodes of same community p.n = 0.3. the probability for edge between two nodes of different communities Pour = 0.001. We use the procedure described in Paliwal et al. (2020) convert the instances from these graph families into DAG. We the same distribution to generate node duration and their machine type as in layered graph for these two graph distributions as well.
The results for 1000-node graphs are shown in the main part of our work, and 500-node graph results in Appendix
D BASELINE ALGORIHTM
D. PRIORITY-BASER LIST SCHEDULING BASELINES
We consider classical list scheduling baselines that use the following values as node priorities Zhang et al. 2020); Wang et al. 202T))
Pritical-Path (CP). The priority for each node is determined by the length of the critical path from the node to the target node in DAG.
Most OPeratioNs Remaining (MOPNR). The priority for each node is computed using the number of operations remaining on the longest path from the node to the target node in a DAG.
Shortest Processing Time (SPT). The priority for each node computed using the inverse of the processing time of the node.
13

---

Published as a conference paper at ICLR 2023
D.2 CONSTRAINT PROGRAMMING BASELINE
Constraint programming models are commonly used to solve scheduling problems the literature. We have developed constraint programming formulations and carried out numerical simulations to provide a baseline comparison for our methods.
1. Constraint programming for computation graph scheduling. The developed constraint pro- gramming problem has start times of jobs and machine assignments as discrete variables. The constraints consist of precedence constraints, correct machine assignment constraints, and con- straints for enforcing no overlap of tasks if they are assigned to the same machine.
Constraint programming for DAG scheduling with TPC-H dataset. The constraint program- ming problem for the TPC-H dataset experiments has only the start times of jobs as variables. The constraints ensure that the precedence relations a DAG are satisfied, and the total resource usage does not exceed the given resource budget. In addition, the task durations for this dataset have been scaled by 1000 and rounded to the nearest integer. The reason for this step is that the CP-SAT solver only works with integer valued data, and the duration values for this dataset are of floating point precision. We have solved the constraint programming problem using the scaled and rounded duration values. We have then divided the resulting makespan by 1000. This rounding step introduces an error of at most 0.0005 seconds per node. Overall, for a 1000-node graph, this leads to an error of most 0.5 seconds in the worst case. This is negligibly small since the makespan values for this dataset are in the order of few thousand seconds.
We have implemented these formulations using Google OR-Tools and solved them using the CP-SAT solver (Perron & Furnon). Differently from other baselines, the constraint programming solver is guaranteed to find the optimal schedule given enough time. Due to practical reasons, we have set a time limit of 24 hours for the solver in the simulations. In most of the experiments, the time limit was reached before finding the optimal solution. In these cases, we have reported the best result obtained prior the moment of time-out. Across all experiments, the solver has found the optimal solution for only 3 graphs in the TPC-H-50 dataset. For these 3 jobs, the amounts of time that the solver has run find the optimal solutions are 2429, 21949, 39274 seconds, which are still significantly larger than our algorithms' run time.
We have observed that the memory requirements of the solver for large graphs could exceed the available memory in our servers. The graphs in TPC-H-150 fall into this category, where the experiments resulted in out-of-memory errors. We have not reported results for these graphs in Table 2
In the experiments, we have provided initial feasible schedules for the solver. This is sometimes referred to as solution hinting or warm start. This speeds up the solver considerably. Initial feasible schedules could be obtained in different ways. One way is to generate a topological order of the DAG and map it to sequential schedule without any parallelization. Another way is to run list scheduling and initialize the solver using the output of the list scheduling algorithm. We have experimented with both options. We have found that using the output of list scheduling leads to better initialization for the solver.
D.3 NEURAL BASELINES
In this section we provide how our neural baselines are used. For all the tasks, we use a machine with a single GPU (Nvidia Tesla V-100) with 32 GB memory that is also used to train and evaluate our model.
L,D. While the original implementation of L.2D uses new set of training graphs for every training iteration, we consider more practical scenario where the number of training graphs is restricted. We use the aforementioned 100 training graphs and 50 test graphs for each JSSP instance. We use environments (that was proposed by Zhang et al. (2020)) which ends up with training 4 graphs per training iteration. We train the L2D model for 3000 iterations over 5 different seeds and report the one that performs the best. Note that tested learning rates 2 10 (the default learning rate in their implementation) and 1 x 10 for L2D and the learning rate x 10 -4 works much better, and thus report the results for this case.
14

---

Published as a conference paper at ICLR 2023
Table 5: Experiment results on synthetic JSSP instances (N₂ 25).
(N,, Nm) (25, 10) Time (sec)
(N,,,
(25, 20)
(N,. Nm
(25, 30)
Makespan
Makespan
Time (sec)
Makespan
Time (sec)
CP
1673.56
0.005
2120.24
0.009
2588.72
0.017
SPT
1807.90
0.001
2265.85
0.002
2739.20
0.005
MOPNR
1656.26
0.007
2115.86
0.012
2625.96
0.023
L2D
1725.58
0.577
2253.94
1.245
2799.00
2.850
Greedy (ours)
1578.58
0.018
2032.70
0.021
2512,40
0.031
S(16) (ours)
1540.92
0,045
1970.98
0.054
2452.64
0,085
S(64) (ours)
1527.40
0.092
1948.76
0.127
2427.30
0.294
S(256) (ours)
1519.16
0.321
1932.42
0.514
2411.68
0.909
Table 6: Experiment results synthetic JSSP instances (N₃ 50).
(N,,
(50, 10) Time (sec)
Nm)
(50, 20)
Makespan
Makespan
Time (sec)
CP
2903.14 14
0.012
3290.80
0.024
SPT
3111.74
0.003
3548.20
0.008
MOPNR
2897.60
0.016
3278.82
0.030
L2D
2964.96
1.084
3452.70
3.711
Greedy (ours)
2814.58
0.024
3108.56
0.049
S(16) (ours)
2800.32
0.064
3032.44
0.138
S(64) (ours)
2797.84
0.165
3009.08
0.469
S(256) (ours)
2796.98
0.644
2997.10
1.527
PPO-BiHlb. For experiment with TPC-H dataset, we use the pretrained mode deployed by the author. We confirmed that the results reported in their paper are reproducible and evaluate the run time in our machine settings. For computation graph scheduling, we train PPO-BiHyb for 40 epochs (where 1 epoch is defined train all training graphs in the dataset once).
E
ADDITIONAL RESULTS
E.
JSSP
We describe empirical results for NN)) (25, 10), (25, 20), (25, 30) in Table and (N) (50, 10), (50, 20) in Table 6 For broader settings, we can see that our neural scheduler outperforms our baselines.
E.2 COMPUTATION GRAPHS
We describe empirical results for both 500 and 1000 node synthetic computation graphs in Table Over all sizes and graph distributions, our algorithm achieves better speedup within short time.
F TRAINING AND MODEL DETAILS
F. TRAINING
In this section we provide the training details of our model on the different scheduling tasks that consider. For all the tasks we trained our model machine with single GPU (Nvidia Tesla V-100) with 32 GB memory. We used Adam optimizer for training on all the tasks.
15

---

Published as a conference paper at ICLR 2023
Table Experiment results synthetic graph datasets are described.
Layered Graph Time (sec)
Prdoos-Renyi
Stoc. Block Model
SpeedUp
SpeedUp
Time (sec)
SpeedUp
Time (sec)
CP SPT
4.314
0.025
4.973
0.015
4.664
0.018
500 Node Graph
4.214
0.004
4.480
0.002
4.416
0.003
MOPNR
4.428
0.033
5.033
0.018
4.715
0.023
Greedy (ours)
4.513
0.042
5.113
0,017
4.835
0.021
S(16) (ours)
4.558
0.129
5.152
0.082
4,901
0.082
S(64) (ours)
4.585
0.269
5.168
0.182
4.935
0.185
S(256) (ours)
4.603
0.803
5.178
0.545
4.956
0.539
1000 Node Graph
CP
4.580
0.058
5.049
0.055
4.701
0.055
SPT
4.526
0.013
4.541
0.008
4.473
0.007
MOPNR
4.745
0.075
5.112
0.068
4.761
0.068
Greedy (ours)
4.819
0.094
5.194
0.046
4.866
0.043
S(16) (ours)
4.848
0.311
5.211
0.214
4.916
0.198
S(64) (ours)
4.872
0.750
5,227
0.542
4.944
0,477
S(256) (ours)
4.889
2.418
5.239
1.799
4.965
1,533
JSSP. In this case we train our model on each JSSP instance with 100 training graphs and test it on 50 unseen graphs. We train our model for 20 epochs with random seeds and pick the best performing model. We use the number of samples N 1000.
TPC-H Dataset. We train our model on each TPC instance for 100 epochs with 10 random seeds and pick the best performing model. We use the number of samples N 2000.
Computation Graphs. For each synthetic graph distribution, we consider the graph size equal to either 500 or 1000 and generate training set of 3000 graphs and 300 unseen test graphs. For the real-world graphs, our dataset consists of 92 train graphs and 23 test graphs. We train for our model for 20 epochs with random seeds and pick the best performing model. We the number of samples N = 1000 for both synthetic and real computation graphs.
We summarize the time and epochs required for convergence in Table g
F.2 MODEL ARCHITECTURE
We use topoformer encoder with the same hyperparameters as described in Gagrani et al. (2022). We
Cogits
0.001 and 0.1 for clipping the denominator in Eq. (13).
F.3
INPUT FEATURES
We use the following input node features Z, for node E V:
Node duration 8(v)
Node resource requirement p(v)
One hot representation of node machine type p(v)
The critical path duration from to the target node and critical path duration from source node to
Note that for tasks where pvv) 1, Ve we ignore p(u) from the input feature. We normalize each entry of the input features across the nodes so that the features are in between 0 and as follows:
Z.
EV L.
In addition, we also augment the node features with the Laplacian positional encodings Dwivedi & Bresson (2000) of dimension 20 by computing it on the undirected version of the DAG. We pass the input feature T, through linear layer to obtain the initial embedding of node v for the topoformer encoder.
16

---

Published as a conference paper at ICLR 2023
Table 8: Training time details for the experiments. We report the rough duration for the convergence time for training and number of epochs to convergence for each of the experiment reported in the paper.
Experiment name
Convergence time
Time per epoch
# epochs to converge
Real- computation graphs
25 min
min
8
Syn. Layered (500)
2.5 hour
1.2 hour
2
Syn. Rerdos-Renyi (500)
hour
1.2 hour
2.5
Syn. Stoc. Block Model (500)
3.5 hour
1.2 hour
3
Syn. Layered (Ik)
hour
2.4 hour
1.25
Syn. Erdos- Erdos-Renyl Renyi (1k)
8 hour
2.4 hour
3.3
Syn. Stoc. Block Model (Ik)
hour
2.4 hour
3.3
TPC-H dataset TPC-50
hour
3.6 min
16.7
TPC-H dataset TPC-100
hour
14.4 min
20.8
TPC-H dataset TPC-150
16 hour
28.8 min
33
JSSP dataset (25, 10)
30 min
2 min
15
JSSP dataset (25, 20)
45 min
4.5 min
10
JSSP dataset- (25, 30)
hour
7.5 min
8
JSSP dataset (50, 10)
45 min
4.5 min
10
JSSP dataset- (50, 20)
hour
12 min
15
G PERFORMANCE IMPROVEMENT RELATIVE TO GAGRANI ET AL. (2022)
We use Topoformer Gagrani et al. 2022 in our method and summarize the key differences between Gagrani al. (2022) method and ours below:
Our method
1. Aims solve scheduling problems
2. Perturbs logits (via i.i.d. Gumbel noise) to sample node priorities (not necessarily satisfying precedence constraints) and convert the priorities into schedules via list scheduling.
3. While using REINFORCE, cost standardization is used.
4. Logits are regularized, i.e., its 1.2 norm is minimized together with REINFORCE loss.
Gagrani et al. (2022)ss 's method
1. Aim to solve peak memory minimization (which is not a scheduling problem).
2. Logits are used sample sequences satisfying the precedence constraints on DAGs.
3. While using REINFORCE, greedy baseline used, motivated by [15].
4. Logits standardized.
For items and 4, our motivations are described as follows. REINFORCE with greedy baseline in Gagrani et al. 2002 uses the policy gradient
Bv~ne
Vo log no (=II)) - -C(Ta
for fixed G, where is the greedy policy and its parameter 0' is copied from 0 when the policy with shows the best performance during training (normally evaluated the end of multiple epochs).
The problems of greedy baseline are:
it may become unstable cost scales differ too much across training graphs,
2. it slows down training since further forward computation with is required,
3. it requires additionally memory to store 6',
all of which can be resolved by the cost standardization technique that used in our paper. Also, we use logit standardization used by Gagrani et al. (2022), we empirically observed performance degradation, which believe due to logit standardization restricting the representation power of the model and motivated us to use logit norm regularization instead (example in Appendix A).
To prove the effectiveness of our technical improvements, did ablation studies for our algorithm with all tasks our submission. Specifically. we compare our method with Gagrani et al. (2022)*s method using both greedy baseline and logit standardization while maintaining the priority sampling
17

---

Published as a conference paper at ICLR 2023
Table Experiment results on synthetic JSSP instances with ablation studies. Greedy and Sampling methods with (-) indicate results from the model trained with greedy baseline and logit standardization Gagrani et 2022).
(Nj, Nm) = (25, 20)
(N,, Nm.
(25,30) 30) Time (sec)
(N N,.)
(50, 20)
Makespan
Time (sec)
Makespan
Makespan
Time (sec)
CP
2120.24
0.009
2588.72
0.017
3290.80
0.024
SPT
2265.88
0.002
2739.20
0.005
3548.20
0.008
MOPNR
2115.86
0.012
2625.96
0.023
3278.82
0.030
L2D
2253.94
1.245
2799.00
2.850
3452.70
3.711
Greedy (-)
2077.60
0.022
2610.86
0.028
3129.16
0,056
S(16) (-)
2004.78
0,055
2493.38
0.083
3055,20
0.142
S(64) (-)
1979.86
0.129
2467.00
0.291
3037.38
0.484
S(256) (-)
1957.76
0.512
2446.04
0.844
3015.78
1.568
Greedy (ours)
2032.70
0.021
2512.40
0.031
3108.56
0.049
S(16) (ours)
1970.98
0.054
2452.64
0.085
3032.44
0.138
S(64) (ours)
1948.76
0.127
2427.30
0.294
3009.08
0.469
S(256) (ours) 1932.42
0.514
2411.68
0.909
2997.10
1.527
via Gumbel and list scheduling (this necessary since we have to fairly compare performances in scheduling domains). We train the model for all scheduling tasks; JSSP, DAG scheduling, real and synthetic computation graphs.
G.1 JSSP
Our method is shown to perform better than Gagrani et al. (2022)s 's method (See Table 9 Also, note that our method always outperform Gagrani et al. (2022)* 's method when they are in the same mode and the same task. Runtimes for both methods are comparable with each other. We observe that Gagrani et al. (0022> 's method can outperform our neural and non-neural baselines for our JSSP instances.
G.2 DAG SCHEDULING ON TPC-H DATASET
Our method outperforms Gagrani et (2022) method (See Table IO]. Runtimes for both methods comparable with each other. For TPC dataset, our method can be regarded to make significant improvement due to the following observations:
1. S (256) of Gagrani et al. (2022)'s 's method (which shows the lowest makespan among the runs with Gagrani et al. (2022)'s method) is always outperformed by Greedy mode of our method (which the highest makespan among our methods).
2. Our neural baseline, -Po-BiHyb] Wang et al. (202T always outperforms Gagrani et al. 2022
method for all TPC tasks.
G.3 SCHEDULING COMPUTATIONAL GRAPHS
G.3.1 SYNTHETIC GRAPHS
Our method outperforms Gagrani et al. (2022) 's method (See Table Runtimes for both methods comparable with each other. For synthetic computation graphs, the following observations make our method more outstanding:
1. S (256) mode Gagrani al. 2022)> 's method (that shows the lowest makespan among the runs using Gagrani et al. (2022)ss 's method) is always outperformed by Greedy mode of our method (which the highest makespan among our methods).
2. Our non-neural baseline, MOPNR, always outperforms Gagrani et al. (002 's Greedy mode and some of sampling modes.
18

---

Published as conference paper at ICLR 2023
Table 10: Experiment results on TPC-H dataset with ablation studies. Greedy and Sampling methods with (-) indicate results from the model trained with greedy baseline and logit standardization in
Gagrani et al. (202).
TPC-H-50
TPC-H-100
TPC-H-150
Makespan
Time (sec)
Makespan
Time (sec)
Makespan
Time (sec)
Const. Prog.
8629.4
19278.3
CP
9821.3
0.008
16914.1
0.027
24429.5
0.048
SPT
12818.4
0.002
19502.7
0.008
27409.4
0.021
MOPNR
11360.1
0.011
17733.1
0.032
24871.2
0.064
PPO-BiHyb
8905.4
66.484
15192.2
149.215
22371.2
571.424
Greedy (-)
9300.4
0.219
16185.5
0.139
23788.9
0.143
S(16) (-)
9079.9
0.322
15974.1
0.340
23477.6
0,425
S(64) (-)
9037.3
0.467
15804.1
0.748
23412.2
1.233
S(256) (-)
8976.2
0.959
15684.9
2.338
23264.4
4.419
Greedy (ours)
8845.6
0.057
14981.2
0.100
22332.7
0,259
S(16) (ours)
8782.4
0.114
14972.0
0.287
22330.2
0.674
S(64) (ours)
8742.5
0.216
14968.1
0.699
22323.0
1.856
S(256) (ours)
8694.4
0.540
14964.7
2.270
22320.8
6.485
Table 11: Experiment results 1000-node synthetic computation graphs with ablation studies.
Greedy and Sampling methods with (-) indicate results from the model trained with greedy baseline and logit standardization in Gagrani al. (2022).
Layered Graph
Hirdos-Reny's
Stoc.
Block Model
SpeedUp
Time (sec)
SpeedUp
Time (sec)
SpeedUp
Time (sec)
CP
4.580
0.058
5.049
0.055
4.701
0.055
SPT
4.526
0,013
4.541
0.008
4.473
0.007
MOPNR
4.745
0.075
5.112
0.068
4.761
0.068
Greedy (-)
4.597
0.051
4.871
0,047
4.657
0,060
S(16) (-)
4.706
0.196
4.989
0.219
4.792
0.217
S(64) (-)
4.740
0.515
5.028
0.529
4.833
0.536
S(256) (-)
4.767
1.730
5.056
1.716
4.863
1.775
Greedy (ours)
4.819
0.094
5.194
0.046
4.866
0.043
S(16) (ours)
4,848
0.311
5.211
0.214
4.916
0.198
S(64) (ours)
4.872
0.750
5.227
0.542
4.944
0.477
S(256) (ours)
4.889
2.418
5.239
1.799
4.965
1.533
G.3.2 REAL-WORLD GRAPHS
Our method outperforms Gagrani et al. (2022) (See Table [2
method and shows comparable runtimes for all tasks
19

---

Published as a conference paper at ICLR 2023
Table 12: Experiment results on real-world computation graphs with ablation studies. Greedy and Sampling methods with (-) indicate results from the model trained with greedy baseline and logit standardization in Gagrani al. (2022)
500 Node Graphs SpeedUp
200
500 700 Node Graphs
700
1000 Node Graphs
Time (sec)
SpeedUp
Time (sec)
SpeedUp
Time (sec)
Const. Prog
3.267
3.183
2.497
3,174
0.007
2.804
0.016
2.739
0,025
SPT
3.107
0.002
2.868
0.005
2.664
0.008
MOPNR
3.181
0.009
2.825
0.020
2.739
0.028
PPO-BiHyb
3,223
17.937
2.965
52.777
2.798
322.793
Greedy (-)
3,227
0.196
3.114
0.028
2.822
0,045
S(16) (+)
3,248
0.243
3,151
0.153
2.843
0,238
S(64) (-)
3,254
0.324
3.179
0.404
2.847
0.656
S(256) (-)
3,258
0.644
3.192
1.317
2.853
2.205
Greedy (ours)
3.245
0.152
3.131
0.098
2.846
0.060
S(16) (ours)
3,271
0.192
3.188
0.245
2.848
0.230
S(64) (ours)
3,278
0.263
3.199
0.456
2.856
0.606
S(256) (ours)
3.286
0.595
3.207
1.309
2.860
2.001
H COMBINING SEARCH ALGORITHM WITH NEURAL SCHEDULER
also conducted ablation studies to study the effect of using more powerful search algorithm with method and analyze the performance-rus time tradeoff against using sampled priorities. We combine our neural scheduler with genetic algorithm (GA) search. The initial population for GA is obtained by sampling 256 node priority vectors using learned model. We let GA run for 100 generations with a population size of 1000. We use the publicly available GA implementaiton of Pymoo (Blank & Deb. 2020) for our experiments.
Computation graph scheduling (synthetic). We observe that for layered graphs and stochastic block model, using GA improves the performance marginally whereas there improvement for the Erdos-RRnyi graphs. Also, the GA search increases the runtime significantly compared to the runtime of our method sampling mode S(256).
Table 13: Experiment results on synthetic computation graphs with GA.
Layered Graph SpeedUp
Erdos-Renyi
Stoc. Block Model
Time (sec)
SpeedUp
Time (sec)
SpeedUp
Time (sec)
S(256)
4.889 4.997
2.418 723.245
5.239
1.799 707.659
4.965
1.533 700.248
S(256) GA
5,239
4.982
Computation graph scheduling (real-world). We found that using GA improves the speedup of schedules, while requiring significantly more computation time. The time required for GA increases the graph size grows.
Table 14: Experiment results on real-world computation graphs with GA.
200 - 500 Node Graphs
500 - 700 Node Graphs
700 1000 Node Graphs SpeedUp
SpeedUp
Time (sec)
SpeedUp
Time (sec)
Time (sec)
S(256)
3.286
0.595
3.207
1.309
2.860
2.001
S(256) GA
3.304
151.331
3.240
403.829
2.882
687.288
20

---

Published as a conference paper at ICLR 2023
DAG Scheduling (TPC dataset). For TPC-H-50, observed that GA search can improve the performance, whereas either no improvement (TPC-H-100) or negligible improvement (TPC-H-150) --50) was observed for other datasets. Similar to computation graph scheduling tasks, significantly longer computation time is required, and the required time increases as the graph size increases.
Table 13: Experimental results on TPC-H datasets with GA.
TPC-H-50
TPC-H-100
TPC-H-150
Makespan
Time (sec)
Makespan
Time (sec)
Makespan
Time (sec)
S(256)
8694.4 8646.2
0.540 246.942
14964.7 14964.7
2.270 965.119
22320.8 22320.7
6.485 2089.547
S(256) GA
21


Reviews:
[{"review_id": "a3AyAdkmeL", "reviewer": "Reviewer_p2nE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to use policy gradients to learn the permutation of node sequences for applications in scheduling operations. Specifically the model takes a pre-defined DAG structure which describes the job dependencies as input, and produces the logits for each job. The permutation is sampling via gumbel-topk trick and then REINFORCE gradient is used to update the policy network. Experiments on some synthetic and real-world scheduling tasks show some improvement over existing approaches.\n", "summary_of_the_review": "Overall this paper presents a simple yet efficient approach for the scheduling problem. There are several aspects that need to be justified, in terms of the 1) experimental quality and thoroughness; 2) technical correctness/improvements; 3) generality to more CO problems that optimize for a permutation.\n", "strength_and_weaknesses": "Strength:\n- The paper is well written and the background and problem is described clearly.\n- The empirical results seem to be good.\n\nWeakness:\n- Technically the paper is relatively straightforward and lacks novelty in terms of the technical contribution. Leveraging policy gradient or RL in general for combinatorial optimization has been studied over the past several years. The main contribution of the paper is to use a variant of the parameterization of the policy, where the encoder is from another recent paper and the gumbel-topk used for decoding from the encoded logits is not new. \n- Continuing with above, the technical part can be potentially improved in following aspects: \n1) It seems that the current approach doesn’t guarantee that the generated permutation satisfies the topological order specified by the DAD (correct me if I’m wrong). It would be nice to make it more compatible with the problem that is being solved, instead of using a generic approach in a straightforward way. \n2) The nice thing about the gumbel trick is that it admits a reparameterization. If the author can build something that allow end to end optimization with gradient propagates from the objective to the graph encoder, then it would be nicer and more efficient than REINFORCE, hopefully.\n\n- The paper claims that the one-shot decoding is beneficial. However it seems that it is a surrogate of the model used in [1], as mentioned by the authors in Sec 3.2. So a natural question would be, how does [1] perform in the experiments presented in the paper? The hypothesis is that [1] would yield better results but worse speed, so a trade-off is expected. More comprehensive comparison with [1] would be needed for the experimental sections.\n\n- The application for the scheduling seems to be specific, but the technique using GNN + gumbel-topk with RL seems to be general. To make it naturally aligned with the technique, there could be several ways to improve:\n1) include other combinatorial optimizations that optimize for a permutation, like the TSP mentioned by the authors;\n2) make the approach more appropriate and adapted to the DAG based optimization, as I mentioned above about the techniques.\n\nReferences:\n[1] Neural topological ordering for computation graphs, Gagrani et.al\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JCpI425XlS", "reviewer": "Reviewer_HjB3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose solving dag scheduling problems by using GNNs to provide node priorities in the list scheduling algorithm. The GNN is trained using good old REINFORCE. The training is made stable using logit renormalization and cost standardization, as introduced by the authors.\n\nThe authors introduce two versions of their Greedy and Sampling. In Sampling, they sample multiple node priorities from an encoded distribution. Greedy is a special case of Sampling, i.e. S(1).'\n\nThe authors evaluate their technique on several baselines, both real and synthetic. Their technique is consistently the best among all relevant baselines.", "summary_of_the_review": "Overall, the paper is interesting, well written and well motivated. However, given that the novelty is incremental, the strength of the paper would lie in its empirical evaluation, hence I encourage the authors to explore more avenues such as how would their technique work using a more advanced search algorithm as proposed above.\n", "strength_and_weaknesses": "The novelty of the technique is mostly incremental - the proposed technique piggy back's on REINFORCE, Gumbel top-k trick and the list scheduling algorithm. The novelty is in how to make training with REINFORCE stable, which can be notoriously hard to get right.\n\nWhat I found most interesting about the work was the authors' decision to create a Sampling version of the algorithm -- One important feature of the proposed technique is that the trained model can be used zero-shot once it has been trained on a dataset, making it extremely efficient compared to other neural-baselines. This is similar to the REGAL technique in Paliwal et al except that they use a genetic algorithm instead of the random search that the authors use in Sampling. Given that Sampling is always stronger (as it can't be worse than Greedy), I ask two things\n- why didn't the authors consider using a more powerful search algorithm instead of just random search in Sampling?\n- why didn't the authors consider using the cost of the of the best solution found by the search algorithm as the reward in REINFORCE instead of just using Greedy?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "6qjPb-0f2By", "reviewer": "Reviewer_zenr", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper considers node/operation scheduling within DAGs, which is previously solved using heuristics but more recently approached with ML-based solutions. ML techniques are typically more costly, but the authors employ a “one-shot” encoder networks to sample node priorities in a single forward pass of the network. These priorities are then converted via list scheduling to the final schedules. This encoder can sample priorities in parallel, which makes it much faster than baselines. Although the method is still a bit slower than heuristic techniques in most cases, the proposed method is shown to actually outperform both heuristic and ML-based techniques. \n", "summary_of_the_review": "Personally, I am not familiar with reinforcement learning. As such, I am completely open to discussion of my comments with other authors/reviewers. This review reflects my initial impression of the paper given that I am somewhat unfamiliar with this area (DAG scheduling/RL). I emphasize my final score will be mostly based upon subsequent discussion that I hope will allow me to gain more perspectives on the work. \n\nI think the paper is well-written, motivated nicely, and novel (to the best of my knowledge). Empirically, the method performs very well (best in almost every case). As such, I feel that this paper provides a nice contribution to the community that is novel (AFAIK), performant, and practical. In its current form, my biggest recommendation to the authors would be to provide more details on training complexity and requirements. To better understand how using your method compares to using heuristic baseline methods, this discussion is needed. Yes, the method performs very well. But, practitioners may not use it if too much effort is required. Currently, I lack from the paper an idea of where your method stands on this spectrum -- is it practical enough to actually use?\n\nNonetheless, I thank the authors for the nicely-written and interesting paper. I think the paper is a good contribution, and I look forward to future discussion!\n\n", "strength_and_weaknesses": "Strength:\n* The problem being addressed by the proposed method is clearly motivated and explained. It also has a big impact, as any scheduling approach must have low runtime to be usable in practical scenarios. The approach considered within this work is novel to the best of my knowledge (i.e., the differences from Gagrani et. al. seem to be significant/non-trivial, as you must find a way to obtain the schedule in a single shot). \n* The use of Gumbel variables and the Gumbel top-k trick is explained well and interesting. \n* The method seems to perform well empirically relative to neural and heuristic baselines.\n\nWeakness:\n* It would be good to provide discussion of how long it takes to train your method, how much data is needed, etc. This is probably an additional latency consideration that would be of interest to anyone trying to practically use your method. I think this discussion is needed to more thoroughly compare your approach to common, heuristic methods on a basis of practicality. \n* Time is still worse compared to heuristic baselines, but it is still close. It seems like to achieve the best performance, more time is needed. However, this is somewhat unavoidable for approaches that leverage neural networks. \n* Outperformed by constraint programming on TPC-H, but this doesn’t seem to be a big issue. It is only in one case and performance of the proposed method is still pretty similar. \n\n\nSmall Comment:\n* Logit norm regularization explanation is well done. I like how you motivate the problem that exists using binary variables, then propose a solution that works well empirically. This explanation makes sense and is non-trivial. \n* In the summary of contributions “(w.r.t to the makespan)”, you can remove the “to”.\n\nQuestions:\n* Is REINFORCE the only way to train this algorithm? Are there potentially other methods (maybe even non-RL) that would be less data hungry or provide a better result? It would be good to understand a little better why REINFORCE is chosen. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster