Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs
Authors: Aditya Paliwal
Felix Gimeno
Vinod Nair
Yujia Li
Miles Lubin
Pushmeet Kohli
Oriol Vinyals
Venue:  (2020)
Keywords: reinforcement learning
learning to optimize
combinatorial optimization
computation graphs
model parallelism
learning for systems

Abstract: We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks. 

Published as a conference paper at ICLR 2020
REINFORCED GENETIC ALGORITHM LEARNING FOR OPTIMIZING COMPUTATION GRAPHS
Aditya Paliwal
Felix Gimeno
Vinod Nair
Google Research adi pa se.oogle.com
DeepMind
DeepMind vi vinair@googles
fgimeno@googles e com
Yujia Li
Miles Lubin
Pushmeet Kohli DeepMind pusimeet@@googlers com
DeepMind
Google Research
yujjaali@gpogle.com
mlubin@google..cm com
Oriol Vinyals
DeepMind
vinyals@google.com
ABSTRACT
We present a deep reinforeement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to extensive of baselines, approach achieves significant improvements over classical and other learning-ban based methods on these two tasks.
INTRODUCTION
Deep Learning frameworks such as MXNet (Chen et al., 2015), PyTorch (Paszke et al.) (017)) and TensorFlow (TensorFlow Authors 2016a) represent neural network models computation graphs. Efficiently executing such graphs requires optimizing discrete decisions about how to map the computations in a graph onto hardware so as to minimize relevant cost metric (e.g., running time, peak memory). Given that execution efficiency is critical for the success of neural networks, there is growing interest in the use of optimizing static compilers for neural network computation graphs, such Glow (Rotem et al. 2018). MLIR (MLIR Authors 2018), TVM (Chen al. 2018a), and XLA (XLA team 2017).
Here consider the model parallelism setting where computation graph can be executed using multiple devices in parallel. Nodes of the graph are computational tasks, and directed edges denote dependencies between them. We consider jointly optimizing over placement, i.e., which nodes executed on which devices, and schedule, i.e., the node execution order on each device. These decisions are typically made in either one two passes in the compiler. We consider two different objectives: 1) minimize running time, subject to not exceeding device memory limits, and 2) minimize peak memory usage. In the optimization literature, such problems are studied under the class of task scheduling, which is known to be hardard in typical settings (Sinnen 2007 Kwok & Ahmad 1999).
As scheduling and placement are just a few of the many complex decisions made in a compiler, is essential in a production setting that solution 1) produce solutions of acceptable quality fast, even on large graphs (e.g., thousands of nodes) and decision spaces, and 2) handle diverse graphs from various types of applications, neural network architectures, and users. In this work we consider
"Google AI Resident

---

