# A Data-Centric Optimization Framework for Machine Learning

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 68 references that are in the map.
- First posted on arXiv 20 Oct 2021.

## Builds on (most cited first, 30 of 40)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [WaveNet: A Generative Model for Raw Audio](https://mlsanity.com/paper/1609.03499/) (2016)
- [YOLOv4: Optimal Speed and Accuracy of Object Detection](https://mlsanity.com/paper/2004.10934/) (2020)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Revisiting Unreasonable Effectiveness of Data in Deep Learning Era](https://mlsanity.com/paper/1707.02968/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Carbon Emissions and Large Neural Network Training](https://mlsanity.com/paper/2104.10350/) (2021)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [cuDNN: Efficient Primitives for Deep Learning](https://mlsanity.com/paper/1410.0759/) (2014)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [The Hardware Lottery](https://mlsanity.com/paper/2009.06489/) (2020)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [Glow: Graph Lowering Compiler Techniques for Neural Networks](https://mlsanity.com/paper/1805.00907/) (2018)
- [MLIR: A Compiler Infrastructure for the End of Moore's Law](https://mlsanity.com/paper/2002.11054/) (2020)
- [FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems](https://mlsanity.com/paper/2008.11359/) (2020)

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

- [Cached Operator Reordering: A Unified View for Fast GNN Training](https://mlsanity.com/paper/2308.12093/) (2023)
- [OpenHLS: High-Level Synthesis for Low-Latency Deep Neural Networks for Experimental Science](https://mlsanity.com/paper/2302.06751/) (2023)
- [DaCe AD: Unifying High-Performance Automatic Differentiation for Machine Learning and Scientific Computing](https://mlsanity.com/paper/2509.02197/) (2025)

## Abstract (arXiv)

Rapid progress in deep learning is leading to a diverse set of quickly changing models, with a dramatically growing demand for compute. However, as frameworks specialize performance optimization to patterns in popular networks, they implicitly constrain novel and diverse models that drive progress in research. We empower deep learning researchers by defining a flexible and user-customizable pipeline for optimizing training of arbitrary deep neural networks, based on data movement minimization. The pipeline begins with standard networks in PyTorch or ONNX and transforms computation through progressive lowering. We define four levels of general-purpose transformations, from local intra-operator optimizations to global data movement reduction. These operate on a data-centric graph intermediate representation that expresses computation and data movement at all levels of abstraction, including expanding basic operators such as convolutions to their underlying computations. Central to the design is the interactive and introspectable nature of the pipeline. Every part is extensible through a Python API, and can be tuned interactively using a GUI. We demonstrate competitive performance or speedups on ten different networks, with interactive optimizations discovering new opportunities in EfficientNet.