# ED-Batch: Efficient Automatic Batching of Dynamic Neural Networks via Learned Finite State Machines

arXiv: https://arxiv.org/abs/2302.03851 · HTML: https://mlsanity.com/paper/2302.03851/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 25 references that are in the map.
- First posted on arXiv 8 Feb 2023.

## Builds on (most cited first, 11 of 11)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Bidirectional LSTM-CRF Models for Sequence Tagging](https://mlsanity.com/paper/1508.01991/) (2015)
- [Improved Semantic Representations From Tree-Structured Long Short-Term Memory Networks](https://mlsanity.com/paper/1503.00075/) (2015)
- [Chinese NER Using Lattice LSTM](https://mlsanity.com/paper/1805.02023/) (2018)
- [DyNet: The Dynamic Neural Network Toolkit](https://mlsanity.com/paper/1701.03980/) (2017)
- [Nimble: Lightweight and Parallel GPU Task Scheduling for Deep Learning](https://mlsanity.com/paper/2012.02732/) (2020)
- [Cortex: A Compiler for Recursive Deep Learning Models](https://mlsanity.com/paper/2011.01383/) (2020)
- [Deep Learning with Dynamic Computation Graphs](https://mlsanity.com/paper/1702.02181/) (2017)
- [Lattice-Based Recurrent Neural Network Encoders for Neural Machine Translation](https://mlsanity.com/paper/1609.07730/) (2016)
- [On-the-fly Operation Batching in Dynamic Computation Graphs](https://mlsanity.com/paper/1705.07860/) (2017)
- [Just-in-Time Dynamic-Batching](https://mlsanity.com/paper/1904.07421/) (2019)

## Built on by (most cited first, 1 of 1)

- [TASTE: Throughput-Aware Batch Size Tuning for On-Device Edge Learning](https://mlsanity.com/paper/2609.07444/) (2026)

## Abstract (arXiv)

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.