Memory Safe Computations with XLA Compiler
Not on the map. Why?
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.
Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Builds on 25 of its 43 references that are in the map. First posted on arXiv 28 Jun 2022.
Builds on
25 of 43 references.
References · 25
- 2017Attention Is All You Needcited by 59,932
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
- 2022Hierarchical Text-Conditional Image Generation with CLIP Latentscited by 6,539
- 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,597
- 2016Geometric deep learning: going beyond Euclidean datacited by 1,042
- 2015Automatic differentiation in machine learning: a surveycited by 846
- 2017Deep Neural Networks as Gaussian Processescited by 617
- 2018Gaussian Process Behaviour in Wide Deep Neural Networkscited by 314
- 2018Mesh-TensorFlow: Deep Learning for Supercomputerscited by 168
- 2021Self-attention Does Not Need O(n2) Memorycited by 152
- 2019Neural Tangents: Fast and Easy Infinite Neural Networks in Pythoncited by 149
13 more
- 2016GPflow: A Gaussian process library using TensorFlowcited by 139
- 2018ANN-Benchmarks: A Benchmarking Tool for Approximate Nearest Neighbor Algorithmscited by 130
- 2019Exact Gaussian Processes on a Million Data Pointscited by 118
- 2020Kernel Operations on the GPU, with Autodiff, without Memory Overflowscited by 101
- 2020Kernel methods through the roof: handling billions of points efficientlycited by 88
- 2018Glow: Graph Lowering Compiler Techniques for Neural Networkscited by 73
- 2016Understanding Probabilistic Sparse Gaussian Process Approximationscited by 68
- 2019Rates of Convergence for Sparse Variational Gaussian Process Regressioncited by 62
- 2020Convergence of Sparse Variational Inference in Gaussian Processes Regressioncited by 53
- 2018Supporting Very Large Models using Automatic Dataflow Graph Partitioningcited by 43
- 2021Tighter Bounds on the Log Marginal Likelihood of Gaussian Process Regression Using Conjugate Gradientscited by 20
- 2014Distributed Variational Inference in Sparse Gaussian Process Regression and Latent Variable Modelscited by 19
- 2018The Generalized Matrix Chain Algorithmcited by 1
Built on by
Citing papers per year: 2022: 1; 2023: 1; 2024: 1.
3 papers.