MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems

Tianqi ChenMu LiYutian LiMin LinNaiyan WangMinjie WangTianjun XiaoBing XuChiyuan ZhangZheng Zhang

article2015arXiv2,343 citations

Introduces the architecture and system implementation of MXNet, showing how unifying declarative symbolic graphs with imperative tensor operations enables scalable, memory-efficient deep learning across devices ranging from mobile phones to distributed GPU clusters.

Listen

As modern machine learning models grow increasingly complex, organizations face difficult trade-offs between programming flexibility, execution speed, memory consumption, and distributed scalability. Existing machine learning software frameworks typically specialize in either declarative approaches—which define entire computation graphs upfront for automated optimization—or imperative approaches, which execute commands step-by-step for greater flexibility and easier debugging.

The article introduces and evaluates MXNet, an open-source machine learning library designed to blend declarative symbolic expressions with imperative tensor computations across heterogeneous hardware ranging from mobile devices to distributed computing clusters.

To evaluate the system, the authors conducted empirical benchmarks comparing MXNet against established deep learning frameworks (Torch7, Caffe, and TensorFlow) on standard convolutional neural network workloads. They evaluated raw processing speed on single graphics processing units (GPUs), analyzed internal memory consumption across multiple memory allocation strategies, and measured multi-machine scaling performance on image classification using Amazon Web Services cloud instances with up to ten multi-GPU machines.

The evaluation produced several key findings regarding speed, memory efficiency, and scalability. First, MXNet matched the raw execution speeds of high-performing dedicated frameworks like Torch7 and Caffe on standard computer vision benchmarks, while operating about twice as fast as early versions of TensorFlow. Second, MXNet's memory optimization strategies—specifically combining in-place memory reuse and non-parallel node memory sharing—reduced internal memory consumption by roughly 50% during model training and by approximately 75% during model prediction. Finally, distributed image classification tests demonstrated near-linear scaling, reducing the time required for a full data pass across 1.3 million images from 14,000 seconds on a single machine down to 1,400 seconds across ten machines.

These findings indicate that teams can achieve high developer productivity and cross-language flexibility (supporting Python, C++, R, Julia, and Go) without incurring runtime performance or memory overhead penalties. By significantly lowering memory footprints and scaling linearly across distributed hardware, the framework enables organizations to train larger neural network architectures on existing compute infrastructure, effectively reducing hardware acquisition and operational cloud costs.

Organizations evaluating or deploying large-scale machine learning workflows should consider MXNet as a versatile, lightweight foundation for both research and deployment pipelines. Teams adopting the framework should leverage its dual-level distributed storage architecture to balance intra-machine and network communication costs based on their infrastructure constraints.

The reported results carry moderate limitations based on the testing scope provided. The empirical benchmarks primarily reflect image classification workloads under specific software library versions and hardware configurations, and the performance differences relative to other tools like TensorFlow may diminish as competitor implementations mature. Nevertheless, the findings demonstrate a sound and highly viable architecture for scalable machine learning systems.

Carnegie Mellon University Massachusetts Institute of Technology Microsoft National University of Singapore New York University Stanford University TuSimple University of Alberta University of Washington gpu libraries mobile devices heterogeneous distributed systems distributed machine learning library parameter server

arXiv: 1512.01274 apache/mxnet

Paper: TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems, Martín Abadi et al. (2016). TensorFlow established the large-scale dataflow graph paradigm that MXNet builds upon and refines.
Paper: Theano: A Python framework for fast computation of mathematical expressions, Rami Al-Rfou et al. (2016). Theano introduced foundational symbolic computation and compilation techniques that inspired MXNet's unified API.
Paper: Large Scale Distributed Deep Networks, Jeffrey Dean et al. (2012). DistBelief pioneered the distributed machine learning architectures that motivated MXNet's heterogeneous cluster support.

Paper: TVM: an automated end-to-end optimizing compiler for deep learning, Tianqi Chen et al. (2018). TVM extends MXNet's foundational execution model by automatically optimizing deep learning operators for diverse hardware targets.
Paper: Training Deep Nets with Sublinear Memory Cost, Tianqi Chen et al. (2016). Sublinear memory training techniques directly build upon MXNet's execution graph to dramatically reduce memory consumption.

Cover for MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems

Abstract

MXNet is a multi-language machine learning (ML) library to ease the development of ML algorithms, especially for deep neural networks. Embedded in the host language, it blends declarative symbolic expression with imperative tensor computation. It offers auto differentiation to derive gradients. MXNet is computation and memory efficient and runs on various heterogeneous systems, ranging from mobile devices to distributed GPU clusters.

