Title: LEARNING TO PROGRESSIVELY PLAN

Abstract: For problem solving, making reactive decisions based on problem description is fast but inaccurate, while search-based planning using heuristics gives better solutions but could be exponentially slow. In this paper, we propose a new approach that improves an existing solution by iteratively picking and rewriting its local components until convergence. The rewriting policy employs a neural network trained with reinforcement learning. We evaluate our approach in two domains: job scheduling and expression simplification. Compared to common effective heuristics, baseline deep models and search algorithms, our approach efficiently gives solutions with higher quality.

1. INTRODUCTION:

In recent years, deep reinforcement learning has achieved strong success in scenarios that are welldefined and can be precisely and efficiently simulated (e.g., games like Atari (Mnih et al., 2013) ). One typical scenario is model-free control: given the current state (e.g., a ball is falling), a decision is made reactively to yield maximal rewards (e.g., move the pedal to bounce the ball back).

In more challenging scenarios, a one-off reactive decision is usually insufficient. It is necessary to plan before a strong decision. As illustrated by many works, especially competitive games like Chess and Go (Silver et al., 2017) , most of the times, search-based policy is stronger than reactive ones. However, the time consumption can also be exponential, if planning is done from scratch. To address this issue, previous works use neural networks to predict (and execute) an entire plan from scratch, given a complete specification of the problem description (Vinyals et al., 2015; Mao et al., 2016; Graves et al., 2014) . While this avoids search, a direct prediction could be very difficult and even impossible in complicated scenarios.

An alternative approach is to start from an existing plan and iteratively improve. When both state and action spaces are continuous, many trajectory optimization techniques have been proposed in the context of robotics and control (MAYNE, 1973; Tassa et al., 2012; Levine & Abbeel, 2014) , which use local gradient information to gradually improve the existing plan. On the other hand, with discrete state and action spaces, such approaches become inapplicable due to the indifferentiablity.

In this paper, we propose a novel approach of progressive planning in the discrete space, by iteratively picking a local part of an existing plan and rewriting it. The rewriting decision is made by a neural network, trained end-to-end with reinforcement learning. During the training, the neural network learns to find common local patterns in the current solution and replace them with a better choice, until there is no more improvement. We apply our approach to two different domains: job scheduling and expression simplification. We show that our rewriting approach is better than strong heuristics of both domains using multiple metrics. For job scheduling, under a controlled setting, we also demonstrate that our model outperforms DeepRM (Mao et al., 2016) , a neural network predicting a holistic scheduling plan, by large margins especially when more heterogeneous resources lead to more complicated scheduling.

2. RELATED WORK:

Trajectory optimization problems in the continuous space have been widely studied in the domain of robotics and control, and many effective techniques have been proposed (MAYNE, 1973; Bradtke et al., 1994; Vrabie et al., 2009; Tassa et al., 2012; Levine & Koltun, 2013; Levine & Abbeel, 2014) . Typically, local gradient information is leveraged to gradually optimize the existing trajectory. However, these techniques are not directly applicable to domains with discrete state and action spaces. In this work, we study approaches of iterative refinement for discrete optimization problems.

Job scheduling and resource management problems are ubiquitous and fundamental in computer systems. Various work have studied these problems from both theoretical and empirical sides (Błażewicz et al., 1996; Grandl et al., 2015; Armbrust et al., 2010; Scully et al., 2017; Terekhov et al., 2014; Mao et al., 2016; Chen et al., 2017) . In particular, recent work use deep reinforcement learning for job scheduling with a simplified setup (Mao et al., 2016; Chen et al., 2017) . Existing work focus on proposing algorithms to construct a schedule from scratch. However, with more complex configurations, it becomes challenging to make an effective arrangement in this way, as indicated in our evaluation. Thus, we make an initial step towards tackling the job scheduling problem through rewriting, and we consider extending our approach to real-world settings as future work.

A recent line of work studies using deep neural networks to discover equivalent expressions (Cai et al., 2018; Allamanis et al., 2017; Zaremba et al., 2014) . In particular, (Cai et al., 2018) trains a deep neural network to rewrite algebraic expressions. However, they use supervised learning to train the model, which requires a collection of ground truth rewriting paths, and lacks the capability to find novel rewriting routines. On the contrary, we use reinforcement learning to train our rewriter for expression simplification, which mitigates these limitations.

Our rewriting formulation is closely related to other applications such as code optimization (Schkufza et al., 2013; Chen et al., 2018) , theorem proving (Huang et al., 2018; Lederman et al., 2018) , and text simplification (Cohn & Lapata, 2009; Paetzold & Specia, 2013; Feblowitz & Kauchak, 2013) . In addition, our rewriting approach could also be extended to classical combinatorial optimization problems (Khalil et al., 2017; Bello et al., 2016; Vinyals et al., 2015; Karp, 1972 ), e.g., Traveling Salesman Problem (Applegate et al., 2006) and Vertex Cover Problem (Bar-Yehuda & Even, 1981) .

3. PROBLEM SETUP:

In this work, we propose to formulate optimization as a rewriting problem, and find the optimal solution through an iterative rewriting process. Before we motivate our formulation, we start with the description of optimization problems of our interest. Typically, an optimization problem can be defined as follows:

Definition 1 (Optimization problem.) Let S be the space of all valid states in the problem domain, and c : S → R is the cost function. The goal is to find arg min s∈S c(s).

When S is continuous, gradient-based algorithms are effective ways of finding the optimal solutions. For example, trajectory optimization techniques such as iLQR leverage the gradient information to gradually improve the current trajectory towards the optimum (MAYNE, 1973) . In contrast, for discrete S, it is not straightforward to apply similar approaches. A long line of works propose algorithms to construct from scratch the optimal solution for combinatorial optimization problems and job scheduling problems (Karp, 1972; Grandl et al., 2015) . However, when dealing with more complicated problems, it is challenging to output an effective solution from the ground up.

In comparison, for many optimization problems, a better alternative approach is to first construct a feasible solution, then make incremental improvement. This is because (1) a feasible solution is often easy to find; (2) an existing solution provides a full context for the improvement, which is not the case if a solution is generated from scratch, (3) many problems, as well as their solutions, have strong local structures that can be utilized when improving incrementally, and (4) different solutions might share a common routine towards the optimal.

For example, in job scheduling, it is usually difficult to decide whether a job with a large resource requirement needs to be postponed; however, from an existing schedule, it is easy to identify what jobs cause a long waiting time for later ones. In this case, re-scheduling them improves the efficiency. Therefore, we propose to solve optimization problems through rewriting, as defined formally below.

Definition 2 (Optimization as a rewriting problem.) Let S be the space of all valid states, A be the rewriting ruleset, c : S → R is the cost function. Suppose s t is the state at rewriting iteration t, we can apply a rewriting rule a t ∈ A to ĝt , where ĝt ⊂ s t , and the rewriting step results in the next state s t+1 = f (s t , ĝt , a t ). Given an initial state s 0 , our goal is to find a sequence of rewriting steps (s 0 , (g 0 , a 0 )), (s 1 , (g 1 , a 1 )), ..., (s T -1 , (g T -1 , a T -1 )), s T that minimizes c(s T ).

To tackle a rewriting problem, rule-based rewriters with manually-designed rewriting routines have been proposed (Halide, 2018) . However, manually designing such routines is not a trivial task. An incomplete set of routines often leads to an inefficient exhaustive search, while a set of kaleidoscopic routines is often cumbersome to design, hard to maintain and lacks flexibility.

In this paper, we propose to train a neural network instead, using reinforcement learning. Recent advance in deep reinforcement learning suggests the potential of well-trained models to discover novel effective policies, such as demonstrated in Computer Go (Silver et al., 2017) and video games (OpenAI, 2018) . In our evaluation, we demonstrate that our approach not only mitigates laborious human efforts, but also enables the model to discover novel rewriting paths from its own exploration.

In the following sections, we discuss the application of our rewriting approach to two different domains: job scheduling (as mentioned above) and expression simplification, in which we minimize the expression length using a well-defined semantics-preserving rewriting ruleset.

3.1. JOB SCHEDULING PROBLEM:

We first study the job scheduling problem, and in particular, we consider a simplified problem setup studied in (Mao et al., 2016) as follows.

Suppose we have a machine with D types of resources. Each job j is specified as g j = (r j1 , r j2 , ..., r jD , S j , T j ), where r j1 , r j2 , ..., r jD denotes the required portion (between 0 and 1) of resources for each type, S j is the arrival timestep, and T j is the duration. We assume that the resource requirement is fixed during the entire job execution, each job must run continuously until finishing, and no preemption is allowed. We adopt an online setting: there is a pending job queue Q which can hold at most L Q jobs. When a new job arrives, it can either be allocated immediately, or be added to Q. If Q is already full, to make space for the new job, at least one job in the Q needs to be scheduled immediately. The goal is to find a time schedule for every job, so that the average waiting time is as short as possible.

For job scheduling, the only type of rewriting is to re-schedule a job g j and allocate it after another job g j finishes or at its arrival time S j . Details of a rewriting step is presented in Appendix B.1. Thus, the size of the rewriting ruleset is |A |= 2L Q , since each job could only switch its scheduling order with at most L Q of its former and latter jobs respectively.

3.2. EXPRESSION SIMPLIFICATION:

We also apply our approach to expression simplification domain. In particular, we consider expressions in Halide, a domain-specific language for high-performance image processing (Ragan-Kelley et al., 2013) , which is widely used at scale in multiple products of Google (e.g., YouTube) and Adobe Photoshop. Simplifying Halide expressions is an important step towards the optimization of the entire code. To this end, a rule-based rewriter is implemented for the expressions, which is carefully tuned with manually-designed heuristics1 . The grammar of the expressions considered in the rewriter is specified in Appendix A. Notice that the grammar includes a more comprehensive operator set than previous works on finding equivalent expressions, which consider only boolean expressions (Allamanis et al., 2017; Evans et al., 2018) or a subset of algorithmic operations (Allamanis et al., 2017) . The rewriter includes hundreds of manually-designed rewriting templates. Given an expression, the rewriter checks the templates in a pre-designed order, and applies those rewriting templates that match any sub-expression of the input.

