Original: R OBUST SCHEDULING WITH GF LOW N ETS
Revised: R OBUST S CHEDULING WITH GF LOW N ETS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Finding the best way to schedule operations in a computation graph is a classical
Revised: Finding the best way to schedule operations in a computation graph is a classical

Original: NP-hard problem which is central to compiler optimization.
Revised: NP-hard problem which is central to compiler optimization.

Original: However, evaluating the goodness of a schedule on the target hardware can be very time-consuming.
Revised: However, evaluating the goodness of a schedule on the target hardware can be very time-consuming.

Original: Traditional approaches as well as previous machine learning ones typically optimize proxy metrics, which are fast to evaluate but can lead to bad schedules when tested on the target hardware.
Revised: Traditional approaches as well as previous machine learning ones typically optimize proxy metrics, which are fast to evaluate but can lead to bad schedules when tested on the target hardware.

Original: In this work, we propose a new approach to scheduling by sampling proportionally to the proxy metric using a novel GFlowNet method.
Revised: In this work, we propose a new approach to scheduling by sampling proportionally to the proxy metric using a novel GFlowNet method.

Original: We introduce a technique to control the trade-off between diversity and goodness of the proposed schedules at inference time and demonstrate empirically that the pure optimization baselines can lead to subpar performance with respect to our approach when tested on a target model.
Revised: We introduce a technique to control the trade-off between diversity and goodness of the proposed schedules at inference time and demonstrate empirically that the pure optimization baselines can lead to subpar performance with respect to our approach when tested on a target model.

Original: Furthermore, we show that conditioning the GFlowNet on the computation graph enables generalization to unseen scheduling problems for both synthetic and real-world compiler datasets.
Revised: Furthermore, we show that conditioning the GFlowNet on the computation graph enables generalization to unseen scheduling problems for both synthetic and real-world compiler datasets.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Efficient execution of computation graphs is paramount to many scientific and industrial applications, with deep learning being a prominent example (Amodei & Hernandez, 2018).
Revised: Efficient execution of computation graphs is paramount to many scientific and industrial applications, with deep learning being a prominent example (Amodei & Hernandez, 2018).

Original: Scheduling is the action of assigning operations to the available compute resources, such as threads, cores, or nodes in a cluster (Hennessy & Patterson, 2011).
Revised: Scheduling is the action of assigning operations to the available compute resources, such as threads, cores, or nodes in a cluster (Kwok & Ahmad, 1999; Hennessy & Patterson, 2011; Pinedo, 2012).

Original: Unfortunately, finding the schedule with the shortest possible makespan (start-to-end runtime) is in general NP-hard (Papadimitriou & Steiglitz, 1998).
Revised: Unfortunately, finding the schedule with the shortest possible makespan (start-to-end runtime) is in general NP-hard (Papadimitriou & Steiglitz, 1998).

Original: As a result, domain experts have come up with heuristics that are tailored to specific problem instances (Ibarra & Kim, 1977).
Revised: As a result, domain experts have come up with heuristics that are tailored to specific problem instances (Ibarra & Kim, 1977).

Original: Machine learning approaches promise the possibility to automate this process allowing for fast adaptation to new graph distributions (Wang & O’Boyle, 2018).
Revised: Machine learning approaches promise the possibility to automate this process allowing for fast adaptation to new graph distributions (Wang & O’Boyle, 2018; Bengio et al., 2021c).

Original: In this work, we consider the problem of scheduling a set of operations with precedence constraints on a fixed number of homogeneous devices, i.e. devices such that each operation can run on any device andhas thesame runtime  on every device.
Revised: In this work, we consider the problem of scheduling a set of operations with precedence constraints on a fixed number of homogeneous devices, i.e., any operation can run on any device and the runtime is the same on all devices.

Original: Evaluating the makespan of a schedule involves running all of the tasks on some target hardware.
Revised: Evaluating the makespan of a schedule involves running all operations in the computation graph on some target hardware.

Original: This can be very resource intensive, especially when the computation graph includes lengthy operations, the evaluated schedule is inefficient, or the intended target hardware is a cluster with many nodes.
Revised: This can be very resource intensive, especially when the computation graph includes lengthy operations, the evaluated schedule is inefficient, or the intended target hardware is a cluster with many nodes.

Original: Heuristic optimizers, like genetic algorithms (Hou et al., 1994), or machine learning (Mao et al., 2019) approaches further exacerbate this problem because they require many evaluations to converge.
Revised: Heuristic optimizers, like genetic algorithms (Hou et al., 1994), or machine learning (Mao et al., 2019) approaches further exacerbate this problem because they require many evaluations to converge (Chen et al., 2018).

Original: Proxies are a much faster alternative that estimate the makespan using a simplified model of the hardware.
Revised: Proxies are a much faster alternative that estimates the makespan using a simplified model of the hardware.

Original: However, this comes at the cost of discrepancies between the proxy makespan and the one observed on the hardware; as a result, performant solutions on the proxy might ultimately be unsatisfactory once tested on the target.
Revised: However, this comes at the cost of discrepancies between the proxy makespan and the one observed on the hardware; as a result, performant solutions on the proxy might ultimately be unsatisfactory once tested on the target.

Original: Nonetheless, proxies remain a good indicator for most schedules.
Revised: Nonetheless, proxies remain a good indicator for most schedules and are essential due to their efficiency.

Original: We aim to learn a scheduler that can be trained using the proxy, whilst being robust to its inaccuracies.
Revised: We aim to learn a scheduler that can be trained using the proxy, whilst being robust to its inaccuracies.

Original: The common approach to scheduling problems (and combinatorial optimization problems in general) is to look for the single best schedule that minimizes a makespan measure which can be a pure mathematical proxy (Paliwal et al., 2020),or the output of a simulator (Zhou et al., 2020) or even the real makespan on hardware (Khadka et al., 2021).
Revised: The common approach to scheduling problems (and combinatorial optimization problems in general) is to look for the single best schedule that minimizes a makespan measure which can be an analytical proxy (Paliwal et al., 2020), the output of a simulator (Zhou et al., 2020), or even the real makespan on hardware (Khadka et al., 2021).

Original: We propose a different philosophy: generate a set of candidate schedules that have a low makespan according to the proxy and are diverse.
Revised: We propose a different philosophy: generate a set of candidate schedules that have a low makespan according to the proxy and are diverse.

Original: By having multiple good schedules that are significantly different, we can reduce the impact of systematic errors in the proxy, and hope for robust performance on the target.
Revised: By having multiple good schedules that are significantly different, we can reduce the impact of systematic errors in the proxy, and hope for robust performance on the target.

Original: See Figure 1 for an illustration of the proposed scheduling pipeline.

Original: GFlowNetSampling ComputationGraph 1.
Revised: GFlowNetSampling ComputationGraph 1.

Original: N.
Revised: N.

Original: Proxy Ranking
Revised: Proxy Ranking

Original: Top-k
Revised: Top-k

Original: Chain Graphs
Revised: Chain Graphs

Original: Schedule 

Revised: Schedule

Original: Schedule
Revised: Schedule

Original: Target Evaluation  4 

Revised: Target Evaluation  4

Original: Our goal isthen to learn a generative model that assigns higher probability to low-makespan schedules, and importantly can also discover the different modes associated with local optima of the makespan cost.
Revised: Our goal is to learn a generative model that assigns higher probability to low-makespan schedules, and importantly can also discover the different modes associated with local optima of the makespan cost.

Original: Generative flow networks (GFlowNets) havebeen recently  introduced as a method for learning a stochastic policy that can piece-by-piece construct discrete and composite objects, proportional to a given reward (Bengio et al., 2021b).
Revised: Generative Flow Networks (GFlowNets) have recently been introduced as a method for learning a stochastic policy that can piece-by-piece construct discrete and composite objects, proportional to a given reward (Bengio et al., 2021b).

Original: By treating the proxy as the reward  we can use GFlowNets to sample a diverse set of candidate schedules.
Revised: By computing the reward from the proxy-makespan we can use GFlowNets to sample a diverse set of candidate schedules.

Original: Our main contributions are: 
Revised: Our main contributions are: 1.

Original: • We introduce an alternative to the pure proxy optimization viewpoint of scheduling that achieves better robustness to proxy errors, by generating multiple candidate schedules to evaluate directly on the target hardware.
Revised: We introduce an alternative to the pure proxy optimization viewpoint of scheduling that achieves better robustness to proxy errors, by generating multiple candidate schedules to evaluate directly on the target hardware.

Original: • We extend GFlowNets to generate schedules conditioned on a computation graph.
Revised: We extend GFlowNets to generate schedules conditioned on a computation graph.

Original: Additionally, we introduce a method to control diversity and goodness at inference time, without the need for retraining.
Revised: Additionally, we introduce a method to control diversity and goodness at inference time, without the need for retraining.

Original: These contributions may be of general interest, beyond the scheduling problem.
Revised: These contributions may be of general interest, beyond the scheduling problem.

Original: • We empirically demonstrate the robustness of our method to proxy errors and verify the generalization ability on a diverse set of synthetic and real-world computation graphs.
Revised: We empirically demonstrate the robustness of our method to proxy errors and verify the generalization ability on a diverse set of synthetic and real-world computation graphs.

Original:  R OBUST SCHEDULING
Revised:  R OBUST SCHEDULING