Published as a conference paper at ICLR 2020
Input
Graph Neural Pollcy
BRKGA
Output
Evaluation
FIlM
Placement
Peak Mancry Rumiin
Parformnne Model
Schodule
Fig. 1: Overview of our approach. The Biased Random Key Genetic Algorithm (BRKGA) is used to optimize execution decisions for a computation graph (c.g., placement and scheduling of nodes) with respect to cost metric (e.g.. running time, peak memory) computed using the performance model. BRKGA requires proposal distributions for each node in the graph to generate candidate solutions in its search loop. The default choice is agnostic to the input graph: uniform distribution over [0, 1] at all nodes. We use graph neural network policy to predict node-specific non-uniform proposal distribution choices (parameterized as beta distributions over [0, 1]). BRKGA is then run with those choices and outputs the best solution found by its iteration limit. By controlling the non non-uniformity of the distributions, the policy directs how BRKGA's search effort is allocated such that better solution can be found with the same search budget.
learning an optimizer that satisfies these requirements. Crucially, we aim to learn an optimizer that generalizes to broad set of previously unseen computation graphs, without the need for training such graphs, thus allowing it to be fast at test time.
Previous works on learning to optimize model parallelism decisions (Mirhoseini al. 2017:
2018
Addanki et al. 2019) have not considered generalization to a broad set of graphs nor joint optimization of placement and scheduling. In Mirhoseini al. (017;) 2018), learning is done from scratch for each computation graph and for placement decisions only, requiring hours (e.g.. 12 to 27 hours per graph). This is too slow to be broadly useful in general-purpose production compiler. We propose an approach that takes only seconds to optimize similar graphs. In concurrent work to ours, AAddnkii al. (2019) shows generalization to unseen graphs, but they are generated artificially by architecture search for single learning task and dataset. In contrast, we collect real user-defined graphs spanning a broad set of tasks, architectures, and datasets. addition, both Mirhoseini et al. (2017) 2018)) and Addanki et al. (2019) consider only placement decisions and rely on TensorFlow s dynamic scheduler; they do not address the static compiler setting where it is natural to jointly optimize scheduling and placement.
The key idea of our approach (Figure D to learn a neural network that, conditioned on the input graph to be optimized, directs existing optimization algorithm's 's search such that it finds better solution in the same search budget. We choose the Biased Random-Key Genetic Algorithm (BRKGA (Gonçalves & Resende 2011) ) as the optimization algorithm after an extensive evaluation of several choices showed that it gives by far the best vu-qualtty trade-off for our application. BRKGA produces good solutions in just few seconds even for real-world TensorFlow graphs with thousands of nodes, and we use learning to improve the solution quality significantly at similar speed. We train a graph neural network (Battaglia et al. 2018) take computation graph as input and output node-specific proposal distributions to use in the mutant generation step of BRKGA's inner loop. BRKGA then run to completion with those deput-dependent distribution choices, instead of input- agnostic default choices, to compute execution decisions. The distributions are predicted at each node, resulting a high-dimensional prediction problem. There is no explicit supervision available, so use the objective value as a reward signal in contextual bandit approach with REINFORCE (Williams 1992). Our approach, "Reinforced Genetic Algorithm Learning" (REGAL), uses the network's ability to generalize to new graphs to significantly improve the solution quality of the genetic algorithm for the same objective evaluation budget.
We follow the static compiler approach of constructing coarse static cost model to evaluate execution decisions and optimizing them with respect to it, as done in (Addanki et al.) 2018 Ja et al., 2018)) This is in contrast to evaluating the cost by executing the computation graph hardware (Mirhoseini
2

---

Published as a conference paper at ICLR 2020
et al. 2017| 2018 A computationally cheap cost model enables fast optimization. It is also better suited for distributed training of RL policies since a cost model is cheap to replicate in parallel actors, while hardware environments problems, so optimizing it is difficult. In this paper we focus fully on learning to optimize this cost model, leaving integration with compiler for future work.
not. Our cost model corresponds to classical NP-hard scheduling
We structure the neural network's s task as predicting proposal distributions to use in the search over execution decisions, rather than the decisions themselves directly. Empirically have found the direct prediction approach to be too slow at inference time for our application and generalizes poorly. Our approach potentially allows the network to learn more abstract policy not directly tied to detailed decisions that are specific to particular graphs, which may generalize better to new graphs. It can also make the learning task easier as the search may succeed even with sub-optimal proposal distribution predictions, thus smoothening the reward function and allowing the network to incrementally learn better proposals. The node-specific proposal distribution choices provide a rich set of knobs for the network to flexibly direct the search. Combining learning with search algorithm has been shown to be successful (e.g., (Silver et al. 2017 2018). and our work can be seen as an instance of the same high-level idea.
This paper makes several contributions:
We are the first demonstrate learning policy for jointly optimizing placement and scheduling that generalizes to broad set of real-world TensorFlow graphs. REGAL significantly outperforms all baseline algorithms on two separate tasks of minimizing runtime and peak memory usage (section 5.3) on datasets constructed from 372 unique real-world TensorFlow graphs, the largest dataset of its kind in the literature and at least an order of magnitude larger than the ones in previous works (Mirhoseini et al. 2017 2018
Chen et al. 2018b Addanki et al. 2018 2019).
We use a graph neural network to predict mutant sampling distributions of genetic algo- rithm, specifically BRKGA, for the input graph to be optimized. This directs BRKGA's search in an inputddependent way, improving solution quality for the same search budget.
We compare extensively to classical optimization algorithms, such as enumerative search, local search, genetic search, and other heuristics, and analyze room-for-im noom-oorimprrovement in the objective value available to be captured via learning. Both are missing previous works.
RELATED WORK
Learning to optimize computation graphs: AutoTVM (Chen et al. 2018b) applics learning to the very different problem of optimizing low-level implementations of operators in tensor program, while we focus on optimizing higher-level decisions such as placement and scheduling of ops. [Ma et al.) (2019) use graph neural nets and RL to learn scheduling policy for data processing jobs on clusters. These works are conceptually similar to ours in their use of learning, applied to a different domain.
Learning for combinatorial optimization: Our work is instance of applying learning for combinato- rial optimization (Bengio et al., 2018). Previous works on learning graph combinatorial optimization algorithms (c.g.. Li et al. (2018); Khalil et al. (2017) have focused on problems such as Minimum Vertex Cover, Maximum Clique, Maximum Independent Set, etc. The task scheduling problem we consider is significantly different in that the objective value is more complex function on node-level decisions. Also, we focus on large-scale, real-world TensorFlow graphs, while e.g., Khalil et al. (0177) uses small-scale, synthetic graph distributions.
Learning a proposal distribution for stochastic search: Bunel et al. (017) learns policy for pre- dicting instance cetence-dependent proposal distributions to be used in the stochastic optimizer STOKE (Schkufza et al., 2013) for superoptimizing programs. However, it uses handcrafted instance fea- tures and shows results on relatively simple, small programs. In contrast, we automatically learn the instance representations and show results real-world graphs. An earlier work by Paige & Wood (2016) similarly learns a neural network to predict input-dependent proposal distributions for sequential Monte Carlo search for inference in graphical model.
Optimization without learning: Parallel task scheduling (Sinnen 2007 Kwok Ahmad 1999) is a classical problem for scheduling ops in computational graph to minimize runtime. Learning not
3

---

Published as a conference paper at ICLR 2020
traditionally part of the approaches proposed in this literature. Mayer al. (2017) studies greedy task scheduling approaches for TensorFlow. Jia et al. (2018) develops a simulation based optimizer for deep learning computation graphs that uses larger decision space by combining data, model, and attribute parallelism. Our approach can potentially he extended to such larger decisions spaces to achieve even bigger improvements in execution cost.
BACKGROUND
Figure T shows an overview of our approach. Given an input graph to optimize, instead of applying BRKGA directly with the default uniform distribution all nodes, a graph neural network predicts beta distribution choices at each node. BRKGA is run with these choices to optimize placement and scheduling decisions with respect to the objective defined by the performance model. We first explain the performance model and BRKGA in this section, and the learning component in the next.
3.1 PERFORMANCE MODEL
A computation graph has set of ops run. Each op produces zero more tensors and requires zero or more tensors as input. The runtime of each op is known and fixed (e.g.. given by a simulator in Jia et al. (2018)). The memory use of each tensor is known (an assumption that holds in static compilers like XLA). We assume a collection of d homogeneous devices that have separate local memory and can run at most one op a time, An op can run only when its input tensors are present in the local memory. Tensors can be transferred across devices by synchronous (blocking) transfers. Tensors are freed from local memory after all local consumers have run.
In this setting, we consider the problem of finding an assignment of ops to devices and an overall schedule such that each op run once with the objectives of (1) minimizing the peak local memory use across devices (e.g.. to find a feasible way to run large computation graph). or (2) minimizing the runtime subject to a constraint on the peak memory used on any device.
The performance model does not consider rematerialization of tensors, fragmentation when computing memory use, and asynchronous transfers between devices. Despite these simplifications, the model yields slight variants of problems that are known to be NP-hard (Eyraud-Dubois. et al., (0155) and therefore remains challenging setting in which to study how to learn an optimizer. See section A.4 for more details of the model.
3.2 BIASED RANDOM-KEY KEY GENETIC ALGORITHM
Biased random-ke key genetic algorithm (BRKGA) is meta-heuristic framework that has been success- ful in a wide array of applications for solving hard combinatorial optimization problems (Gonçalves) Resende 2011) J. In BRKGA, chromosomes in a population are encoded n-dimensional vectors with entries in [0, 1] for some fixed n. This random- key encoding decouples the application from the genetic algorithm, specifically the crossover and mutant generation procedures (Bean) 1994).
The BRKGA variant we use is specified by (1) fitness evaluation function f: [0, 1|" + R. (2) scalar integer parameters æ, Tes and n. representing the population size, number of elites, and number of children, resp., (3) an elite bias p E [0.5. 1.0), and (4) a mutant generation distribution D over [0, 1] The procedure aims to find a chromosome that maximizes f.
The initial population (a collection of chromosomes) is created by sampling from D. (Known good solutions may also be used to initialize a population.) One evolution step is completed as follows.
1. Sort the chromosomes in order of decreasing fitness using f. Denote the first म. chromo- somes as elites and the remaining chromosomes as nonelites.
2. Construct the next generation from three different sources of chromosomes: (a) Copy the elite chromosomes unmodified from the last generation. (b) For each of the मं new children, select two parent chromosomes uniformly at random, one from the nonelites and one from the elites. Apply the crossover procedure (described below) to generate a new chromosome given the two parents. (c) Generate the remaining ne n. by sampling from D.
4

---

Published as a conference paper at ICLR 2020
We continue the evolution procedure for fixed number of evaluations, i... calls to f. Given an elite and nonelite chromosome a, b E [0, 1]" (resp.), the crossover procedure produces a child chromosome by independently combining entries from the parents. Specifically, for each index i E ,,...,, n independently, let C; a, with probability P and C; = b, with probability 1 -
Our use of BRKGA is standard except for the mutant-sampling distribution D, which is usually fixed to the uniform distribution. We generalize BRKGA for instance-specific learning by sampling from independent beta distributions, whose parameters can vary by index. Beta flexibly admits non-uniform distribution choices and also subsumes the uniform choice.
Given computation graph, let d be the number of devices, o the number of ops, and the number of tensors. We define the chromosome encoding a scheduling solution to have three distinct parts: (I) x d entries specifying op-to-device affinities; (2) entries specifying scheduling priorities for each op: (3) x d entries specifying censor-to-device priorities for transfers that may be needed. Given a chromosome, op placements are picked by maximum affinity. Transfer ops are created as implied by the placements. We then obtain schedule by performing topological sort over the ops given their tensor dependencies, breaking ties by using the corresponding node priorities. Once a schedule is constructed, the performance model is used to evaluate its peak memory and/or runtime. When enforcing a memory constraint, the fitness of a schedule is encoded such that all feasiblee schedules have better fitness than infeasible schedules. An example is provided in section
REGAL
4.1 CONTEXTUAL BANDIT FORMULATION
We train a contextual bandit policy that predicts beta distribution choices for each of the nodes of computation graph to be used by BRKGA to optimize it. For each round in the bandit setting, first the context is observed by drawing a computation graph G as an i.i.d. sample from a distribution DD (e.g.. a distribution over TensorFlow graphs). G has a set of nodes V and set of edges E, with features associated with the nodes and edges. A policy p(a G) is applied to make a set of decisions at each node. These decisions, denoted a, for each a = {avev) J. One decision in a corresponds to playing one arm of the bandit, and specifying the entire a corresponds to playing several arms together in single round. This can be viewed combinatorial multi-armed bandit problem (Chen et al. 2013).
V, across all nodes form action
The action a specifies all the node-specifie beta distributions BRKGA needs to optimize placement and scheduling decisions for G. To enable a policy over discrete choices, we quantize the mean and variance parameters of the beta distribution. The environment then runs BRKGA with those distribution choices with a fixed iteration limit. The final objective value is used to compute the reward. To make the reward values comparable across different graphs. we divide the objective value on
o,(() (G) achieved on graph G with action by the objective value o, (G) achieved by standard BRKGA using uniform distributions. Since we want to minimize the objective (e.g., runtime or peak memory), we define the reward as r(a, G) =
(C) So a reward > -1 corresponds to an action
ea
that achieves better objective value than standard BRKGA on a graph.
We maximize the expected reward L EG E ()r(a, G)), where Eg is expectation over graphs in our training set. Learning is done by REINFORCE (Williams 1992). We added scalar baseline b(G) to reduce the variance of the gradient estimates.
4.2 GRAPH NEURAL NETWORK POLICY
From computation graphs, we derive multigraphs with attributed nodes and directed edges. Denote a multigraph G = (V, E). In our setup, the nodes V correspond 1:1 to the ops. An edge E E exists from u to v for each tensor that op e requires that is produced by op u. As a tensor can be required by multiple ops, the correspondence from edges to tensors may be many to one. Each node v E V and edge e E E' has an attribute vector I, and E.. The attributes contain respective features, e.g., sizes of the tensors.
We learn a model that predicts good mutant sampling distributions for BRKGA given this multigraph. Each node has d + independent beta distributions, corresponding to device affinities and scheduling
5

---

Published as a conference paper at ICLR 2020
priorities, whose parameters are represented as vector d,. These the model's actions in RL terms, and our model specifies a distribution over actions a = vaelee for each graph, p(a|G). Note the action space is different from graph to graph.
We use Graph Neural Networks (GNNs) (Scarselli et al. 2009 Li ct al. 2015 Gilmer et al. 2017 Batttglia al., 2018) to learn representations for computation graphs. Given (multi)graph G, GNN computes representation vectors h, for each node through an iterative message passing process
follows:
h (0) MLP,t n (xx);
hg = (x))
m (1) MLP usg (he her hel),
(6) (1)
m (1) MLP
(0) hel) [
(0)
h ((+1) MMP [n Dec.=p
me
(4)
Ee:e.=v m
where E, the source node of edge and e, is the target node. In our formulation, MLP, n and MLP. are multilayer perceptrons (MLPs) that encode node and edge attributes, MLP msg and MLP msg
compute messages along the edges in the edge direction (m (0) and the opposite direction (m. (t) ) MLP node updates node representations and [-] represents flat vector concatenation. After rounds of message passing, the representation for each node h. hee
(T) will contain information from the
T Thoop neighborhood around v in the graph. Given the h, 's we produce a, 's through conditionally independent predictions, where the prediction for one node does not depend on the predictions of other nodes given the computed representations:
G = II p(a, GG) II pla,aa(MP.... (M...(hh))))
(2)
MLP, is shared across all nodes for predicting the parameters of the output distributions. In experiments, we quantize the continuous beta distribution parameters and use a discrete action space. The outputs are therefore categorical, and softmax distributions. More details are included in section section[A.6] The baseline is computed using a separate GNN, where after obtained the node representations h., we aggregate across nodes and
the MLP to compute the logits of the corresponding
compute b(G) as b(G) = MLP, (M E. MLP, y y(h.))
EXPERIMENTAL RESULTS
5.1 TASKS AND DATASETS
We consider two tasks, one is minimizing peak memory and the other is minimizing running time, both on two homogeneous devices with 16 GiB of memory each and synchronous tensor transfers with zero cost (zero latency and infinite bandwidth). We train separate neural network for each task-dataset pair for the case of two devices.
We have collected dataset of 372 topologically-une uuiqque real-worke TensorFlow graphs by mining machine learning jobs on Google's internal compute cluster (..2). These jobs are from wide range of production and research use cases. The dataset is split into [train, valid, test) sets containing (60%, 20%. 20%) graphs, respectively. These sets are disjoint with respect to graph topology, so at test time the policy needs to generalize new topologies.
We augment the dataset by applying multiplicative noise to tensor sizes and op running times to create several variants per graph. Even though the variants of the same graph share the same topology, they represent different optimization problem instances. We create separate datasets for minimizing runtime and peak memory. The TF runtime dataset has 16329 training. 5470 validation, and 5266 test graphs. The TF peak memory datasel has 22400 training, 7400 validation, and 7400 test graphs.
For reproducibility, we have released synthetic dataset of computation graphs with 10000 training, 1000 validation, and 1000 test cases. The graph topologies are generated from several classical random graph models, and the op running times and tensor sizes are sampled from Gaussian distributions (see A.1.4) On this dataset we minimize running time without a memory constraint (c.g., on two homogeneous devices with infinite memory).
ht pttps wttps////vvvigeeeee wwo.g le wwt.l//e.ooooooo..oeoeeeee ve/ ders/ 1xx TocsMu- www.deooooeeeeeeeeeeeeeeeee

---

Published as a conference paper at ICLR 2020
5.2 BASELINES
Graph Parlitioning + Depth First Search (GP+DFS): Combines graph partitioning (GP) baseline for device placement to minimize communication across devies and a Depth-First Search heuristic similar to the one implemented in XIAA (TensorFlow Authors 20T6b) to compute per-device schedules given placements. This is representative of the XLA compiler's 's solution for model parallclism.
Local Search: The method starts with a random placement and schedule and greedily improves it by moving an op across devices or changing an op's order in the current schedule.
Graph-As-Sequenc Model (GAS): Like Mirhoseini et al. (2017) 2018), we convert the graph into a sequence using a topological sort and apply a recurrent neural network to predict node-level distributions to be used by BRKGA. This comparison measures the usefulness of graph structure for learning.
BRKGA XK: Run BRKGA for X thousand iterations with uniform sampling distributions using default hyperparameters consistent with Gonçalves & Resende (201T). This comparison measures the performance of the default version of BRKGA.
Tuned BRKGA: Apply grid search to BRKGA's hyperparameters on the training set and pick the best. This represents how well BRKGA performs by customizing it to the distribution of computation graphs, but without instance-depender customization.
Instance-dependent Random Search (IDRS): Same as REGAL, but BRKGA is replaced with random search. This is done by running BRKGA for only one generation using the proposal distributions computed by the neural network.
Additionally, we use Constraint Programming (CP) approach with the CP-SAT solver of Google OR- tools (Google 2019) establish a provably global optimum for each computation graph optimization problem instance by running for up to 24 hours. As an enumerative algorithm, it is generally not competitive when run only for seconds.
For fair comparison, we fix the number of performance model evaluations allowed per graph to be the same across algorithms. (Except GP+DFS, which does not allow fixing it.) Given typical TensorFlow graph sizes and compiler running time constraints, we estimate that a budget of 5000 evaluations is feasible in practice, so we use that in the experiments.
Learning to directly predict a solution: We have explored two more approaches for training graph neural network to predict placement and scheduling solutions directly, without BRKGA. We used supervised learning to train network to predict BRKGA's solutions. The best accuracy was achieved by predicting the solution autoregressively. one variable at a time conditioned on previously predicted variables. We also used RL to learn a policy with IMPALA (Espeholt et al. 2018) to optimize the objective value by incrementally predicting the solution one variable at a time, and once complete, iteratively improving it with learned local search policy. The inference cost for both approaches is quadratic in the number of nodes (the graph net is applied a linear number of times, each with linear cost), while REGAL's inference cost linear, making them orders of magnitude slower than REGAL at test time, An evaluation on a test set of small graphs showed that neither approach improves on BRKGASK. Improving the scalability and the generalization of these approaches is left as future work, and we do not present their results here.
5.3 COMPARISON TO BASELINE ALGORITHIMS
We use two metrics to compare algorithms. 1) Average percent improvement over BRKGA 5K: For a given graph, compute the percent improvement in the objective achieved by an algorithm relative to BRKGA with evaluation limit set to 5000. BRKGA 5K a natural reference for measuring the effect of learning approaches that predict proposal distributions for it. 2) Average percent gap from best known solution: Compute the best known objective value among all the algorithms. (This will be found by CP-SAT if it finishes within the time limit.) Compute the percent difference between an algorithm's 's solution and the best known objective value. We report averages over test set graphs. Training set results similar and reported in section A.3]
Table compares REGAL to other algorithms on the two TensorFlow test sets and the synthetic dataset. REGAL outperforms all the baselines on all three tasks. It gives 1.9 and 4.4x bigger improvements
7

