ED-Batch: Efficient Automatic Batching of Dynamic Neural Networks via
  Learned Finite State Machines
Batching has a fundamental influence on the efficiency of deep neural network
(DNN) execution. However, for dynamic DNNs, efficient batching is particularly
challenging as the dataflow graph varies per input instance. As a result,
state-of-the-art frameworks use heuristics that result in suboptimal batching
decisions. Further, batching puts strict restrictions on memory adjacency and
can lead to high data movement costs. In this paper, we provide an approach for
batching dynamic DNNs based on finite state machines, which enables the
automatic discovery of batching policies specialized for each DNN via
reinforcement learning. Moreover, we find that memory planning that is aware of
the batching policy can save significant data movement overheads, which is
automated by a PQ tree-based algorithm we introduce. Experimental results show
that our framework speeds up state-of-the-art frameworks by on average 1.15x,
1.39x, and 2.45x for chain-based, tree-based, and lattice-based DNNs across CPU
and GPU.
Siyuan Chen, Pratik Fegade, Tianqi Chen, Phillip B. Gibbons, Todd C. Mowry
# ED-Batch: Efficient Automatic Batching of Dynamic Neural Networks

###### Abstract

Batching has a fundamental influence on the efficiency of deep neural network (DNN) execution. However, for dynamic DNNs, efficient batching is particularly challenging as the dataflow graph varies per input instance. As a result, state-of-the-art frameworks use heuristics that result in suboptimal batching decisions. Further, batching puts strict restrictions on memory adjacency and can lead to high data movement costs. In this paper, we provide an approach for batching dynamic DNNs based on finite state machines, which enables the automatic discovery of batching policies specialized for each DNN via reinforcement learning. Moreover, we find that memory planning that is aware of the batching policy can save significant data movement overheads, which is automated by a PQ tree-based algorithm we introduce. Experimental results show that our framework speeds up state-of-the-art frameworks by on average 1.15x, 1.39x, and 2.45x for chain-based, tree-based, and lattice-based DNNs across CPU and GPU.

Machine Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Deep Learning, Learning, Deep Learning, Deep Learning, Deep
training of RL, we design a reward function inspired by a sufficient condition for the optimal batching policy.

For the static subgraphs of the dynamic DNN, we take a general approach to optimize it by memory-efficient batching. Our key insight is that the memory operations can be significantly minimized by better planning the inter-tensor memory layouts after batching, which we perform using a novel PQ tree-based (Booth & Lueker, 1976) algorithm that we have designed.

In summary, this paper makes the following contributions:

* We propose an FSM-based batching algorithm to batch dynamic DNNs that finds near-optimal batching policy.
* We design a PQ tree-based algorithm with almost linear complexity to reduce memory copies introduced by dynamic batching.
* We compare the performance of ED-Batch with state-of-the-art dynamic DNNs frameworks on eight workloads and achieved on average 1.15x, 1.39x, and 2.45x speedup for chain-based, tree-based, and lattice-based networks across CPU and GPU. We will make the source code for ED-Batch publicly available.

## 2 FSM-based Algorithm for Dynamic Batching

In this section, we identify the shortcomings of current batching techniques, propose a new FSM-based dynamic batching algorithm, and the mechanism to learn it by RL.

### Problem Characterization

_Dynamic batching_ was initially proposed in TensorFlow Fold (Looks et al., 2017) and DyNet (Neubig et al., 2017b) to enable batched execution of operations for dynamic DNNs. Specifically, given a mini-batch of input instances, dataflow graphs are generated for each of the input instances in the mini-batch and each operation is given a type (indicating operation class, tensor shape, etc.). Upon execution, the runtime identifies batching opportunities within the dataflow graphs by executing operations of the same type together. Therefore, the batching algorithm cannot have high complexity to avoid severe runtime overhead. However, the problem of minimizing the number of launched (batched) kernels is an NP-hard problem with no constant approximation algorithm1, making the batching problem extremely challenging. As a result, the heuristics used for dynamic batching in the current systems often find a suboptimal policy. As shown later (Fig.9), the number of batches executed by current frameworks can be cut by up to 3.27 times.

Footnote 1: Proved by reducing from the _shortest common supersequence_ problem (Raihá & Ukkonen, 1981) in Appendix.A.1.