Original: In this section, we first provide a definition of the scheduling problem we consider in this work.
Revised: In this section, we first provide a definition of the scheduling problem we consider in this work.

Original: Then, we discuss how a proxy simulates the schedule execution as well as the difficulties of specifying a reliable proxy.
Revised: Then, we discuss how a proxy simulates the schedule execution as well as the difficulties of specifying a reliable proxy.

Original: Finally, we describe our proposed generative scheduling framework.
Revised: Finally, we describe our proposed generative scheduling framework.

Original: 2. P ROBLEM DEFINITION
Revised: 2. P ROBLEM DEFINITION

Original: In scheduling, we are given a computation graph G C
Revised: In scheduling, we are given a computation graph G C

Original: =
Revised: =

Original: ( O, P ) that is a Direct Acyclic Graph (DAG) consisting of operations (nodes) o ∈ O and precedence constraints (edges) p ∈ P that encode a partial order in which the operations need to be executed.
Revised: ( O, P ) that is a direct acyclic graph (DAG) consisting of operations (nodes) o ∈ O and precedence constraints (edges) p ∈ P that encode a partial order in which the operations need to be executed.

Original: In particular, the edge p ij encodes that operation
Revised: In particular, the edge p ij encodes that operation

Original: o
Revised: o

Original: i needs to finish before o j can start, for example because o j requires the output of o i as input.
Revised: i needs to finish before o j can start, for example because o j requires the output of o i as input.

Original: Our task is to run all operations on a set of devices D = { δ 1 , . . .
Revised: Our task is to run all operations on a set of devices D = { d 1 , . . .

Original: , δ m } , without violating the precedence constraints.
Revised: , d m } , without violating the precedence constraints.

Original: In addition to the precedence constraints, the devices can only run one operation at a time.
Revised: In addition to the precedence constraints, the devices can only run one operation at a time.

Original: We can then view scheduling as performing two distinct tasks: assign a device to each operation, and determine a (complete) order among all operations on the same device that is compatible with the precedence constraints encoded in G C .
Revised: We can then view scheduling as performing two distinct tasks: assign a device to each operation, and determine a (complete) order among all operations on the same device that is compatible with the precedence constraints encoded in G C .

Original: We can model the schedule as a chain of operations for each device, where the chain denotes the order in which the operations run on that device.
Revised: We can model the schedule as a chain of operations for each device, where the chain denotes the order in which the operations run on that device.

Original: See Figure 1 for a visual example of the chain graphs.
Revised: See Figure 1 for a visual example of the chain graphs.

Original: Our aim is to find the schedule with the lowest makespan for some target hardware.
Revised: Our aim is to find the schedule with the lowest makespan for some target hardware.

Original: 2. T ARGET MODEL VS  PROXIES
Revised: 2. T ARGET MODEL VS . PROXIES

Original: The makespan on any schedule can be evaluated on the target hardware by running all the operations in the specified order and on the specified devices.
Revised: The makespan of any schedule can be evaluated on the target hardware by running all the operations in the specified order and on the specified devices.

Original: However, this can take up significant time and compute resources when the computation graph is large, has costly operations, or the target hardware is a cluster with many nodes.
Revised: However, this can take up significant time and compute resources when the computation graph is large, has costly operations, or the target hardware is a cluster with many nodes.

Original: In addition to this, when optimizing the makespan one needs to evaluate many different schedules, further exacerbating the resource requirements.
Revised: In addition to this, when optimizing the makespan one needs to evaluate many different schedules, further exacerbating the resource requirements.

Original: A proxy is any tool that allows one to estimate the makespan of a given schedule, without having to run the schedule on the target hardware.
Revised: A proxy is any tool that allows one to estimate the makespan of a given schedule, without having to run the schedule on the target hardware.

Original: Proxies come with significant speed advantages, which remedy the problems mentioned above.
Revised: Proxies come with significant speed advantages, which remedy the problems mentioned above.

Original: However, this comes at the cost of possible mistakes in the estimation of the makespan and relative comparison of schedules.
Revised: However, this comes at the cost of possible mistakes in the estimation of the makespan and relative comparison of schedules.

Original: Mistakes can occur for example when the task durations are not accurately profiled, memory movements are too complex to fully model, or additional hardware-specific features are changed.
Revised: Mistakes can occur for example when the task durations are not accurately profiled, memory movements are too complex to fully model, or additional hardware-specific features are changed.

Original: Ideally, we would like to rely on a proxy for the majority of the schedule evaluations, and only evaluate a small fraction of promising schedules on the target hardware.
Revised: Ideally, we would like to rely on a proxy for the majority of the schedule evaluations, and only evaluate a small fraction of promising schedules on the target hardware.

Original: This approach differs from previous works, that either evaluate every schedule on the target (Khadka et al., 2021), leading to very long optimization times or evaluate everything on the proxy (Paliwal et al., 2020), which is susceptible to modeling failures.
Revised: This approach differs from previous works, that either evaluate every schedule on the target (Khadka et al., 2021), leading to very long optimization times, or evaluate everything on the proxy (Paliwal et al., 2020), which is susceptible to modeling failures.

Original: Next, we describe how the proxy we use in this work assigns start times to each operation given a schedule, and from those produces a makespan estimate.
Revised: Next, we describe how the proxy we use in this work assigns start times to each operation given a schedule and estimates the makespan based on those.

Original: We recall that a schedule is an order of operations for each device, which can be represented by one chain graph per device.
Revised: We recall that a schedule is an order of operations for each device, which can be represented by one chain graph per device.

Original: For each of these, let us denote with C δ , δ ∈ D the set of edges of the chain graph associated to device δ and with D ≡ (cid:83) mk
Revised: For each of these, let us denote with C d , d ∈ D the set of edges of the chain graph for device d and with D := (cid:83) m k

Original: =1 C δ k the set of all device constraints.
Revised: =1 C d k the set of all device constraints.

Original: The operations correspond to graph nodes and are labeled in the same way as in G C .
Revised: The operations correspond to graph nodes and are labeled in the same way as in G C .

Original: During the runtime or duration ρ i of operationo i , no other operation can run on the same device.
Revised: No other operation can run on the same device during the runtime or duration ρ i of operation

Original:  In practice, ρ i is estimated directly on the hardware in a profiling stage that precedes scheduling.
Revised: o i In practice, ρ i is estimated directly on the hardware in a profiling stage that precedes scheduling.

Original: We denote the start time of o i as τs i and can thus express the precedence constraints as:
Revised: We denote the start time of o i as τ i and can thus express the precedence constraints as:

Original: [Equation]
Revised: [Equation]

Original: An operation cannot start unless all of those that produce its inputs and all of those that precede it on its assigned device have finished first.
Revised: An operation cannot start unless all of those that produce its inputs and all of those that precede it on its assigned device have finished first.

Original: One can then simply assign each operation
Revised: To ensure that these constraints are satisfied the proxy assigns each operation

Original: o i the start time
Revised: o i the start time

Original: [Equation]
Revised: [Equation]

Original: which can be done as long as start times τ s k of all of its parents have been assigned previously.
Revised: If a node has no parents in P ∪ D the proxy assigns the start time τ i = 0 . The start times of all o i ∈

Original: It is easy to convince oneself that, if the graph ( O, P ∪ D ) isin turn a DAG, these assignments can be done recursively by starting from the source node (i.e. the one node with no parents) of G C .
Revised: O can be computed by assigning a start time to a node whenever it has no parents or all its parents have an assigned start time. If the graph ( O, P ∪ D ) is a DAG, then this algorithm is guaranteed to assign start times that satisfy Equation 2.

Original: We can then define the makespan T of the schedule x as:
Revised: The proxy then estimates the makespan T of the schedule x as:

Original: [Equation]
Revised: [Equation]

Original: This is the estimate of the makespan produced by our proxy given a schedule and a set of node durations.

Original: Optimizing this cost over the set of all possible schedules is already a very rich problem, and yet, we made significant simplifying assumptions.
Revised: Optimizing this cost over the set of all possible schedules is already a very rich problem, and yet, we made significant simplifying assumptions in the construction of the proxy.

Original: In particular, we assumed the runtimes to have been estimated with perfect accuracy, and in equation (2) we effectively assumed that an operation can start as soon as all of the operations producing its inputs finish, meaning that data can be moved between devices instantenously (zero latency) independently of their size (infinite bandwidth).
Revised: In particular, we assume perfectly estimated runtimes, and in Equation 2 we effectively assume that an operation can start as soon as all of the operations producing its inputs finish, meaning that data can be moved between devices instantaneously (zero latency) independently of their size (infinite bandwidth).

Original: In principle, neither of these applies to real hardware, as we shall later discuss.
Revised: These assumptions might be unrealistic, depending on the specific target devices (Valiant, 1990).

Original: 2. G ENERATIVE SCHEDULING
Revised: 2. G ENERATIVE SCHEDULING

Original: Our aim is to come up with a good set of candidate schedules to be tested on the target hardware while relying only on the proxy for generating this set.
Revised: Our aim is to come up with a good set of candidate schedules to be tested on the target hardware while relying only on the proxy for generating this set.

Original: While the proxy is imperfect, it still offers good guidance for most schedules; therefore, we would like to include schedules that perform well according to the proxy.
Revised: While the proxy is imperfect, it still offers good guidance for most schedules; thus, we would like to include schedules that perform well according to the proxy.

