Optimizing Automatic Differentiation with Deep Reinforcement Learning

• URL: http://arxiv.org/abs/2406.05027v3
• Date: Mon, 27 Jan 2025 10:21:04 GMT
• Title: Optimizing Automatic Differentiation with Deep Reinforcement Learning
• Authors: Jamie Lohoff, Emre Neftci,
• Abstract summary: We present a novel method to optimize the number of necessary multiplications for Jacobian computation by leveraging deep reinforcement learning (RL) We show that this method achieves up to 33% improvements over state-of-the-art methods on several relevant tasks taken from diverse domains.
• Score: 0.9353041869660692
• License:
• Abstract: Computing Jacobians with automatic differentiation is ubiquitous in many scientific domains such as machine learning, computational fluid dynamics, robotics and finance. Even small savings in the number of computations or memory usage in Jacobian computations can already incur massive savings in energy consumption and runtime. While there exist many methods that allow for such savings, they generally trade computational efficiency for approximations of the exact Jacobian. In this paper, we present a novel method to optimize the number of necessary multiplications for Jacobian computation by leveraging deep reinforcement learning (RL) and a concept called cross-country elimination while still computing the exact Jacobian. Cross-country elimination is a framework for automatic differentiation that phrases Jacobian accumulation as ordered elimination of all vertices on the computational graph where every elimination incurs a certain computational cost. We formulate the search for the optimal elimination order that minimizes the number of necessary multiplications as a single player game which is played by an RL agent. We demonstrate that this method achieves up to 33% improvements over state-of-the-art methods on several relevant tasks taken from diverse domains. Furthermore, we show that these theoretical gains translate into actual runtime improvements by providing a cross-country elimination interpreter in JAX that can efficiently execute the obtained elimination orders.

Related papers

• ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning [54.08906841213777]
  Asynchronous methods are fundamental for parallelizing computations in distributed machine learning. We propose ATA (Adaptive Task Allocation), a method that adapts to heterogeneous and random distributions of computation times. We show that ATA identifies the optimal task allocation and performs comparably to methods with prior knowledge of computation times.
  arXiv  Detail & Related papers  (2025-02-02T12:22:26Z)
• Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians [0.0]
  Jacobians and Hessians have many potential use cases in Machine Learning (ML), but conventional wisdom views them as computationally prohibitive. This paper presents advances in Automatic Sparse Differentiation (ASD), starting with a new perspective on sparsity detection. We also describe a novel ASD pipeline in Julia, consisting of independent software packages for sparsity detection, matrix coloring, and differentiation.
  arXiv  Detail & Related papers  (2025-01-29T16:21:54Z)
• An Efficient Algorithm for Clustered Multi-Task Compressive Sensing [60.70532293880842]
  Clustered multi-task compressive sensing is a hierarchical model that solves multiple compressive sensing tasks. The existing inference algorithm for this model is computationally expensive and does not scale well in high dimensions. We propose a new algorithm that substantially accelerates model inference by avoiding the need to explicitly compute these covariance matrices.
  arXiv  Detail & Related papers  (2023-09-30T15:57:14Z)
• One-step differentiation of iterative algorithms [7.9495796547433395]
  We study one-step differentiation, also known as Jacobian-free backpropagation, a method as easy as automatic differentiation. We provide a complete theoretical approximation analysis with specific examples along with its consequences in bilevel optimization.
  arXiv  Detail & Related papers  (2023-05-23T07:32:37Z)
• Gaussian Elimination versus Greedy Methods for the Synthesis of Linear Reversible Circuits [0.0]
  reversible circuits represent a subclass of reversible circuits with many applications in quantum computing. We propose a new algorithm for any linear reversible operator by using an optimized version of the Gaussian elimination algorithm and a tuned LU factorization. Overall, our algorithms improve the state-of-the-art methods for specific ranges of problem sizes.
  arXiv  Detail & Related papers  (2022-01-17T16:31:42Z)
• Efficient GPU implementation of randomized SVD and its applications [17.71779625877989]
  Matrix decompositions are ubiquitous in machine learning, including applications in dimensionality data compression and deep learning algorithms. Typical solutions for matrix decompositions have complexity which significantly increases their computational cost and time. We leverage efficient processing operations that can be run in parallel on modern Graphical Processing Units (GPUs) to reduce the computational burden of computing matrix decompositions.
  arXiv  Detail & Related papers  (2021-10-05T07:42:41Z)
• Self Normalizing Flows [65.73510214694987]
  We propose a flexible framework for training normalizing flows by replacing expensive terms in the gradient by learned approximate inverses at each layer. This reduces the computational complexity of each layer's exact update from $mathcalO(D3)$ to $mathcalO(D2)$. We show experimentally that such models are remarkably stable and optimize to similar data likelihood values as their exact gradient counterparts.
  arXiv  Detail & Related papers  (2020-11-14T09:51:51Z)
• Efficient Learning of Generative Models via Finite-Difference Score Matching [111.55998083406134]
  We present a generic strategy to efficiently approximate any-order directional derivative with finite difference. Our approximation only involves function evaluations, which can be executed in parallel, and no gradient computations.
  arXiv  Detail & Related papers  (2020-07-07T10:05:01Z)
• Global Optimization of Gaussian processes [52.77024349608834]
  We propose a reduced-space formulation with trained Gaussian processes trained on few data points. The approach also leads to significantly smaller and computationally cheaper sub solver for lower bounding. In total, we reduce time convergence by orders of orders of the proposed method.
  arXiv  Detail & Related papers  (2020-05-21T20:59:11Z)
• Accelerating Feedforward Computation via Parallel Nonlinear Equation Solving [106.63673243937492]
  Feedforward computation, such as evaluating a neural network or sampling from an autoregressive model, is ubiquitous in machine learning. We frame the task of feedforward computation as solving a system of nonlinear equations. We then propose to find the solution using a Jacobi or Gauss-Seidel fixed-point method, as well as hybrid methods of both. Our method is guaranteed to give exactly the same values as the original feedforward computation with a reduced (or equal) number of parallelizable iterations, and hence reduced time given sufficient parallel computing power.
  arXiv  Detail & Related papers  (2020-02-10T10:11:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.