---

Published as a conference paper at ICLR 2020
Table 1: Comparison of methods the TensorFlow and Synthetic test Results averages over test set graphs. Higher is better for % Improvement over BRKGASK, and lower is better for % Gap from best known. Note: CP-SAT, an enumerative algorithm, is run for up to 24 hours only to establish provably global optima (if possible) for evaluation purposes.
TF Runtime test
TF Peak Memory test set Gap from best known
Symthetic Runtime test set
% Improv.
% Gap
Improv.
% Improv.
% Gap from best known
Algorithm
over BRKGA 5K
from best known
over
over BRKGA 5K
BRKGA 5K
CP-SAT 24hr
13.85%
1.00%
-1.48%
8.06%
19.50%
0.00%
GP + DFS
-37.32%
66.98%
-6.51%
14.77%
-55,8%
93.66%
Local Search
-1.66%
22.63%
0.63%
7.24%
0.08%
24.60%
BRKGA 5k
0.00%
20.19%
0.00%
7.98%
0.00%
24.63%
Tuned BRKGA
3.20%
16.40%
0.80%
7.11%
3.11%
20.76%
GAS
3.79%
15.24%
0.16%
7.67%
0.80%
23.48%
IDRS
-6.87%
28.60%
-3.16%
12.39%
-12.12%
39.72%
REGAL 7.09%
11.04%
3.56%
4.44%
4.81%
18.57%
than the next best algorithm runtime and peak memory minimization tasks, respectively. The percent improvement over GP + DFS is 44.4% and 10.1% 1 for runtime and peak memory, respectively. REGAL reduces the average percent gap from the best known solution by about 1.8x with respect BRKGA 5K on both TensorFlow test sets, and by about 6x and 3.3x with respect to GP + DFS on the TensorFlow Runtime and Peak Memory test sets, respectively. (For an XLA user, GP DFS is the current, albeit weak, stae--of-tee- algorithm.) The synthetic test set shows similar results. The learned policy successfully generalizes to previously unseen graphs, the extent that a large fraction of the estimated room for improvement over BRKGA 5K is captured using the same evaluation limit.
To further test the limits of generalization for the policies learned using REGAL, we evaluate them on XLA graphs from production compiler team's 's internal performance benchmark. XLA uses different set of ops from TensorFlow, and the benchmark graphs on average have about an order of magnitude more nodes and edges than the TensorFlow graphs in our training set, this is difficult generalization challenge. REGAL achieves 0.58% average runtime improvement over BRKGA 5K on 94 graphs, and 3.7 3.74% average peak memory improvement on 32 graphs. It is promising that any improvements are possible at all despite training only on TensorFlow graphs, and points to the possibility of bigger improvements by training directly XLA graphs.
Optimizer running times: BRKGA 5K takes on average 0.89 seconds on the TensorFlow Peak Memory test set to optimize computation graph, while REGAL takes 1.04 seconds. (The times are similar on the Runtime test set.) Instead of taking hours to compute a solution per graph (e.g., Mirhoseini al. (2017 2018)), REGAL produces solutions orders of magnitude less time, while still being better than all the baselines.
5.4 COMPARING REGAL vs. BRKGA
Figure 2 shows histograms of percent improvements runtime (left) and peak memory (right) achieved by REGAL over BRKGA 5K on the test sets. Green bars correspond to graphs on which REGAL improved over BRKGA 5K, while red bars correspond graphs on which REGAL was worse. (Ties have been omitted for clarity.) REGAL matches or beats BRKGA 5K on 87 4.4% of the runtime set, and 88.9% of the peak memory test set. The highest improvement 26.0% for run time and 54.3% for peak memory, while the worst regression is 24.0% for run time and 17.9% for peak memory.
To assess whether the improvements provided by REGAL's policy generalize to evaluation limits other than the one for which it was trained, we varied the evaluation limit used by both BRKGA and REGAL at test time. The results are shown in figure REGAL's performance improves with more evaluations, confirming that the policy generalizes to higher evaluation limits. other words, there exist nove-level choices for the distributions used in BRKGA that perform well regardless of the evaluation limit, and REGAL learns to predict those choices. This is particularly useful cases where the actual evaluation limit will be known only test time, that the same policy be
8