This paper describes both the API design and the system implementation of MXNet, and explains how embedding of both symbolic expression and tensor operation is handled in a unified fashion. Our preliminary experiments reveal promising results on large scale deep neural network applications using multiple GPU machines.

Topics

gpu libraries mobile devices ranging deep learning convolutional neural networks distributed environments

Keywords

heterogeneous distributed systems distributed machine learning library parameter server dataflow graph tensor computation declarative programming data synchronization GPU clusters lazy evaluation

1 Introduction
2 Programming Interface
2.1 Symbol: Declarative Symbolic Expressions
2.2 NDArray: Imperative Tensor Computation
2.3 KVStore: Data Synchronization Over Devices
2.4 Other Modules
3 Implementation
3.1 Computation Graph
3.2 Dependency Engine
3.3 Data Communication
4 Evaluation
5 Conclusion
References

Knowls

Knowl 1 — Blending Declarative Symbolic Graphs and Imperative Tensor Computation
model/method
MXNet unifies declarative symbolic computation graphs with imperative tensor computations within a single execution backend.

The framework provides two complementary programming abstractions:
- Declarative Symbolic Expressions (Symbol): Computational dataflow graphs are declared using multi-output symbolic expressions. Operators define neural network layers or mathematical primitives, and variables can be unbound free variables or intermediate operator outputs. Symbolic graph declarations give the runtime global graph visibility, enabling auto-differentiation, memory reuse planning, and graph optimizations prior to execution.
- Imperative Tensor Computation (NDArray): Multi-dimensional tensor arrays are manipulated imperatively on CPU or GPU devices, matching the programming feel of standard numerical libraries.
Both abstractions execute through the same backend engine via lazy evaluation. When an imperative tensor update (such as $w \leftarrow w - \eta g$ ) is interleaved with forward-backward graph passes, the runtime dependency engine automatically resolves data dependencies across the two abstractions, enabling dynamic parameter updates and interactive debugging without sacrificing the optimizations of static computation graphs.
Knowl 2 — Dynamic Dependency Engine with Explicit Resource Mutation Tracking
model/method
MXNet uses a multi-threaded dynamic dependency engine to schedule and parallelize operations across heterogeneous resources, including CPUs, GPUs, and memory/PCIe buses. Resources (such as NDArray instances, random number generator states, or temporary memory allocations) are registered with unique resource tags.