After investigating into the rewriting templates in the rule-based rewriter, we find that a large number of rewriting templates enumerate specific cases for an uphill rule, which lengthens the expression first and shortens it later (e.g., "min/max" expansion). Similar to momentum terms in gradient descent for continuous optimization, such rules are used to escape a local optimum. However, they should only be applied when the initial expression satisfies certain pre-conditions, which is traditionally specified by manual design, a cumbersome process that is hard to generalize.

Observing these limitations, we hypothesize that a neural network model has the potential of doing a better job than the rule-based rewriter. In particular, we propose to only keep the core rewriting rules in the ruleset, remove all unnecessary pre-conditions, and let the neural network decide which and when to apply each rewriting rule. In this way, the neural rewriter has a better flexibility than the rule-based rewriter, because it can learn such rewriting decisions from data, and has the ability of discovering novel rewriting patterns that are not included in the rule-based rewriter.

s t Sample ! " # ∼ %& " # , " # ⊂ ) # Sample * # ∼ +% ! " # ) #,-= /() # , ! " # , * # )

Input Encoder Score Predictor Rule Selector:

Figure 1 : The framework of the neural rewriter architecture. Here, the score predictor computes SP (g t ), which is the rewriting score; and the rule selector predicts RS(ĝ t ), which is the probability distribution of applying each rewriting rule a ∈ A.

In the following, we present the design of our rewriting model, i.e., Neural Rewriter. We first provide an overview of our model framework, then present the design details for different applications.

4.1. MODEL OVERVIEW:

Figure 1 illustrates the overall framework of our neural rewriter, and we describe the two key components for rewriting as follows. More details can be found in Appendix C.

Score predictor. Given the state s t , the score predictor computes a score SP (g t ) for every g t ⊂ s t , which measures the benefit of rewriting g t . A high score indicates that rewriting g t could be desirable.

Rule selector. Given ĝt ⊂ s t to be rewritten, the rule selector predicts a probability distribution RS(ĝ t ) over the entire ruleset A, and selects a rule a t ∈ A to apply accordingly.

4.2. MODEL DETAILS FOR JOB SCHEDULING PROBLEM:

Figure 2a demonstrates the model architecture for job scheduling, and we discuss the details below.

Input embedding.:

As described in Section 3, each job is specified by g j = (r j1 , r j2 , ..., r jD , S j , T j ). In addition, we define A j as the schedule time, and C j = A j + T j as the completion time. We embed each job into a (D × (T max + 1) + 1)-dimensional vector e j , where T max is the maximal duration of a job. This vector encodes the information of its attributes and the machine status during its execution, and we defer the embedding details to Appendix C.1.

We represent each schedule as a directed acyclic graph (DAG), which describes the dependency among the schedule time of different jobs. Specifically, we denote each job g j as a node in the graph, and we add an additional node g 0 to represent the machine, which has a zero embedding vector e 0 . If a job g j is scheduled at its arrival time S j , then we add a directed edge g 0 , g j in the graph. Otherwise, there must exist at least one job g j such that C j = A j . We add an edge g j , g j for every such job g j to the graph. Figure 3 illustrates an example of the graph construction.

To encode the graphs, we extend the Child-Sum Tree-LSTM architecture in (Tai et al., 2015) , which is similar to the DAG-structured LSTM in (Zhu et al., 2016) . Specifically, for a job g j , suppose (h 1 , c 1 ), (h 2 , c 2 ), ..., (h p , c p ) are the LSTM states of all parents of g j , then its LSTM state is

EQUATION

For each node, the d-dimensional hidden state h is used as the embedding for other two components.

Score predictor. This component is a L P -layer fully connected neural network with a hidden size of N p , and the input to the predictor of job g j is h j .

Rule selector. The rule selector is a L S -layer fully connected neural network with a hidden size of N S . The rewriting options described in Section 3 is equivalent to moving the current job g j to be a child of another job g j or g 0 in the graph, which means allocating job g j after job g j finishes or at its arrival time S j . Thus, the input to the rule selector not only includes h j , but also h j of all other Each circle with index greater than 0 represents a job node, and node 0 is an additional one representing the machine. Edges in the graph reflect job dependencies. The score predictor selects a job ĝt ∼ SP (g t ) from all job nodes to re-schedule. The rule selector chooses a moving action a t ∈ A for ĝt , then modifies s t to get a new dependency graph s t+1 . In (b), s t is the expression parse tree, where each square represents a node in the tree. The set of g t ⊂ s t includes every sub-tree rooted at a non-terminal node, from which the score predictor selects ĝt to rewrite. Afterwards, the rule selector predicts a rewriting rule a t , then rewrites the sub-tree ĝt to get the new tree s t+1 .

g j that could be used for rewriting. The output layer is an |A|-dimensional softmax layer, where |A |= 2L Q as discussed in Section 3. More details can be found in Appendix C.1.