---

Published as a conference paper at ICLR 2020
30.4%
18%
26.6%
16%
Percensagge 18t
22.8%
Parsonage of lest sel
12%
19.0%
10%
15.2%
8%
4%
3.8%
2%
0.0%
0%
20
110
10
Percentage improvement over BRKGA 5K
Percenlago improvvmment over BRKGA 5K
Fig. 2: Histogram of percent improvements in objective value on the TensorFlow runtime (left) and peak memory (right) datasets for test graphs on which REGAL is better (green) and worse (red) than BRKGA. (Ties are omitted from the figure for clarity but are included in the histogram percentage calculation.)
Percentage Improvement BRKGA 5K
Patcentage impporeeeent BRKGA 5K
12
6-
10
BRKGA REGAL
BRKGA
REGAL
10
20
30
40
50
60
10
20
30
40
50
60
Evaluation Limit (x 10')
Brahution Limit (x (x10')
Fig. 3: Average percent improvement over BRKGA 5K given by REGAL and BRKGA the TensorFlow test set for running time (left) and peak memory (right) as the evaluation limit is increased.
applied without re-trinn Interestingly, even with 50K evaluations, BRKGA is not able match REGAL's performance with just 5K evaluations!
5.5 GRAPH-DEPENDENT POLICY
The RL agent's actions are instance dependent. The agent that performs the best on the TF Runtime dataset has a choice of 16 different node placement actions for each node in graph. For each graph in the TF Runtime test set, we compute the entropy of the distribution of the node placement actions taken by the agent and plot histogram these entropies in Figure (a). The mean this distribution is 1.71 nats which implies that the actions are neither uniform random, nor constant, and vary from graph to graph.
Furthermore, the agent's performance overall gets better with more graph message passing iterations T. Figure shows the peak validation reward reached within a hyperparameter sweep for each T for the TF runtime optimization task. Models that utilize the GNN with message passing (T 0) reach higher performance than T = (i.e., ignoring the graph structure).
CONCLUSIONS AND FUTURE WORK
By training graph neural network policy to predict craph-conditional node-level distributions for BRKGA, REGAL successfully generalizes to new graphs, significantly outperforms all baselines in solution quality, and computes solutions in about second on average per TensorFlow test set

---

Published as a conference paper at ICLR 2020
2500
-0.921
-0.926
2000
Validation reward
-0.928
Frequency
1500
-09300
-0932
1000
-0.9344
500
-0.936
-0.98
-0.40
1.0
1.5
2.0
16
Entropy (nats)
Number message passing layers
(b)
Fig. 4: The agent learns to utilize the graph structure. (a) The TF runtime agent picks a diverse set of actions. This plot shows the histogram the entropy the agent's s actions across graphs in the dataset. (b) This plot shows the best validation reward achieved within a sweep of hyperparameters for each number of graph message passing rounds T. The performance gets overall better T increases, and models with T > 0 perform better than T = 0, which does not utilize the structure.
graph. REGAL's speed and generalization make a strong choice for use production compiler that needs to handle diverse set of graphs under limited time budget.
We foresee several extensions. Integrating REGAL into neural network compiler would allow us to evaluate the end-to-end gains due to better placement and scheduling decisions. To further improve REGAL's own performance, one could use Mixture of Experts architecture. Given the diversity of graphs, mixture model can train specialized sub-models convolutional networks, recurrent networks, etc.). Another is to replace BRKGA with alternatives, e.g., combining learned neural policies with local search.
different types of graphs (e.g..
ACKNOWLEDGMENTS
The authors would like to thank Ross Anderson, David Applegate, and Peter Hawkins for a significant amount of infrastructure on which this work builds, including the BRKGA and CP baselines, and Peter Ma, HyoukJoong Lee, and Peter Dolan for help with data collection.
REFERENCES
Ravichandra Addanki, Shaileshh Venkatakrishnan, Shreyan Gupta, Hongzi Mao, and Mohammad Alizadeh. Placeto: Efficient progressive device placement optimization. Workshop on fo. Systems at NeurIPS 2018, 2018
Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan, Shreyan Gupta, Hongzi Mao, and Mo- hammad Alizadch. Placeto: Learning generalizable device placement algorithms for distributed machine learning. In Proceedings of NeurlPS '9, 2019.
Albert-Láászóó Barabási and Réka Albert. Emergence of scaling in random networks. science, 286 (5439):509-512, 1999.
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinícius Flo- res Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Çaglar Gülçehre, Francis Song, Andrew J. Ballard, Justin Gilmer, George E. Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matthew Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pas- canu. Relational inductive biases, deep learning, and graph networks. CoRR, 2018. URL
James C. Bean. Genetic algorithms and random keys for sequencing and optimization. ORSA Journal on Computing. 6(2):154-160, 1994. doi: 10. 2.877//c......1...... URL https: /ttttp:://dooooo. 10.1287/1 0...7111...6.2111.
10

---

Published as a conference paper at ICLR 2020
Yoshua Bengio, Andrea Lodi, and Antoine Prouvost. Machine learning for combinatorial optimization: a methodological tour d'horizon. CoRR, abs/1811.06128,, 2018. URL htttp://axxiv.og//
abs/1811.06128
Rudy Bunel, Alban Desmaison, M. Pawan Kumar, Philip H. S. Torr, and Pushmeet Kohli. Learning to superoptimize programs. In International Conference on Learning Representations, 2017. URL
htttps://ppenrevieew..o
forum? id=r1:z6U5
Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. Mxnet: A flexible and efficient machine learning library for heterogeneous distributed systems. arXiv preprint arXiv:1512.0122: 2015.
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. Tvm: automated end-to-end optimizing compiler for deep learning. In OSDI 2018, 2018a. URL
lttpps://arxivorrg/// 1800.00 799
Tianqi Chen, Lianmin Zheng, Eddic Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. Learning to optimize tensor programs. In Neural Information /ttpaa/wwww..ta/as////..0.....1.
Processing Systems 2018, 5 2018b.
Wei Chen, Yajun Wang. and Yang Yuan. Combinatorial multi-armed bandit: General framework and applications. In Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp- 151-159, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. TtRL:attee//prrceeiiisssaaappppp 1ttpeesss/22/chee13......
Paul Erdos and Alfréd Rényi. On the evolution of random graphs. Publ. Math. Inst. Hung. Acad. Sci, 5(1):17-60, 1960.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In Proceedings of the International Conference on Machine Learning (ICMI), 2018.
Lionel Eyraud-Dubois, Loris Marchal, Oliver Sinnen, and Frédéric Vivien. Parallel scheduling of task trees with limited memory. ACM Truns. Parallel Comput., 112):13: June 2015. ISSN 2329-4949. doi: 10.1145/277905.. /tt/:ttt..//bo.......orr////11/1
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. arXiv preprint
2017.
José Fernando Gonçalves and Mauricio G. Resende. Biased random-key genetic algorithms for combinatorial optimization. Journal of Heuristics, 17(5):487-225, October 2011. ISSN 1381- 1231. doi: 10. 10.10077/110732--000994533
URL http://s///...ooooo///0000/////0
010 9143-1
Google. CP-SAT solver. http:////errrlpprrsooooeoeoo//e2 [p__olver] 2019. [Online; accessed 21 21-January-2011)
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Paul Holland, Kathryn Blackmond Laskey, and Samuel Leinhardt. Stochastic blockmodels: First steps. Social networks, 5(2): 1983.
Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. CoRR, 2018. URL
G. Karypis and V. Kumar. A fast and high quality multilevel scheme for partitioning irregular graphs SIAM Journal on Scientific Computing, 20(1):359-392, 1998. doi:
B. W. Kernighan and S. Lin. An efficient heuristic procedure for partitioning graphs. The Bell System Technical Journal, 49(2):291-30) Feb 1970. ISSN 0005-8580. 7305. 1970.tb01170...
doi: 10.1002/j.
11

---

Published as a conference paper at ICLR 2020
Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial optimization algorithms over graphs. In I. Guyon, V. Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems 30,
pp. 6348-6358. 2017.
Yu- Ywnnw Kwok and Ishfaq Ahmad. Static scheduling algorithms for allocating directed task graphs to multiprocessors. ACM Comput. Surv., 31(4):406-471 December 1999. ISSN 0360-0300.
Yujia Li, Daniel Tarlow, Mare Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. arXiv preprint arXiv: 151.04931 2015.
Zhuwen Li, Qifeng Chen, and Vladlen Koltun. Combinatorial optimization with graph con- volutional networks and guided tree search. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Infor- mation Processing Systems 31, pp. 539-548. Curran Associates, Inc., 2018. URL http://poorssieeeooo/o//scccaaaa oot.:aaoo-ooioiiiiaaoooooo-- htg.::/wwww.aa.a...a.aaaaaoooaa-
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, SIGCOMM 19, pp. 270-288, New York, NY, USA, 2019. ACM. ISBN 978- 678-1--500--9956-6 doi: 00.1145/3341322.344000
URL
Ruben Mayer, Christian Mayer, and Larissa Laich. The tensorflow partitioning and scheduling problem: It's the critical path! In Proceedings of the 1st Workshop on Distributed Infras- tructures for Deep Learning, DIDL '17, pp. 1-6, New York, NY, USA. 2017. ACM. ISBN /tLL:tttppp//aaa..oz///////////
778-144003-516900 doi: 10. 00.1145/114844.315544
0154842.3154847
Azalia Mirhoseini, Hieu Pham, Quoc V. Le. Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. Device placement optimization with reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-1 August 2017, pp. 2430-24333 2017. URL 17a.eeiii1....mm
httt:://proceedinne
r.press/ vv00
Azalia Mirhoseini, Anna Goldie, Hicu Pham, Benoit Steiner, Quoe V. Le, and Jeff Dean. A hierarchi- cal model for device placement. In International Conference on Learning Representations, 2018.
URL
MLIR Authors. Multi-level intermediate representation. mlir 2018. Accessed: 2019-05-22.
Brooks Paige and Frank Wood. Inference networks for sequential monte carlo in graphical models. In Proceedings of the 33rd International Conference on International Conference on Machine Learnnng Volume 48, 2016.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. International Conference on Machine Learning, pp. 1310-1318, 2013.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Nadav Rotem, Jordan Fix, Saleem Abdulrasool, Summer Deng. Roman Dzhabarov, James Hegeman, Roman Levenstein, Bert Maher, Nadathur Satish, Jakob Olesen, Jongsoo Park, Artem Rakhov, and Misha Smelyanskiy. Glow: Graph lowering compiler techniques for neural networks. CoRR, 2018. URL http:/ttpp//axx..orrg////// http:///rgggabb//...00000
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions Neural Networks, 20(1):61-80, 2009.
Eric Schkufza, Rahul Sharma, and Alex Aiken. Stochastic superoptimization. 2013.
12

---

Published as a conference paper at ICLR 2020
David Silver, Julian Schrittwieser, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap. Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of go without human knowledge. Nature, 550, 2017.
David Silver, Thomas Hubert, Julian Schrittwiseer, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Si- monyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6411)) 2018. ISSN 0036-8075. doi: 10. hetp://cienne.aar-4.1. 6419/1140
O. Sinnen. Task Scheduling for Parallel Systems. Wiley Series on Parallel and Distributed Computing. Wiley, 2007. ISBN 9780470121160.
TensorFlow Authors. Tensorflow: Large-scale machine learning on heterogeneous distributed systems.
CoRR, abs/1603.04467, 2016a. ULL/Ltt.///xxx10r//////000000050
TensorFlow Authors. hlo memory memoorysccccuuuerr...
tensorflow wet.roosww...c...eeeaaeeeeeeeeee Accessed: 2019-01-25.
2016b.
Duncan Watts and Steven H Strogatz. Collective dynamics of 'small-world' networks.... nature, 393 (6684):440, 1998.
Ronald Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach. Learn., 8(3-4):229-22 256,, 1992.
XI.A team. blog.
tensorflow compiled. post in the google developers http:///saacceeoo///////00010010
Xla
aet.00/:..000llloo000/0///0.a.e0 2017.
Yuan Yu, Martín Ahadi, Paul Barham, Fugene Brevdo, Mike Burrows, Andy Davis, Jeff Dean, Sanjay Ghemawat, Tim Harley, Peter Hawkins, Michael Isard, Manjunath Kudlur, Rajat Monga, Derek Gordon Murray, and Xiaoqiang Zheng. Dynamic control flow in large-scale machine learning, CoRR, 2018. URL htt.:::://aarawooo/////////....
13