Specifically, previous state-of-the-art algorithms use heuristics depending on aggregated graph statistics to guide batching. The _depth-based algorithm_ in Tensorflow Fold (Looks et al., 2017) batches operations with the same type at the same topological depth (the input operation to the network has depth 0). And the _agenda-based algorithm_ in DyNet (Neubig et al., 2017b) executes operations of the type with minimal average topological depth iteratively. However, topological depth cannot always capture the regularity of the dataflow graph and result in sub-optimal batching choices. Fig. 1(a) shows a dataflow graph of the tree-based network, which builds upon the parse tree of a sentence with three types of operations: internal nodes (\(I\)), output nodes (\(O\)), and reduction nodes (\(R\)). The ideal batching policy executes all \(O\) nodes in one batch. However, the depth-based algorithm in Fig. 1(b) executes the \(O\) nodes in four batches because they have different topological depths. For the agenda-based algorithm, consider the scenario when it is deciding the next batch after batching the \(I\) nodes first (Fig. 1(c)). Because the \(O\) nodes have a lower average depth (\(\overline{Depth}=(1+1+1+1+2+3+4)/7=1.85\)) than the \(I\) nodes (\(\overline{Depth}=(1+2+3)/3=2\)), the algorithm will pick the \(O\) nodes for the next batch, resulting in an extra batch.

### FSM-based Dynamic Batching

To fully overcome the limitation of specific graph statistics, we found that an FSM-based approach (1) offers the opportunity to specialize for network structure under a rich design space of potential batching policies and (2) can generalize to any number of input instances, as long as they share the same regularity in topology.

Shown in Alg.1, the FSM-based dynamic batching approach is an iterative process of choosing the operation type for the next batch. The process differs from the agenda-based

Figure 1: Example on current dynamic batching algorithms.
algorithm only in how it computes the next type for batching (line 3). During each iteration, the next type is decided by first encoding the current dataflow graph \(G\) into a state \(S=E(G)\), and then using a policy \(\pi\) to map the state into an operation type \(t=\pi(S)\). Then, the operations of type \(t\) on the frontier of the \(G\) form the next batch. After they are executed, these operations are removed from \(G\) and the next iteration begins.

For the model in Fig. 1(a), an optimal FSM-based batching policy is shown in Fig. 2(a), where we encode the dataflow graph by the set of types on the frontier. Fig. 2(b) shows the batching process. From iterations 1 to 3, the dataflow graph is encoded into \(S_{2}=\{I,O\}\), thus the policy continues to batch nodes of type \(I=\pi(S_{2})\), avoiding batching the \(O\) nodes as past heuristics would do. At the same time, it is not hard to see that this FSM-based batching policy can be applied to batch multiple input instances of different parse trees.

### Using RL to Learn the FSM

As the FSM provides us with the design space for potential batching policies, we need an algorithm to come up with the best batching policy specialized for a given network structure. In ED-Batch, we adopt an RL-based approach for the design-space-exploration and learn the best FSM by a reward function inspired by a sufficient condition for optimal batching.

In RL, an agent learns to maximize the accumulated reward by exploring the environment. At time \(t\), the environment is encoded into a state \(S_{t}\), and the agent takes action \(a_{t}=\pi(S_{t})\) following the policy \(\pi\) and receives a reward \(r_{t}=R(S_{t},a)\). After this, the environment transforms to the next state \(S_{t+1}\). This results in a sequence of states, actions, and rewards: \((S_{0},a_{0},r_{0},S_{1},a_{1},r_{1},...,S_{N-1},a_{N-1},r_{N-1},S_{N})\), where \(N\) is the number of time steps and \(S_{N}\) is an end state. The agent aims to maximize the accumulated reward \(\Sigma_{t}r_{t}\) by updating the policy \(\pi\). For FSM-based dynamic batching, the environment is the dataflow graph, which is encoded into states by the encoding function \(E\). For every iteration, the agent decides on the type for the next batch, receives a reward on that decision, and the environment gets updated according to Alg. 1. Now we elaborate on the state encoding, reward design, and training respectively.

Before that, we explain some notations. For a dataflow graph \(G\), \(G_{t}\) refers to its status at step \(t\), \(G^{a}\) refers to the extracted subgraph of \(G\) composed solely of type \(a\) operations (See example in Fig.2(c)), \(Frontier(G)\) refers to the set of ready-to-execute operations, \(Frontier_{a}(G)\) refers to the subset of \(Frontier(G)\) with type \(a\).