4.3. MODEL DETAILS FOR EXPRESSION SIMPLIFICATION:

We present the instantiation of our neural rewriter framework for expression simplification in Fig- ure 2b . We mainly discuss the design choices different from the model for job scheduling below.

Input embedding. We use expression parse trees as the input, and employ the N-ary Tree-LSTM designed in (Tai et al., 2015) as the input encoder to compute the embedding for each node in the tree. Notice that in this problem, each non-terminal has at most 3 children. Thus, let x be the embedding of a non-terminal, (h L , c L ), (h M , c M ), (h R , c R ) be the LSTM states maintained by its children nodes, the LSTM state of the non-terminal node is computed as

EQUATION

) Where [a; b] denotes the concatenation of vectors a and b. For non-terminals with less than 3 children, the corresponding LSTM states are set to be zero.

Score predictor. The score predictor is a L P -layer fully connected neural network with a hidden size of N P . For each sub-tree g i , its input to the score predictor is represented as a 2d-dimensional vector [h 0 ; h i ], where h 0 embeds the entire tree. More details can be found in Appendix C.2. Rule selector. The rule selector is a L S -layer fully connected neural network with a hidden size of N S , and its input format is the same as the score predictor.

4.4. TRAINING DETAILS:

Let (s 0 , (ĝ 0 , a 0 )), (s 1 , (ĝ 1 , a 1 )), ..., (s T -1 , (ĝ T -1 , a T -1 )), s T be the rewriting sequence in the forward pass, we define the reward function r(s t , (ĝ t , a t )) as follows. For job scheduling problem, r(s t , (ĝ t , a t )) = c(s t ) -c(s t+1 ). For expression simplification problem, r(s t , (ĝ t , a t )) = max t+1≤k≤T (γ k-t (c(s t ) -c(s k )))/c(ĝ t ), where γ is a discount factor. This reward function is designed to credit uphill rules that could lead to a simplified expression in the end. More discussion about the forward pass algorithm and reward design can be found in Appendix D.

The reward function r(s t , (ĝ t , a t )) is also used as the training target of SP (ĝ t ). Specifically, let θ be the model parameters, then the loss function of the score predictor is

EQUATION

To train the rule selector, we employ the Advantage Actor-Critic algorithm (Sutton et al., 1998) , and we use SP (ĝ t ) as the critic. In particular, let the advantage function A(s t , (ĝ t , a t )) = r(s t , (ĝ t , a t ))-SP (ĝ t ), then the loss function of the rule selector is

EQUATION

The overall loss function is L(θ) = L RS (θ) + αL SP (θ), where α is a hyper-parameter.

5. EXPERIMENTS:

In this section, we present the evaluation results of both job scheduling and expression simplification problems. All neural network models in our evaluation are implemented in PyTorch (Paszke et al., 2017) . To calculate the inference time during testing, we run all algorithms on the same server equipped with 2 Quadro GP100 GPUs and 80 CPU cores. Only 1 GPU is used when evaluating neural network models, and 4 CPU cores are used for search algorithms . For both tasks, we set the timeout of search algorithms to be 10 times as long as the timeout of our neural rewriter. 1 : Experimental results of the job scheduling problem on the test set. For each approach, we report the average slowdown of the jobs with different number of resource types D.



5.1.1. EVALUATION SETUP:

We randomly generate 100K job sequences, and use 80K for training, 10K for validation, and 10K for testing. We use an online setting where jobs arrive on the fly with a pending job queue of length L Q = 10. When the number of resource types D = 2, we follow the same setup as in (Mao et al., 2016) . The maximal job duration T max = 15, and the latest job arrival time is S max = 50. With larger D, except changing the resource requirement of each job to include more resource types, other configurations stay the same. We use average job slowdown as our evaluation metric, which is computed by (C j -S j )/T j . This metric is also used in (Mao et al., 2016) .

5.1.2. MODEL CONFIGURATION:

For our neural rewriter model, we provide an initial schedule that allocates jobs in a first-come-firstserve manner for every job sequence, then feed it to the neural rewriter for refinement. Such an initial schedule is intuitive to design, easy to compute with a negligible overhead, while is much less effective than the optimal solution, as we will demonstrate in Table 1 .

We compare our neural rewriter model with two kinds of baselines. The first kind of baselines use manually designed heuristics: Shortest Job First (SJF) always allocates the shortest job in the pending job queue at each timestep, also used as a baseline in (Mao et al., 2016) . Shortest First Search searches over the shortest k jobs to schedule at each timestep, and returns the optimal one. We find that other heuristic-based baselines used in (Mao et al., 2016) generally perform worse than SJF, especially with large D. Thus, we omit the comparison.

The second kind of baselines use a deep neural network to construct the job schedule from scratch. We re-implement DeepRM (Mao et al., 2016) , a neural network trained with reinforcement learning, and test it on larger D. For a fair comparison, we tune the hyper-parameters of DeepRM for best performance.