---

Published as a conference paper at ICLR 2020
Number of nodes
Number of edges
25
TF Memory
15
TF Runtime
30.
is
10.
Synthetic
25
30
10
100
1000
10000
10
100
1000
10000
Fig. Histograms of number of nodes (left) and edges (right) for the different datasets. The y-axis shows the percentage of graphs.
45
40
35
800
25
20
15
10
9.0
10.0
50.0
graph diameters
Fig. 6: Histogram of diameters of the graphs in the TF Runtime and TF Memory datasets.
A APPENDIX
A.1 DATASETS
A.1.1 DATASET STATISTICS
Figure and figure give statisties for the number of nodes and edges in the datasets. The broad range of graph sizes indicates the diversity of the datasets.
A.1.2 TENSORFLOOW DATASET
We collected a dataset by mining TensorFlow jobs running in shared production cluster and extract- ing computation graphs in the MetaGraphDDef format. As many computation graphs were repeated due to device/machine/joh replicas, we de-duplicate the dataset by graph topology (specifically node in-degree sequence). We have not applied any other kind of filtering to restrict the dataset in any way (e.g., by architecture, input modality, learning task, dataset, etc.). Since the graphs were collected from large and diverse set of production and research use cases across input modalities, learning types, and datasets, we strongly believe our dataset is representative of broad, real-world distribution of TeosorFlow graphs.
14

---

Published as a conference paper at ICLR 2020
Computational costs for these computation graphs are simulated with an in-house simulator (based on Grappler that outputs memory and running time profiled information in the CostGraphDet format. The simulator TensorFlow Op coverage didn't it include custom kernels complicated control flow Yu et al. 2018) like cycles (e.g. t.. wh le_loop).
The vrain-validation-ee
split is made by selecting set of graphs from a list of graphs sorted by number of nodes, and splitting them as 3-1-1 across the three sets, respectively. This ensures that the distribution of the number of nodes is similar for the three sets.
For each graph in the train/validationntestt sets, we make 99 copies of it and multiply each tensor size and each Op running time cost with uniform sampled number in the interval (0.5, 1.5) (one sample per tensor size copy plus one sample per TF Op per copy). The modified copies are added back to the respective set so that the graph topologies in train/validation/tess do not overlap.
Graphs with no relevant cost information or no room for improvement (e.g. graph with a single node, a chain in the minimizing running time task) are filtered. This results in two different datasets for the two objective functions, one for runtime and another for peak memory.
The encoding is described in A.2
1.1.3 XLA DATASET
We also collected dataset by extracting CostGraphDefss during the XLA compilation of several benchmark graphs and extracted the greatest-size control-flow-free subgraph. After deduplication 94 graphs remained.
The encoding is described in A.2
A.1.4 SYNTHETIC DATASET
We sample synthetic graphs from a set of classic random graph models, including the Erdos-Renyi model Erdos & Rényil (1960), the Barabasi-Albes Strogatz model Watts & Strogatz (998) and the stochastic block model Holland et al. (1983). The parameters we used for each of the random graph models are listed in Table 2
model |Barabási & Albert (1999), the Watts-
Model Type
Parameters
Redos-Renyi
Edge probability 0.05.
Barabasi-Albert
Each node connected to m previous nodes.
Watts-Strogatz
Each node connected to k 4 neighbors initially, with probability
p = 0.3 to swap an edge with a random edge.
Stochastic Block Model
Number of blocks k = 4, within block edge probability p 0.3, cross block edge probability q = 0.01.
Table Parameters for the random graph models used to generate synthetic graphs.
The number of nodes in a graph is sampled uniformly from range [50, 200). Note that all the above random graph models generate undirected graphs. To convert the graphs into directed acyclic graphs, for each graph we sample a random ordering of nodes æ, and then set the direction of all edges (f. j) such that क (i) < T()), i.e. node i comes before node in the ordering directions, we locate all the head nodes that don't 't have any predecessors and all the tail nodes that don't have any successors, and then create a source node and a sink node, and create edges from the source node to all the head nodes, and from all the tail nodes to the sink node. Real TensorFlow graphs all contain one source and one sink node. Examples of synthetic graphs generated from the random graph models are shown in Figure
After setting all the edge
grappler
http:////ttbooooeerrrrrrrrrr/oor Tramexorkk/osss_rrapppppo
15