Unlike traditional dataflow engines that assume functional immutability and track only read dependencies, MXNet explicitly tracks both read and write (mutation) tags for each scheduled operation:
- Read Dependencies: Operations that only read from a resource can be scheduled concurrently across multiple threads and execution devices.
- Write/Mutation Dependencies: Operations that mutate an existing resource are serialized relative to preceding reads and writes on that resource tag, and block subsequent accesses until mutation finishes.
Explicit mutation tracking allows direct in-place parameter array updates ( $w \leftarrow w - \eta g$ ) without creating temporary array copies, enables seamless scheduling of imperative tensor operations, and serializes operations sharing state (such as random number generators using an identical seed) to guarantee reproducible execution.
Knowl 3 — Linear-Time Memory Allocation Heuristics for Computation Graphs
model/method
Finding an optimal memory reuse layout among non-overlapping intermediate variables in an arbitrary computation graph has an $O(n^2)$ time complexity for $n$ variables. MXNet employs two linear-time $O(n)$ heuristic strategies to recycle internal intermediate memory buffers:
1. In-place Reuse (inplace): During a simulated topological graph traversal, the allocator tracks a reference counter for each intermediate node representing the number of pending consumer operations. Once an operation consumes an intermediate tensor and the reference counter reaches zero, its memory buffer is immediately recycled and made available for allocation to subsequent operations.
2. Co-sharing Reuse (co-share): Two variables can share the same physical memory allocation if and only if their active lifespans do not overlap and they cannot be executed concurrently in parallel. The allocator identifies the longest pending execution path in the graph upon scheduling and imposes sequential dependency constraints along that path to safely reuse memory blocks across distinct nodes.
Knowl 4 — Two-Level Hierarchical Parameter Server and Distributed KVStore
model/method
MXNet implements a distributed key-value store (KVStore) based on a two-level parameter server architecture to synchronize model parameters across multiple devices and machines:
- $\text{push}(k, v)$ : Pushes a key-value pair consisting of key $k$ and tensor $v$ (such as computed gradients) from a worker device to the store, merging them via a user-defined updater function.
- $\text{pull}(k, d)$ : Pulls the updated aggregated value associated with key $k$ from the store into a local destination tensor $d$ .
The synchronization hierarchy consists of two levels:
1. Level-1 Server (Intra-machine): Coordinates synchronization among local devices (such as multiple GPUs) within a single physical host. Outbound updates are aggregated locally on the machine before being transmitted over the network.
2. Level-2 Server (Inter-node): Manages communication and parameter synchronization across different physical hosts in the cluster.
The system supports multiple consistency models across the hierarchy (such as sequential consistency for local intra-machine communication and eventual consistency for inter-machine synchronization). All push and pull operations are scheduled as asynchronous tasks within MXNet's dependency engine, allowing communication to overlap automatically with computation via lazy evaluation.
Knowl 5 — Computation Graph Transformations and Optimizations in MXNet
model/method
Before evaluating a bound symbolic computation graph, MXNet applies static graph-level optimizations:
1. Dead Subgraph Elimination and Pruning: The engine analyzes the output variables requested during binding and trims unneeded subgraphs. During inference, backward-pass gradient nodes are discarded; during intermediate feature extraction, trailing layers beyond the target activation are skipped.
2. Operator Fusion: Multiple consecutive elementwise operations (such as linear scaling and shift $a \times b + 1$ ) are merged into a single BLAS or GPU kernel launch, reducing memory bandwidth pressure and kernel invocation overhead.
3. Optimized Composite Operations: High-level structural components, such as convolution and activation layers, are mapped directly to tuned, monolithic library implementations.
Knowl 6 — Memory Footprint Reduction via In-Place Reuse and Co-Sharing
empirical result
The effectiveness of MXNet's inplace and co-share memory allocation strategies was evaluated by measuring internal intermediate variable memory usage (excluding model outputs) on AlexNet, GoogLeNet, and VGG with a batch size of 64:
- Model Training (Forward and Backward Passes): Combining inplace and co-share heuristics achieves approximately a $2\times$ reduction in internal memory footprint across all three network architectures compared to naive allocation without memory sharing.
- Model Inference (Forward Pass Only): Combining both heuristics achieves approximately a $4\times$ reduction in internal memory footprint compared to naive allocation.
- Absolute Memory Overhead: For the VGG network during training, the combined allocation strategy requires less than $16\text{ MB}$ of extra internal intermediate buffer memory.
Knowl 7 — Single-GPU Raw Performance Comparison across Deep Learning Frameworks
empirical result
Single-GPU execution time for combined forward and backward passes was evaluated using the convnet-benchmarks suite on an Nvidia GTX 980 GPU with a batch size of 32 across AlexNet, GoogLeNet, and VGG architectures:
- MXNet, Torch7, and Caffe (all compiled with CUDA 7.5 and cuDNN 3) demonstrated comparable execution times across all three model architectures, as computation time was dominated by optimized CUDA/cuDNN kernel calls.
- TensorFlow (compiled with CUDA 7.0 and cuDNN 2) was approximately $2\times$ slower across all three networks, primarily due to the older underlying cuDNN version.
Knowl 8 — Distributed Training Scalability and Convergence of GoogLeNet on ILSVRC12
empirical result
Distributed scalability was evaluated by training GoogLeNet with Batch Normalization on the ImageNet ILSVRC12 dataset ( $1.3$ million images, $1{,}000$ classes) across Amazon EC2 g2.8x instances (each containing 4 Nvidia GK104 GPUs and 10 Gbps Ethernet). Training used stochastic gradient descent with learning rate $\eta = 0.05$ , momentum $0.9$ , weight decay $10^{-4}$ , and a mini-batch size of 36 images per GPU (144 images per host).

Key performance results comparing 1 machine (4 GPUs) and 10 machines (40 GPUs):
- Pass Execution Time: Average time per full data pass was $14{,}000\text{ s}$ on 1 machine and $1{,}400\text{ s}$ on 10 machines, representing an exact $10\times$ linear speedup in throughput.
- Convergence: While the 10-machine setup converged slightly slower in the initial few data passes due to the larger effective batch size, it caught up and surpassed the single-machine test accuracy after 10 data passes, reaching over $60\%$ top-1 test accuracy within 20 data passes.

Coverage note — Omitted standard auxiliary engineering utilities, such as the compact binary record packing tool, multi-threaded image data iterator prefetchers, and high-level training loop wrapper APIs.

References