Original: Nevertheless, we also know that systematic errors in the proxy can cause it to incorrectly predict  low makespan for some schedules.
Revised: Nevertheless, we also know that systematic errors in the proxy can cause it to incorrectly predict a low makespan for some schedules.

Original: Therefore, we would like the set of candidate schedules to be diverse, while still high-quality from the point of view of the proxy.
Revised: Therefore, we would like the set of candidate schedules to be diverse, while still high-quality from the point of view of the proxy.

Original: If we had a ranking over all the schedules according to the proxy, we could just go through the list top-to-bottom, and add a schedule to the batch whenever it is significantly different from the previous ones.
Revised: If we had a ranking over all the schedules according to the proxy, we could just go through the list top-to-bottom, and add a schedule to the batch whenever it is significantly different from the previous ones.

Original: A full ranking like this is infeasible to construct, but we can instead learn a generative model that samples higher ranked schedules with higher probability.
Revised: A full ranking like this is infeasible to construct, but we can instead learn a generative model that samples higher ranked schedules with higher probability.

Original: When generating schedules, we need to satisfy the precedence and device constraint outlined in Section 2.1.
Revised: When generating schedules we need to satisfy the precedence and device constraint outlined in Section 2.1.

Original: To avoid generating invalid schedules, we picture this generation as a step-by-step process, where an empty schedule corresponds to the initial state s 0 ,a full schedule to a final state s n (with n = | O | ), and  each  state st with 0 < t < n to a partial schedule.
Revised: To avoid generating invalid schedules we construct a schedule in a step-by-step process: start with an empty schedule at the initial state s 0 , and at each step add an operation to the partial schedule until the schedule contains all operations at the terminal state s n .

Original: At each intermediate state s t , an action a t consists in picking an operation and assigning it to one of the devices, leading to a new state s t +1 .
Revised: At each intermediate state s t , an action a t consists in picking an operation and assigning it to one of the devices, leading to a new state s t +1 .

Original: We define the set of valid actions at every step t in a way such that the constraint in Equation 1 is automatically satisfied.
Revised: We define the set of valid actions at every step t in a way such that the precedence constraints are automatically satisfied.

Original: In particular, the agent can pick an operation o t  if and only if ∀ k : ( k, t ) ∈ P , o k is already in the partial schedule at state s t .
Revised: In particular, adding an operation o t is a valid action if and only if ∀ k : ( k, t ) ∈ P , o k is already in the partial schedule at state s t .

Original: This is a sufficient condition for the final “schedule” graph ( O, P ∪ D ) to be a DAG, implying that the constructed schedule is feasible.
Revised: This is a sufficient condition for the final “schedule” graph ( O, P ∪ D ) to be a DAG, implying that the constructed schedule is feasible.

Original: The final states represent full schedules x , for which we can compute the makespan T ( x ) with the proxy, given the runtimes { ρ i } ni =1 .
Revised: The final states represent full schedules x , for which we can compute the makespan T ( x ) with the proxy, given the runtimes { ρ i } ni =1 .

Original: From the makespan, we compute its speedup with respect to the makespan on a single device as U ( x )=
Revised: We compute the relative speedup compared to the makespan on a single device as U ( x )=

Original: (cid:80) i ρ i /T
Revised: (cid:80) i ρ i /T

Original: ( x ) , from which we compute the reward as we present in the next section.
Revised: ( x ) , from which we compute the reward as we present in the next section.

Original:  G ENERATIVE F LOW N ETWORKS FOR SCHEDULING
Revised:  G ENERATIVE F LOW N ETWORKS FOR SCHEDULING

Original: GFlowNets (Bengio et al., 2021a;b) are a method for training a stochastic policy to sample discrete and composite objects proportionally to a given unnormalized density.
Revised: GFlowNets (Bengio et al., 2021a;b) are methods for training a stochastic policy to sample discrete and composite objects proportionally to a given reward.

Original: Each object x is generated incrementally by a sequence of actions.
Revised: Each object x is generated incrementally by a sequence of actions.

Original: In Section 2.3, we have presented how valid schedules can be iteratively constructed by limiting the action space appropriately.
Revised: In the previous section, we discussed how to limit the action space to guarantee that we sample valid schedules.

Original: After a brief introduction of GFlowNets, the following sections will present a variety of novel extensionsof GFlowNets that underpin our method.
Revised: After a brief introduction to GFlowNets, the following sections will present our proposed extensions that include a new training objective that is suitable for learning conditional GFlowNets and a method for controlling the selectivity of samples at inference time.

Original: 3. B ACKGROUND
Revised: 3. B ACKGROUND

Original: We start by introducing some notation.
Revised: We start by introducing some notation.

Original: We denote by s =( s 0 , s 1 , . . .
Revised: We denote by s =( s 0 , s 1 , . . .

Original: , s n ) a trajectory made up ofintermediate states s t .
Revised: , s n ) a trajectory that consists of a sequence of states s t .

Original: In the case of scheduling, these correspond to partial schedules, withs 0 being the empty scheduleand s n a complete schedule.
Revised: In the case of scheduling, trajectories start with the empty schedule s 0 , followed by partial schedules, and end with a complete schedule s n .

Original: Wealso denote by T the set of all such trajectories, and by T x the set of trajectories that end at x .
Revised: We denote by T the set of all such trajectories and by T x the set of trajectories that end at x .

Original: We then define a flow function F : T → R +, and its associated normalized probability distribution P ( s ) =
Revised: Based on this, we define a flow function F : T → R + and its associated normalized probability distribution P ( s ) =

Original: F ( s ) /Z
Revised: F ( s ) /Z

Original: , Z = (cid:80) s ∈T F ( s ) .
Revised: , Z = (cid:80) s ∈T F ( s ) .

Original: By matching the trajectory flow of every terminal state x and its reward, i.e., R ( x ) = (cid:80) s ∈T x F ( s ) , the total trajectory distribution for all x becomes proportional to the reward, P (s ) ∝ R ( x ) for s ∈ T x and
Revised: A flow function that fulfills the condition: R ( x ) = (cid:80) s ∈T x F ( s ) (every terminal state has a total flow matching its reward), results in a probability over schedules P ( x ) =

Original: Z = (cid:80)  x R ( x ) .
Revised: (cid:80) s ∈T x F ( s ) /Z

Original: The training objectives considered in previous works aim to achieve a consistent flow (Bengio et al., 2021b; Malkin et al., 2022). While there are multiple ways todefine such an objective, we parametrize the two state transition probabilities P F and P B ,which we will call the forward and backward-sampling distributions respectively.
Revised: that is proportional to the reward P ( x ) ∝ R ( x ) , and further entails that Z = (cid:80) x R ( x ) .

Original: For any Markovian flow, we can decompose the probability of a trajectory in terms of the forward probability:
Revised: For any Markovian flow, we can decompose the probability of a trajectory in terms of the forward probability:

Original: [Equation]
Revised: [Equation]

Original: This way, we can generate trajectories s by sampling a sequence of actions starting from s 0 .
Revised: This way, we can generate trajectories s by sampling a sequence of actions starting from s 0 .

Revised: In Section 2.3 we described how to limit the action space appropriately to guarantee that every sampled schedule is valid.

Original: Similarly, the backward probability P B  factorizes the trajectory probability conditioned on a terminal state:
Revised: Similarly, we can define a backward probability P B that factorizes the trajectory probability conditioned on a terminal state:

Original: [Equation]
Revised: [Equation]

Revised: The training objectives considered in previous works aim to achieve a consistent flow (Bengio et al., 2021b; Malkin et al., 2022), where consistency means that the flow estimated for the forward direction should equal the flow for the backward direction.

Original: A consistent flow F ( s ) for trajectories s ∈ T x can then be written in terms of P F or P B and these two expressions must be equal:
Revised: A consistent flow F ( s ) for trajectories s ∈ T x can then be written in terms of P F and P B and has to fulfill the equality:

Original: [Equation]
Revised: [Equation]

Original: Based on this equation, Malkin et al.
Revised: Based on this equation, Malkin et al.

Original: (2022) propose to estimate Z , P F , and P B by optimizing the trajectory balance loss which is the squared difference between the logarithms of the l.h.s.
Revised: (2022) propose to estimate Z , P F , and P B by optimizing the trajectory balance loss which is the squared difference between the logarithms of the l.h.s.

Original: and the r.h.s.
Revised: and the r.h.s.

Original: of (6).
Revised: of Equation 6.

Original: 3. L OG - PARTITION VARIANCE LOSS
Revised: [Equation] - PARTITION VARIANCE LOSS

Original: Previous works on GFlowNets do not consider the conditional case, where the Z might differ depending on the context, and simply treat Z as a scalar parameter to optimize. Unfortunately, conditioning the flow network on the computation graph G C when using the trajectory balance loss described above requires learning a regression model to estimate log Z conditioned on G C . 
Revised: In order to apply the trajectory balance loss in the conditional case, we would need to learn an additional regression model that estimates the log-partition function log Z conditioned on G C . Training such a network accurately is difficult but important for learning the probabilities P F . In particular, a wrong estimation of log Z can incorrectly change the direction of the gradients of the loss function.

Original: Training such a network accurately is difficult and in  practice, we found this  to lead to non-convergent training behaviors.
Revised: We explain why this occurs in Appendix B. In practice, we found this approach to perform poorly when different computation graphs had large differences in their log Z value.