---

Published as a conference paper at ICLR 2020
Frdos-Reny'
Barabasi-Alber
Watts-Strogaal
Stochastic Block Model
Fig. 7: Example synthetic graphs.
16

---

Published as a conference paper at ICLR 2020
Each edge (i.j) in the graph represents either a control dependency (op j can only when op i is finished) or a data dependency (op i can run only when some output tensor(s) produced by op j is available). We assign probability of 0.1, 0.8 and 0.1 for each op to produce O, 1 or 2 output tensors. When op é produces () output tensors, any other op j that depends on can only be through a control dependency. Otherwise, op will have a probability of 0.2 to make the dependency a control dependency, and probability of 0,8 to make this data dependency, in which case an output tensor is picked according an uniform distribution.
We fill in the memory cost for each tensor by sampling from a Gaussian distribution with mean 50 and standard deviation 10. The time cost for each op is computed as the sum of all the input and output tensor memory costs plus a random noise that is fraction r of the total memory cost, with sampled from Gaussian with 0 mean and standard deviation 0.1. The source and sink nodes do not have any memory or time costs.
To make sure the generated synthetic graphs are interesting, we apply an additional filtering step by running BRKGA IK and BRKGA 10K, and keeping only the graphs whose runtime improved by at least 18%. This gives us dataset of graphs that on average improve runtime by 20% from running BRKGA IK to 10K.
The synthetic data in CostGraphDer format is available at d httpa///wws.eeeeeee////tteterr/e
The encoding is described in A.2
A.2 DATA ENCODING
We consider "control dependencies" as being tensors of size zero
Each triple (op producer, tensor, op consumer) in the CostGraphDef is encoded as a separate edge. Each of these edges e is associated with three features denoted by ze: the size of the tensor that is associated with the edge, one-hot feature indicating if the edge is control edge, and normalized index of the hyperedge to which the edge belongs.
This means that the graph neural network's 's input is directed graph with multiple edges. (There alternative encodings like a bipartite graph where both TF ops and TF tensors are nodes, and edges exist only between ops and tensors when the op consumes produces the tensor.)
Each node e in the computation graph is associated with node features T.. There are 11 node features in total, which can be grouped into three categories: memory-basede runtime-based (not used in the peak memory task) and BRKGA- based (task dependent).
As memory-based based node features, we use the sum of input tensor sizes, the sum of output tensor sizes, the extra internal memory of the TensorFlow op, and a one-hot indicator of whether the op is the one which uses the greatest memory in the graph.
As muntime-based node features, we use the sum of direct predecessor nodes' running times, the sum of direct successor nodes' running times, the running time cost of the op itself, and one-hot indicator of whether the is the one with the greatest runtime cost in the graph.
As BRKGA- node features, we have a node aggregation (the expectation of the placement per device and the schedule order for each node) of the chromosomes found by BRKGA (minimizing peak memory for the peak memory dataset and minimizing runtime for the runtime dataset) running for 400 evaluations with uniform random distributions. To make comparisons fair, REGAL with K fitness evaluations means 400 evaluations to compute features, and K 400 fitness evaluations for BRKGA using the instance-specific distributions.
For each graph, all node or edge features relating to memory size are normalized by the greatest memory size number in that graph and all features relating to runtime greatest op runtime cost.
normalized by the
To break symmetry and reduce single degree of freedom without loss of performance, fix the placement of the node with highest memory for the memory task and runtime for the runtime task to the first device.
17

---