**State Encoding:** The design of state encoding should be simple enough to avoid runtime overhead. In practice, we experimented with three ways of encoding: (1) \(E_{base}(G)=\{v.type|v\in Frontier(G)\}\) is the set of operation types on the frontier, (2) \(E_{max}(G)=(E_{base}(G),argmax_{t\in T}|Frontier_{t}(G)|)\) is \(E_{base}(G)\) plus the most common type on the frontier and (3) \(E_{sort}(G)=sort(\{v.type\in T|v\in Frontier(G)\},t:|Frontier_{t}(G)|)\) is \(E_{1}(G)\) sorted by the number of occurrences on the frontier. Empirically, we found that \(E_{sort}\) was the best among the three (SS5.3).

**Reward:** We design the reward to minimize the number of batches, thus increasing the parallelism exploited. The reward function is defined as

\[r(S_{t},a_{t})=-1+\alpha*\frac{|Frontier(G_{t}^{a_{t}})|}{|Frontier_{a_{t}}(G _{t})|} \tag{1}\]

where \(\alpha\) is a positive hyper parameter and \(S_{t}=E(G_{t})\). The constant \(-1\) in the reward penalizes every additional batch, thereby helping us minimize the number of batches.

**Lemma 1** (Sufficient Condition for Optimal Batching).: 2 _If \(\frac{|Frontier(G_{t}^{a_{t}})|}{|Frontier_{a_{t}}(G_{t})|}=1\), then there exists a shortest batching sequence starting with \(a_{t}\)._

Footnote 2: Proof in the Appendix A.2

The second term is inspired by a sufficient condition for optimal batching (Lemma 1) to prioritize the type that all operations in the frontier of the subgraph of this type are ready to execute. For the tree-based network, this term prioritizes the batching choice made by the optimal batching policy in Fig. 2(a). For example, at iteration \(2\), this term is \(\frac{5}{7}\) and \(\frac{1}{1}\) for the \(O\) and \(I\) node respectively and the \(I\) node is given higher priority for batching. For other networks, like

Figure 2: Dynamic Batching Policy by FSM
the chained-based networks (Fig. 9), this sufficient condition continues to hold.

**Training:** We adopt the tabular-based Q-learning (Watkins and Dayan, 1992) algorithm to learn the policy. And an N-step bootstrapping mechanism is used to increase the current batching choice's influence on longer previous steps. During the training phase, the algorithm learns a \(Q\) function, which maps each state and action pair to a real number indicating its score, for every new topology. We found that this simple algorithm learns the policy \(\pi\) within hundreds of trials (Table 3). During inference, at each state \(S\), we select the operation type with the highest \(Q\) value for the next batch, i.e. \(\pi(S)=argmax_{a}Q(S,a)\). This step is done by a lookup into stored \(Q\) functions in constant time.

## 3 Memory-efficient Batching for Static Subgraphs

### Background and Motivation

In order to invoke a batched tensor operator in a vendor library, the source and result operands in each batch are usually required to be contiguous in memory (as per the vendor library specifications). Current batching frameworks such as Cavs and DyNet ensure this by performing explicit memory gather and/or scatter operations, leading to high data movement. On the other hand, as mentioned above, Cortex relies on specialized, hand-optimized batched kernels instead of relying on vendor libraries. This approach however, is unable to reuse the highly performant optimizations available as part of vendor libraries.

In ED-Batch, we take a different approach to fit the memory layout into the batching policy, where operations in the source and result operands for batched execution are already contiguous in memory.

We illustrate the approach by the example. Fig. 3(a) shows an example code for a static subgraph and Fig. 3(b) shows its batched version. In Fig. 3(c), we compare two memory layouts. On the left, we directly allocate memory according to the variable's label, then two memory gather for \([x_{1},x_{3}],[x_{2},x_{1}]\) and one scatter for \([x_{8},x_{6},x_{7}]\) is performed because they are either not contiguous or aligned in memory. We say an operand of a batch is aligned in memory if the order of its operations matches with the one in memory. Now, consider the memory allocation on the right, which allocates memory following \((x_{2},x_{1},x_{3},x_{4},x_{5},x_{8},x_{6},x_{7})\). Then, every source and result operand of the batched execution is already contiguous and aligned in memory, saving us from extra memory copies.