Original: Therefore, we propose to instead use a variance loss on the log-partition function. Namely, we can rewrite Equation 6 as:
Revised: Instead, we can rewrite Equation 6 to implicitly estimate log Z based on the forward and backward flows of a single trajectory s , where P F and P B are neural networks with parameters θ :

Original: [Equation]
Revised: [Equation]

Original:  ζ ( s  ) is a constant (equal to  log Z ) for all s associated to a single computation graph G C .
Revised: In the optimal case, ζ ( s ; θ ) is equal to the true log Z which is the same for all trajectories corresponding to the same computation graph G C .

Original: Our optimization goalnow turns into minimizing the variance of ζ ( s  ) over different trajectories s with the loss
Revised: Thus, our optimization goal turns into minimizing the variance of ζ ( s ; θ ) over different trajectories s with the loss

Original: [Equation]
Revised: [Equation]

Original: In practice, we use the training distribution to estimate E s [ ζ ( s )] with a mini-batch of  trajectories.
Revised: In practice, we use the training distribution to estimate E s [ ζ ( s )] with a mini-batch of sampled trajectories.

Original: We note that by optimizing the log-partition variance loss (8), one only needs to parametrize the forward and backward probabilities P F and P B .
Revised: We note that by optimizing the log-partition variance loss in Equation 8, one only needs to parametrize the forward and backward probabilities P F and P B .

Original: This is similar to the non-forward trajectory loss mentioned in the appendix by Malkin et al.
Revised: This is similar to the non-forward trajectory loss mentioned in the appendix by Malkin et al.

Original: (2022), which  does not involve learning any state flows, including the initial flow Z .
Revised: (2022), which also does not involve learning any state flows, including the initial flow Z .

Original: However, our loss does not mix forward and backward steps from different trajectories and directly optimizes the consistency of the total flow Z for each trajectory associated with a given computation graph G C .
Revised: However, our loss does not mix forward and backward steps from different trajectories and directly optimizes the consistency of the total flow Z for each trajectory associated with a given computation graph G C .

Original: 3. T EMPERATURE C ONDITIONED T OPOFORMER
Revised: 3. T EMPERATURE - CONDITIONED T OPOFORMER

Original: Reward Temperature.
Revised: Reward temperature.

Original: We compute the reward as a function of the speedup.
Revised: We compute the reward as a function of the speedup.

Original: In particular, we choose log R ( x ; m, σ ) =
Revised: In particular, we choose log R ( x ; m, σ ) =

Original: ( U ( x ) − m ) /σ where U ( x ) is the speedup of the schedule x , m is the number of devices, and σ ∈ R + plays the role of a temperature.
Revised: ( U ( x ) − m ) /σ where U ( x ) is the speedup of the schedule x , m is the number of devices, and σ ∈ R + plays the role of a temperature.

Original: The temperature allows us to concentrate the distribution on the modes and control the selectivity of the generator.
Revised: The temperature allows us to concentrate the distribution on the modes and control the selectivity of the generator.

Original: This is useful since there can be many more schedules with low speedup when compared to good ones.
Revised: This is useful since there can be many more schedules with low speedup when compared to good ones.

Original: For example, when simply setting the rewardR ( x ) equal to the speedup of x , we observed that finding schedules with high speedup requires a prohibitively large amount of samples.
Revised: For example, when simply setting the reward equal to the speedup, we observed that finding schedules with high speedup requires a prohibitively large amount of samples.

Original: We expect this temperature term to allow trade-offs between diversity and shifting the mean of the sampling distribution towards better schedules.
Revised: We expect this temperature term to allow trade-offs between diversity and shifting the mean of the sampling distribution towards better schedules.

Original: Previous works on GFlowNets apply a constant temperature value during training and at inference time (Bengio et al., 2021a; Jain et al., 2022).
Revised: Previous works on GFlowNets apply a constant temperature value during training and at inference time (Bengio et al., 2021a; Jain et al., 2022).

Original: This can lead to low performance (when set too high), and low diversity or unstable training (when set too low).
Revised: This can lead to low performance (when set too high), and low diversity or unstable training (when set too low).

Original: Furthermore, different computation graphs can have different ideal temperature values, making this approach less suitable when learning conditional GFlowNets.
Revised: Furthermore, different computation graphs can have different ideal temperature values, making this approach less suitable when learning conditional GFlowNets.

Original: Instead, we propose to learn a single model for multiple different reward functions R ( x ; m, σ ) , by conditioning the policy network on the temperature σ .
Revised: Instead, we propose to learn a single model for multiple different reward functions R ( x ; m, σ ) , by conditioning the policy networks ( P F and P B ) on the temperature σ .

Original: Approximating the  policy with a neural network is feasible because flows for a given reward can be continuously morphed into flows for any other reward function.
Revised: Approximating the temperature-conditioned policy with a neural network is feasible because flows for a given temperature can be continuously morphed into flows for any other temperature.

Original: Since our reward R ( x ; m, σ ) is continuous with respect to the temperature σ , we expect the change of flow for different temperatures to be learnable by a neural network.
Revised: Since our reward R ( x ; m, σ ) is continuous with respect to the temperature σ , we expect the change of flow for different temperatures to be learnable by a neural network.

Original: We provide a proof for the following theorem in Appendix A.
Revised: We provide a proof for the following theorem in Appendix A.

Original: Theorem 1 (Flow Continuity) .
Revised: Theorem 1 (Flow Continuity) .

Original: Let { R i } ∞
Revised: Let { R i } ∞

Original: i
Revised: i

Original: =1 be a sequence of positive reward functions such that for all terminal states x , ( R i ( x ) /R
Revised: =1 be a sequence of positive reward functions such that for all terminal states x , ( R i ( x ) /R

Original: ( x ))
Revised: ( x ))

Original: → 1 as i → ∞ .
Revised: → 1 as i → ∞ .

Original: Then, for any flow F R with reward R , there exists a sequence of flow functions { F R i } ∞ i =1 with F R i ( τ ) → F R ( τ ) for all τ ∈ T .
Revised: Then, for any flow F R with reward R , there exists a sequence of flow functions { F R i } ∞ i =1 with F R i ( s ) → F R ( s ) for all s ∈ T .

Original: The output policy changes more rapidly as a function of the temperature forvery values close tothan for larger values.
Revised: The output policy changes more rapidly as a function of the temperature for values close to 0 than for larger values.

Original: To account for this, we parametrize the temperature in terms of its logarithm.
Revised: To account for this, we use the logarithm of the temperature as input to the policy instead.

Original: During training, we sample temperatures from the log-uniform distribution with support between [log σ min , log σ max ] , where σ min is a minimum temperature that is necessary for numerical stability.
Revised: During training, we sample temperatures from the log-uniform distribution with support between [log σ min , log σ max ] , where σ min is a minimum temperature that is necessary for numerical stability.

Original: In comparison to sampling from U ( σ min , σ max ) , this avoids oversampling from high temperature regions that have little difference in the resulting network flow.
Revised: In comparison to sampling from U ( σ min , σ max ) , this avoids oversampling from high temperature regions that have little difference in the resulting flow network.

Original: At inference time, we choose how close the samples are to the mode by adjusting the σ .
Revised: At inference time, we choose how close the samples are to the mode by adjusting the σ .

Original: Finally, we note that sampling multiple temperatures at training time indeed helps recover good schedules at inference time, as opposed to training on a single fixed temperature.

Original: This is illustrated in Figure 4.

Original: Topoformer Architecture.
Revised: Topoformer architecture.

Original: For the  architecture of our GFlownet, we use the Topoformer (Gagrani et al., 2022), which has been recently introduced for learning topological orderings of computation graphs.
Revised: For the neural network architecture of our policy, we use the Topoformer (Gagrani et al., 2022), which has been recently introduced for learning topological orderings of computation graphs.

Original: It builds on the Transformer encoder architecture (Vaswani et al., 2017) and additionally masks the multi-head attention depending on the topology of the computation graph.
Revised: It builds on the Transformer encoder architecture (Vaswani et al., 2017) and additionally masks the multi-head attention depending on the topology of the computation graph.

Original: Both forward and backward policies use separate MLP heads on top of a shared Topoformer encoder.
Revised: Both forward and backward policies use separate MLP heads on top of a shared Topoformer encoder.

Original: Taking inspiration from the successful use of time conditioning in diffusion models (Song et al., 2020; Ho et al., 2020), wealso add temperature conditioning by first embedding the temperature using an MLP to produce e σ , and then reuse the embedding in every first linear layer block of the Topoformer:[Equation]
Revised: Taking inspiration from the successful use of time conditioning in diffusion models (Song et al., 2020; Ho et al., 2020), we add temperature conditioning by first embedding the temperature using an MLP to produce e σ , and then reuse the embedding in every first linear layer block of the Topoformer:

Original: [Equation]
Revised: [Equation]

Original: Here lin scale and lin shift are linear layers and ⊙ is the elementwise multiplication (Perez et al., 2018).
Revised: Here lin scale and lin shift are linear layers and ⊙ is the elementwise multiplication (Perez et al., 2018).

Original: In contrast to diffusion models, we observe better performance on large temperature ranges with the ReLU  activation function.
Revised: In contrast to diffusion models, we observe better performance on large temperature ranges with the ReLU (Nair & Hinton, 2010) activation function.

