L EARNING TO PROGRESSIVELY PLAN
Anonymous authors
Paper under double-blind review
A BSTRACT
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 simpliﬁcation.
Compared to common effective heuristics, baseline deep models and search algorithms, our approach efﬁciently gives solutions with higher quality.
I NTRODUCTION
In recent years, deep reinforcement learning has achieved strong success in scenarios that are welldeﬁned and can be precisely and efﬁciently 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 insufﬁcient.
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 speciﬁcation 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 difﬁcult 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 ﬁnd 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 simpliﬁcation.
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.
R ELATED W ORK
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 reﬁnement 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˙zewicz 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 simpliﬁed 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 conﬁgurations, 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 ﬁnd novel rewriting routines.
On the contrary, we use reinforcement learning to train our rewriter for expression simpliﬁcation, 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 simpliﬁcation (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).
P ROBLEM S ETUP
In this work, we propose to formulate optimization as a rewriting problem, and ﬁnd 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 deﬁned as follows:
Deﬁnition 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 ﬁnd arg min s ∈S c ( s ) .
When S is continuous, gradient-based algorithms are effective ways of ﬁnding 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 ﬁrst construct a feasible solution, then make incremental improvement.
This is because (1) a feasible solution is often easy to ﬁnd; (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 difﬁcult 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 efﬁciency.
Therefore, we propose to solve optimization problems through rewriting , as deﬁned formally below.
Deﬁnition 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 ˆ g t , where ˆ g t ⊂ s t , and the rewriting step results in the next state s t +1 =
f ( s t , ˆ g t , a t ) .
Given an initial state s 0 , our goal is to ﬁnd 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 inefﬁcient exhaustive search, while a set of kaleidoscopic routines is often cumbersome to design, hard to maintain and lacks ﬂexibility.
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 simpliﬁcation, in which we minimize the expression length using a well-deﬁned semantics-preserving rewriting ruleset.
3. J OB S CHEDULING P ROBLEM
We ﬁrst study the job scheduling problem, and in particular, we consider a simpliﬁed problem setup studied in (Mao et al., 2016) as follows.
Suppose we have a machine with D types of resources.
Each job j is speciﬁed as g j = ( r j 1 , r j 2 , ..., r jD , S j , T j ) , where r j 1 , r j 2 , ..., 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 ﬁxed during the entire job execution, each job must run continuously until ﬁnishing, 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 ﬁnd 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 (cid:48) ﬁnishes 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
| = 2 L Q , since each job could only switch its scheduling order with at most L Q of its former and latter jobs respectively.
3. E XPRESSION S IMPLIFICATION
We also apply our approach to expression simpliﬁcation domain.
In particular, we consider expressions in Halide, a domain-speciﬁc 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 heuristics 1 .
The grammar of the expressions considered in the rewriter is speciﬁed in Appendix A. Notice that the grammar includes a more comprehensive operator set than previous works on ﬁnding 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 ﬁnd that a large number of rewriting templates enumerate speciﬁc cases for an uphill rule , which lengthens the expression ﬁrst 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 satisﬁes certain pre-conditions , which is traditionally speciﬁed 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 ﬂexibility 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.
N EURAL R EWRITER M ODEL
[Figure] [Equation] [Figure] [Equation]
[Figure] [Equation] [Figure] [Equation]
[Figure]
Score Predictor In the following, we present the design of our rewriting model, i.e., Neural Rewriter . We ﬁrst provide an overview of our model framework, then present the design details for different applications.
4. M ODEL O VERVIEW
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 beneﬁt of rewriting g t .
A high score indicates that rewriting g t could be desirable.
Rule selector.
Given ˆ g t ⊂ s t to be rewritten, the rule selector predicts a probability distribution
RS (ˆ g t ) over the entire ruleset A , and selects a rule a t ∈ A to apply accordingly.
4. M ODEL D ETAILS FOR J OB S CHEDULING P ROBLEM
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 speciﬁed by g j =
( r j 1 , r j 2 , ..., r jD , S j , T j ) .
In addition, we deﬁne 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.
Speciﬁcally, 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 (cid:104) g 0 , g j (cid:105) in the graph.
Otherwise, there must exist at least one job g j (cid:48) such that C j (cid:48)
=
A j .
We add an edge (cid:104) g j (cid:48) , g j (cid:105) for every such job g j (cid:48) 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).
Speciﬁcally, 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 (cid:48) or g 0 in the graph, which means allocating job g j after job g j (cid:48) ﬁnishes or at its arrival time S j .
Thus, the input to the rule selector not only includes h j , but also h j (cid:48) of all other g 2 0.
g 4 0.
g 5 0.
Input Encoder
Score Predictor
Rule Selector 0.
DAG-LSTM
Embedding
FC FC
FC
Softmax
S t a t
S t+     4-0.
#    3(a) a 1 0 … a 5 1 (ConstantReduction) …a
InputEncoder
ScorePredictor
Rule Selector
Embedding
Tree-LSTM
FC0.
FC
Softmax
S t
S t+ 𝐠" 𝐭 a t <= minv v v v <
= minv v v v <=
min v v (b) g j (cid:48) that could be used for rewriting. The output layer is an |A| -dimensional softmax layer, where |A | = 2 L Q as discussed in Section 3. More details can be found in Appendix C.1.
4. M ODEL D ETAILS FOR E XPRESSION S IMPLIFICATION
We present the instantiation of our neural rewriter framework for expression simpliﬁcation in Figure 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 2 d -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. T RAINING D ETAILS
Let ( s 0 , (ˆ g 0 , a 0 )) , ( s 1 , (ˆ g 1 , a 1 )) , ..., ( s T − 1 , (ˆ g T − 1 , a T − 1 )) , s T be the rewriting sequence in the for- ward pass, we deﬁne the reward function r ( s t , (ˆ g t , a t )) as follows.
For job scheduling problem, r ( s t , (ˆ g t , a t ))
=
c ( s t ) − c
( s t +1 ) .
For expression simpliﬁcation problem, r ( s t , (ˆ g t , a t ))
=
max t +1 ≤ k ≤ T ( γ k − t ( c ( s t ) − c ( s k ))) /c
(ˆ g t ) , where γ is a discount factor.
This reward function is designed to credit uphill rules that could lead to a simpliﬁed 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 , (ˆ g t , a t )) is also used as the training target of SP (ˆ g t ) .
Speciﬁcally, 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 (ˆ g t ) as the critic.
In particular, let the advantage function A ( s t , (ˆ g t , a t ))
=
r ( s t , (ˆ g t , a t )) −
SP (ˆ g 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.
E XPERIMENTS
In this section, we present the evaluation results of both job scheduling and expression simpliﬁcation 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.
[Table] [Table] [Table]  E VALUATION S ETUP
We randomly generate 100 K job sequences, and use 80 K for training, 10 K for validation, and 10 K for testing.
We use an online setting where jobs arrive on the ﬂy 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 conﬁgurations 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. M ODEL C ONFIGURATION
For our neural rewriter model, we provide an initial schedule that allocates jobs in a ﬁrst-come-ﬁrstserve manner for every job sequence, then feed it to the neural rewriter for reﬁnement.
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 ﬁrst 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 ﬁnd 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 ofﬂine 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. R ESULTS
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 efﬁciency, then progressively reﬁne 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
[Table]
[Table]
5. E XPRESSION S IMPLIFICATION 5.2. E VALUATION S ETUP
To construct the dataset
[Table]
we ﬁrst generate random pipelines using the generator in the Halide repository 2 , then extract expressions from them.
We ﬁlter 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
[Table]
and more details can be found in Appendix A.
We measure the following metrics in our evaluation: (1) Average expression length reduction
[Table]
which is the length reduced from the initial expression to the rewritten one
[Table]
and the length is deﬁned 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.
[Table]
[Table]
M ODEL C ONFIGURATION
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 simpliﬁed, 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 ﬁrst 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 ﬁrst 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 simpliﬁer works by traversing each sub-formula in the input expression and invoking the solver to ﬁnd a simpler equivalent one to replace it.
Therefore, the simpliﬁcation 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 beseconds for each input expression, and we ﬁnd that the results are not signiﬁcantly better with a longer timeout.
5.2. R ESULTS
Table 3 presents the main results of our expression simpliﬁcation 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( v 0 , 3) + 3 , max( v 1 , v 2)) , which could be reduced to True by expanding max(max( v 0 , 3) + 3 , max( v 1 , v 2)) and max( v 0 , 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 ﬁnd the correct order of expanding the right hand size of the expression when more “min/max” are included, which would make the search less efﬁcient.
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.
C ONCLUSION
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 simpliﬁcation 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 efﬁciency of our rewriting approach and extending it to more complicated scenarios as future work.
A G RAMMAR OF THE H ALIDE E XPRESSIONS
Figure 4 presents the grammar of Halide expressions in our evaluation.
B M ORE D ETAILS ON THE R EWRITING R ULESET B.
M ORE D ETAILS FOR J OB S CHEDULING P ROBLEM
Algorithm 1 describes a single rewriting step for job scheduling problem.
B. M ORE D ETAILS FOR E XPRESSION S IMPLIFICATION P ROBLEM
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 simpliﬁed.
Notice that although applying this rewriting rule has no beneﬁt in most cases, since it will increase the expression
[Figure]
e
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
e
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
length, it is necessary to include it in the ruleset, because when either a < c or b < c is always true, expanding the “min” term could reduce the entire expression to a tautology, which ends up simplifying the entire expression.
Figure 5 shows an example of the rewriting process using uphill rules properly.
C M ORE D ETAILS ON M ODEL A RCHITECTURES C. M ORE D ETAILS FOR J OB S CHEDULING P ROBLEM
Job embedding.
We describe the details of job embedding as follows.
Consider a job g j = ( r j 1 , r j 2 , ..., r jD , S j , T j ) .
We denote the amount of resources occupied by all jobs at each timestep t as r (cid:48) t
= ( r (cid:48) t 1 , r (cid:48)
t 2 , ..., r (cid:48) tD ) .
Each job g j is represented as a ( D × ( T max + 1) + 1) -dimensional vector, where the ﬁrst D dimensions of the vector are ( r j 1 , r j 2 , ..., r jD ) , representing its resource requirement.
The following D × T j dimensions of the vector are the concatenation of r (cid:48)
A j , r (cid:48)
A j +1 , ..., r (cid:48)
A j + T j − 1 , which describes the machine usage during the execution of the job g j .
When T j < T max , the following D × ( T max − T j ) dimensions are zero.
The last dimension of the embedding vector is the slowdown of the job in current schedule.
The deﬁnition of the slowdown is the same as in (Mao et al., 2016), which is computed by ( C j − S j ) /T
j , where C j =
A j + T j is the completion time.
We denote the embedding of each job g j as e j .
Figure 6 shows an example of our job embedding approach.
Require: initial state s 0 , hyper-parameters β, (cid:15), p c , T iter , T sp , T rs 1: for t =
→ T iter − 1 do 2: for i =
→ T sp do 3:
Sample ˆ g t ∼ exp( β · SP ( g t )) , where g t ⊂ s t 4: if SP (ˆ g t ) <
(cid:15) then 5:
Re-sample ˆ g (cid:48) t ∼ exp( β · SP ( g t )) with a probability of 1 − p c 6: if Re-sampling is not performed then break ﬁ 7: else 8: break 9: end if 10: end for 11: for i =
→ T rs do 12:
Sample a t ∼ RS (ˆ g t ) 13: if a t can be applied to ( s t , ˆ g t ) then break ﬁ 14: end for 15: if a t does not applied to ( s t , ˆ g t ) then break ﬁ 16: s t
+1 = f
( s t , ˆ g t , a t ) 17: end for
Rule selector.
The rule selector has two modules.
The ﬁrst 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 (cid:48) k } denotes the set of such jobs.
For each g j (cid:48) k , the input is [ h j ; h j (cid:48) k ] , and this module computes a d -dimensional vector h (cid:48)
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 (cid:48) 1 ; h (cid:48) 2 ; ... ; h (cid:48)|A| ] , where |A | = 2 L Q .
When N j < |A| , h (cid:48) N j +1 , h (cid:48)
N j +2 , ..., h (cid:48)|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. M ODEL D ETAILS FOR E XPRESSION S IMPLIFICATION
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 2 d -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 insufﬁcient to determine whether it is beneﬁcial 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 simpliﬁcation, since the resulted expression max ( a + 2 , b + 2) < a + 2 could be reduced to false .
We have tried other approaches of combining the parent information into the input, but we ﬁnd that including the embedding of the entire tree is the most efﬁcient way.
C. M ODEL HYPER - PARAMETERS
For both tasks, L S =
L P = 1 , N S = N P = 256 , d = 512 .
D M ORE D ETAILS ON T RAINING
Algorithm 2 presents the details of the forward pass during training.
The forward pass during evaluation is similar, except that we compute ˆ g t and a t as ˆ g t = arg max g t SP ( g t ) and a t = arg max a ( RS (ˆ g t )) , and the inference immediately terminates when SP (ˆ g t ) < (cid:15) or a t does not apply.
Reward design.
For job scheduling problem, we simply deﬁne the reward function as r ( s t , (ˆ g t , a t )) =
c ( s t ) − c
( s t +1 ) .
The same reward function does not apply to the expression simpliﬁcation problem, since it would always give a negative reward for uphill rewriting rules such as expanding a “min/max”.
Thus, for expression simpliﬁcation problem, we modify the reward function to r ( s t , (ˆ g t , a t )) =
max t +1 ≤ k ≤ T ( γ k − t ( c ( s t ) − c ( s k ))) /c
(ˆ g t ) , where γ is a discount factor.
This
[Table]
design of reward function would assign a positive value to an uphill rule if it is included in a path that results in a simpliﬁed expression in the end.
We normalize the reward by c (ˆ g t ) , so that the reward function is bounded.
In our evaluation, we set γ = 0 .
Hyper-parameters.
In Algorithm 2, β = 10 . 0 , (cid:15) = 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 1 e − 4 , and is decayed by 0 .
9 for every 1000 timesteps.
Batch size is 128.
Gradients with L 2 norm larger than 5. 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 .
E M ORE R ESULTS FOR J OB S CHEDULING P ROBLEM
To examine how the initial schedules affect the ﬁnal results, besides earliest-job-ﬁrst schedules discussed in Section 5, we also evaluate initial schedules with different average slow down.
Speciﬁcally, 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 ﬁnal 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 M ORE R ESULTS FOR E XPRESSION S IMPLIFICATION
In Figures 7 and 8, we present some success cases of expression simpliﬁcation, where we can simplify better than the Halide rule-based rewriter.
[Figure] [Figure]