To measure the optimality of these algorithms, we also compute the following empirical bounds. Earliest Job First (EJF) schedules each job by their arrival time. This provides an upper bound of the average job slowdown. We also use this algorithm to generate initial schedules for the neural rewriter. Optim assumes an offline setting: the entire job sequence is available before scheduling. This leads to a better algorithm that sorts jobs' duration at the earliest time runnable on the machine. It gives a lower bound of the average job slowdown in an idealized setting.

5.1.3. RESULTS:

Table 1 presents the results for the job scheduling problem. Our neural rewriter model outperforms both heuristic algorithms and the baseline neural network DeepRM. In particular, we observe that while the performance of DeepRM and our neural rewriter are similar when D = 2, with larger D, DeepRM starts to perform worse than heuristic-based algorithms, which is consistent with our hypothesis that it becomes challenging to design a schedule from scratch when the environment becomes more complex. On the other hand, our neural rewriter could capture the bottleneck of an existing schedule that limits its efficiency, then progressively refine it to obtain a better one. Meanwhile, our results are also closer to the empirical lower bound computed by the Optim algorithm, which further demonstrates the effectiveness of our rewriting approach. More results and discussion can be found in Appendix E.

Number of expressions in the dataset Length of expressions Size of expression parse trees:

Total 



To construct the dataset, we first generate random pipelines using the generator in the Halide repository2 , then extract expressions from them. We filter out those expressions that can not be further rewritten, then split the rest into 8/1/1 for training/validation/test set respectively. We summarize the statistics of the dataset in Table 2 , and more details can be found in Appendix A.

We measure the following metrics in our evaluation: (1) Average expression length reduction, which is the length reduced from the initial expression to the rewritten one, and the length is defined as the number of characters in the expression; (2) Average tree size reduction, which is the number of nodes decreased from the initial expression parse tree to the rewritten one.

We discuss the ruleset design of our neural rewriter as follows. We look into the Halide rewriting ruleset with hundreds of templates, and for those templates that can not be further simplified, e.g., reducing v -v into 0 as in Figure 2b , we simply include them in our ruleset. As discussed in Section 3, a large number of templates are enumerating pre-conditions to apply uphill rules. For these templates, we remove the manually-designed pre-conditions, and only include the uphill rules themselves in the ruleset. In this way, we manually build a ruleset with |A |= 19 categories of rewriting rules. More details about the ruleset and rewriting process can be found in Appendix B.2.

We examine the effectiveness of our neural rewriter against two kinds of baselines. The first kind of baselines are heuristic-based rewriting approaches, including the rule-based Halide rewriter (Section 3) and a heuristic search algorithm using our ruleset only. In each rewriting iteration, the search algorithm first enumerates all combinations of rewriting node and rule, then selects the top k shortest resulted expressions to be rewritten in the next iteration.

In addition, we evaluate Z3 solver on our dataset (De Moura & Bjørner, 2008) . Z3 solver is a highperformance theorem prover developed by Microsoft Research. Its simplifier works by traversing each sub-formula in the input expression and invoking the solver to find a simpler equivalent one to replace it. Therefore, the simplification steps performed by this solver may not be included in the Halide ruleset, which makes it a strong baseline to compare with. We set the timeout to be 10 seconds for each input expression, and we find that the results are not significantly better with a longer timeout.

Table 3 presents the main results of our expression simplification problem. We can observe that our neural rewriter model outperforms both the rule-based rewriter and the heuristic search by a large margin. In particular, our neural rewriter could reduce the expression length and parse tree size by around 45% on average; meanwhile, compared to the rule-based rewriter, our model further reduces the average expression length and tree size by 30% and 40% respectively. We observe that the main performance gain comes from learning to apply uphill rules appropriately in ways that are not included in the manually-designed templates. For example, consider the expression 5 ≤ max(max(v0, 3) + 3, max(v1, v2)), which could be reduced to T rue by expanding max(max(v0, 3) + 3, max(v1, v2)) and max(v0, 3). Using a rule-based rewriter would require the need of specifying the pre-conditions recursively, which becomes prohibitive when the expressions become more complex. On the other hand, heuristic search may not be able to find the correct order of expanding the right hand size of the expression when more "min/max" are included, which would make the search less efficient. Meanwhile, the performance of our neural rewriter is also much closer to the performance of Z3 solver, which could perform rewriting steps that are not included in the Halide ruleset. More results can be found in Appendix F.

6. CONCLUSION:

In this work, we propose to formulate optimization as a rewriting problem, and solve the problem by iteratively rewriting an existing solution towards the optimum. We utilize deep reinforcement learning to train our neural rewriter. In our evaluation, we demonstrate the effectiveness of our neural rewriter on job scheduling and expression simplification problems, where our model outperforms heuristic-based algorithms and baseline deep neural networks that generate an entire solution directly.

Meanwhile, we observe that since our approach is based on local rewriting, it could become timeconsuming when large changes are needed. In extreme cases where each rewriting step needs to change the global structure, starting from scratch becomes preferrable. We consider improving the efficiency of our rewriting approach and extending it to more complicated scenarios as future work. if j = 0 then A j = C j else A j = S j fi 6:

C j = A j + T j 7: 8:

//Resolve potential resource occupation overflow within [A j , C j ] 9:

J = all jobs in s t except g j that are scheduled within [A j , C j ] 10:

Sort J in the topological order 11:

for g i ∈ J do 12:

A i = the earliest time that job g i can be scheduled 13:

C i = A i + T i 14:

end for 15:

For

g i ∈ J, A i = A i , C i = C i 16: s t+1 = {(A i , C i )} 17:

return s t+1 18: end function 

A GRAMMAR OF THE HALIDE EXPRESSIONS:



B.2 MORE DETAILS FOR EXPRESSION SIMPLIFICATION PROBLEM:

More discussions about the uphill rules. A commonly used type of uphill rules is "min/max" expansion, e.g., min(a, b) < c → a < c||b < c. Dozens of templates in the ruleset of the Halide rewriter are describing conditions when a "min/max" expression could be simplified. Notice that although applying this rewriting rule has no benefit in most cases, since it will increase the expression s t+1 = f (s t , ĝt , a t ) 17: end for Rule selector. The rule selector has two modules. The first module is a L S -layer fully connected neural network with a hidden size of N S . For each job g j , let N j be the number of jobs that could be the parent of g j , and {g j k } denotes the set of such jobs. For each g j k , the input is [h j ; h j k ], and this module computes a d-dimensional vector h k to encode such a pair of jobs. The second module of the rule selector is another L S -layer fully connected neural network with a hidden size of N S . For this module, the input is an (|A| × d)-dimensional vector [h 1 ; h 2 ; ...; h |A| ], where |A |= 2L Q . When N j < |A|, h Nj +1 , h Nj +2 , ..., h |A| are set to be zero. The output layer of this module is an |A|-dimensional softmax layer, which predicts the probability of each different move of g j .

C.2 MODEL DETAILS FOR EXPRESSION SIMPLIFICATION:

Input representation. As discussed in Section 4, for each sub-tree g i , its input to both the score predictor and the rule selector is represented as a 2d-dimensional vector [h 0 ; h i ], where h 0 is the embedding of the root node encoding the entire tree. The reason why we include h 0 in the input is that looking at the sub-tree itself is sometimes insufficient to determine whether it is beneficial to perform the rewriting. For example, consider the expression max(a, b) + 2 < a + 2, by looking at the sub-expression max(a, b) + 2 itself, it does not seem necessary to rewrite it as max(a + 2, b + 2). However, given the entire expression, we can observe that this rewriting is an important step towards the simplification, since the resulted expression max(a + 2, b + 2) < a + 2 could be reduced to f alse. We have tried other approaches of combining the parent information into the input, but we find that including the embedding of the entire tree is the most efficient way.

C.3 MODEL HYPER-PARAMETERS:

For both tasks, L S = L P = 1, N S = N P = 256, d = 512.

D MORE DETAILS ON TRAINING:

Algorithm 2 presents the details of the forward pass during training. The forward pass during evaluation is similar, except that we compute ĝt and a t as ĝt = arg max gt SP (g t ) and a t = arg max a (RS(ĝ t )), and the inference immediately terminates when SP (ĝ t ) < or a t does not apply.

Reward design. For job scheduling problem, we simply define the reward function as r(s t , (ĝ t , a t )) = c(s t ) -c(s t+1 ). The same reward function does not apply to the expression simplification problem, since it would always give a negative reward for uphill rewriting rules such as expanding a "min/max". Thus, for expression simplification problem, we modify the reward function to r(s t , (ĝ t , a t )) = max t+1≤k≤T (γ k-t (c(s t ) -c(s k )))/c(ĝ t ), where γ is a discount factor. This design of reward function would assign a positive value to an uphill rule if it is included in a path that results in a simplified expression in the end. We normalize the reward by c(ĝ t ), so that the reward function is bounded. In our evaluation, we set γ = 0.9.

Hyper-parameters. In Algorithm 2, β = 10.0, = 0.0, T sp = 10, T rs = 10, T iter = 50. p c is initialized with 0.5, and is decayed by 0.8 for every 1000 timesteps until p c = 0.01, where it is not decayed anymore. In the training loss function,α = 10.0. The initial learning rate is 1e -4, and is decayed by 0.9 for every 1000 timesteps. Batch size is 128. Gradients with L 2 norm larger than 5.0 are scaled down to have the norm of 5.0. The model is trained using Adam optimizer. All weights are initialized uniformly randomly in [-0.1, 0.1].

E MORE RESULTS FOR JOB SCHEDULING PROBLEM:

To examine how the initial schedules affect the final results, besides earliest-job-first schedules discussed in Section 5, we also evaluate initial schedules with different average slow down. Specifically, for each job sequence, we generate different initial schedules by randomly allocating one job at a time.

In Table 4 , we present the results with D = 20 types of resources. For each job sequence, we randomly generate 10 different initial schedules. We can observe that although the effectiveness of initial schedules affect the final schedules, the performance is consistently better than other baseline approaches, which demonstrates that our neural rewriter is able to substantially improve the initial solution regardless of its quality.