Original: We hypothesize that this is connected to the monotonicity of the underlying policy function with respect to decreasing temperatures (see the appendix, Corollary 1) and the propensity for linear extrapolation of ReLU MLPs (Xu et al., 2020).
Revised: We hypothesize that this is connected to the monotonicity of the underlying policy function with respect to decreasing temperatures (see the appendix, Corollary 1) and the propensity for linear extrapolation of ReLU MLPs (Xu et al., 2020).

Original: Sub-graph training.
Revised: Sub-graph training.

Original: Training with a full computation graph might not always be necessary and we hypothesize that learning on sub-graphs can lead to policies that generalize to the full computation graph.
Revised: Training with a full computation graph might not always be necessary and we hypothesize that learning on sub-graphs can lead to policies that generalize to the full computation graph.

Original: This can be seen as a form of data augmentation and increases the amount of training data.
Revised: This can be seen as a form of data augmentation and increases the amount of training data, while simultaneously improving the training time.

Original: We shall use sub-graph training for the larger graphs that we study in this paper.
Revised: We shall use sub-graph training for the larger graphs that we study in this work.

Original:  R ELATED WORK 
Revised:  R ELATED WORK Reinforcement learning for scheduling.

Original: Reinforcement learning for scheduling.

Original: Reinforcement learning has been the predominant machine learning approach to optimize the makespan for computation graph schedules (Addanki et al., 2019; Paliwal et al., 2020; Zhang et al., 2020).
Revised: Reinforcement learning has been the predominant machine learning approach to optimize the makespan for computation graph schedules (Addanki et al., 2019; Paliwal et al., 2020; Zhang et al., 2020).

Original: The rewards used include simple  proxies of the makespan (Paliwal et al., 2020;
Revised: The rewards used include simple analytical proxies of the makespan (Paliwal et al., 2020;

Original: Zhou et al., 2020), but also more refined proxies which incorporate modeling of memory movements (Addanki et al., 2019).
Revised: Zhou et al., 2020), but also more refined proxies which incorporate modeling of memory movements (Addanki et al., 2019).

Original: Khadka et al.
Revised: Khadka et al.

Original: (2021) directly train on the target hardware, but consider only  few computation graphs, and do not show generalization to unseen ones.
Revised: (2021) directly train on the target hardware, but consider only a few computation graphs, and do not show generalization to unseen ones.

Original: Addanki et al.
Revised: Addanki et al.

Original: (2019) use a sophisticated simulator of the makespan which is customized to the target hardware.
Revised: (2019) use a sophisticated simulator of the makespan which is customized to the target hardware.

Original: Similar to our work, Zhang et al.
Revised: Similar to our work, Zhang et al.

Original: (2020) also construct the schedule pieceby-piece.
Revised: (2020) also construct the schedule piece-by-piece.

Original: Instead of finding a single (local) mode of the proxy, our work proposes to learn the full distribution over the proxy to improve the robustness against inaccuracies in the proxy.
Revised: Instead of finding a single (local) mode of the proxy, our work proposes to learn the full distribution over the proxy to improve the robustness against inaccuracies in the proxy.

Original: Generative Flow Networks.
Revised: Generative Flow Networks.

Original: GFlowNets have been applied to generating small molecules (Bengio et al., 2021a), Bayesian networks (Deleu et al., 2022), discrete images (Zhang et al., 2022), and biological sequences (Jain et al., 2022).
Revised: GFlowNets have been applied to generating small molecules (Bengio et al., 2021a), Bayesian networks (Deleu et al., 2022), discrete images (Zhang et al., 2022), and biological sequences (Jain et al., 2022).

Original: We extend its application to scheduling, a classical combinatorial optimization problem.
Revised: We extend its application to scheduling, a classical combinatorial optimization problem.

Revised: Similar to previous works, our state-action space is also a DAG, hence training the policy with maximum entropy reinforcement learning methods is inadequate (Bengio et al., 2021a).

Revised: Our robust scheduling approach shares the same motivation as methods in drug discovery which leverage cheap proxies to generate multiple candidates to be evaluated in the true environment (Bengio et al., 2021a; Jain et al., 2022).

Original: Conditional GFlowNets have previously only been theoretically discussed by (Bengio et al., 2021b).
Revised: Conditional GFlowNets have previously only been theoretically discussed by Bengio et al. (2021b).

Original: We demonstrate them in practice with the log-partition variance loss, and empirically demonstrate generalization to unseen computation graphs.
Revised: We enable training conditional GFlowNets with our proposed log-partition variance loss and empirically demonstrate generalization to unseen computation graphs.

Revised: Note that this differs from previous work that tests the generalization of

Original:  To control the selectiveness of the generator, previous works augment the reward with a fixed temperature (Bengio et al., 2021a; Deleu et al., 2022; Jain et al., 2022).
Revised: GFlowNets to unseen data (Nica et al., 2022). To control the selectiveness of the generator, previous works augment the reward with a fixed temperature (Bengio et al., 2021a; Deleu et al., 2022; Jain et al., 2022).

Original: Instead, we condition the policy neural network on the temperature term which allows us to tune the degree of diversity of the generated schedules at inference time.
Revised: Instead, we condition the policy neural network on the temperature term which allows us to tune the selectiveness of the generator at inference time.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we evaluate different aspects of our generative scheduling approach by incrementally adding complexity to the computation graph dataset.
Revised: In this section, we evaluate different aspects of our generative scheduling approach by incrementally adding complexity to the computation graph dataset.

Original: First, we restrict training and evaluation to a single computation graph, which corresponds to the same unconditional setting considered by previous works on GFlowNets (Bengio et al., 2021a; Deleu et al., 2022; Jain et al., 2022).
Revised: First, we restrict training and evaluation to a single computation graph, which corresponds to the same unconditional setting considered by previous works on GFlowNets (Bengio et al., 2021a; Deleu et al., 2022; Jain et al., 2022).

Original: Next, we train with multiple computation graphs and evaluate on unseen ones.
Revised: Next, we train with multiple computation graphs and evaluate on unseen ones.

Original: To the best of our knowledge, this is the first time that  generalization of  GFlowNets  is tested empirically.
Revised: To the best of our knowledge, this is the first time that the generalization of conditional GFlowNets to unseen conditioning is tested empirically.

Original: Finally, we verify the generalization ability on real-world computation graphs of neural networks that are being used in a diverse set of AI products.
Revised: Finally, we verify the generalization ability on real-world computation graphs of neural networks that are being used in a diverse set of AI products.

Original: Experimental setup.
Revised: Experimental setup.

Original: In all experiments, we only use the node time duration as a feature of the computation graph.
Revised: In all experiments, we only use the node time duration as a feature of the computation graph.

Original: For simplicity and ease of reproducibility, we avoid any complicated heuristics to add extra features.
Revised: For simplicity and ease of reproducibility, we avoid any complicated heuristics to add extra features.

Original: All our experiments are based on four homogenous devices, i.e. any device can execute any operation.
Revised: All our experiments are based on four homogenous devices.

Original: Candidate sampler.
Revised: Candidate sampler.

Original: We consider two heuristic and two neural methods for generating candidate schedules.
Revised: We consider two heuristic and two neural methods for generating candidate schedules.

Original: The first is our GFlowNet approach described in Section 3 from which we generatesamples at temperature σ = 0 .
Revised: The first is our GFlowNet approach described in Section 3 from which we generatesamples at temperature σ = 0 .

Original: 005 and take the top 100 following the proxy; the other three are:
Revised: 005 and take the top 100 following the proxy; the other three are:

Original: • Critical path-based list scheduling, a greedy algorithm for scheduling on homogeneous devices (Micheli, 1994).
Revised: • Critical path-based list scheduling, a heuristic algorithm for scheduling on homogeneous devices (Micheli, 1994).

Original: The algorithm takes as input an ordered list of jobs, which can be generated by any convenient sorting method.
Revised: List scheduling first forms a topological order of the operations and then assigns them in that order one by one to a free device.

Original: In our implementation, we use the Critical Path method (Micheli, 1994), which assigns the highest priority to jobs that are on the DAG path (from source to sink) with the highest sum of runtimes.
Revised: In our implementation, we use the Critical Path method (Micheli, 1994) to order the operations. It ensures that operations on the time critical path are scheduled first.

Original: This method produces a single schedule.
Revised: This method produces a single schedule.

Original: • Biased Random Key Genetic Algorithm (BRKGA) (Gonc¸alves & Resende, 2011), a genetic algorithm that has previously shown good performance on scheduling tasks (Paliwal et al., 2020).
Revised: • Biased Random Key Genetic Algorithm (BRKGA) (Gonc¸alves & Resende, 2011), a genetic algorithm that has previously shown good performance on scheduling tasks (Paliwal et al., 2020).

Original: We use the top-100 schedules from the final population as the candidate schedules.
Revised: We use the top 100 schedules from the final population as the candidate schedules.

Original: • Proximal Policy Optimization (PPO) (Schulman et al., 2017), a deep reinforcement learning method that has been successfully applied to scheduling problems Zhou et al.
Revised: • Proximal Policy Optimization (PPO) (Schulman et al., 2017), a deep reinforcement learning method that has been successfully applied to scheduling problems (Zhou et al., 2020).