### PQ tree-based memory allocation

To find the ideal layout, we designed an almost linear complexity memory allocation algorithm based on PQ tree, which is a tree-based structure used to solve the consecutive one property (Meidanis et al., 1998) and is previously applied to Biology for DNA-related analysis (Landau et al., 2005).

We define the _ideal memory layout_ as a sequence of variables satisfying two constraints:

* **Adjacency Constraint:** Result and source operands in every batch should be adjacent in the sequence. E.g. \(\{x_{4},x_{5}\},\{x_{1},x_{3}\},\{x_{2},x_{1}\}\) for \(B1\), \(\{x_{6},x_{7},x_{8}\},\{x_{4},x_{3},x_{5}\}\) for \(B2\) are adjacent in the sequence.
* **Alignment Constraint:** The order of the result and source operands should be aligned in a batch. E.g. for \(B1\), \(x_{4}\prec x_{5}\iff x_{1}\prec x_{3}\iff x_{2}\prec x_{1}\) in the sequence.

The adjacency constraint is satisfied by PQ tree (Booth and Lueker, 1976) algorithm. Given several subsets of a set \(S\), the PQ tree algorithm returns a data structure in linear time called PQ tree, representing potential permutations of \(S\) that elements in each subset are consecutive. Fig.4(a) shows the PQ tree for the example code. The tree has three kinds of nodes: P-node, Q-node, and leaf node. Leaf nodes represent the variables; P-nodes have more than two children, whose appearance in the sequence is contiguous but could be permuted; Q-nodes have more than one child, whose appearance in the sequence follows an order but could be reversed. A depth-first traversal of the leaf nodes gives the sequence. For example, there is one P-nodes and three Q-nodes in Fig.4(a). \(Q_{2}\) indicates the order should only be \((x_{2},x_{1},x_{3},Q_{3})\) or \((Q_{3},x_{3},x_{1},x_{2})\), while \(P_{2}\) indicates that one permutation of \(\{x_{6},x_{7},x_{8}\}\) appears in the sequence. The adjacency of \(\{x_{4},x_{5}\}\) is embedded in \(Q_{3}\), \(\{x_{1},x_{3}\},\{x_{2},x_{1}\},\{x_{4},x_{3},x_{5}\}\) in \(Q_{2}\), and \(\{x_{6},x_{7},x_{8}\}\) in \(P_{1}\). A possible sequence is \((x_{2},x_{1},x_{3},x_{4},x_{5},x_{6},x_{7},x_{8})\).

To satisfy the alignment constraint, we annotate each node on the PQ tree with an order. An annotated PQ tree is shown in Fig.4(c), where a direction mark is attached to every Q-node, indicating its traversal order. As a result, any leaf node sequence of legal traversal on this annotated PQ tree indicates a memory allocation order satisfying the alignment constraint.

Figure 3: Example on memory allocation. \(\alpha\), \(\sigma\) represent operators.
Two passes obtain the order annotation to the PQ tree (Fig.4, Alg.2). The first pass, BroadcastConstraint, makes the tree structure of each batch's operands isomorphic. For \(B2\)'s operands, \(\{x_{3},x_{4},x_{5}\}\)'s tree structure is \(Q_{2}=(...,x_{3},Q_{3}=(x_{4},x_{5}))\), and \(\{x_{6},x_{7},x_{8}\}\)'s tree structure is \(P_{1}=(x_{6},x_{7},x_{8})\). They are made isomorphic in this pass. The second pass, DecideNodeOrder, derives the equivalent class of node and order pairs among P and Q-nodes and their orders and searches for a compatible solution for the direction of the Q-node and a permutation choice of the P-node complied with the equivalence relationship.

We walk through the algorithm in the example. At first, the PQ tree is constructed by the standard algorithm to satisfy the adjacency constraint (Fig.4(a)). After that, in BroadcastConstraint, for \(B2\), we parse the adjacency constraint of it (line 11) by 1) parsing the adjacency constraint for each operand, e.g. \(\{x_{3},x_{4},x_{5}\}\) and \(\{x_{4},x_{5}\}\) for the source operand, \(\{x_{6},x_{7},x_{8}\}\) for the result operand, and 2) transforming them across operands by alignment information, e.g. \(\{x_{4},x_{5}\}\) is transformed into \(\{x_{6},x_{7}\}\). After that, \(\{x_{6},x_{7}\}\) is applied to the PQ tree3 and we keep records on the batch whose tree structure has changed (line 12). Now tree structures for \(B2\)'s operands are isomorphic, and we apply this process to other batches in a breadth-first search until no update on the tree structure happens.