F MORE RESULTS FOR EXPRESSION SIMPLIFICATION:

In Figures 7 and 8 , we present some success cases of expression simplification, where we can simplify better than the Halide rule-based rewriter. 

:

The code is released in their public repository here: https://github.com/halide/Halide.

https://github.com/halide/Halide/tree/new_autoschedule_with_new_ simplifier/apps/random_pipeline.

Figures and Tables:

Figure 2: The instantiation of neural rewriter architectures for different domains: (a) job scheduling; (b) expression simplification. In (a), s t is the dependency graph representation of the job schedule.Each circle with index greater than 0 represents a job node, and node 0 is an additional one representing the machine. Edges in the graph reflect job dependencies. The score predictor selects a job ĝt ∼ SP (g t ) from all job nodes to re-schedule. The rule selector chooses a moving action a t ∈ A for ĝt , then modifies s t to get a new dependency graph s t+1 . In (b), s t is the expression parse tree, where each square represents a node in the tree. The set of g t ⊂ s t includes every sub-tree rooted at a non-terminal node, from which the score predictor selects ĝt to rewrite. Afterwards, the rule selector predicts a rewriting rule a t , then rewrites the sub-tree ĝt to get the new tree s t+1 .

Figure 3: An example to illustrate the graph construction approach for the job scheduling problem. Node 0 is an additional node representing the machine.

Figure4presents the grammar of Halide expressions in our evaluation.

Figure 7: The rewriting process that simplifies the expression ((v0 -v1 + 18)/35 * 35 + 35) ≤ v0 -v1 + 119 to 34 ≤ (v0 -v1 + 13)%35.

Figure 8: The rewriting process that simplifies the expression ((v0 -v1 + 12)/137 * 137 + 137) ≤ min((v0 -v1 + 149)/137 * 137, v0 -v1 + 13) to 136 ≤ (v0 -v1 + 12)%137.

Statistics of the dataset for expression simplification.

Table content: <table><tr><td>: 1.36M</td><td>Average: 106.84</td><td>Average: 27.39</td></tr><tr><td>Training/Val/Test: 1.09M/136K/136K</td><td>Min/Max: 10/579</td><td>Min/Max:3/100</td></tr><tr><td colspan="3">Average expression length reduction Average tree size reduction</td></tr><tr><td>Halide Rule-based Rewriter</td><td>36.13</td><td>9.68</td></tr><tr><td>Heuristic Search</td><td>43.27</td><td>12.09</td></tr><tr><td>Neural Rewriter (Ours)</td><td>46.98</td><td>13.53</td></tr><tr><td>Z3 Solver</td><td>50.81</td><td>15.82</td></tr></table>

Experimental results of the Halide expression simplification task on the test set. Note that the Z3 solver can perform rewriting steps that are not included in the Halide ruleset.

Table content: <table/>

Figure4: Grammar of the Halide expressions in our evaluation. "select (c, e1, e2)" means that when the condition c is satisfied, this term is equal to e1, otherwise is equal to e2. In our dataset, all constants are integers ranging in[-1024, 1024], and variables are from the set {v0, v1, ..., v12}. Algorithm of a Single Rewriting Step for Job Scheduling Problem 1: function REWRITE(g j , g j , s t ) 2: if C j < S j or C j == A j then

Table content: <table><tr><td/><td colspan="2">&lt;Expr&gt; ::= &lt;AlgExpr&gt; | &lt;BoolExpr&gt;</td></tr><tr><td/><td colspan="2">&lt;BoolExpr&gt; ::= &lt;AlgExpr&gt; &lt;&lt;AlgExpr&gt;</td></tr><tr><td/><td>|</td><td>&lt;AlgExpr&gt; &lt;= &lt;AlgExpr&gt;</td></tr><tr><td/><td>|</td><td>&lt;AlgExpr&gt; == &lt;AlgExpr&gt;</td></tr><tr><td/><td>|</td><td>(!&lt;BoolExpr&gt;)</td></tr><tr><td/><td>|</td><td>(&lt;BoolExpr&gt; &amp;&amp; &lt;BoolExpr&gt;)</td></tr><tr><td/><td>|</td><td>(&lt;BoolExpr&gt; || &lt;BoolExpr&gt;)</td></tr><tr><td/><td colspan="2">&lt;AlgExpr&gt; ::= &lt;Term&gt;</td></tr><tr><td/><td>|</td><td>(&lt;AlgExpr&gt; + &lt;Term&gt;)</td></tr><tr><td/><td>|</td><td>(&lt;AlgExpr&gt; -&lt;Term&gt;)</td></tr><tr><td/><td>|</td><td>(&lt;AlgExpr&gt; * &lt;Term&gt;)</td></tr><tr><td/><td>|</td><td>(&lt;AlgExpr&gt; / &lt;Term&gt;)</td></tr><tr><td/><td>|</td><td>(&lt;AlgExpr&gt; % &lt;Term&gt;)</td></tr><tr><td/><td colspan="2">&lt;Term&gt; ::= &lt;Var&gt; | &lt;Const&gt;</td></tr><tr><td/><td>|</td><td>max(&lt;AlgExpr&gt;, &lt;AlgExpr&gt;)</td></tr><tr><td/><td>|</td><td>min(&lt;AlgExpr&gt;, &lt;AlgExpr&gt;)</td></tr><tr><td/><td>|</td><td>select(&lt;BoolExpr&gt;, &lt;AlgExpr&gt;, &lt;AlgExpr&gt;)</td></tr><tr><td colspan="2">Algorithm 1 3: return S</td></tr><tr><td>4:</td><td>end if</td></tr><tr><td>5:</td><td/></tr></table>