Original: PPO also trains a stochastic policy, which makes it a natural choice for comparison with GFlowNets (Bengio et al., 2021a).
Revised: PPO also trains a stochastic policy, which makes it a natural choice for comparison with GFlowNets (Bengio et al., 2021a).

Original: We employ the same definitions of states, actions
Revised: We employ the same definitions of states, actions

Original: and reward function (with temperature σ = 0 .
Revised: and reward function (with temperature σ = 0 .

Original: 25 ; lower was not beneficial) as the GFlowNet approach.
Revised: 25 ; lower was not beneficial) as the GFlowNet approach.

Original: To ensure that PPO keeps exploring even after finding a local optimum we employ entropy regularization, anddecay both the entropy.
Revised: To ensure that PPO keeps exploring even after finding a local optimum we employ entropy regularization and

Original:  coefficient (Ahmed et al., 2019) and the learning rate to ensure convergence to a good solution.
Revised: decay both the entropy coefficient (Ahmed et al., 2019) and the learning rate to ensure convergence to a good solution.

Original: Same as for GFlowNets, we sample 1000 schedules and pick the top 100 as the candidate schedules Metrics.
Revised: Same as for GFlowNets, we sample 1000 schedules and pick the top 100 as the candidate schedules.

Original: We measure
Revised: Metrics.

Original:  the performance in terms of the speedup U ( x ) .
Revised: We measure the performance in terms of the speedup U ( x ) .

Original: For the diversity, we report three different measures: graph-edit distance (GED), the L2 distance between the proxy start-times ( d inv ), and the L2 distance between the proxy start-times concatenated with the device placement ( d sen ).
Revised: For the diversity, we report three different measures: graph-edit distance (GED), the L2 distance between the proxy start-times ( d inv ), and the L2 distance between the proxy start-times concatenated with the device placement ( d sen ).

Original: See Appendix B.2 for more details on distance computation.
Revised: See Appendix C.2 for more details on diversity measures.

Original: 5. P ROXY ERRORS : DIVERSITY FOR ROBUST SCHEDULING
Revised: 5. P ROXY ERRORS : DIVERSITY FOR ROBUST SCHEDULING

Original: [Figure]
Revised: [Figure]

Original: We examine how differences between the proxy and the target performance model can affect the final runtime.
Revised: We examine how differences between the proxy and the target performance model can affect the final runtime.

Original: To do so, we first focus on a single computation graph that is used both for training and testing, to avoid possible confounding factors that may happen in the generalization setting.
Revised: To do so, we first focus on a single computation graph that is used both for training and testing to avoid possible confounding factors that may happen in the generalization setting.

Original: Based on the possible reasons for proxy errors discussed in Section 2.2 we design three different target models that each reflect a different setting.
Revised: Based on the possible reasons for proxy errors discussed in Section 2.2 we design three different target models that each reflect a different setting.

Original: In the first setting node durations are incorrectly profiled (Noisy Runtimes).
Revised: In the first setting node durations are incorrectly profiled (Noisy Runtimes).

Original: In the second and third settings, the target models the memory movement across devices with a linear model (Culler et al., 1999), which can be either bottlenecked by limited bandwidth (Bandwidth Limited), or by high latency (Latency Limited). 
Revised: In the second and third settings, the target models the memory movement across devices with a linear model (Valiant, 1990), which can be either bottlenecked by limited bandwidth (Bandwidth Limited), or by high latency (Latency Limited). The linear model has been shown to be a good makespan estimator for certain devices (Hockney, 1994; Culler et al., 1993). We refer to Appendix C.3 for more details.

Original: In Figure 2, we show the correlation between the proxy and the different  target environments.
Revised: In Figure 2, we show the correlation between the proxy and the different [Table] target environments.

Original: For all three targets, the proxy is highly correlated but can have target speedups that differ by a factor of up to × 2 for the schedules with high proxy speedups.
Revised: For all three targets, the proxy is highly correlated but can have target speedups that differ by a factor of up to × 2 for the schedules with high proxy speedups.

Original: We report the speedups and diversity measures in Table 1.
Revised: We report the speedups and diversity measures in Table 1.

Original: The results highlight that any method that can generate multiple good candidate schedules achieves higher speedups than list scheduling,[Table] which only produces a single unique schedule.
Revised: The results highlight that any method that can generate multiple good candidate schedules achieves higher speedups than list scheduling, which only produces a single unique schedule.

Original: Furthermore, if the candidate schedules are diverse, as is the case for GFlowNet [Table] the target performance is also better on average.
Revised: Furthermore, if the candidate schedules are more diverse — as is the case for GFlowNet — the target performance is also better on average.

Original: The results confirm our hypothesis that a diverse set of candidate schedules can improve robustness towards a misspecified proxy [Table]
Revised: The results confirm our hypothesis that a diverse set of candidate schedules can improve robustness towards a misspecified proxy.

Original: 5. G ENERALIZING TO UNSEEN COMPUTATION GRAPHS
Revised: 5. G ENERALIZING TO UNSEEN COMPUTATION GRAPHS

Original: Next, we evaluate how well our conditional GFlowNet can generalize to unseen computation graphs.
Revised: Next, we evaluate how well our conditional GFlowNet can generalize to unseen computation graphs.

Original: We train and evaluate on a diverse set of synthetic computation graphs sampled from different random graph distributions.
Revised: We train and evaluate on a diverse set of synthetic computation graphs sampled from different random graph distributions.

Original: In particular, we train on graphs of size 50 sampled from the random graph distributions (a) Erd ˝ os–R ´ enyi (Erd ˝ os et al., 1960), and (b) Layered Graphs (Gagrani et al., 2022) and evaluate, in addition to (a) and (b), on stochastic block model (Holland et al., 1983), WattsStrogatz (Watts & Strogatz, 1998), and Barab ´ asi–Albert (Albert & Barab ´ asi, 2002).
Revised: In particular, we train on graphs of size 50 sampled from the random graph distributions (a) Erd ˝ os–R ´ enyi (Erd ˝ os et al., 1960), and (b) Layered Graphs (Gagrani et al., 2022) and evaluate, in addition to (a) and (b), on stochastic block model (Holland et al., 1983), WattsStrogatz (Watts & Strogatz, 1998), and Barab ´ asi–Albert (Albert & Barab ´ asi, 2002).

Original: The generative process of the computation graphs can be found in Appendix B.4.
Revised: For details on the generative process of the computation graphs, we refer to Appendix C.5.

Original: In Table 2, we demonstrate that the conditional GFlowNet is able to generalize to previously unseen computation graphs, regardless of whether they originate from the same random graph distribution.
Revised: In Table 2, we demonstrate that the conditional GFlowNet is able to generalize to previously unseen computation graphs, regardless of whether they originate from the same random graph distribution.

Original: Furthermore, our proposed temperature conditioning method allows us to trade off better average performance for more diversity.
Revised: Next, we ablate our proposed temperature conditioning method by generating 1000 samples at different temperature values.

Original: In Figure 3, we observe that decreasing the temperature does indeed shift the sample distribution to the right and also sharpens it when the temperature approaches zero.
Revised: In Figure 3, we observe that decreasing the temperature does indeed shift the sample distribution to the right and also sharpens it when the temperature approaches zero.

Original: Notably, the temperature σ = 0 .
Revised: Notably, the temperature σ = 0 .

Original: 005 is not in the training distribution, which demonstrates that the model can extrapolate to temperature values outside of the training range.
Revised: 005 is not in the training distribution, which demonstrates that the model can extrapolate to temperature values outside of the training range.

Original: Finally, we have observed that training on variable temperatures (as opposed to a fixed one) indeed improves performance at inference time, as can be seen, in Figure 4.
Revised: Surprisingly, we observe that training with a variable temperature can improve the performance further than is possible with a fixed temperature, which we demonstrate in Appendix D. [Table] [Table] [Figure]

Original: 5. R EAL WORLD COMPUTATION GRAPHS
Revised: 5. R EAL WORLD COMPUTATION GRAPHS

Original: Finally, we verify the generalization ability on a small set of real-world computation graphs used for the commercial development of our artificial intelligence hardware and software products (see Appendix B.5 for details on their number and size).
Revised: Finally [Figure] we verify the generalization ability on a small set of real-world computation graphs used for the commercial development of our artificial intelligence hardware and software products (see Appendix C.6 for details).

Original: This time we also report the speedup on the same target models used in Section 5.1 to assess robustness on a real graph set.
Revised: We report the speedup on the same target models used in Section 5.1 to assess robustness on unseen real-world computation graphs.

Original: To speed up training,
Revised: To speed up training

Original: we

Original: [Table]
Revised: [Figure]

Original: [Table] [Table] [Table] [Table]
Revised: we apply the graph subsampling strategy presented in Section 3.3 to randomly pick between 25 to 75 nodes at every training step.

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table] [Table] [Table] [Table]
Revised: In Table

Original: [Figure]
Revised: [Figure]

Original: apply the graph subsampling strategy presented in Section 3.3 to randomly pick between 25 tonodes for training.

Original: In Table 3, we observe that the conditional GFlowNet retains the benefits of high diversity and robustness to misspecifications in the proxy even when applied to graphs not seen during training and of larger sizes.
Revised: we observe that the conditional GFlowNet retains the benefits of high diversity and robustness to misspecifications in the proxy even when applied to graphs not seen during training and of larger sizes.