Footnote 3: Perform by standard Reduce step in the Vanilla PQ tree algorithm to satisfy adjacency constraint by restructuring the tree.

In DecideNodeOrder (line 18), we assign directions for the Q-node and the permutation for the P-node. We start by parsing the equivalence relationship (line 22) among Q-node and direction pairs or the P-node and permutation pairs from the isomorphic tree structures after the first pass, e.g. \(<Q_{2},\leftarrow>\Longleftrightarrow<Q_{3},\leftarrow>\) for \(B1\), \(<Q_{3},\leftarrow>\Longleftrightarrow<Q_{4},\leftarrow>\) and \(<Q_{2},\leftarrow>\Longleftrightarrow<Q_{5},\rightarrow>\) for \(B2\). After that, we spread the equivalence relationship across batches with the support of _union-find data structure_. Shown in Fig.5, a graph carrying the equivalence relationship is constructed by iteratively Union equivalent relationship (line 23), i.e. the node and order pairs. When processing \(<Q_{2},\leftarrow>\Longleftrightarrow<Q_{5},\rightarrow>\), a \(R\) edge \(<Q_{2},Q_{5}>\) is added to the graph, indicating \(Q_{2}\)'s direction is determined by the reverse of \(Q_{5}\)'s direction. When processing the \(<Q_{2},\leftarrow>\Longleftrightarrow<Q_{3},\leftarrow>\), we first find the decider of their order, i.e. \(Q_{5}\) for \(Q_{2}\) and \(Q_{4}\) for \(Q_{3}\), and add an \(R\) edge between them. In this way, \(Q_{2}\) and \(Q_{3}\) always have the same order. Finally, the deciders in the graph are assigned arbitrary directions, which spread across the graph following the relationship on the edge (Fig.5(b)).

The PQ tree memory allocation algorithm's time complexity is given in lemma 2, showing the PQ tree algorithm is linear to the problem size if the operations on a single batch are limited to a constant.

**Lemma 2**.: _PQ tree memory allocation algorithm's time complexity is \(O(\Sigma_{b\inbatches}|b|\max_{b\inbatches}^{2}|b|)\), where \(|.|\) counts the operation in a batch._

Right now, PQ tree memory optimization is applied to the static subgraph because its execution time still does not fit

Figure 4: Example for PQ tree-based algorithm. \(x_{1}\)-\(x_{8}\) are variables.

Figure 5: Illustration for order decision in Pass2.
into the high runtime constraint for dynamic DNNs. But the algorithm itself and the idea of better memory planning for batching is applicable to any batching problem. A detailed explanation of the algorithm is in AppendixB.

## 4 Implementation

The optimizations in ED-Batch are fully automated and implemented as a runtime extension to DyNet in 5k lines of C++ code. The user can optionally enable the batching optimization by passing a flag when launching the application, and enable the static subgraph optimization by defining subgraphs in DyNet's language with a few annotations. Before execution, the RL algorithm learns the batching policy and ED-Batch optimizes the static subgraph by the approach in SS3. Upon execution, ED-Batch calls DyNet's executor for batched execution, which is supported by vendor libraries.

## 5 Evaluation

### Experiment Setup

We evaluate our framework against DyNet and Cavs, two state-of-the-art runtimes for dynamic DNNs, which are shown to be faster than traditional frameworks like Pytorch and Tensorflow (Xu et al., 2018; Neubig et al., 2017a). Cavs' open-sourced version has worse performance than DyNet, as stated in Fegade, Chen, Gibbons, and Mowry (2021) because certain optimizations are not included. To make a fair comparison with Cavs, we use an extended version of DyNet with the main optimizations in Cavs enabled as a reference for Cavs' performance (referred to as Cavs DyNet). Namely, the static subgraphs in the network are pre-defined and batching optimization is applied to them. ED-Batch is implemented on top of this extended version, with the RL-based dynamic batching algorithm (\(E_{sort}\) for state encoding) and memory optimization on the static subgraph by PQ Tree. On the other side, the agenda-based algorithm and the depth-based algorithm are used for dynamic batching on Vanilla/Cavs DyNet. Depending on the workload and configuration, a better-performing algorithm is chosen for Vanilla/Cavs DyNet in the evaluation.