1.Frederic Bastien, Pascal Lamblin, Razvan Pascanu, James Bergstra, Ian Goodfellow, Arnaud Bergeron, Nicolas Bouchard, David Warde-Farley, and Yoshua Bengio. Theano: new features and speed improvements. arXiv preprint arXiv:1211.5590, 2012.
2.Soumith Chintala. Easy benchmarking of all public open-source implementations of convnets, 2015. https://github.com/soumith/convnet-benchmarks.
3.Ronan Collobert, Koray Kavukcuoglu, and Clement Farabet. Torch7: A matlab-like environment for machine learning. In BigLearn, NIPS Workshop, number EPFL-CONF-192376, 2011.
4.J. Dean, G. Corrado, R. Monga, K. Chen, M. Devin, Q. Le, M. Mao, M. Ranzato, A. Senior, P. Tucker, K. Yang, and A. Ng. Large scale distributed deep networks. In Neural Information Processing Systems, 2012.
5.Chainer Developers. Chainer: A powerful, flexible, and intuitive framework of neural networks, 2015. http://chainer.org/.
6.Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167, 2015.
7.Yangqing Jia, Evan Shelhamer, Jeff Donahue, Sergey Karayev, Jonathan Long, Ross Girshick, Sergio Guadarrama, and Trevor Darrell. Caffe: Convolutional architecture for fast feature embedding. In Proceedings of the ACM International Conference on Multimedia, pages 675–678. ACM, 2014.
8.M. Li, D. G. Andersen, J. Park, A. J. Smola, A. Amhed, V. Josifovski, J. Long, E. Shekita, and B. Y. Su. Scaling distributed machine learning with the parameter server. In OSDI, 2014.
9.M. Li, D. G. Andersen, A. J. Smola, and K. Yu. Communication efficient distributed machine learning with the parameter server. In Neural Information Processing Systems, 2014.
10.Min Lin, Shuo Li, Xuan Luo, and Shuicheng Yan. Purine: A bi-graph based deep learning framework. arXiv preprint arXiv:1412.6249, 2014.
11.Abadi Martn, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mane, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viegas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: Large-scale machine learning on heterogeneous systems. 2015.
12.Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211–252, 2015.
13.Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, pages 1–42, 2014.
14.Minjie Wang, Tianjun Xiao, Jianpeng Li, Jiaxing Zhang, Chuntao Hong, and Zheng Zhang. Minerva: A scalable and highly efficient training platform for deep learning, 2014.

Citation

MLA

Chen, T., et al. “MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems”. arXiv, 2015, http://arxiv.org/abs/1512.01274v1.

APA

Chen, T., Li, M., Li, Y., Lin, M., Wang, N., Wang, M., Xiao, T., Xu, B., Zhang, C., & Zhang, Z. (2015). MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems. arXiv. http://arxiv.org/abs/1512.01274v1

Chicago

Chen, T., M. Li, Y. Li, et al. 2015. “MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems”. arXiv. http://arxiv.org/abs/1512.01274v1.

Harvard

Chen, T. et al. (2015) “MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems”, arXiv [Preprint]. Available at: http://arxiv.org/abs/1512.01274v1.

Vancouver

1. Chen T, Li M, Li Y, Lin M, Wang N, Wang M, Xiao T, Xu B, Zhang C, Zhang Z (2015) MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems. arXiv

BibTeX

@article{chen2015mxnet,
  title = {MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems},
  author = {Chen, Tianqi and Li, Mu and Li, Yutian and Lin, Min and Wang, Naiyan and Wang, Minjie and Xiao, Tianjun and Xu, Bing and Zhang, Chiyuan and Zhang, Zheng},
  year = {2015},
  journal = {arXiv},
  url = {http://arxiv.org/abs/1512.01274v1},
  eprint = {1512.01274}
}

Metadata:arXiv

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Blending Declarative Symbolic Graphs and Imperative Tensor Computation

Knowl 2 — Dynamic Dependency Engine with Explicit Resource Mutation Tracking

Knowl 3 — Linear-Time Memory Allocation Heuristics for Computation Graphs

Knowl 4 — Two-Level Hierarchical Parameter Server and Distributed KVStore

Knowl 5 — Computation Graph Transformations and Optimizations in MXNet

Knowl 6 — Memory Footprint Reduction via In-Place Reuse and Co-Sharing

Knowl 7 — Single-GPU Raw Performance Comparison across Deep Learning Frameworks

Knowl 8 — Distributed Training Scalability and Convergence of GoogLeNet on ILSVRC12

References

Citation

BibTeX

Source Code

Access the Paper