Published as a conference paper at ICLR 2020
-0/9
Resradd carea coldag
-0.00
-000
2195
-0.94
-008
0.6
1.00
-100
3000
10300
50000
50030
100000
-1.0
300000
10000
600000
500000 1000000
Fig. 8: Average reward achieved on a minibatch of training graphs during training for runtime minimization (left) and peak memory minimization (right)
A.3 TRAINING SET RESULTS
Figure X shows the reward curves on the training set for runtime minimization (left) and peak memory minimization (right). Each point in the curves is the average reward achieved on a minibatch of training graphs at the corresponding training step. The graph neural network policy trained using TensorFlow 10 cores of a CPU machine (no GPUs were used), with multi multi-threading used by BRKGA for evaluating chromosomes and by TensorFlow. training run takes approximately 2-3 days to be completed. The final average percent improvement over BRKGASK on the training set similar to that of the test set for both tasks: 7.25% on the training set vs. 7.09% on the test set for runtime minimization, and 4.36% on the training set 3.56% on the test set for peak memory minimization. The small gap between train and test results shows that the policy is able to generalize successfully to unseen graphs test time.
A.4 PERFORMANCE MODEL
The scheduling problem is specified by a set of devices D and a computation graph. The computation graph has the list of ops j E N and tensors E T, the tensors produced by each op I(j) S T. the tensors consumed by each op C(j) C T, the memory used by each tensor m,, and the execution time of each op rj. A tensor is produced by exactly one op but can be consumed by many.
Solutions to the scheduling problem are constructed as follows. A placement is an assignment pl + D from ops to devices. Given placement we define Not as the set of ops N extended with synchronous inter-device transfer operations. A transfer operation consumes tensor on the device where it was created and produces on a device where it is needed.
Given a placement pi, a schedule is total ordering s Ñ (1, [N on ops Npt. We say that op j runs at simulation time step s(j). We model the schedule execution as follows. At each simulation time step t f1.2..... [N]}. each device d has list l,,0 of tensors currently in memory. tensor is added to the list when produced by an op that runs on the device or by a transfer op that receives the tensor the device. A tensor is removed immediately after all of its consumers the device have run. A schedule is valid if for each op j. all the input tensors are available on the corresponding device at simulation time step s(j). See Section f.r example schedule.
The memory used on device at simulation time step / is the sum of the memory used by each tensor that is in memory, i.e., E my. The peak memory of a schedule is the maximum value of the memory used at any time and on any device.
The runtime of schedule is computed by stepping through the simulation time steps in order and accounting for the execution time of each op on each device. Synchronous transfers block until both the sender and the receiver have completed their preceding tasks, and their execution time depends on a known bandwidth between devices.
A.5 BRKGA CHROMOSOME ENCODING
Let the input graph G contain o ops and t tensors which must be placed over devices. Then, the BRKGA chromosome for this graph is vector
1 1|oxd+otxxd composed of the following parts
18

---

the node placement and scheduling decisions. The Beta distribution Published as a conference paper at ICLR 2020
Pricttire
Thester
2en
Dev-1
D
Der
Dev2
Der!
Der2
Dei1
Den2
Der1
Fig. 9: A computation graph and example of chromosome encoding.
1. The first o x entries in e. represent the node-device affinities, one value for each (node, device). Each node is assigned to the device for which it has the highest value in the chromosome.
2. The next o entries represent the node scheduling priorities. A valid schedule of the computa- tion graph is obtained by performing topological sort over the nodes of the graph, breaking ties using the node scheduling priorities. Nodes with higher priority are scheduled first.
3. The final x d entries represent the tensor transfer priorities, one entry for each (tensor, device) pair. These priorities determine the order in which tensors are transferred across devices.
An example of chromosome encoding is shown in Figure 9 for a graph with o = 3 nodes, / = 3 tensors and placed on device 2. The scheduling order over the nodes is 1, 2, 3. Since nodes and 2 are placed on different devices, tensors and C must be transferred from device 1, where they are produced, to their consumer, node 2, which is on device 2. As per the tensor transfer priorities, tensor C is transferred before tensor A since tensor C has a higher priority to get transferred to device 2.
devices. As per the example, nodes and are placed on device while node is
A.6 ACTION DEQUANTIZATION
Each real number in a BRKGA chromosome is sampled from its own Beta distribution, which is parameterized by two real numbers a and 8. To be more precise, if we denote the chromosome by E R2,, then 6; Daa,, 3.) A < iis L where D(a,, 3,) is Beta distribution with parameters a; and B,. To be able to rum BRKGA, REGAL must propose the values for and B3 for each i.
As described in A.5 the BRKGA chromosome consists three parts. In REGAL, we optimize the RL agents over choices of a; and 3; for the first two parts of the chromosome, i.c., the parts of the chromosome corresponding parameters for the tensor transfer priorities are fixed to B, = which correspond to the uniform random distribution. Thus, for a graph with o ops and with devices, the RL agent must propose (d + 1) x values for each of the ops in the graph.
To make the learning task easier, rather than directly predicting values of a and B, we quantize the output space of the RL agent's 's actions such that each action uniquely maps to Beta distribution. This mapping is done as follows:
For each node in the graph 1 s and for each entry S S (d 1) correponding to node i in the BRKGA chromosome, the agent performs the set of actions
a: (mi,1, 1 ma,2 02,,,..,, Mi.j, Uijj...,, 1i,++1 0-
miy and Vij (0, 1, . ko 1} where is some fixed constant greater than 1, and these represent the quantized mean Pyy and variance oī of the Beta distribution which are related to each other as follows:
Py = Maj +
oi = Mij x (1 pig)
Uyy kij
kij +
+
Piy and of can be mapped to ai and Big for Beta distribution as follows:
(1
Mij
Pij
Biz Pij
1 pijz Oaj
Bij
Psj
19

---

Published as a conference paper at ICLR 2020
The values ma and bij sampled from Categorical distribution whose logits determined by MLPa(hap (k+ * (j - 1) k j].
We use a similar quantization strategy for the BRKGA erossover probabilities. For every crossover probability, we sample an integer E 10. k l} from a Categorical distribution for some fixed integer constant k and the dequantized crossover probability is given by 0.5 (1 + +
A.7 EXTRA MODEL DETAILS
MLPs Multi-layer perceptrons, multi-layer fully connected neural networks are models that map input vectors to oulput vectors through layers of linear transformations and nonlinear activation functions, like the following:
h = (=) = WVo1-1(..02(W..1) (Wix = + bi) + b2)) bi,
where æ is an input vector, (W; b:) are the parameters for the ith layer, and h is the output vector. a a nonlinear scalar function applied element-wise to the input vectors. Typical choices include the logistic sigmoid function (xx) = TTe tanh function (xr <i+e== or (z) = x).
--- and the ReLU function
RNNs and LSTMs Recurrent neural networks (RNNs) are good sequence models. Typical RNNs contains recurrent memory e that is updated recursively by taking some input each step the following:
C = RNNCell(e.-1, 22),
(4)
where is the input at step f. The simplest RNN cell has the following form
C σ (W(e (, b),
(5)
where W, are the parameters and a is a nonlinearity.
Long-short term memory (LSTM) models are type of RNNs that uses explicit gating control the access to the memory. LSTMs distinguish the memory e₁ and the output of the LSTM h, as two sets of vectors, and compute the update step /
sigmoid W; [1,-1, b:)
(6)
f = sigmoid W [h;-1, [3] + by)
(7)
sigmoid W (,, 2t] b.)
(8)
tanh( W, 2; [ + by)
(9)
C f @ C--1
(10)
h; = o @ tcn)
(11)
Here i, f. o are the input, forget and output gates and C is element-wise multiplication. The carefully designed memory access control through gating makes LSTMs better modeling long-term dependencies.
A.8 AUTOREGRESSIVE PREDICTION MODELS
We can use an autoregressive model to capture structure in the outputs. Given the node representations hy for each of the nodes from the GNN, we can utilize an ordering of the nodes, e.g. from topological sort, and treat the node representations Schmidhuber 1997) to predict the outputs yo sequentially.
sequence, and then use an LSTM (Hochreiter
We tried this approach but found that using an LSTM on top of the h.'s to predict yo's did not perform as well as the conditionally independent model. The reasons for this might be: (1) the autoregressive approach relies sequential ordering of the nodes, and this ordering might not be reliable nor consistent across graphs; (2) the number of nodes in the computation graphs can be large, and learning recurrent models long sequences is known to be challenging (Pascanu al. 2013); (3) the noisy training signal in our REINFORCE, based training setup makes this model even more difficult train.
20

---

Published as a conference paper at ICLR 2020
A.9 PLACEMENT AND SCHEDULING EXAMPLE
Figure TO illustrates a computation graph, valid schedule, and how we account for which tensors in memory at given time under the model presented in sec. 3.1
Op
In
In
mem.
mem.
(dev1)
(dev2)
Run
A, B
0
Transfer B
A, B
B
Run
A. C
B
Run
C
B,
Run 4
C, E
D
Transfer
E, D
Run 5
D, E
0
Fig. 10: An example computation graph and execution schedule across two devices. Op 3 is assigned to device 2 while all others are assigned to device 1.
A.10 BASELINES
Here provide supplementary information about baselines.
The multi-device peak memory minimization problem formulated for the CP solver using a model of the operation execution order, tensor lifetimes, and cumulative constraints evaluate peak memory usage. The solver is guaranteed to find the globally optimal solution given sufficient time.
CP SAT:
Graph partition + DFS: The graph partitioning objective is to minimize data transferred across devices. We use the modified implementation of the Kernighan-Lin algorithm (Kernighan & Lin 1970) used by XLA for device placement in some settings. This im- plementation generally slower than heuristics implemented popular libraries like METIS (Karypis & Kumnr (998) although it tends to find better quality solutions.
The initial schedule is topological sort order of the ops, and the initial placement selects devices uniformly randomly. A local move either changes the device assignment of the op, changes the op order in the current schedule. The hyperparameters (e.g.. number of random restarts) are set to values that perform the best on sample of 10,000 graphs in the training
Local Search:
found by grid search.
Tuned BRKGA:
The following hyperparameters of BRKGA using grid search: the beta distribution parameters (two scalars), and the number of chromosomes, elites, mutants, and populations. The grid search tries 648 hyperparameter settings and picks the best as evaluated on 10,000 training set graphs.
The performance of REGAL is stochastic both because the actions are stochastic and because BRKGA itself depends on random samples for mutation and crossover. We estimated the standard deviation of the percent improvement statistics with respect to these sources of randomness as below 0.1%, which small compared the differences we observe. Hence we have omitted the enror bars from figures
REGAL:
A.11 RUNNING TIME COMPARISON FOR ALGORITHMS
Table 3 shows the average running times of the various algorithms on the TensorFlow test set and the XLA dataset, as measured on an Intel Xeon E5- E6-1650 3. 3.60GHz machine. The times are averaged over the unaugmented graphs in the test set. REGAL provides both fast running time well as high solution quality. For a slightly higher running time than BRKGA 5K, REGAL improves the solution quality significantly. Almost all of the added running time is due extra cost of sampling beta distributions by REGAL compared to uniform distributions by BRKGA. This can be seen from the nearly identical running times of REGAL and Tuned BRKGA, which also uses beta distributions,
21

---

Published as a conference paper at ICLR 2020
Table 3: Average running times for all methods (time of running the algorithms, to be confused with solution quality).
Algorithm SAT
TF Peak Memory test ~2 hours
+ DFS
144 sec
Local Search
122 sec
BRKGA 5K
0.89 sec
Tuned BRKGA
1.04 sec
GAS REGAL
1.04 sec
1.04 sec
Table 4: Performance of REGAL on peak memory with various subsets of actions.
Placement
Scheduling
Valid
Test
Yes
No
-0.4%
-0.2%
No
Yes
4.4%
3.65% %
Yes
Yes
4.67%
3.56%
but without the neural network policy. The local search heuristic slowly because it implemented efficiently, e.g., with incremental evaluations of the objective; we show its timing results for completeness only.
A.12 ABLATION ANALYSIS OF AGENT ACTION TYPES
REGAL can train a policy to generate any subset of the following actions for BRKGA: 1) Actions for node placement priorities, and 2) Actions for node scheduling priorities. We train REGAL with various subsets of these actions and compare their performance against each other in table
We observe that on the validation set, REGAL performs best when it has to learn actions for both placement and scheduling compared just scheduling or placement alone.
A.13 ANALYSIS OF THE STRUCTURE OF PLACEMENT AND SCHEDULING DECISIONS,
Our best model trained on the TF peak memory dataset capable of generating 16 different kinds of actions for node placement decisions and 4 different kind of actions for node scheduling decisions. Each these actions determine the shape of the Beta distributions from which we sample the node-device affinities and the node scheduling priorities. In this section attempt to gain insights into the structure these actions.
We divide our node placement decisions into three categories:
Actions that give node higher probability to be placed on device
Actions that give node higher probability to be placed on device
Actions that give equal preference to the two devices.
Similarly, we divide our node scheduling decisions in two categories:
Actions that give nodes a "high" scheduling priority.
Actions that give node "low" scheduling priority.
Finally, we aggregate the average relative memory consumption of all nodes that were assigned the same set of actions, where memory consumption of node is defined the sum of the memory uses of all its input and output tensors. The relative memory usage is the memory usage normalized by the largest memory usage of node in the graph.
22

---

Published as a conference paper at ICLR 2020
54
0.135
Device
48
Device 2
0.120
42
Device affinity
Device D affinity
0.105
36
0.090
30
Device
Device
24
0.075
18
0.060
12
Same
Same
-0.455
6
-0.030
Low
High
Low
High
Node scheduling priority
Node scheduling priority
Fig. 11: The agent picks diverse set of actions. These plots show the frequency of actions chosen (in percent, left) and the average relative weights of the nodes that the actions are applied to (right)
We plot this data in Figure On the right, each cell represents the average relative memory consumption of the nodes that were assigned a particular placement and scheduling decision (darker cells indicate nodes with higher average relative memory consumption). On the left, each cell represents the frequency of the nodes that were assigned a particular placement and scheduling decision (darker cells indicate higher frequency). With these we can make the following observations:
On average, nodes with higher normalized memory consumption are assigned lower schedul- ing priorities.
Most of the nodes with the highest relative memory consumption have no affinity for either of the two devices.
For nodes with the lowest relative memory consumption, most of them have an affinity be placed on device 2, while a smaller but still significant number of them prefer device 1.
This implies that the node placement strategy is more complicated than trivial "place lighter nodes on device 2" strategy and REGAL's actions are non-trivially dependent on the input.
A.14 HYPERPARAMETERS OF THE BEST AGENT FOR PEAK MEMORY TF
The graph neural network had state size of 32 for each node and edge, 16 propagations, all networks MLP,, MLPP MLP node MLP usng MIP' msg being two layers of size 32, the aggregation used was mean pooling. For faster training, the reward of the training set was made with 1000 fitness evaluations for REGAL and BRKGA (4600 for REGAL and 5000 for BRKGA for the validation and test sets). Training lasted 100000 gradient steps with each step having mini-batch of size 4 and with gradient clipping by Le norm with value 10. The baseline mean squared error term't 's contribution to the overall loss was weighted by 0.0001 The optimizer was Adam with betal 0.9 beta2 0.999 epsilon 1e and learning rate 0.0001 The number of devices (for the memory model) was 2.
A.15 HYPERPARAMETERS THE BEST AGENT FOR RUNTIME TF
The graph neural network had a state size of 32 for each node and edge, 16 residual graph propagations,
all networks MLP,, MLP, MLP node MLP mag MLP' msg being two layers of size 32, the aggregation used was sum. Training lasted 100000 gradient steps with each step having mini-batch of size 4 and with gradient clipping by L2 norm with value 10. The baseline mean squared error term's 's contribution the overall loss weighted by 0.0001 The optimizer Adam with betal 0.9 beta2 0.999 epsilon 1e 8 and learning rate 0.0001 With 16 for scheduling and = 2 for placement (k being the quantization level defined in A.6
23

---

Published as a conference paper at ICLR 2020
Reward (higher is better)
-0.8
-1.0
-1.2
Test graphs, by unique topology
Fig. 12: A box plot rewards on the TF Runtime test set by unique graph topology. There are 100 graphs for each topology, 99 of which generated by data augmentation. A reward greater than -1 implics that REGAL finds better solution than BRKGA. Box plots visualize the 25th, 50th, and 75th percentiles and the range of a set of points.
A.16 HYPERPARAMETERS OF THE BEST AGENT FOR RUNTIME SYNTHETIC
Same A.15 but with 2 graph propagations and GRU node updates and aggregation using mean.
A.17 PERFORMANCE BY GRAPH TOPOLOGY
Figure 12 shows the distribution of rewards on the TF Runtime test set broken down by graph topology. As described Section for each unique topology we augment the dataset by perturbing the tensor sizes and op running times. This generates a distribution of 100 rewards per topology. The variance for a fixed topology is typically relatively small.
24


Reviews:
[{"review_id": "ryxXud915S", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this work the authors propose a deep RL approach to minimize the makespan and the peak memory usage of a computation graph as produced by MXNet/PyTorch/TensorFlow.  This is an increasingly important problem as distributed deep learning is necessary in many cases. The authors aim to minimize the execution time and/or the peak memory usage. For this purpose they generate a training dataset out of a real-world dataset of various TensorFlow computation graphs using simulation software. The proposed RL approach consists of two steps. First  a GNN is used to derive representations for computation graphs. Then the authors use a heuristic BRKGA to learn a policy for the placement of computation graphs, that actually works on unseen graphs. Overall this paper is well-written, deals with an important practical problem. While it is not immediately clear to the reader the effect of BRKGA on the mapping of the graph to the resource network and why it works so well, the results are convincing (but still there is space for improvement).  That is why I rate it as a \"weak accept\". \n\n- Can you explain why the beta distribution choices at each node may have a negative impact on the makespan in certain cases? Have you looked into them? \n- To what extent are the simulations realistic? Can you please comment more on this aspect? \n- Have you tried Scotch? https://www.labri.fr/perso/pelegrin/scotch/. Since the software aims to achieve a different objective, it serves as a baseline. \n- Can you obtain insights with respect how you could cluster the TensorFlow computation graphs? \n- Can you improve the discussion on BRKGA? Since it is a vital component of the proposed framework, it would be informative to read few more self-contained details on how it works in section 2.\n- Once you obtain a mapping, can you comment on any insights concerning the structure of the partition and schedule? ", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SJlyOpyNcH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n\nThis paper proposes an ML-based method to optimize TensorFlow Graph execution. Specifically, it combines graph neural networks (GNNs) and BRKGA (a genetic algorithm) to search over the joint space of TF node-device placement and scheduling. The core claims on the advantages of this method are that (1) it co-searches placement and scheduling space, (2) the trained model can generalize to different graphs and inference cost is very  small. The experimental results show that REGEL can outperform a few baseline methods on this problem.\n\nWriting\n- The paper is well-written and I enjoyed reading the paper.\n- Some more descriptions about the BRKGA algorithm could be added in.\n\n\nMethod and Results\n\nSome confusion if the authors could answer:\n- I am very confused by one of the claims that “ the first work on learning a policy for jointly optimizing placement and scheduling”. I don’t see much evidence in the result section about showing the co-searching the joint space yield advantages? I am fairly familiar with the line of work on only optimizing device placement, but it would be good to see some ablation studies showing search over the joint space is advantageous. \n\n- The model is trained with standard REINFORCE -- how many training time and resources are needed to train a REGEL model for a task? How’re the training dynamics looking like (variance, convergence, etc?)? \n\n- In terms of the generalization ability of REGEL, the paper has clearly shown that REGEL is able to generalize to differently shaped graphs, with acceptable cost, but I am wondering for the same dataflow graph, how REGEL generalizes to different input data configurations (size, modality, etc.)? E.g. if the batch size of the input data is changed, the execution time of each kernel and their memory usage (in general, the system treatment) would change; Can a trained REGEL model on a data config A generalize to B? How would this affect the performance of REGEL?\n\n- It seems the method and assumptions about graphs or training data are pretty coupled with TensorFlow and graph-mode execution, how could the method be generalized to other ML frameworks (e.g. frameworks with eager execution)\n\n- Could the authors clarify why the two methods mentioned in “Learning to directly predict a solution” has quadratic complexity w.r.t. # of nodes and whereas REGEL is linear? \n\n- Confusion on Figure 4(b): Could some more critical statistics about the graphs in the training/test dataset be reported? e.g. what’s the average depth of the training graphs? When there are 32 MP layers a node’s feature will be passed across its 32-hop neighborhood, which seems surprising as it is common to observe GNN starts degenerating with increased depth (because all node features become similar during message passing)\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "SJlqiP0S5H", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors proposed a framework to generating a task scheduling for a compiler to reduce the execution cost of neural networks. A computation graph is first fed into a GNN to produce a beta distribution, which is then fed into the BRKGA algorithm to yield the encoded solutions. The motivation is interesting, and the proposed method is technically reasonable. The details are also included in the appendix. To improve the quality, the following concerns may be considered:\n\n1. The detailed explanations of o_a(G) and o_s(G) should be included.\n\n2. How were the attribute vectors  x_v and x_e defined in your experiments?\n\n3. The baseline (GP+DFS) may not be strong enough, since it is designed to reduce the communication cost. With the information of the input size and time complexity of ops, a better greedy algorithm can be designed. Moreover, the performance of Local Search and BRKGA 5K are similar, and REGAL is just slightly better than BRKGA 5K. Hence, the improvement over the best efficient greedy algorithm seems small.\n\nOverall, the studied topic is interesting, and this paper is also intriguing.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission presents an approach that leverages machine learning to optimize the placement and scheduling of computation graphs (such as TensorFlow graphs) by a compiler. The work is interesting and well-executed. All reviewers recommend accepting the paper.
Decision: Accept (Poster)