We test the framework on 8 workloads, shown in Table 1. They follow an increase in dynamism, from chains to trees and graphs. Except for lattice-based networks, all workloads appeared as the benchmark for past works.

We run our experiments on a Linux server with an Intel Xeon E2590 CPU (28-physical cores) and an Nvidia V100 GPU. The machine runs CentOS 7.7, CUDA 11.1, and cuDNN 8.0. We use DyNet's latest version (Aug 2022, commit c418b09) for evaluation.

### Overall Performance

We compare ED-Batch's end-to-end inference throughput against Vanilla/Cavs DyNet. We follow past work to evaluate different batch sizes (1, 8, 32, 64, 128, 256) and model sizes (32, 64, 128, 256, 512), which is the size for the hidden vector length and the embedding size. The throughput is calculated as the maximum throughput among all bath size choices. For all cases, ED-Batch outperforms Vanilla DyNet significantly due to the reduction in graph construction and runtime overhead by pre-definition of the static subgraph.

We now discuss the comparison with Cavs DyNet. For chain-based models, the BiLSTM-tagger and LSTM-NMT, ED-Batch achieved on average 1.20x, 1.11x speedup on CPU and 1.20x, 1.12x on GPU. Because the network structure is basically made up of chains, shown in Fig.9, both the agenda-based algorithm and the FSM-based batching algorithm find the optimal batching policy. On the other hand, the LSTMCell is 1.54x faster with the PQ-tree optimization compared to the one with the DyNet's memory allocation, which explains the speedup.

For the tree-based model, compared to agenda/depth-based batching heuristic ED-Batch reduces the number of batches by 37%. This is because the FSM-based algorithm executes the output nodes in one batch (Fig.1). For TreeLSTM and TreeGRU, ED-Batch achieved on average 1.63x, 1.46x speedup on CPU and 1.23x, 1.29x speedup on GPU. ED-Batch's performance is close to Cavs DyNet on MVRNN because the execution is bounded by matrix-matrix multiplications which can hardly benefit from extra batch parallelism and the reduction in runtime overhead.

For the lattice-based models, the LatticeLSTM and LatticeGRU, ED-Batch increases DyNet Cavs's throughput significantly by 1.32-2.97x on CPU and 2.54-3.71x on GPU, which is attributed to both the better dynamic batching and static subgraph optimization. For the lattice-based models'

