# MONET: Modeling and Optimization of neural NEtwork Training from Edge to Data Centers

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 41 references that are in the map.
- First posted on arXiv 16 Mar 2026.

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

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization](https://mlsanity.com/paper/1603.06560/) (2016)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Estimating the Carbon Footprint of BLOOM, a 176B Parameter Language Model](https://mlsanity.com/paper/2211.02001/) (2022)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [On-Device Training Under 256KB Memory](https://mlsanity.com/paper/2206.15472/) (2022)
- [The rising costs of training frontier AI models](https://mlsanity.com/paper/2405.21015/) (2024)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [Stream: Design Space Exploration of Layer-Fused DNNs on Heterogeneous Dataflow Accelerators](https://mlsanity.com/paper/2212.10612/) (2022)
- [FuseMax: Leveraging Extended Einsums to Optimize Attention Accelerator Design](https://mlsanity.com/paper/2406.10491/) (2024)
- [DaCe AD: Unifying High-Performance Automatic Differentiation for Machine Learning and Scientific Computing](https://mlsanity.com/paper/2509.02197/) (2025)

## Abstract (arXiv)

While hardware-software co-design has significantly improved the efficiency of neural network inference, modeling the training phase remains a critical yet underexplored challenge. Training workloads impose distinct constraints, particularly regarding memory footprint and backpropagation complexity, which existing inference-focused tools fail to capture. This paper introduces MONET, a framework designed to model the training of neural networks on heterogeneous dataflow accelerators. MONET builds upon Stream, an experimentally verified framework that that models the inference of neural networks on heterogeneous dataflow accelerators with layer fusion. Using MONET, we explore the design space of ResNet-18 and a small GPT-2, demonstrating the framework's capability to model training workflows and find better hardware architectures. We then further examine problems that become more complex in neural network training due to the larger design space, such as determining the best layer-fusion configuration. Additionally, we use our framework to find interesting trade-offs in activation checkpointing, with the help of a genetic algorithm. Our findings highlight the importance of a holistic approach to hardware-software co-design for scalable and efficient deep learning deployment.