Original: PPO shows unstable training behavior and the reward training curve does not converge, despite using the same hyperparameters that worked for the previous two experiments.
Revised: PPO shows unstable training behavior and the reward training curve does not converge, despite using the same hyperparameters that worked for the previous two experiments.

Original: We conjecture that this is due to the inhomogeneous maximum possible speedup of the training graphs that lead to different reward scales per training graph.
Revised: We conjecture that this is due to the inhomogeneous maximum possible speedup of the training graphs that lead to different reward scales per training graph.

Original: In comparison, GFlowNet still trains as before without any changes to the hyperparameters.
Revised: In comparison, GFlowNet still converges as before without any changes to the hyperparameters.

Original: [Figure]
Revised: We ablate our proposed log-partition variance loss by comparing it against the trajectory balance loss that uses a Topoformer to predict log Z given a computation graph. Learning such a model is difficult due to large differences in the output space of different computation graphs that arise from the differences in the number of nodes, which in turn impedes the training progress of the policy network. We confirm in Appendix C.6 that our proposed loss function remedies the slow start problem of the baseline and achieves a higher speedup in the end.

Original: [Table]
Revised: [Table]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We have empirically demonstrated how the conventional optimization approach to scheduling, which optimizes a proxy of the real makespan, is brittle to modeling failures in the proxy itself.
Revised: We have empirically demonstrated how the conventional optimization approach to scheduling, which optimizes a proxy of the real makespan, is brittle to modeling failures in the proxy itself.

Original: Our proposed approach evaluates multiple schedules on the target and thereby achieves more robustness to discrepancies between the proxy and the target.
Revised: Our proposed approach evaluates multiple schedules on the target and thereby achieves more robustness to discrepancies between the proxy and the target.

Original: We demonstrated that GFlownets can sample a diverse set of candidate schedules that achieve better target performance than alternative methods which achieve lower diversity.
Revised: We demonstrated that GFlownets can sample a diverse set of candidate schedules that achieve better target performance than alternative methods which achieve lower diversity.

Original: Further, we have shown that conditioning on temperature allows a trade-off between diversity and proxy performance, and that conditional GFlowNets can generalize to unseen computation graphs.
Revised: Further, we showed that conditioning on temperature allows a trade-off between diversity and proxy performance, and that conditional GFlowNets can generalize to unseen computation graphs.

Original: Interesting future directions include scaling up our method to larger graphs and integrating scheduling heuristics to speed up training.
Revised: Interesting future directions include scaling up our method to larger graphs and integrating scheduling heuristics to speed up training.

Original: E THICS S TATEMENT
Revised: E THICS S TATEMENT

Original: Our work is concerned with the optimization of compute schedules on hardware devices.
Revised: Our work is concerned with the optimization of compute schedules on hardware devices.

Original: Since this process is at the moment carried out via tailored heuristics and not by human operators, we cannot foresee any negative societal impact of our work.
Revised: Since this process is at the moment carried out via tailored heuristics and not by human operators, we cannot foresee any negative societal impact of our work.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: In the main text and appendix, we provide comprehensive information about our GFlowNet scheme, the NN architecture employed therein[Table] and the training hyperparameters we use.
Revised: In the main text and appendix we provide comprehensive information about our GFlowNet scheme, the NN architecture employed therein and the training hyperparameters we use.

Original: We believe that this allows any interested reader to reproduce the results in Table 1 and Table 2.
Revised: We believe that this allows any interested reader to reproduce the results in Table 1 and Table 2.

Original: We shall also release our code if the paper is accepted.
Revised: We shall also release our code if the paper is accepted

Original: A P ROOF FOR T HEOREM
Revised: A P ROOF FOR T HEOREM

Original: In the following, we will denote F R to be the flow corresponding to a flow function F : T → R + with corresponding reward R .
Revised: In the following, we will denote F R to be the flow corresponding to a flow function F : T → R ≥with corresponding reward R .

Original: That is, F R : T → R + is a flow function, such that for any terminal state x∈ S we have  F R (  x  )
Revised: That is, F R : T → R ≥ 0 is a flow function such that for any terminal state x we have R ( x ) = (cid:80) s ∈T x F R ( s ) , where (cid:80) s ∈T x F R ( s )

Original: =
Revised: =:

Original:  R ( x ) , where F ( s ′ )
Revised: F R ( x ) is the total flow at the terminal state x .

Original: = (cid:80) s → s ′ ∈ τ F ( τ ) is the sum of trajectory flow values over trajectories containing the transition from an arbitrary state s into s ′ .

Original: Theorem (Flow Continuity) .
Revised: Theorem (Flow Continuity) .

Original: Let { R i } ∞
Revised: Let { R i } ∞

Original: i
Revised: i

Original: =1 be a sequence of positive reward functions such that  ( R i ( x ) /R( x ))
Revised: =1 be a sequence of positive reward functions such that for all terminal states x , ( R i ( x ) /R

Original:  → 1 for all (terminal) states x ∈ S .
Revised: ( x )) → 1 as i → ∞ .

Original: Then, for any flow F R with reward R , there exists a sequence of flow functions { F R i } ∞ i =1 with F R i ( τ ) → F R ( τ ) for all τ ∈ T .
Revised: Then, for any flow F R with reward R , there exists a sequence of flow functions { F R i } ∞ i =1 with F R i ( s ) → F R ( s ) for all s ∈ T .

Original: Proof.
Revised: Proof.

Original: Let { x i } Mi =1 be the set of terminal states and define
Revised: Let { x i } Mi =1 be the set of terminal states and define

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: To see that F R i is a valid flow for R i , we note that F R i ≥ 0 and
Revised: To see that F R i is a valid flow for R i , we note that F R i ≥ 0 and

Original: [Equation]
Revised: [Equation]

Original: R i ( x [Equation] (11) for any terminal state x∈ S .
Revised: 11) for any terminal state x .

Original: Finally, for any τ ∈ T with terminal state x we have
Revised: Finally, for any s ∈ T with terminal state x we have

Original: [Equation] (12) ■
Revised: [Equation] [Equation]

Original: Corollary 1.
Revised: Corollary 1.

Original: Let R σ i :=
Revised: Let R σ i :=

Original: log R ( x ; m, σ i ) =
Revised: log R ( x ; m, σ i ) =

Original: ( U ( x ) − m ) /σ
Revised: ( U ( x ) − m ) /σ

Original: i be a sequence of temperature conditioned reward functions with σ i ↘ σ 0 .
Revised: i be a sequence of temperature conditioned reward functions with σ i ↘ σ 0 .

Original: Then, for any ϵ > 0 and flow F R σ 0 there exists a neighborhood ( σ 0 , σ 0 + δ ) containing flow functions F R σ with F R σ ( τ ) − F R σ 0 ( τ ) <
Revised: Then, for any ϵ > 0 and flow F R σ 0 there exists a neighborhood ( σ 0 , σ 0 + δ ) containing flow functions F R σ with F R σ ( s ) − F R σ 0 ( s ) <

Original: ϵ ∀ τ ∈ T .
Revised: ϵ for all s ∈ T .

Original: Furthermore, we can approximate F R σ 0 with monotonically changing flow functions F R σi .
Revised: Furthermore, we can approximate F R σ 0 with monotonically changing flow functions F R σi .

Revised: B F AILURE MODES WHEN ESTIMATING THE LOG - PARTITION FUNCTION

Revised: In this section, we present the different cases in which an incorrect estimate of log Z can lead to gradients that point in the opposite direction of the “true gradients”.

Revised: We start by rewriting the trajectory balance loss (Malkin et al., 2022) as:

Revised: [Equation]

Original: B E XPERIMENT DETAILS B. C ANDIDATE SAMPLERS
Revised: For better readability, we do not explicitly write out the dependence on the computation graph G C for the trajectories, their probabilities, and the reward. We can rearrange the terms in Equation 13, such that: [Equation] In words, φ F is the log probability of the trajectory s as computed by the forward-sampling distribution, and φ RBZ can also be viewed as the log probability of s but instead estimated by a combination of the reward, log-partition function, and the backward-sampling distribution.

Revised: To simplify the argument, we assume that the backward-sampling distribution is independent of θ , which can be realized for example by setting it equal to the distribution that assigns equal probability to all parents.

Revised: Thus, we denote the true log-partition function as log Z ∗ ( G C ) and by assuming access to it when computing the loss we can rewrite part of Equation 14 as:

Revised: [Equation]

Revised: We denote the loss using φ RBZ ∗ as L ∗ TB ( s ; θ ) .

Revised: There exist two cases in which L ∗ TB ( s ; θ ) can be non-zero and in both cases, inaccurate estimation of the log-partition function can lead to wrong gradients.

Revised: In the first case, we assume that the neural network that estimates P F is overestimating the probability of s :

Revised: [Equation]

Revised: Furthermore, we consider the scenario in which the regression model makes an error when predicting the log-partition value.

Revised: In particular, we assume that the neural network is underestimating the true value by an amount that is greater than the error φ F − φ RBZ ∗ , i.e., we assume:

Revised: [Equation]

Revised: The implications of this are:

Revised: [Equation]

Revised: Equation 16 and Equation 21 imply that −∇ θ L TB ( s ; θ ) is an ascent direction for L ∗ TB ( s ; θ ) , instead of the desired descent direction.