Under review as a conference paper at ICLR 2019 Algorithm 2 Forward Pass Algorithm for the Neural Rewriter during Training Require: initial states 0 , hyper-parameters β, , p c , T iter , T sp , T rs 1: for t = 0 → T iter -1 do Sample ĝt ∼ exp(β • SP (g t )), where g t ⊂ s t if a t does not applied to (s t , ĝt ) then break fi

Table content: <table><tr><td>2:</td><td>for i = 1 → T sp do</td></tr><tr><td>3:</td><td/></tr><tr><td>4:</td><td>if SP (ĝ t ) &lt; then</td></tr><tr><td>5: 6:</td><td>Re-sample ĝ t ∼ exp(β • SP (g t )) with a probability of 1 -p c if Re-sampling is not performed then break fi</td></tr><tr><td>7:</td><td>else</td></tr><tr><td>8:</td><td>break</td></tr><tr><td>9:</td><td>end if</td></tr><tr><td>10:</td><td>end for</td></tr><tr><td>11:</td><td>for i = 1 → T rs do</td></tr><tr><td>12:</td><td>Sample a t ∼ RS(ĝ t )</td></tr><tr><td>13:</td><td>if a t can be applied to (s t , ĝt ) then break fi</td></tr><tr><td>14:</td><td>end for</td></tr><tr><td>15:</td><td/></tr><tr><td>16:</td><td/></tr></table>

Experimental results of the job scheduling problem using initial schedules with different average slow down. The number of resource types D = 20. We also include the results from Table1for reference.

Table content: <table><tr><td>Initial average slow down</td><td colspan="3">≤ 10 10 -25 &gt; 25</td></tr><tr><td>Final average slow down</td><td>4.49</td><td>4.61</td><td>4.78</td></tr><tr><td colspan="2">Results from Table 1</td><td/><td/></tr><tr><td>Shortest Job First</td><td/><td>5.00</td><td/></tr><tr><td>Shortest First Search</td><td/><td>4.98</td><td/></tr><tr><td>DeepRM (Mao et al., 2016)</td><td/><td>10.18</td><td/></tr><tr><td>Neural Rewriter (Ours)</td><td/><td>4.63</td><td/></tr><tr><td>Optim (Lower bound)</td><td/><td>4.26</td><td/></tr><tr><td>Earliest Job First (Upper bound)</td><td/><td>24.23</td><td/></tr></table>

Peer review (iclr2019):
 This paper addresses the challenges of prediction-based, progressive planning on discrete state and action spaces. Their proposed method applies existing DAG-LSTM/Tree-LSTM architectures to iteratively refine local sections in the existing plan that could be improved until convergence. These models are then evaluated on a simulated job scheduling dataset and Halide expression simplification. While this paper presents an interesting approach to the above two problems, its presentation and overall contribution was pretty unclear to me. A few points: 1. Ambiguous model setup: It may have been more advantageous to cut a large portion of Section 3 (Problem Setup), where the authors provide an extensive definition of an optimization problem, in favor of providing more critical details about the model setup. For example, how exactly should we view the job scheduling problem from an RL perspective? How are the state transitions characterized, how is the network actually trained (REINFORCE? something else?), is it episodic (if so, what constitutes an episode?), what is the exploration strategy, etc. It was hard for me to contextualize what exactly was going on 2. Weak experimental section: The authors mention that they compare their Neural Rewriter against DeepRM using a simplified problem setup from the original baseline. I wonder how their method would have fared against a task that was comparable in difficulty to the original method -- this doesn’t feel like a fair comparison. And although their expression simplification results were nice, I would also like to know why the authors chose to evaluate their method on the Halide repository specifically. Since they do not compare their method against any other baselines, it’s hard for me to gauge the significance of their results. 3. Variance across initializations: It would have been nice to see an experiment on how various initializations of schedules/expressions affect the policies learned. I would imagine that poor initializations could lead to poor results, but it would be interesting if the Neural Rewriter was robust to the quality of the initial policy. Since this is not addressed in the paper, it is difficult to gauge whether the authors’ model performed well due to an unfair advantage. Additionally, how much computational overhead is there to providing these (reasonable) initial policies as opposed to learning from scratch? 4. Unclear notation: As previously addressed by other reviewers, key definitions such as the predicted score SP(.) are missing from the text.