\begin{table}
\begin{tabular}{p{113.8pt}|p{113.8pt}|p{113.8pt}} \hline \hline
**Model** & **Short name** & **Dataset** \\ \hline A bi-directional LSTM Named Entity Tagger (Huang et al., 2015) & BiLSTM-Tagger & WikiNER English Corpus(Nothman et al., 2013) \\ \hline An LSTM-based encoder-decoder model for neural machine translation. & LSTM-NMT & IWSLT 2015 En-Vi \\ \hline N-ary TreeLSTM (Tai et al., 2015a) & TreeLSTM & Penn treebank (Marbank et al., 1994) \\ \hline An extension to TreeLSTM that contains two types of internal nodes, each with 50\% probability & TreeLSTM-2Type \\ \hline A lattice-based LSTM network for Chinese NER (Zhang & Yang, 2018) & Lattices generated based on Chinese Chinese Dataset\({}^{4}\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Models and datasets used in evaluation
network structure in Fig.7, the FSM-based algorithm prioritizes the execution of the character cell and delays the execution of the word cell, whereas the depth/agenda-based algorithms batch the character cell and word cell more arbitrarily. As a result, the number of batches is reduced by up to 3.27 times (Fig.9). For the static subgraph, the used LSTMCell and GRUCell's latency is cut by 34% and 35%, which adds to the speedup.

### Analysis

**Where does ED-Batch's speedup come from?** Shown in Fig.8, we decompose the inference pass into the construction time, scheduling time, and execution time. Construction time is the time to define the dataflow graph. Scheduling time is the time for dynamic batching analysis. Execution time is the left time of the forward pass, mainly composed of the execution of operations. While having similar construction/scheduling time, ED-Batch speeds up Cavs DyNet in the great cutdown in execution time benefited from better batching and fewer kernels for data movement.

**Does the algorithm find a good enough batching policy?** Shown in Fig.9, compared to the agenda/depth-based batching algorithm, ED-Batch's FSM-based batching algorithm uniformly executes fewer batches. Among three state encoding choices, \(E_{sort}\) is slightly better because of the stronger expressiveness, finding the optimal batching policy on BiLSTM-daggerger, LSTM-NMT, and Tree-based models and executing 23% and 44% more batches on TreeLSTM-2Type, Lattice-Based models.

To demonstrate the efficiency of the reward function, we measure the number of batches executed by a sufficient-condition-guided heuristic, which selects the type for the next batch that maximizes the second term in Eq.1. Shown in Fig. 9, this heuristic executes batches paramount to the best FSM-based algorithm. However, this heuristic has higher time complexity and adds to unacceptable runtime overhead. Thus, on the evaluated workloads, the FSM-based algorithm can be treated as a time-efficient distiller of this heuristic.

**Ablation Study of the Static Subgraph Optimization.** In

Figure 8: Cavs DyNet v.s. ED-Batch: Time Decomposition when model size is 128 and batch size is 64.

Figure 7: Lattice Network for Chinese NER. The topology for one input sentence is a chain of character cells with jump links of word cells. Agenda/Depth-based batching algorithm fails to batch the word cells together.

Figure 6: ED-Batch v.s. Vanilla/Cavs DyNet: Inference Throughput

Figure 9: The number of batches for different batching algorithms. FSM-base/sort/max refers to the FSM based algorithm with different state encodings.
Table 2, we evaluate ED-Batch's memory layout optimization on the static subgraphs. For all evaluated cases, the PQ-tree algorithm finds the _ideal memory allocation order_ (Remained data transfer is caused by broadcast that cannot be optimized by better memory layout). Compared to the baseline, ED-Batch reduces the latency of the static subgraph by up to 1.6x, memory kernels by up to 4x, and memory transfer amount by up to 66x. This significant reduction in memory transfer can be attributed to the better arrangement of the weight parameters. For example, there are four gates in the LSTM cell that perform feed-forward arithmetic \(y_{i}=W_{i}x_{i}+b_{i}\), which are executed in a batch. The memory arrangement in ED-Batch makes sure the inputs, parameters, and intermediate results of batched kernels are contiguous in the memory, which is not considered by DyNet's policy. Since the weight matrix occupies memory relative to the square of the problem size, this leads to a huge reduction in memory transfer.

**Comparison with a more specialized framework.** Cortex (Fegade et al., 2021) is highly specialized for optimizing a class of recursive neural networks and it requires the user to not only express the tensor computation, but also specify low-level optimizations specific to underlying hardware, both through TVM's domain-specific language (Chen et al., 2018). We compare ED-Batch with Cortex on TreeLSTM and TreeGRU. To make more of an apples-to-apples comparison in terms of the user's effort in developing the application, we enabled Cortex's automated optimizations like _linearization_ and _auto-batching_ and used simple policies on optional user-given (manual) optimizations like kernel fusion and loop transformation (details in SSC). As shown in Table 5, ED-Batch can speed up Cortex by up to 3.98x.

**The compilation overhead.** We trained the RL for up to 1000 trials and stopped early if the number of batches reaches the lower bound (check every 50 iterations). On the static subgraph, batching is performed as a grid search and the PQ tree optimization is applied afterward. As shown in Table 3 and Table 4, it takes tens of milliseconds to optimize the static subgraph and up to 22 seconds to learn the batching policy for tested workloads.

## 6 Related Work

There are a variety of frameworks specialized for dynamic neural network training (Neubig et al., 2017; Looks et al., 2017; Xu et al., 2018), inference (Fegade et al., 2021; Feagade, 2023; Zha et al., 2019), and serving (Gao et al., 2018). Concerning batching for the dynamic neural networks, DyNet (Neubig et al., 2017) and TFFold (Looks et al., 2017) laid the system and algorithm foundation to support _dynamic batching_. However, their batching heuristics are often sub-optimal as we saw above. Nevertheless, their algorithms have been used in other frameworks, like Cavs (Xu et al., 2018) and Acrobat (Fegade, 2023). Apart from batching, another major direction of optimization is to extract the static information from the dynamic DNN and optimize them during compile time. Cavs (Xu et al., 2018) proposed the idea of predefining the static subgraphs, which was later extended in (Zha et al., 2019) to batch on different granularities. ED-Batch adopts this multi-granularity batching idea to perform batching on both the graph level and the subgraph level. For static subgraphs, traditional techniques are used for optimization, like the kernel fusion in Cavs and Cortex (Fegade et al., 2021), the AoT compilation (Kwon et al., 2020), and specialized kernel generation in ACRo-Bat (Fegade, 2023). However, though with high efficiency, these optimizations can hardly be automated because either the developer or the user needs to optimize each subgraph manually. In ED-Batch, fully automated runtime optimizations are used instead to enable both efficient execution and generalization.

## 7 Conclusion

In ED-Batch, we designed an FSM-based algorithm for batched execution for dynamic DNNs. Also, we mitigated the memory transfer introduced by batching through the memory layout optimization based on the PQ-tree algorithm. The experimental results showed that our approach achieved significant speedup compared to current frameworks by the reduction in the number of batches and data movement.

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline  & \multicolumn{2}{c}{TreeGRU} & \multicolumn{2}{c}{TreeLSTM} \\ batch\_size & model\_size & Cortex & Ours & Cortex & Ours \\ \hline
10 & 256 & 2.30 & **2.27** & **2.244** & 2.78 \\  & 512 & 5.60 & **3.04** & 8.500 & **4.70** \\
20 & 256 & 3.73 & **3.03** & **3.460** & 3.52 \\  & 512 & 11.70 & **3.70** & 19.210 & **4.82** \\ \hline \hline \end{tabular}
\end{table}
Table 5: ED-Batch v.s. Cortex: Inference Latency (ms).

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline  & \multicolumn{2}{c}{Latency (ms)} & \multicolumn{2}{c}{Mem Kernel/Subgraph} & \multicolumn{2}{c}{Memory Amount (kB)} \\ Subgraph & value & ratio & value & ratio & value & ratio \\ \hline GRUCell & 0.11/**0.07** & 1.54 & 6/2 & 3.0 & 666.0/ **1.40** & 47.57 \\ LSTMCell & 0.2/**0.13** & 1.52 & **4/**1 & 4.0 & 1054.0/ **1.60** & 65.87 \\ MVCell & **0.20/**0.08** & 0.96 & 2.2 & 1.0 & 260.0/ **2.00** & 1.0 \\ TreeGRU-Internal & 0.2/**0.15** & 1.6 & **8/2** & 4.0 & 552.0/ **1.60** & 34.5 \\ TreeGRU-Leaf & 0.09/**0.07** & 1.4 & **4/2** & 2.0 & 268.0/ **8** & 33.5 \\ TreeLSTM-Internal & 0.19/**0.12** & 1.61 & 7/3 & 2.33 & 1054.0/ **2.20** & 48.36 \\ TreeLSTM-Leaf & 0.12/**0.09** & 1.27 & 3/**1** & 3.0 & 396.0/ **0.60** & 66.0 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Batching with DyNet’s memory allocation (left) v.s. Batching with PQ tree-based memory allocation on static subgraphs. batch size = 8, model size = 64.

\begin{table}
\begin{tabular}{l c c} \hline \hline  & \multicolumn{2}{c}{TreeGRU} & \multicolumn{2}{c}{TreeLSTM} \\ batch\_size & model\_size & Cortex & Ours \\ \hline
10 & 256 & 2.30 & **2.27** & **2.244** & 2.78 \\  & 512 & 5.60 & **3.04** & 8.500 & **4.70** \\
20 & 256 & 3.73 & **3.03** & **3.460** & 3.52 \\  & 512 & 11.70 & **3.70** & 19.210 & **4.82** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Batching with DyNet’s memory allocation (left) v.s. Batching with PQ tree-based memory allocation on static subgraphs. batch size = 8, model size = 64.
## Acknowledgements

We thank Peking University's supercomputing team for providing the hardware to run the experiments and thank Kevin Huang of CMU for his help in early exploration of this research space.