Original: We use the popular open-source library pymoo (Blank & Deb, 2020) implementation.
Revised: In detail, for non-zero gradients the inner product is positive: [Equation] In the second case, φ F − φ RBZ ∗ < 0 , we assume that the neural network overestimates the logpartition function. Analog to the first case we can show that the gradients point in the opposite direction of the true gradients.

Original: We implement the list scheduling algorithm with critical path priorities.
Revised: These two cases happen more frequently when log Z ( G C ; θ ) is more challenging to regress accurately. For example, when there are computation graphs with drastically different node numbers, their log Z ( G C ; θ ) values can have significant differences (e.g., 10s vs. 1000s). Training a regression model on such large ranges is notoriously difficult (Pohlen et al., 2018). Our proposed log-partition variance loss in Equation 8 avoids this regression task completely, and we confirm its benefits empirically in Appendix C.6.

Revised: C E XPERIMENT DETAILS C. C ANDIDATE SAMPLERS

Revised: We use the popular open-source library pymoo (Blank & Deb, 2020) to implement the BRKGA candidate sampler.

Original: Our PPO implementation is based on algorithm 1 at https://spinningup.openai.com/en/latest/algorithms/ppo.html , and we follow Schulman et al.
Revised: Our PPO implementation is based on algorithm 1 at https://spinningup. openai.com/en/latest/algorithms/ppo.html , and we follow Schulman et al.

Original: (2017) to implement the entropy regularisation by adding the entropy term directly to the PPO-clip loss.
Revised: (2017) to implement the entropy regularisation by adding the entropy term directly to the PPO-clip loss.

Revised: We decay this entropy term during training similar to Ahmed et al.

Original: We use the same learning rate for both the actor and the critic, and we decay it with an exponential schedule.
Revised: We use the same learning rate for both the actor and the critic, and we decay it with an exponential schedule.

Original: We train GFlowNets conditioned on a temperature randomly sampled between 0.01 and 1.
Revised: We train GFlowNets conditioned on a temperature randomly sampled between 0.01 and 1.

Original: At inference, we use 0.005 for the temperature in all experiments.
Revised: At inference, we use 0.005 for the temperature in all experiments.

Revised: We use the Adam optimizer with default hyperparameters to optimize the parameters.

Revised: We compute the gradients at each update step based on a minibatch that consists of 100 sampled trajectories for a single computation graph and use a single temperature value to compute their rewards.

Revised: We observed no benefits in our initial experiments that used multiple computation graphs or temperatures in a single minibatch.

Revised: We represent the computation graph as a directed graph with a single node feature (runtime of the operation).

Revised: In addition, we represent different states by concatenating a state feature vector to each node, which consists of:

Original: B. M ETRICS
Revised: • Start time assigned by the proxy (default: − 1 ) • Device placement as a one-hot vector (default: 0 vector) • Binary indicator: 1 if adding it to the schedule is a valid action, else• Binary indicator: 1 if removing it is a valid backward-action, else• Binary indicator: 1 if the node is already part of the schedule, else The combined computation graph and state features are treated as a single graph by the Topoformer neural network. C. M ETRICS

Original: The graph-edit distance (GED) compares two schedules in their chain-graph form.
Revised: The graph-edit distance (GED) compares two schedules in their chain-graphs form.

Original: In particular, we can model a schedule for a computation graph, by constructing a chain graph for each device that specifies the additional precedence constraints we introduce to complete the order in which the operations are run on each device.
Revised: In particular, we can model a schedule for a computation graph, by constructing a chain graph for each device that specifies the additional precedence constraints we introduce to complete the order in which the operations are run on each device.

Original: The GED is then computed simply by taking the difference between the adjacency matrices and normalizing it by the total number of edges.
Revised: The GED is then computed simply by taking the difference between the adjacency matrices and normalizing it by the total number of edges.

Original: The L2 distance between the start ( d inv ) times simply takes the start times as assigned by the proxy model and computes the L2 norm of the difference.
Revised: The L2 distance between the start ( d inv ) times simply takes the start times as assigned by the proxy model and computes the L2 norm of the difference.

Original: The L2 distance including the device assignment ( d sen ) additionally concatenates the device placement to the times.
Revised: The L2 distance including the device assignment ( d sen ) additionally concatenates the device placement to the times.

Revised: C. T ARGET ENVIRONMENT WITH LINEAR MEMORY MODEL

Revised: The linear memory model (Valiant, 1990) computes the delay as a linear function f ( m )

Revised: = am + b of the memory m with a modeling the amount of data that can be transferred per time and b modeling the startup delay.

Revised: In the Bandwidth Limited setting the a term dominates the delay, while in the

Revised: Latency Limited setting b has a greater effect.

Original: B. P ROXY ERRORS : D IVERSITY FOR ROBUST SCHEDULING
Revised: C. P ROXY ERRORS : D IVERSITY FOR ROBUST SCHEDULING

Original: In this experiment, we consider a single real-world graph that has around 78 nodes.
Revised: In this experiment, we consider a single real-world graph that has around 78 nodes.

Original: B. G ENERALIZATION TO UNSEEN COMPUTATION GRAPHS
Revised: C. G ENERALIZATION TO UNSEEN COMPUTATION GRAPHS

Original: We generate the synthetic graph dataset from random graph distributions over undirected graphs.
Revised: We generate the synthetic graph dataset from random graph distributions over undirected graphs.

Original: To get DAGs from these graphs, we randomly choose a direction for every edge in a way that produces no cycles.
Revised: To get DAGs from these graphs, we randomly choose a direction for every edge in a way that produces no cycles.

Original: We sample the runtimes for each node from the uniform distribution U (0 , 1) .
Revised: We sample the runtimes for each node from the uniform distribution U (0 , 1) .

Original: For training, we use 1000 different computation graphs, with equally many sampled from the two random graph distributions: Erd ˝ os–R ´ enyi (Erd ˝ os et al., 1960), and Layered Graphs (Gagrani et al., 2022).
Revised: For training, we use 1000 different computation graphs, with equally many sampled from the two random graph distributions: Erd˝os–R´enyi (Erd˝os et al., 1960), and Layered Graphs (Gagrani et al., 2022).

Original: We report test performances on 50 different computation graphs with equally many sampled from the five different random graph distributions: Erd ˝ os–R ´ enyi (Erd ˝ os et al., 1960), Layered
Revised: We report test performances on 50 different computation graphs with equally many sampled from the five different random graph distributions: Erd˝os–R´enyi (Erd˝os et al., 1960), Layered

Original: Graphs (Gagrani et al., 2022), stochastic block model (Holland et al., 1983), Watts-Strogatz (Watts & Strogatz, 1998), and Barab ´ asi–Albert (Albert & Barab ´ asi, 2002). 
Revised: Graphs (Gagrani et al., 2022), stochastic block model (Holland et al., 1983), Watts-Strogatz (Watts & Strogatz, 1998), and Barab´asi–Albert (Albert & Barab´asi, 2002). [Figure] A v e r a g e Sp ee d [Figure]

Original: B. R EAL W ORLD C OMPUTATION G RAPHS
Revised: C. R EAL - WORLD COMPUTATION GRAPHS

Revised: The computation graphs in this dataset originate from a diverse set of neural network architectures with different applications, including for example classification and denoising.

Original: We train on 8 real-world computation graphs of sizes below 100 nodes and evaluate on 4 different computation graphs of sizes between a dozen andbelow 128 nodes.
Revised: We train on 8 realworld computation graphs of sizes below 100 nodes and evaluate on 4 different computation graphs of sizes between a dozen and 128 nodes.

Revised: In Figure 4, we report the average speedup of the schedules sampled during training, at varying training steps.

Revised: We compare our proposed log-partition variance loss against the trajectory balance loss (Malkin et al., 2022) that uses a Topoformer (Gagrani et al., 2022) to regress the log-partition function.

Revised: The results demonstrate that the log-partition variance loss starts to improve the average speedup much earlier on during training and achieves a better performance in the end.

Original: C T EMPERATURE CONDITIONING ABLATION
Revised: D T EMPERATURE CONDITIONING ABLATION

Original: In order to increase the likelihood of sampling good schedules, one could introduce a fixed temperature throughout training and inference.
Revised: In order to increase the likelihood of sampling good schedules, one could introduce a fixed temperature throughout training and inference.

Original: However, we have observed that this procedure is unreliable for small temperatures.
Revised: However, we have observed that this procedure is unreliable for small temperatures.

Original: Figure 4 shows generalization performance during training on the synthetic graph experiment of Section 5.2.
Revised: Figure 5 shows generalization performance during training on the synthetic graph experiment of Section 5.2.

Original: As can be seen, choosing a temperature of 0 .
Revised: As can be seen, choosing a temperature of 0 .

Original: 01 results in a smaller maximum reward as opposed to training on a higher temperature of 0 .
Revised: 01 results in a smaller maximum reward as opposed to training on a higher temperature of 0 .

Original: On the other hand, sampling a range of temperatures between 0 .
Revised: On the other hand, sampling a range of temperatures between 0 .

Original: 01 and 1 and evaluating on 0 .
Revised: 01 and 1 and evaluating on 0 .

Original: 01 samples the best performing schedules on unseen computation graphs.
Revised: 01 samples the best performing schedules on unseen computation graphs.

Original: [Figure]
Revised: [Figure]