Robust Scheduling with GFlowNets
Authors: David W Zhang
Corrado Rainone
Markus Peschl
Roberto Bondesan
Venue: ICLR 2023 poster (2023)
Keywords: Scheduling
GFlowNets
Combinatorial Optimization

Abstract: Finding the best way to schedule operations in a computation graph is a classical NP-hard problem which is central to compiler optimization. However, evaluating the goodness of a schedule on the target hardware can be very time-consuming. Traditional approaches as well as previous machine learning ones typically optimize proxy metrics, which are fast to evaluate but can lead to bad schedules when tested on the target hardware. In this work, we propose a new approach to scheduling by sampling proportionally to the proxy metric using a novel GFlowNet method. We introduce a technique to control the trade-off between diversity and goodness of the proposed schedules at inference time and demonstrate empirically that the pure optimization baselines can lead to subpar performance with respect to our approach when tested on a target model. Furthermore, we show that conditioning the GFlowNet on the computation graph enables generalization to unseen scheduling problems for both synthetic and real-world compiler datasets.

Published as a conference paper at ICLR 2023
ROBUST SCHEDULING WITH GFLOWNETS
David W. Zhang
Corrado Rainone Markus Peschl
Roberto Bondesan
University of Amsterdam
Qualcomm AT Research
ABSTRACT
Finding the best way to schedule operations in a computation graph is a classical NP-hard problem which is central to compiler optimization. However, evaluating the goodness of a schedule on the target hardware can be very time-consuming. Traditional approaches as well as previous machine learning mize proxy metrics, which are fast to evaluate but can lead to bad schedules when tested on the target hardware. In this work, we propose a new approach to schedul- ing by sampling proportionally to the proxy metric using a novel GFlowNet method. We introduce a technique to control the trade-off between diversity and goodness of the proposed schedules at inference time and demonstrate empirically that the pure optimization baselines can lead to subpar performance with respect to our approach when tested on a target model. Furthermore, we show that condi- tioning the GFlowNet on the computation graph enables generalization to unseen scheduling problems for both synthetic and real-world compiler datasets.
typically opti-
INTRODUCTION
Efficient execution of computation graphs is paramount to many scientific and industrial applica- tions, with deep learning being prominent example (Amodei & Hernandez, 2018). Scheduling the action of assigning operations to the available compute resources, such threads, cores, or nodes in a cluster (Kwok& & Ahmad, 1999; Hennessy & Patterson, 2011; Pinedo, 2012). Unfortu- nately, finding the schedule with the shortest possible makespan (start-to-end runtime) in general NP-hard (Papadimitriou & Steiglitz, 1998). As result, domain experts have come up with heuristics that are tailored to specific problem instances (Ibarra & Kim, 1977). Machine learning approaches promise the possibility to automate this process allowing for fast adaptation to new graph distri- butions (Wang & O' Boyle, 2018; Bengio et al., 2021c). In this work, we consider the problem of scheduling a set of operations with precedence constraints devices, i.e., any operation can run on any device and the runtime is the same on all devices.
fixed number of homogeneous
Evaluating the makespan of schedule involves running all operations in the computation graph on some target hardware. This can be very resource intensive, especially when the computation graph includes lengthy operations, the evaluated schedule inefficient, or the intended target hardware is cluster with many nodes. Heuristic optimizers. like genetic algorithms (Hou et al., 1994), or machine learning (Mao et al., 2019) approaches further exacerbate this problem because they require many evaluations to converge (Chen et al., 2018). Proxies are a much faster alternative that estimates the makespan using a simplified model of the hardware. However, this comes the cost of discrepancies between the proxy makespan and the one observed on the hardware; as a result, performant solutions on the proxy might ultimately be unsatisfactory once tested on the target. Nonetheless, proxies remain a good indicator for most schedules and are essential due to their efficiency. We aim to learn a scheduler that can be trained using the proxy, whilst being robust to its inaccuracies.
The common approach to scheduling problems (and combinatorial optimization problems in gen- eral) is to look for the single best schedule that minimizes a makespan measure which can be an analytical proxy (Paliwal et al., 2020), the output of simulator (Zhou et al., 2020). or even the real makespan on hardware (Khadka et al., 2021). We propose a different philosophy: generate a set of candidate schedules that have low makespan according the proxy and are diverse. By hav-
'Work completed during internship at Qualcomm Technologies Netherlands B. V.; Qualcomm Research is an initiative of Qualcomm Technologies, Inc. This work was done while the author worked Qualcomm Technologies, Inc.

---

Published as a conference paper at ICLR 2023
Computation Graph
GFlowNet Sampling
Proxy Ranking
Target Evaluation
Schedule
2.
Top-k
Schedule
dg
2
5
dg
3
Chain Graphs
Full pipeline of our generative scheduling approach. Conditioned on the computation graph we generate multiple candidate schedules using GFlowNet, filter for the best k with the proxy and pick the best performing one out of the k that we check on the target. Here we illustrate the pipeline for k = 2 and two devices, d,, d2-
Figure 1:
ing multiple good schedules that are significantly different, we can reduce the impact of systematic errors in the proxy, and hope for robust performance on the target.
Our goal is to learn a generative model that assigns higher probability to low-makespan schedules, and importantly can also discover the different modes associated with local optima of the makespan cost. Generative Flow Networks (GFlowNets) have recently been introduced as method for learn- ing a stochastic policy that can piece-by by-pie construct discrete and composite objects. proportional to given reward (Bengio et al., 2021b). By computing the reward from the proxy-makespan we can use GFlowNets to sample diverse set of candidate schedules.
Our main contributions are: 1. We introduce an alternative to the pure proxy optimization view- point of scheduling that achieves better robustness to proxy errors, by generating multiple candidate schedules to evaluate directly on the target hardware. 2. We extend GFlowNets to generate sched- ules conditioned on computation graph. Additionally, we introduce a method to control diversity and goodness at inference time, without the need for retraining. These contributions may be of general interest, beyond the scheduling problem. 3. We empirically demonstrate the robustness of method to proxy errors and verify the generalization ability diverse set of synthetic and real-world computation graphs.
2 ROBUST SCHEDULING
In this section, we first provide definition of the scheduling problem we consider in this work. Then, we discuss how a proxy simulates the schedule execution well the difficulties speci- fying a reliable proxy. Finally, we describe our proposed generative scheduling framework.
2.1 PROBLEM DEFINITION
In scheduling, we are given computation graph Gc = (O, that is a direct acyclic graph (DAG) consisting of operations (nodes) O and precedence constraints (edges) p E that encode partial order in which the operations need to be executed. In particular, the edge pij encodes that operation o; needs finish before 0j can start, for example because 0j requires the output of o as input. Our task is to run all operations the precedence constraints. In addition to the precedence constraints, the devices can only run one operation at time. We can then view scheduling as performing two distinct tasks: assign device to each operation, and determine a (complete) order among all operations on the same device that is chain
a set of devices 'D = {di,... dm), without violating
compatible with the precedence constraints encoded in Gc. We model the schedule of operations for each device, where the chain denotes the order in which the operations run on that device. See Figure for visual example of the chain graphs Our aim to find the schedule with the lowest makespan for some target hardware.
2.2 TARGET MODEL VS. PROXTES
The makespan of any schedule can be evaluated on the target hardware by running all the operations in the specified order and the specified devices. However, this can take up significant time
2

---

Published as a conference paper at ICLR 2023
and compute resources when the computation graph is large, has costly operations, the target hardware is a cluster with many nodes. In addition to this, when optimizing the makespan one needs to evaluate many different schedules, further exacerbating the resource requirements.
A proxy is any tool that allows one to estimate the makespan of a given schedule, without having to run the schedule the target hardware. Proxies come with significant speed advantages, which remedy the problems mentioned above. However, this comes at the cost of possible mistakes in the estimation of the makespan and relative comparison of schedules. Mistakes can occur for example when the task durations are not accurately profiled, memory movements are too complex to fully model, or additional hardware-specific features are changed. Ideally, we would like to rely on proxy for the majority of the schedule evaluations, and only evaluate a small fraction of promising schedules on the target hardware. This approach differs from previous works, that either evaluate every schedule on the target (Khadka et al., 2021), leading to very long optimization times, or evaluate everything on the proxy (Paliwal et al., 2020), which susceptible to modeling failures.
Next, we describe how the proxy we use in this work assigns start times to each operation given a schedule and estimates the makespan based on those. We recall that a schedule an order of operations for each device, which can be represented by one chain graph per device. For each of these, let us denote with Ca, d E D the set of edges of the chain graph for device d and with D D==U== Ca, the set of all device constraints. The operations correspond to graph nodes and are labeled in the same way as in Gc- No other operation can run on the same device during the runtime or duration p: of operation o; In practice, P; is estimated directly on the hardware in profiling stage that precedes scheduling. We denote the start time of o as T; and can thus express the precedence constraints as:
(1) An operation cannot start unless all of those that produce its inputs and all of those that precede on its assigned device have finished first. To ensure that these constraints are satisfied the proxy assigns each operation o, the start time
Tj 727 T3 Psr Vi.j) E PUD
T3 (mx (k) E DJ
If a node has no parents in PUD the proxy assigns the start time T, = 0. The start times of all o, E O can be computed by assigning a start time to node whenever it has no parents or all its parents have an assigned start time. If the graph (O, PU D) is a DAG, then this algorithm guaranteed assign start times that satisfy Equation 2. The proxy then estimates the makespan T of the schedule x as:
T(x) == max(T: + ps) - min(T:)
Optimizing this cost over the set of all possible schedules is already a very rich problem, and yet, we made significant simplifying assumptions in the construction of the proxy. In particular, we as- sume perfectly estimated runtimes, and in Equation we effectively assume that an operation can start as soon all of the operations producing its inputs finish, meaning that data can be moved be- tween devices instantaneously (zero latency) independently of their size (infinite bandwidth). These assumptions might be unrealistic, depending on the specific target devices (Valiant, 1990).
2.3 GENERATIVE SCHEDULING
Our aim is to come up with a good set of candidate schedules to be tested on the target hardware while relying only on the proxy for generating this set. While the proxy is imperfect, it still of- fers good guidance for most schedules; thus, we would like to include schedules that perform well according to the proxy. Nevertheless, we also know that systematic errors in the proxy can cause it to incorrectly predict low makespan for some schedules. Therefore, we would like the set of candidate schedules to be diverse, while still high-quality from the point of view of the proxy.
If we had ranking over all the schedules according to the proxy, we could just go through the list top-to-bottom, and add a schedule to the batch whenever it is significantly different from the previous ones. A full ranking like this is infeasible to construct, but we can instead learn generative model that samples higher ranked schedules with higher probability. When generating schedules we need to satisfy the precedence and device constraint outlined in Section 2.1. To avoid generating invalid schedules we construct a schedule in a step-by-step process: start with an empty schedule at the initial state so, and at each step add an operation to the partial schedule until the schedule contains
3

---

Published as a conference paper at ICLR 2023
all operations at the terminal state Sn- At each intermediate state st, action de consists in picking an operation and assigning it to one of the devices, leading to a new state 81+1- We define the set of valid actions at every step f in a way such that the precedence constraints are automatically satisfied. In particular, adding an operation or is a valid action if and only if vk (k, 1) E P. og is already in the partial schedule state 8t. This is sufficient condition for the final "schedule" graph (O, PU D) to be a DAG, implying that the constructed schedule is feasible. The final states represent full schedules x, for which can compute the makespan T T(x) with the proxy, given the runtimes {pi ) We compute the relative speedup compared to the makespan on single device as (zr)= P.P:TT(x), from which we compute the reward as we present in the next section.
GENERATIVE FLOW NETWORKS FOR SCHEDULING
GFlowNets (Bengio et al., 2021a;b) are methods for training a stochastic policy to sample discrete and composite objects proportionally to a given reward. Each object x is generated incrementally by sequence of actions. In the previous section, guarantee that we sample valid schedules. After a brief introduction to GFlowNets, the following sections will present our proposed extensions that include new training objective that is suitable for learning conditional GFlowNets and a method for controlling the selectivity of samples at inference time.
discussed how to limit the action space to
3.1 BACKGROUND
We start by introducing some notation. We denote by s=( .1.-..,8, 8x) of a sequence of states st. In the case of scheduling, trajectories start with the empty schedule so. followed by partial schedules, and end with a complete schedule Sz. We denote by T the set of all such trajectories and by Ta the of trajectories that end x. Based on this, we define a flow function F: ' - R+ and its associated normalized probability distribution P(s) = F(s)/Z. Z ΣXET F(s). flow function that fulfills the condition: has a total flow matching reward), results in a probability over schedules P(x) EmET. F(s)/Z
a trajectory that consists
R(x)= = E.ET, F(s) (every terminal state
that is proportional to the reward P(x) x R(x), and further entails that Z E., R(x).
For any Markovian flow, we can decompose the probability of trajectory in terms of the forward probability:
P(s) II Pr(s: 11) 1=1
(4)
This way, generate trajectories s by sampling sequence of actions starting from 8g. In Sec- tion 2.3 we described how to limit the action space appropriately to guarantee that every sampled schedule is valid. Similarly, probability conditioned on a terminal state:
define a backward probability Pa that factorizes the trajectory
12
P(s|., = II (r(
The training objectives considered in previous works aim to achieve a consistent flow (Bengio et al., 2021b; Malkin et al., 2022), where consistency means that the flow estimated for the forward direction should equal the flow for the backward direction. A consistent flow F(s) for trajectories 8 E Tn can then be written in terms of Pr and Pn and has to fulfill the equality:
n
II PF R(x) II PB(s 1=1
(6)
Based on this equation, Malkin et al. (2022) propose to estimate Z, PF, and PB by optimizing the trajectory balance loss which is the squared difference between the logarithms of the hhs.s. and the r.h.s. of Equation 6.
3.2 LOG-PARTITION VARIANCE LOSS
In order to apply the trajectory balance loss in the conditional case, we would need to learn an addi- tional regression model that estimates the log-partition function log Z conditioned Gg. Training
4

---

Published as a conference paper at ICLR 2023
such a network accurately is difficult but crucial for learning the probabilities Pr. In partieular, wrong estimation of log Z can incorrectly change the direction of the gradients of the loss func- tion. We explain why this occurs in Appendix B. In practice, poorly when different computation graphs had large differences in their log Z value. Instead, we can rewrite Equation to implicitly estimate log Z based on the forward and backward flows of single trajectory 8, where PF and PB are neural networks with parameters O:
found this approach to perform
{(s; 0) = log R(x)+ + Σ log PB( 1(; 06 Σ log (r((((((-1:0)
In the optimal case, (ss; 0) is equal to the true log Z which is the same for all trajectories corre- sponding to the same computation graph Ge. Thus, our optimization goal turns into minimizing the variance of ((s; 0) over different trajectories with the loss
Lv(s; 0) = (C(s; 0) - EE [[s;
(8)
In practice, we use the training distribution to estimate E, [S()]] with mini- batch of sampled trajectories. For more details on the training process, we refer to Appendix C.
We note that by optimizing the log-partition variance loss in Equation 8, one only needs to parametrize the forward and backward probabilities PF and PB. This similar the non-forward trajectory loss mentioned in the appendix by Malkin et al. (2022), which also does not involve learning any state flows, including the initial flow Z. However, our loss does not mix forward and backward steps from different trajectories and directly optimizes the consistency of the total flow Z for each trajectory associated with a given computation graph Gc.
3.3 TEMPERATURE-CONDIIIN
TOPOFORMER
Reward temperature. We compute the reward as a function of the speedup. In particular, we choose log R(x; m. a) (U mi/a where (xx) is the speedup of the schedule æ, m is the number of devices, and a E IR+ plays the role of temperature. The temperature allows us to concentrate the distribution on the modes and control the selectivity of the generator. This is useful since there can be many more schedules with low speedup when compared to good ones, For example, when simply setting the reward equal to the speedup, we observed that finding schedules with high speedup requires a prohibitively large amount of samples. We expect this temperature term to allow trade-offs between diversity and shifting the mean of the sampling distribution towards better schedules.
Previous works on GFlowNets apply a constant temperature value during training and at inference time (Bengio et al., 2021a; Jain et al., 2022). This can lead to low performance (when set too high), and low diversity or unstable training (when set too low). Furthermore, different computation graphs can have different ideal temperature values, making this approach less suitable when learning conditional GFlowNets. Instead, we propose to learn a single model for multiple different reward functions R(x; m. o). by conditioning the policy networks (PF and PB) on the temperature a. Ap- proximating the cnnpeerauue-cooditionn given temperature
policy with a neural network is feasible because flows for be continuously morphed into flows for any other temperature. Since our reward R(I; m, o) is continuous with respect to the temperature a, we expect the change of flow for different temperatures be learnable by neural network. We provide proof for the following theorem in Appendix A.
Theorem 1 (Flow Continuity). Let [Re be a sequence of non-negative reward functions such that for all terminal states x, (x) + R(x) as + cc. Then, for any flow F with reward R, there exists sequence of flow functions { (R.) Jiz1 with FRi (s) PR () for all E T.
The output policy changes more rapidly function of the temperature for values close to 0 than for larger values. To account for this, we use the logarithm of the temperature as input to the policy instead. During training, we sample temperatures from the log-uniform distribution with support between [log min, log maz]). where miin a minimum temperature that is necessary for numerical stability. In comparison sampling from (d Tmaa). ). this avoids oversampling from high temperature regions that have little difference in the resulting flow network. At inference time, we choose how close the samples are the mode by adjusting the o.
5

---

Published as a conference paper at ICLR 2023
Topoformer architecture. For the neural network architecture of policy, use the Topo- former (Gagrani et al., 2022), which has been recently introduced for learning topological orderings of computation graphs. It builds on the Transformer encoder (Vaswani et al., 2017) and additionally masks the multi-head attention depending on the topology of the computation graph. Both forward and backward policies use separate ML.P heads on top of shared Topoformer encoder. Taking inspiration from the successful use of time conditioning in diffusion models (Song et al., 2020; Ho al., 2020), we add temperature conditioning by first embedding the temperature using an MLP to produce Ca. and then reuse the embedding in every first linear layer block of the Topoformer:
1in(h, c.) 11 (en @ 1in(h) + 1in,hift
Here 1 Disallel and 1 Nigift are linear layers and @ is the elementwise multiplication (Perez et al., 2018). In contrast to diffusion models, observe better performance on large temperature ranges with the ReLU (Nair& & Hinton, 2010) activation function. We hypothesize that this is connected to the monotonicity of the underlying policy function with respect to decreasing temperatures (see Corollary in the Appendix) and the propensity for linear extrapolation of ReLU MLPs (Xu et al., 2020). For detailed description of the neural network architecture, refer to Appendix D.
Sub-graph training. Training with full computation graph might not always be necessary and we hypothesize that learning on sub-gras sub-graphs can lead to policies that generalize to the full computa- tion graph. This can be seen data, while simultaneously improving the training time. We shall use sub-graph training for the larger graphs that we study in this work.
form of data augmentation and increases the amount of training
4 RELATED WORK
Reinforcement learning for scheduling. Reinforcement learning has been the predominant ma- chine learning approach to optimize the makespan for computation graph schedules (Addanki al., 2019; Paliwal et al., 2020; Zhang et al., 2020). The rewards used include simple analytical proxies of the makespan (Paliwal et al., 2020; Zhou et al., 2020), but also more refined proxies which incor- porate modeling of memory movements (Addanki et al., 2019). Khadka et al. (2021) directly train on the target hardware, but consider only a few computation graphs, and do not show generalization to unseen ones. Addanki et al. (2019) use sophisticated simulator of the makespan which is cus- tomized to the target hardware. Similar to our work, Zhang et al. (2020) also construct the schedule piece-by-piece. Instead of finding single (local) mode of the proxy, work proposes to learn the full distribution over the proxy to improve the robustness against inaccuracies in the proxy.
Generative Flow Networks. GFlowNets have been applied to generating small molecules (Ben- gio et al., 2021a), Bayesian networks (Deleu et al., 2022), discrete images (Zhang et al., 2022), and biological sequences (Jain et al., 2022). We extend its application to scheduling. a classical combina- torial optimization problem. Similar to previous works, our state-action space is also a DAG, hence training the policy with maximum entropy reinforcement learning methods is inadequate (Bengio et al., 2021a). Our robust scheduling approach shares the same motivation methods in drug dis- covery which leverage cheap proxies to generate multiple candidates to be evaluated in the true environment (Bengio et al., 2021a; Jain al., 2022). Conditional GFlowNets have previously only been theoretically discussed by Bengio et al. (2021b). We enable training conditional GFlowNets with our proposed log-partition variance loss and empirically demonstrate generalization un- seen computation graphs. Note that this differs from previous work that tests the generalization of GFlowNets to unseen data (Nica et al., 2022). To control the selectiveness of the generator, previous works augment the reward with fixed temperature (Bengio et al., 2021a; Deleu et al., 2022; Jain et al., 2022). Instead, we condition the policy neural network on the temperature term which allows us to tune the selectiveness of the generator at inference time.
5 EXPERIMENTS
In this section, we evaluate different aspects of our generative scheduling approach by incrementally adding complexity to the computation graph dataset. First, we restrict training and evaluation to single computation graph, which corresponds to the same unconditional setting considered by previous works on GFlowNets (Bengio et al., 2021a; Deleu et al., 2022; Jain et al., 2022). Next, we train with multiple computation graphs and evaluate on unseen ones. To the best of our knowledge,

---

Published as a conference paper at ICLR 2023
this is the first time that the generalization of conditional GFlowNets unseen conditioning is tested empirically. Finally, we verify the generalization ability on weal-world computation graphs of neural networks that are being used in a diverse set of Al products.
Experimental setup. In all experiments, we only use the node time duration as a feature of the computation graph. For simplicity and ease of reproducibility, to add extra features. All our experiments are based on four homogenous devices, which implies that the speedup is upper bounded by 4. In practice, most computation graphs have lower maximal possible speedup due to their precedence constraints.
avoid any complicated heuristics
We consider two heuristic and two neural methods for generating candidate schedules. The first is our GFlowNet approach described in Section from which we generate 1000 samples at temperature o = 0.005 and take the top 100 following the proxy; the other three are:
Candidate sampler.
Critical path- list scheduling, a heuristic algorithm for scheduling on homogeneous devices (Micheli, 1994). List scheduling first forms topological order of the operations and then assigns them that order one by one to free device. our implementation, we use the Critical Path method (Micheli, 1994) to order the operations. It ensures that operations on the time critical path are scheduled first. This method produces a single schedule.
Biased Random Key Genetic Algorithm (BRKGA) (Gonçalves & Resende, 2011), ge- netic algorithm that has previously shown good performance on scheduling tasks (Paliwal et al., 2020). We use the top 100 schedules from the final population as the candidate schedules.
Proximal Policy Optimization (PPO) (Schulman al., 2017), deep reinforcement learn- ing method that has been successfully applied to scheduling problems (Zhou et al., 2020). PPO also trains a stochastic policy, which makes it a natural choice for comparison with GelowNets (Bengio et al., 2021a). We employ the same definitions of states, actions, and reward function (with temperature a = 0.25; lower was not beneficial) as the GFlowNet approach. To ensure that PPO keeps exploring even after finding local optimum we em- ploy entropy regularization and decay both the entropy coefficient (Ahmed et al., 2019) and the learning rate to ensure convergence to good solution. Same as for GFlowNets, we sample 1000 schedules and pick the top 100 as the candidate schedules.
Metrics. We measure the performance in terms of the speedup Uxx). For the diversity, we report three different measures: graph-edit distance (GED), the L2 distance between the proxy start-times (dim),, and the L2 distance between the proxy start-times concatenated with the device placement (dyen). For diversity, we report the average pairwise distances over the top 100 candidate schedules. See Appendix E.2 for more details on diversity measures.
5.1 PROXY ERRORS: DIVERSITY FOR ROBUST SCHEDULING
We examine how differences between the proxy and the target performance model can affect the final runtime. To do so, first focus on single computation graph that is used both for training and testing to avoid possible confound- ing factors that may happen in the generaliza- tion setting. Based on the possible reasons for proxy errors discussed in Section 2.2 we design three different target models that each reflect different setting. In the first setting node dura- tions are incorrectly profiled (Noisy Runtimes). In the second and third settings, the target mod- els the memory movement across devices with linear model (Valiant, 1990), which can be ei- ther bottlenecked by limited bandwidth (Band- width Limited), by high latency (Latency Limited). The linear model has been shown be
Latwecy Limited Sandwilth
Nojay Runtimes.
Proxy Spentup
Figure 2: Correlation between proxy and target speedup for different target environments. Modes with varying performance can be observed for fixed proxy speedup.
7

---

Published as a conference paper at ICLR 2023
Table 1: Robustness results on single computation graph. We compare different methods for gen- erating candidate schedules. Higher diversity correlates with better robustness against a mismatch of the proxy and the target, with GFlowNet achieving the best diversity and the best target performance on average.
Speedup
Diversity
Noisy Runtimes
Bandwidth Limited
Latency Limited
Proxy
GED
dim
dye
List scheduling
3.23:10.00
2.75:10.00
1.02.10.00
1.74.10.00
BRKGA
3.22+0.00
2.86+0.15
1.29++0.4
1.80+0.34
55.92+2.56
22.83++.339
16.21+1.50
PPO
3.28
3.07:10.09
0.38
1.87:10.88
85.08
31.71
105.64 0.00
GFlowNer 3.21+0002
3.05=0.04
1.78#0.03
2.11=0003
94.79=0.15
42.08=0.33
115.98=0.09
good makespan estimator for certain devices (Hockney, 1994; Culler al., 1993). We refer to Ap- pendix E.3 for more details. In Figure 2, we show the correlation between the proxy and the different target environments. For all three targets, the proxy is highly correlated but can have target speedups that differ by factor of up to x2 for the schedules with high proxy speedups.
We report the speedups and diversity measures in Table 1. The results highlight that any method that can generate multiple good candidate schedules achieves higher speedups on the target environments than list scheduling, which only produces a single unique schedule. Furthermore, if the candidate schedules are more diverse as is the case for GFlowNet the target performance is also better on average. PPO and BRKGA exhibit high variability in performance between different runs, where few runs end up with high speedups speedups. In contrast, the GFlowNet model is consistent over the different random seeds, both in terms of diversity and speedup. The results confirm hypothesis that diverse set of candidate schedules with high average proxy speedups can improve robustness towards a misspecified proxy.
some targets, and other runs result in much lower target
5.2 GENERALIZING TO UNSEEN COMPUTATION GRAPHS
Next, we evaluate how well our conditional GelowNet generalize to unseen computation graphs We train and evaluate on a diverse set of synthetic computation graphs sampled from different ran- dom graph distributions. In particular,
train graphs of size 50 sampled from the random graph distributions (a) Rédós-Rényí (Erdós et al., 1960), and (b) Layered Graphs (Gagrani et al., 2022) and evaluate, in addition (a) and (b), stochastic block model (Holland al., 1983), Watts- Strogatz (Watts & Strogatz, 1998), and Barabási-Albert (Albert & Barabási, 2002). For details on the generative process of the computation graphs, we refer to Appendix E.5.
In Table 2, we demonstrate that both PPO and the conditional GFlowNet are able to generalize to previously unseen computation graphs, regardless of whether they originate from the same random graph distribution. Next, we ablate our proposed temperature conditioning method by generating 1000 samples at different temperature values. Figure 3, we observe that decreasing the tempera- ture does indeed shift the sample distribution to the right and also sharpens it when the temperature approaches zero. Notably, the temperature
0.005 is not in the training distribution, which demonstrates that the model can extrapolate to temperature values outside of the training range. Surprisingly, we observe that training with a variable temperature can improve the performance further than is possible with fixed temperature, which we demonstrate in Appendix F.
Table 2: Generalization to different random graph distributions. We report the speedup and diversity for the top 100 schedules. PPO and GFlowNet are trained on graphs from the Erdos-Renyi and I Lay- ered Graph distribution and report the average performance over all random graph distributions. The Speedup Proxy 100 column reports the average proxy speedup over the top 100 schedules.
Speedup
Diversity
Proxy
Proxy 100
GED
dim
dan
List scheduling
3.44
3.44
0
0
BRKGA
3.46
3.45
46.59
12.75
40.11
PPO
3.48
3.46
69.54
13.45
80.84
GelowNet
3.46
3.41
92.02
24.27
90.17
8

---

Published as a conference paper at ICLR 2023
Temperoture:
Temperature: 0.25
Temperature: 0.05
Temperature: 0.005
Probabbiltty
6.4
1.0
0:0
0.0
150 1.0
203
2.2
130
223
230
1.0
208
1.5
2.25
1.7
2.20
2.23
325
2000
Speedup
Spoedup
Speedup
Spoedup
Figure 3: Empirical reward distribution on 25-node Erdos-Reng graph for different inference temperatures in the conditional GFlowNet. Lower temperatures allocate more probability mass to better schedules.
5.3 REAL WORLD COMPUTATION GRAPHS
Finally, we verify the generalization ability on a small set of real-world computation graphs used for the commercial development of our artificial intelligence hardware and software products (see Ap- pendix E.6 for details). We report the speedup on the same target models used in Section 5.1 to assess robustness on unseen real-world computation graphs. To speed up training, we apply the graph subsampling strategy presented in Section 3.3 randomly pick between 25 to 75 nodes at every training step.
In Table 3, observe that the conditional GFlowNet retains the benefits of high diversity and robustness to misspecifications in the proxy even when applied to graphs not seen during training and of larger sizes. PPO shows unstable training behavior and the reward training curve does not converge, despite using the same hyperparameters that worked for the previous two experiments. We conjecture that this is due the inhomogeneous maximum possible speedup of the training graphs that lead to different reward scales per training graph. In comparison, GFlow still converges as before without any changes the hyperparameters. Note that while PPO exhibits higher diversity than compared to BRKGA, it still underperforms BRKGA due to low average proxy speedups. This highlights that high diversity alone is not sufficient, otherwise, a uniform distribution as the forward policy would already suffice.
We ablate our proposed log-partition variance loss by comparing it against the trajectory balance loss that uses Topoformer predict log Z given a computation graph. Learning such model difficult due to large differences in the output space of different computation graphs that arise from the differences in the number of nodes, which in turn impedes the training progress of the policy network. We confirm in Appendix E.6 that our proposed loss function remedies the slow start problem of the baseline and achieves higher speedup in the end.
Table 3: Generalization on real-world graphs. We train on a small set of real-world graphs and evaluate on unseen ones. GFlowNet retains high diversity and exhibits consistently better per- formances than the baselines on the target models. PPO uses the same hyperparameters in the previous experiments but does not manage to converge on this dataset.
Speedup
Diversity
Noisy Runtimes
Bandwidth Limited
Latency Limited
Proxy
Proxy 100
GFD
Aur
deea
scheduling
2.74+0.00
2.74-0.00
2.51++0.00
0.89+00.00
1.43=00.0
BRKGA
2.59.10.18
2.58
2.46.10.16
1.511
1.80
$2.32 21.59 17.1418.17
42.64.12.23
PPO
2.1100.00
2.23:0:27
2.2810.266
0.1120.20
1.43-0.10
42.70113.4
64.92114.08
GFlowNet
2..110.03
2.66
2.71
1.73.10.01
1.95
87.95.10.13
26.56.10.56
91.33
6 CONCLUSION
We have empirically demonstrated how the conventional optimization approach to scheduling, which optimizes a proxy of the real makespan, is brittle to modeling failures in the proxy itself. Our proposed approach evaluates multiple schedules on the target and thereby achieves more ro- bustness to discrepancies between the proxy and the target. We demonstrated that GFlownets can sample diverse set of candidate schedules that achieve better target performance than alternative methods which achieve lower diversity. Further, we showed that conditioning on temperature allows trade-off between diversity and proxy performance, and that conditional GFlowNets can gener- alize to unseen computation graphs. Interesting future directions include scaling up our method to larger graphs and integrating scheduling heuristics speed up training.

---

Published as a conference paper at ICLR 2023
REFERENCES
Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan, Shreyan Gupta, Hongzi Mao, and Mo- hammad Alizadeh. Placeto: Learning generalizable device placement algorithms for distributed machine learning. Advances in Neural Information Processing Systems 32 (NIPS 20/9), 2019.
URL https::///// /par nsf. https/..ooobbbiooooiieeeee.
Zafarali Ahmed, Nicolas Le Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy on policy optimization. In International conference on machine learning, pp. 151-160. PMLR, 2019.
Réka Albert and Albert-László Barabási. Statistical mechanics of complex networks. Reviews of modern physics, 74(1):47,
Dario Amodei and Danny Hernandez. Ai and compute.
ai bal-and-compute/ 2018. Accessed: 2022-09-29.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. In Advances in Neural Information Processing Systems, volume 34, pp. 27381-27394, 2021a.
Yoshua Bengio, Tristan Deleu, Edward Hu, Salem Lahlou, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. arXiv preprint arXiv:2111.0922 2021b.
Yoshua Bengio, Andrea Lodi, and Antoine Prouvost. Machine learning for combinatorial opti- mization: a methodological tour d'horizon. European Journal of Operational Research, 290(2): 405-421, 2021c.
J. Blank and K. Deb. pymoo: Multi-objective optimization in python. Access, 8:89497-89509, 2020.
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang. Yuwei Hu, Luis Ceze, et al. [TVM}: An automated (End-to-End) op- timizing compiler for deep learning. 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pp. 578-594, 2018.
David Culler, Richard Karp, David Patterson, Abhijit Sahay, Klaus Erik Schauser, Eunice Santos, Ramesh Subramonian, and Thorsten Von Eicken. Logp: Towards realistic model of parallel computation. In Proceedings of the fourth ACM SIGPLAN symposium on Principles and practice of parallel programming, pp. 1-12, 1993.
Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. arXiv preprint arXiv:2202.1390) 2022.
Paul Erdös, Alfréd Rényi, et al. On the evolution of random graphs Publ, Math. Inst. Hung, Acad. Sci, 5(1):17-60, 1960.
Mukul Gagrani, Corrado Rainone, Yang Yang. Harris Teague, Wonseok Jeon, Herke Van Hoof, Weiliang Will Zeng, Piero Zappi, Christopher Lott, and Roberto Bondesan. Neural topological ordering for computation graphs. arXiv preprint arXiv:2207.05880 2022.
José Fernando Gonçalves and Mauricio GC Resende. Biased random key genetic algorithms for combinatorial optimization. Journal of Heuristics, 17(5):487-525, 2011.
John I. Hennessy and David A Patterson. Computer architecture: quantitative approach. Elsevier, 2011.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Roger W Hockney. The communication challenge for mpp: Intel paragon and meiko cs-2. Parallel
computing, 20(3):389-399 1994.
10

---

Published as a conference paper at ICLR 2023
Paul W Holland, Kathryn Blackmond Laskey, and Samuel Leinhardt. Stochastic blockmodels: First
steps. Social networks, 5(2):109-137, 1983.
Edwin SH Hou, Nirwan Ansari, and Hong Ren. A genetic algorithm for multiprocessor scheduling. IEEE Transactions on Purallel and Distributed systems, 5(2): 13-120, 1994.
Oscar H Ibarra and Chul E Kim. Heuristic algorithms for scheduling independent tasks on noniden- tical processors. Journal of the ACM (JACM), 24(2):280-289 1977.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garchs Jarrid Rector-Brooks Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gftownets. In International Conference on Machine Learning, pp. 9786-9801. PMLR, 2022.
Shauharda Khadka, Estelle Aflalo, Mattias Mardar, Avrech Ben-David, Santiago Miret, Shie Man- nor, Tamir Hazan, Hanlin Tang, and Somdeb Majumdar. Optimizing memory placement using evolutionary graph reinforcement learning. In International Conference on Learning Representa- tions, 2021. URL httttp:://ppeneeveee
net feet//formm?d==== 6v9 2aceedd-vv____________ 4Kf2C.
Yu-Kwong Kwok and Ishfaq Ahmad. Static scheduling algorithms for allocating directed task graphs to multiprocesssos, ACM Computing Surveys (CSUR), 31(4):406-471, 1999.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. arXiv preprint arXiv:2201.13254 2022.
Hongzi Mao, Malte Schwarzkopf, Shaileshh Bojja Venkatakrishnan, Zili Meng, and Mohammad Alizadeh. Learning scheduling algorithms for data processing clusters. In Proceedings of the ACM special interest group data communication, pp. 270-288 ACM, 2019.
Giovanni De Micheli. Synthesis and optimization of digital circuits. McGraw-Hil Higher Education, 1994.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Icml, 2010.
Andrei Cristian Nica, Moksh Jain, Emmanuel Bengio, Cheng-Hao Liu, Maksym Korablyov, Michael M Bronstein, and Yoshua Bengio. Evaluating generalization in gflownets for molecule design. In ICLR2022 Machine Learning for Drug Discovery, 2022.
Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs. In Interna- tional Conference on Learning Representations, 2020. URL https::/openrevai http///pppneeeeeee/eeee/ee .netttewwwweeeeeeeeeee......ee.e Corum?id=rkxDo.l.
Christos H Papadimitriou and Kenneth Steiglitz. Combinatorial optimization: algorithms and com- plexity. Courier Corporation, 1998.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.
Michael L Pinedo. Scheduling, volume 29. Springer, 2012.
Tobias Pohlen, Bilal Piot, Todd Hester, Mohammad Gheshlaghi Azar, Dan Horgan. David Budden, Gabriel Maron, Hado Van Hasselt, John Quan, Mel Večerík, et al. Observe and look further: Achieving consistent performance on atari. arXiv preprint arXiv: 1805. 1593, 2018.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 17Xiv:1707.06347, 2017.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1. 13466, 2020.
11

---

Published as a conference paper at ICLR 2023
Leslie G Valiant. A bridging model for parallel computation. Communications of the ACM, 33(8): 103-111, 1990.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Zheng Wang and Michael O' Boyle. Machine learning in compiler optimization. Proceedings of the IEEE, 106(1 1879 199-1901, 2018.
Duncan Watts and Steven H Strogatz. Collective dynamics of 'small-world networks. nature, 393 (6684):440 442, 1998.
Keyulu Xu, Mozhi Zhang, Jingling Li, Simon $ Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint UrXiv:2009.11866 2020.
Cong Zhang, Wen Song, Zhiguang Cao, Jie Zhang, Puay Siew Tan, and Xu Chi. Learning to dispatch for job shop scheduling via deep reinforcement learning. Advances in Neural Information Processing Systems, 33:1621-16322 2020.
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. arXiv preprint 2rXiv:220..01361, 2022.
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 ml compilers. Advances in Neural Information Processing Systems, 33:13844- 118844-13555, 13855, 2020.
12

---

Published as a conference paper at ICLR 2023
A PROOF FOR THEOREM 1
In the following, we will denote FR to be the flow corresponding to a flow function F : -> Rg0 with corresponding reward R. That is, F8 + -R20 is a flow function such that for any terminal state x we have R(x)) E.ET. (s),, where ESET. (R(s) = (x) the total flow at the
terminal state
Let { [R,] be sequence of nonnegative reward functions such that for all terminal states x, R₂(z) - R(x) as i + co. Then, for any flow FR with reward R, there exists sequence flow functions {F j with (s) - (()(s) for all T.
Theorem (Flow Continuity).
Proof. Let {2x )
be the set of terminal states and define
( R.
if R(x) > 0,
F (s (80- (e.....,))) x))
R(x)
(10)
R((a)
else.
17.1
To see that FR, is a valid flow for Rs we note that Fm, and for any terminal state x with R(x) we have
R1(x)
(x) Σ (
Σ (s)
R(x)
BET,
NET,
(11)
(z)
R(x) R.(x)).
R(x)
Using similar reasoning we arrive at F (x) when R(x) 0. Finally, for any with terminal state and R(x) > 0 we have
(x)x) (s((s (R)
F (s) R(x)
(12)
In the case of R(x)= 0, we note that (s)(s) and
R² (x)
R(x)
FRi (s) =
O,
(13)
[Ta
TT
thus proving convergence of F.,, to FR.
Corollary 1. Let log Rs, log R(x; m, (U(s)-m)/a m)/a; conditioned log reward functions with ' Og. Then, for any e >> 0 and flow F Reo there exists a neighborhood (oo, on 5) containing flow functions R with F.. (s) (R) for all E T. Furthermore, FR., (s) monotonically decreases to pReo (8) for all T as i goes to O.
be a sequence temperature
The above corollary suggests that is feasible to use single neural network - that can approximate arbitrary continuous functions to learn all flow functions for different temperature values.
B FAILURE MODES WHEN ESTIMATING THE LOG-PARTITION FUNCTION
In this section, we present the different cases in which an incorrect estimate of log Z can lead to gradients that point in the opposite direction of the "true gradients". We start by rewriting the trajectory balance loss (Malkin et al., 2022) as:
(
CTB(s; 0)
log ((Gc:0z) +
log Pp(sc
1=1
(14)
- log R(x) -
Σ
log Pr 1 0p)
13

---

Published as a conference paper at ICLR 2023
For better readability,, we do not explicitly write out the dependence the computation graph Ge for the trajectories, their probabilities, and the reward. The parameters of the neural networks are denoted as 0p and 0 respectively. We can rearrange the terms in Equation 14, such that:
(T 0) -
Σ log (=(()(((-1:0p)
(15)
log R(x) +
log PB (B((st11s;;0p)) log Z(G 0z)
In words, ypp is the log probability of the trajectory as computed by the forward distribution, and PRBZ can also be viewed as the log probability of but instead estimated by combination of the reward, log-partition function, and the backward distribution. To simplify the argument, we assume that the backward distribution is independent of 0, which can be realized for example by setting equal to the distribution that assigns equal probability to all parents. Thus, we denote the true log-partition function log Z (Gc) and by assuming access to it when computing the loss can rewrite part of Equation 15 as:
PRBZ log R(x)
log Pn(st-1(88) - -log 7' (Gc)
(16)
We denote the loss using PRB2 as (s;
There exist two cases in which (s; 8; 0) can be non-zero and both cases, inaceurate estimation of the log- partition function can lead to wrong gradients for the policy neural network. In the first case, we assume that the neural network that estimates Pp is overestimating the probability of s:
PF PRBZ 0.
(17)
Furthermore, we consider the scenario in which the regression model makes error when predicting the log-partition value. In particular, we assume that the neural network is underestimating the true value by amount that is greater than the error f FRBZ* i.e., we assume:
fr PRBZ < log 7- (Gc) - log Z(Gc; 0z)
(18)
The implications of this are:
PP PRB7 < log 7* (Gc) - log Z(Gc; Oz)
(19)
PF < log Z (Gc) - log Z (Gc; 0z) + PRBz
(20)
PF < RRBZ
(21)
PFF PRBZ <
(22)
Equation 17 and Equation 22 imply that -V,, CTB is ascent direction for Lip (s: 0),
instead of the desired descent direction. In detail, for non-zero gradients the inner product is positive:
-V, 0), Ver (s:
(23)
(-(Pp PRZZ) V Pp. (PF PRBZ Vor
(24)
0
(25)
In the second case, fFF PRBz* O, we assume that the neural network overestimates the log- partition function. Analog to the first case we can show that the gradients point in the opposite direction of the true gradients.
These two cases happen more frequently when log (GG; 0z) is more challenging to regress ac- curately. For example, when there computation graphs with drastically different node numbers, their log (GG; 02) values can have significant differences (e.g.. 10s vs. 1000s). Training a re- gression model on such large ranges is notoriously difficult (Pohlen et al., 2018). Our proposed log partition variance loss in Equation 8 avoids this regression task completely, and we confirm its benefits empirically in Appendix E.6.
14

---

Published as a conference paper at ICLR 2023
C GFLOWNET TRAINING PROCESS
In each parameter update step, we first sample a single computation graph from the training dataset, next we sample single reward temperature from the log-uniform distribution, then sample tra- jectories we refer to these as a mini-batch - for that computation graph and reward temperature, and finally, we compute the loss based on the mini-batch of trajectories. The trajectories are sampled using the current forward policy Pp conditioned on the temperature. Empirically we observe that the temperature conditioning method allows us to forgo using a special exploratory policy that mixes Pp and uniform distribution over the allowed actions, which was used by previous works.
D TEMPERATURE-CONDIIIONN TOPOFORMER
The Topoformer (Gagrani et al., 2022) has the same structure as the Transformer (Vaswani et al., 2017) encoder, that is it stacks layers and each layer (for s L) consists of two steps:
7 = (1-1) + (Tay h ( (4) MLP(LLayeNorm(( ()
(26)
(27)
where Moooooorrerrr denotes the Topoformer version of the multi-head attention, and (() is the input to the Topoformer, which is the output of a linear layer (i.e., affine transformation with no activation function) applied on the computation graph and state features. Topoformer uses the same MLP in the original Transformer:
MLP(h")) 1 1ing(Re..((11n
(28)
We inject the temperature by replacing lin, with (110
Co = RelJ/Linn(Re1(((1nn.))))
(29)
1ing (h (.)) linseale o lin linaneeee)
(30)
Note that e, is the same for all layers / and Equation 30 corresponds to Equation 9 in the main paper.
E EXPERIMENT DETAILS
E.1 CANDIDATE SAMPLERS
We use the popular open-source library pymoo (Blank & Deb, 2020) to implement the BRKGA candidate sampler. Our PPO implementation is based algorithm 1 at ht / //tpnnningnnnpppp./ openai. cwmmee/////cccccccccccc/.ccct en eatest/ / algorr thms ppo. html, and we follow Schulman et al. (2017) implement the entropy regularisation by adding the entropy term directly the PPO-clip loss. We decay this entropy term during training similar to Ahmed et al. (2019). We use the same learning rate for both the actor and the critic, and we decay it with an exponential schedule.
We train Neto conditioned on a temperature randomly sampled between 0.01 and 1. At infer- ence, we use 0.005 for the temperature all experiments. We use the Adam optimizer with default hyperparameters to optimize the parameters. We compute the gradients at each update step based on minibatch that consists of 100 sampled trajectories for single computation graph and use single temperature value to compute their rewards. We observed no benefits in our initial experiments that used multiple computation graphs or temperatures in a single minibatch.
We represent the computation graph as directed graph with single node feature (runtime of the operation). In addition, we represent different states by concatenating a state feature vector to each node, which consists of:
Start time assigned by the proxy (default: -1)
Device placement as a one-hot vector (default: vector)
Binary indicator: adding the schedule a valid action, else 0
Binary indicator: 1 if removing it is valid backward-action, else 0
15

---

Published as a conference paper at ICLR 2023
Binary indicator: 1 if the node is already part of the schedule, else 0
The combined computation graph and state features are treated as single graph by the Topoformer neural network.
E.2 METRICS
The graph-edit distance (GED) compares two schedules in their chain-graphe form. In particular, can model schedule for a computation graph, by constructing chain graph for each device that specifies the additional precedence constraints we introduce to complete the order in which the operations are rum on each device. The GED is then computed simply by taking the difference between the adjacency matrices and normalizing it by the total number of edges.
The L2 distance between the start (dmw) times simply takes the start times as assigned by the proxy model and computes the 1.2 norm of the difference.
The L.2 distance including the device assignment (dhen) additionally concatenates the device place- ment to the times.
E.3 TARGET ENVIRONMENT WITH LINEAR MEMORY MODEL
The linear memory model (Valiant, 1990) computes the delay as linear function f(m) am am of the memory m with a modeling the amount of data that can be transferred per time and b the startup delay. In the Bandwidth Limited setting the dominates the delay, while in the Latency Limited setting has a greater effect.
modeling
E.4 PROXY ERRORS: DIVERSITY FOR ROBUST SCHEDULING
In this experiment, consider single real-world graph that has around 78 nodes.
E.5 GENERALIZATION TO UNSEEN COMPUTATION GRAPHS
We generate the synthetic graph dataset from random graph distributions over undirected graphs To get DAGs from these graphs, we randomly choose direction for every edge in way that produces no cycles. We closely follow the setup described Appendix A.1.4 by Paliwal et (2020), and the setup described in Appendix A by Gagrani et al. (2022).
We sample the runtimes for each node from uniform distribution 2( 1).
For training, we use 1000 different computation graphs, with equally many sampled from the two random graph distributions: Rrdös-Rényí (Erdös et al., 1960), and Layered Graphs (Gagrani al., 2022). We report test performances on 50 different computation graphs with equally many sam- pled from the five different random graph distributions: Erdös-Rényi (Erdōs et al., 1960), Layered Graphs (Gagrani et al., 2022), stochastic block model (Holland et al., 1983), Watts-Strogatz (Watts & Strogatz, 1998), and Barabási-Albert (Albert & Barabási, 2002).
E.6 WORLD COMPUTATION GRAPHS
The computation graphs in this dataset originate from a diverse set of neural network architectures with different applications, including for example classification and denoising. We train on 8 real- world computation graphs of sizes below 100 nodes and evaluate on 4 different computation graphs of sizes between dozen and 128 nodes. Note that the maximal achievable speedup is on average lower for the real-world computation graphs compared to the synthetic ones. Furthermore, the synthetic graphs are also more homogenous in terms of the maximal achievable speedup, with a lower limit of at least 3. In contrast, some real-world graphs could not exceed speedup of 1.5 and others went beyond 3.
In Figure 4, report the average speedup of the schedules sampled during training, varying training steps. We compare our proposed log- partitioo variance loss against the trajectory balance loss (Malkin et al., 2022) that uses Topoformer (Gagrani al., 2022) to regress the log-partition
16

---

Published as a conference paper at ICLR 2023
2.2
Average Speadup
2.1
2.0
1.7
1.6
variance
Trajectory balance
100000
120000
Training Steps
Figure 4: Average speedup of schedules sampled during training. We compare the log-partition variance loss against the trajectory balance loss that models the log-partition function with neural network. The nog-partition variance loss starts improving earlier and achieves much higher final speedup.
The results demonstrate that the log-partition variance loss starts to improve the average speedup much earlier during training and achieves better performance in the end.
function.
F TEMPERATURE CONDITIONING ABLATION
In order increase the likelihood of sampling good schedules, could introduce fixed tempera- ture throughout training and inference. However, we have observed that this procedure is unreliable for small temperatures. Figure 5 shows generalization performance during training on the syn- thetic graph experiment of Section 5.2. As can be seen, choosing a temperature of 0.01 results in a smaller maximum reward as opposed to training on higher temperature of 0.03. On the other hand, sampling range of temperatures between 0.01 and I and evaluating on 0.01 samples the best performing schedules on unseen computation graphs.
3.4
3.2
Speedup
3.0
Ammnn
2.0
WMAL
Training Temperature
2.6
[0.01,
0.01
0.03
2.4
20000
40000
60000
80000
100000
120000
Training Steps
Figure 5: The impact of different temperature regimes on top-1 generalization performance. Train- ing on single temperatures prevents learning when set too low (orange). On the other hand, training on a range of different temperatures (blue) results in better performance when performing inference with the minimum training temperature.
17


Reviews:
[{"review_id": "mbBp6eaiwly", "reviewer": "Reviewer_uJuK", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 examines the computation scheduling problem, recently tackled by RL methods, and attempts to tackle it with the GFlowNet framework. The application is fairly straightforward, but the authors introduce additional interesting aspects that improve performance.\n\nThis paper has several contributions:\n- applying GFlowNets to the scheduling problem\n- training temperature-conditional GFNs (with temperatures sampled log-uniformly)\n- training Trajectory Balance-style GFNs without having to estimate $Z(G_C)$ using a variance loss trick\n- training on subgraphs as a data augmentation", "summary_of_the_review": "While I am unfortunately unable to judge the impact of this work on scheduling problems, the proposed improvements to the GFlowNet framework seem interesting. I think the paper could do much more in terms of details, empirical validation of the proposed ideas, and potentially tackle larger scale benchmarks.", "strength_and_weaknesses": "**Strengths**:\n- clear writing, clear exposition, clear methods\n- brings novel insights into GFlowNet research\n- good use of the diversity-as-a-strength feature of GFlowNets to improve scheduling\n\n**Weaknesses**:\n- there are some details lacking, for example, I'm not sure what the architecture is to condition on a particular computation graph, or how temperature is fed to the MLP producing $e_\\sigma$.\n- there are some convincing results for the temperature-conditional GFN method proposed, but not for the log-partition variance loss. The authors claim that learning $Z_\\theta(G_C)$ is divergent, but no evidence is provided to help understand the reader why this is the case, nor if the learned flows converge to the correct quantities.\n- similarly, the effect of training on subgraphs vs full graphs is not very deeply discussed nor empirically justified.\n\n> To the best of our knowledge, this is the first time that generalization of GFlowNets is tested empirically.\n\nIf I understand prior work correctly, generalization to unseen _data_ has been tested (see e.g. [1]). Perhaps here the authors are instead referring to generalization to unseen conditionals (in the present case, unseen computation graphs to be scheduled)?\n\n> In practice, we use the training distribution to estimate $\\mathbb{E}_s [\\zeta(s)]$ with a mini-batch of trajectories\n\nHow does this interface with (temperature and graph) conditioning? $Z$ is effectively a function of $G_C$ and $\\sigma$. If I understand correctly this means that for a given $G_C$ (and $\\sigma$) many samples are taken. Does this require a \"combinatorial\" number of samples? i.e. for a \"minibatch\" of size $n$ do we need $n=n_b n_G n_\\sigma$ samples?\n\n[1] Evaluating Generalization in GFlowNets for Molecule Design, Andrei Cristian Nica, Moksh Jain, Emmanuel Bengio, Cheng-Hao Liu, Maksym Korablyov, Michael M. Bronstein, Yoshua Bengio, MLDD 2022", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "f36Q-YwBB6s", "reviewer": "Reviewer_umWR", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Scheduling operations in a computation graph on parallel hardware is NP-hard. This makes ML approaches such as RL attractive. However, the high cost of evaluation necessitates the use of proxy reward models, which can be overfitted to with reward maximization, resulting in poor performance on actual hardware. This is analogous to drug discovery where a proxy model is used in place of wet-lab experiments.\n\nThe solution is to seek not just the best mode in the reward distribution but all modes proportional to their goodness. This can be done by training an amortized sampler which samples proportional to the Boltzmann distribution defined by the EBM (the proxy model). The GFlowNet framework is used to train such a sampler with temperature conditioning. Empirical results show that GFlowNets achieve a higher degree of diversity, which correlates with better on-hardware performance.\n", "summary_of_the_review": "Well-motivated paper with adequate theoretical and empirical contributions.", "strength_and_weaknesses": "This paper applies the right tool, GFlowNets, to an important problem, operation scheduling, and obtains good results. Note that it might be worth mentioning that the proposed state space is a DAG, which necessitates the GFlowNet machinery; otherwise, max-entropy RL suffices.\n\nAlong the way, it also empirically evaluates the idea of conditional GFlowNets and utilizes a log-partition variance loss, which are not done in prior work.\n\nThe presented empirical evidence is comprehensive and includes speedup on target hardware and real-world computation graphs.\n\nApproximating E_s[\\zeta(s)] with the mini-batch average of \\zeta(s) reminds me of variance reduction via baseline subtraction where the baseline is estimated with the batch reward. I encourage the authors to explore how the proposed loss is related to established variance reduction methods in RL and beyond.\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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "daaVW3u73MF", "reviewer": "Reviewer_b75p", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 paper works on scheduling of a computational graph which is an NP-hard problem.\nThe paper considers a problem of scheduling a computational graph on a fixed no. of homogeneous devices.\nThe paper claims that the previous approaches take large number of evaluations for convergence and suggests proxies as a faster alternative.\nInstead of trying to minimize the makespan measure w.r.t the proxy, this paper creates a set of candidates and learns a generative model to assign higher probability to the low-makespan schedules.\nOverall, the paper is a case study of GFlowNets on scheduling computational graphs.\nMain difference over GFlowNet baseline is (1) using a variance loss on the log-partition function, (2) learning a single model for multiple different reward functions by conditioning the policy network on a temperature.\n\n ", "summary_of_the_review": "The paper works on scheduling of a computational graph which is an NP-hard problem.\nThe paper considers a problem of scheduling a computational graph on a fixed no. of homogeneous devices.\nThe paper combines GFlowNet with Topoformer with appropriate changes to the formulation.\n\nWhile it may look incremental from a cursory look, the paper makes a very solid contribution.\nIn fact, its experiments also support its significant benefits.\n\nHowever, it seems that the results are limited to simulation.\nIt would be great if we can see more real system experiments.", "strength_and_weaknesses": "+ High impact problem, considering that scheduling is a very generic problem in computer systems.\n+ Very solid contribution on how to use GFlowNets in scheduling\n+ Solid experiments with good baselines (BRKGA, List Scheduling, GFlowNet)\n\n- Seems to be only simulation (would be great if we can see some real system experiments)\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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "C4I9i7veJDK", "reviewer": "Reviewer_9S3Z", "summary": "", "questions": "", "limitations": "", "rating": 6, "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 apply GFlowNets to the problem of scheduling operations in a computational graph on homogeneous parallel hardware. The evaluate conditional GFlowNets for the first time with the computational graph being the conditioning input and the schedule being the output. They also innovate by conditioning on the reward temperature in order to find a good trade-off between diversity of the generated samples and their average reward. Diversity is important in this application because the generative policy is trained using an imperfect but cheap to compute proxy and only a few of the generated schedules are then evaluated on the ground truth hardware. GFlowNets are interesting in this context because their training objective makes them sample with probability proportional to the reward (rather than maximizing the reward as in other RL approaches). Experiments show improvements against previous methods, including standard RL methods.", "summary_of_the_review": "This paper introduced a new form of application for GFlowNets, mapping a computational graph to a distribution over efficient schedules for it. To achieve this, the demonstrate the feasibility and success of the conditional GFlowNet formalism and also show how to use it to control the trade-off between diversity and average reward by conditioning on the reward temperature (exponent of the reward).\nI consider the paper should be published if they fix their claim about conditional GFlowNets (that it was not considered before, see above).", "strength_and_weaknesses": "GFlowNets are a new RL variant (less than a year old) with only a few applications yet, so this paper introduces a new type of application, in the realm of scheduling, and an approach which may be applicable more broadly. The successfully show empirical success on two GFlowNet novelties: conditioning on a rich input (the computational graph) and on temperature (to control the diversity-reward trade-off).\n\nIn page 5, they claim that previous works on GFlowNets did not consider the conditional case, but this is not true and they themselves mention in section 4 that Bengio et al 2021b introduced the theoretical framework for conditional GFlowNets. What is true is that may be the first (or among the first, with concurrent work) to validate experimentally a conditional GFlowNet and show its usefulness.\n\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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster