# Memory Safe Computations with XLA Compiler

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 43 references that are in the map.
- First posted on arXiv 28 Jun 2022.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Geometric deep learning: going beyond Euclidean data](https://mlsanity.com/paper/1611.08097/) (2016)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Deep Neural Networks as Gaussian Processes](https://mlsanity.com/paper/1711.00165/) (2017)
- [Gaussian Process Behaviour in Wide Deep Neural Networks](https://mlsanity.com/paper/1804.11271/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [Neural Tangents: Fast and Easy Infinite Neural Networks in Python](https://mlsanity.com/paper/1912.02803/) (2019)
- [GPflow: A Gaussian process library using TensorFlow](https://mlsanity.com/paper/1610.08733/) (2016)
- [ANN-Benchmarks: A Benchmarking Tool for Approximate Nearest Neighbor Algorithms](https://mlsanity.com/paper/1807.05614/) (2018)
- [Exact Gaussian Processes on a Million Data Points](https://mlsanity.com/paper/1903.08114/) (2019)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [Kernel methods through the roof: handling billions of points efficiently](https://mlsanity.com/paper/2006.10350/) (2020)
- [Glow: Graph Lowering Compiler Techniques for Neural Networks](https://mlsanity.com/paper/1805.00907/) (2018)
- [Understanding Probabilistic Sparse Gaussian Process Approximations](https://mlsanity.com/paper/1606.04820/) (2016)
- [Rates of Convergence for Sparse Variational Gaussian Process Regression](https://mlsanity.com/paper/1903.03571/) (2019)
- [Convergence of Sparse Variational Inference in Gaussian Processes Regression](https://mlsanity.com/paper/2008.00323/) (2020)
- [Supporting Very Large Models using Automatic Dataflow Graph Partitioning](https://mlsanity.com/paper/1807.08887/) (2018)
- [Tighter Bounds on the Log Marginal Likelihood of Gaussian Process Regression Using Conjugate Gradients](https://mlsanity.com/paper/2102.08314/) (2021)
- [Distributed Variational Inference in Sparse Gaussian Process Regression and Latent Variable Models](https://mlsanity.com/paper/1402.1389/) (2014)
- [The Generalized Matrix Chain Algorithm](https://mlsanity.com/paper/1804.04021/) (2018)

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

- [Pex: Memory-efficient Microcontroller Deep Learning through Partial Execution](https://mlsanity.com/paper/2211.17246/) (2022)
- [A differentiable brain simulator bridging brain simulation and brain-inspired computing](https://mlsanity.com/paper/2311.05106/) (2023)
- [Recommendations for Baselines and Benchmarking Approximate Gaussian Processes](https://mlsanity.com/paper/2402.09849/) (2024)
- [(A)iSpy: Parasitic Trojans for Machine Learning Infrastructure](https://mlsanity.com/paper/2607.17550/) (2026)

## Abstract (arXiv)

Software packages like TensorFlow and PyTorch are designed to support linear algebra operations, and their speed and usability determine their success. However, by prioritising speed, they often neglect memory requirements. As a consequence, the implementations of memory-intensive algorithms that are convenient in terms of software design can often not be run for large problems due to memory overflows. Memory-efficient solutions require complex programming approaches with significant logic outside the computational framework. This impairs the adoption and use of such algorithms. To address this, we developed an XLA compiler extension that adjusts the computational data-flow representation of an algorithm according to a user-specified memory limit. We show that k-nearest neighbour and sparse Gaussian process regression methods can be run at a much larger scale on a single device, where standard implementations would have failed. Our approach leads to better use of hardware resources. We believe that further focus on removing memory constraints at a compiler level will widen the range of machine learning methods that can be developed in the future.