Optimizing Automatic Differentiation with Deep Reinforcement Learning
Authors: Jamie Lohoff
Emre Neftci
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Reinforcement Learning
Automatic Differentiation
AlphaZero
Deep Learning
Primary area: reinforcement_learning
TL;DR: We trained a neural network to produce new automatic differentiation algorithms using Deep RL.

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.
Finding the optimal elimination order that minimizes the number of necessary multiplications can be seen as a single player game which in our case 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 relevant domains.
Furthermore, we show that these theoretical gains translate into actual runtime improvements by providing a cross-country elimination interpreter in JAX that can execute the obtained elimination orders.

Optimizing Automatic Differentiation with Deep Reinforcement Learning
Jamie Lohoff
Emre Neftci
Peter Grünberg Institute
Peter Grünberg Institute
Forschungszentruum Jülich & RWTH Aachen ja.Johoff@fz-juelichhcoo
Forschungszentrum Jülich & RWTH Aachen e. meftci@ff-juelicce
Abstract
Computing Jacobians with automatic diffcrentiation 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 (RI.) and a concept called cross-country elimination while still computing the exact Jacobian. Cross-country elimination is framework for automatic differentiation that phrases Jacobian accumulation as ordered elimination of all vertices on the computational graph where every elimination incurs certain computational cost. We formulate the search for the optimal elimination order that minimizes the number of necessary multiplications an RL agent. We demonstrate that this method achieves up to 33% improvements over state-of-the of-the-art art methods several relevant tasks taken from diverse domains. Furthermore, we show that these theoretical gains translate into actual runtime improvements by providing cross-country elimination interpreter in JAX that can efficiently execute the obtained elimination orders.
single player game which is played by
Introduction
Automatic Differentiation (AD) is widely utilized for computing gradients and Jacobians across diverse domains including machine learning (ML), computational fluid dynamics (CFD), robotics, differential rendering, and finance [Baydin
al. 2018 Margossian, 2018 Forth al.) 2004a
Tadjouddine et al., 2002b Giftthaler et al. 2017 Kato et al. 2020 [Schiid et al. Caprioti and Giles 2011 Savine and Andreasen 2021 To many researchers in the machine learning community, AD is synonymous with the backpropagation algorithm [Linnainmaa] 1976) Scchiidhuber [014]. However, backpropagation is just one particular way of algorithmically computing the Jacobian that very efficient in terms of computations for "funnel-like" functions, i.e. with many inputs and a single scalar output such in neural networks. In many other domains, do not have this particular property and thus backpropagation might not be optimal for computing the respective Jacobian [Albrecht et al. 2003 Capriotti and Giles 2011 Naumann, 2020]. In fact, there exists a wide variety of AD algorithms, each of them coming with its own ad advantages and drawbacks regarding computational cost and memory consumption depending
may find functions that
the function they are applied to. Many of these AD algorithms can be viewed as special cases of cross-country elimination [Griewank] and Walther 2008). Cross- Cross-County Elimination frames AD ordered vertex climination problem on the computational graph with the goal of reducing the required number of multiplications and additions. However, finding the optimal elimination procedure is a NP-complete problem [Naumann
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

Computational Graph
Reinforcement Learning
Optimal Autmatic Differentiation
elimination action sequence
Function Python JAX
State
fix, y):
Python code Graphax
x-y
V2 sin[v]]
return Log(v2), v1.v2
Graph Repr.
Policy
jitjoee/ andenender. angease 11..1, 2.1
AlphaZero
Optimal AD algorithm as "jaxpr" for execution in JAX
Compulational graph and adjacency matrix representation
a11111 krttt)
action eliminate vertex (= application of chain rule) reward -computational cost
Figure
Summary of the AlphaGrad pipeline. We trained neural network to produce new Automatic Differentiation (AD) algorithms using Deep RL that can be used in JAX. The resulting algorithms significantly outperform the current state of the art.
2008 Inspired by recent advances in finding optimal matrix-min paarri--muttpliiaation and sorting algorithms
Fawzi et al., 2022 Mankowitz et al. 2023 we demonstrate that deep RL successfully finds efficient elimination orders which translate into new automatic differentiation algorithms and practical runtime gains (figure I. Cross-country elimination is particularly amenable for automatization since it provably yields the exact Jacobian for every elimination order. The solution we seek thus reduces to only finding optimal elimination order, without the need to evaluate the quality of Jacobian approximations (such as in Neural Architecture Search). An important body of prior work aimed to find more efficient climination techniques through hcuristics, simulated annealing dynamic programming and minimizing related quantities such as fill-in [Naumann] [999] 2020J. However, none of these works successful in optimizing with respect to relevant quantities such number of multiplications or memory consumption. We set up our optimization problem by formulating cross- single player RL game called VertexGame. At each step of VertexGame, the
country elimination agent selects a vertex to eliminate from the computational graph according to a certain scheme called vertex elimination. The reward is equal to the negative of the number of multiplications incurred by the particular choice of vertex. VertexGame
agent (Silver et al. 2017 Schrittwieser et al. 2019 Danihelka et al., 2022] with policy and value functions modeled with transformer architecture that processes the graph representation and predicts the next vertex to eliminate, thereby incrementally building the AD algorithm.
played by an AlphaZero-base
Our approach discovers from scratch new vertex elimination orders, i.e. new AD algorithms that are tailored to specific functions and improve over the established methods such as minimal Markowitz degree. We further demonstrate the efficacy the discovered algorithms on real world tasks by including Graphax, a novel sparse AD package which builds on JAX [Bradbury et al., 2018 and enables the user to differentiate Python code with cross-country elimination. Our main contributions are summarized as follows:
We demonstrate that optimizing elimination order can be phrased as a reinforcement learning game by leveraging the graph view of AD,
We show that deep RL agent finds new, tailored AD algorithms that improve the state-of- the-art on several relevant tasks,
We investigate how the discovered novel elimination procedures translate into actual runtime improvements by implementing Graphax, cross-country elimination interpreter JAX allowing the efficient execution of newly found climination orders.
Related Work
RL for Algorithm Research: AlphaTensor and AlphaDev successfully demonstrated that model- based deep RL finds new and improved matrix-min muliiix-ullilliaation and sorting algorithms (Fawzi et al.) 2022 Mankowitz et al. 2023 In particular, AlphaTensor used extension of the AlphaZero agent to search for new matrix- multiilulllpiiiiinnns algorithms that require fewer multiplication operations by directly using this quantity reward. The key insight is that different matrix-min muttiiiiuulipiaiiionn

---

algorithms have a common, simple representation through the three-dimension matrix-mulipllcce tensor which can be manipulated by taking different actions, resulting in algorithms of varying efficiency. Feeding this tensor into the RL agent, they successfully improved on matrix cnttrix-muliplicattion algorithms for 4x4 matrices by beating Strassen's 's algorithm, the current state-of-the-art, with an improvement from 49 to 47 multiplications.
In similar vein, AlphaDev improved simple sorting algorithms by representing the sorting algorithm as a series of CPU instructions which then have to be arranged in the correct way to achieve the correct sorting output. Instead of using the number of CPU operations agent was trained on actual execution times. Our work follows in these footsteps by tackling the difficult problem of finding new and improved AD algorithms for arbitrary functions and hence we termed our method AlphaGrad.
an optimization target, the
RL for Compiler Optimization: A number of works have tackled the complex issue of optimizing the compilation of various computational graphs with deep RL. Knossos [Jinnai et al. 2019] leverages the A algorithm optimize the compilation of simple neural networks. It employs a model to estimate computational cost and utilizes expression rewriting techniques to enhance performance. While Knossos is hardware agnostic, it needs to be trained from scratch for every new computational graph. GO and REGAL both improve on this shortcoming and generalize to new, unseen graphs at
the cost of losing the hardware-agnostic property [Paliwal et al. 2019] Zhou et al.. [2020]
REGAL learns a graph neural network-based policy using a RHINFORCE- optimize the scheduling of the individual operations of a graph to the set of available devices, thereby successfully reducing peak memory usage for different deep learning workloads. Only GO is directly trained actual wall time and handles all relevant optimizations jointly, including device placement, operation fusion, and operation scheduling. GO learns a policy based on graph neural networks and recurrent attention using PPO and successfully demonstrates improvements over Tensorflow's 's default compilation strategy. While our work also makes use of the computational graph, the goal is to find novel AD algorithms instead of optimizing compilation itself, although these problems are related since the new AD algorithm is compiled before execution.
genetic algorithm to
Optimization of AD While no prior work directly aims at improving AD with deep RL. several studies aimed at enhancing AD by other methods. This includes Enzyme [Moses and Churavy/ 2020), which presents reverse-mode AD package that operates on the intermediate representation level using LLVM. Enzyme is thus distinct from other AD packages because it can synthesize gradients for many different high-level languages. Another related work is LAGrad (Peng and Dubach| 2023), , a source-to-source AD package written in Julia which introduces a set of new static optimizations to accelerate reverse-mode AD. It leverages high-level MLIR information, such as sparsity structure and control flow semantics of the computational graph to produce more efficient differentiated code. While LAGrad can improve the performance of AD workloads by orders of magnitude, it is currently limited to the use of reverse-mode mode AD which can be suboptimal for certain tasks. both works, a suboptimal choice of algorithm might nullify the benefits gained from careful engineering. Our work aims to close this gap by additionally providing a novel way of finding the optimal AD algorithm using Deep RI.
The closest related work is [Naumann] 1999 where simulated annealing was applied to reduce the number of multiplications necessary for Jacobian accumulation. The algorithms struggled to significantly outperform state-of-the-ee even when was initialized with a reasonably good elimination order. In a similar manner, [Naumann| 2020) directly optimized the elimination order with dynamic programming albeit with respect to a different optimization target called fill-in on randomly generated graphs that do not necessary represent well-defined, executable functions.t Our work directly optimizes for the number of multiplications required accumulate the Jacobian on real-world problems. Another approach described in [Chen et al., 2012) utilized integer linear programming to find optimal climination orders with respect to number of multiplications, but only dealt with very small problems with up to twenty intermediate vertices. Our approach successfully finds new AD algorithms from scratch for complex problems with hundreds of intermediate vertices.
2 Automatic Differentiation and Cross-Counte Elimination
AD is systematic approach to computing the derivatives of dependent variables y = fxx R with respect to the independent variables x E R" utilizing the chain rule. AD enables the precise and efficient calculation of gradients, Jacobians, Hessians, and higher-order derivatives [Linnainmaa 1976). Unlike methods that rely on finite differences or symbolic differentiation, AD offers a systematic way to compute derivatives up to machine precision, making it an indispensable tool in
3

---

31
log
X1
log
sin
COS D
sin
13
X2
X2
(a) Computational graph of f(11,
(b) Adding partial derivatives.
Xany'
Fs' cos
log
log
cos
x
sin
Fz0el cos
1-1 cDs f)
12
(c) Elimination vertex 2.
(d) Graph after eliminating vertex 1.
Figure 2: Step-by-step description of cross-country elimination with the simple example function [x1 x22 = (log sin 22), 21 sin (#1=))) (a) Initial computational graph. (b) The partial derivatives are added the edges of the computational graph. The intermediate variables and t2 are defined through i1 = 22 and U/ = sin U1. O Elimination vertex 2 associated with the sin operation. The dotted red lines represent the edges that are deleted. {d) Final bipartite graph after both intermediate vertices have been eliminated. All remaining edges contain entries of the Jacobian.
many numerical scientific problems and machine learning [Baydin et al. 2018 Griewank and Walther 2008]. AD leverages the fact that most computer programs can be broken down into a sequence of simple elemental operations, for example additions, multiplications and trigonometrie functions. Partial derivatives of these elemental operations are coded into the AD software and the Jacobian accumulated by recursively applying the chain rule the evaluation procedure. Since the partial derivatives are known up to machine precision, AD gives the Jacobian up to machine precision.
2.1 Graph View and Vertex Elimination
We take the graph view of AD where a function is defined through its computational graph G = (V, E) with its vertices being the elemental operations Φj and directed edges & that describe the data dependencies between the operations (figure(a). The relation i < states that vertex i has an edge connecting it with vertex j. meaning that the output v; of di is an input of dj.. The partial derivatives of the elemental operations with respect to their dependents are assigned to the connecting edges (figure 2b We then identify the edges of the graph with their respective partial derivatives Ojz Sv. The cross-country elimination algorithm computes the Jacobian by a procedure called vertex elimination.
Definition [Griewank and Walther 2008 For a computational graph G = (V, E) with partial derivatives Ctj, veriex elimination of vertex j is defined the update
Cys += CkjCzs V(i, k) V x V where < and j < k:
(1)
and then setting Cjt = Ckj for all involved vertices. The += operator creates edge there no edge Cks and otherwise adds the new value to the existing value.
Intuitively, vertex elimination be understood the local application of the chain rule single vertex in the graph since the multiplication CyCfk is exactly the result of applying the chain rule to (d& (w.)) If vertex has multiple incoming and outgoing edges, all combinations of incoming and outgoing edges are resolved to create new edges. If an edge already exists, we add the result of the product to it, in accordance with the rules for total derivatives. After the new edges added to the graph, we delete all edges connected to the climinated vertex since all the derivative information is now contained in the edges (figure 2c). Note that the climination no longer directly represents the data dependencies of the function since the climinated vertex is now disconnected. Furthermore, the application of vertex climination as described above
graph resulting from vertex

---

requires the computational graph to be static and precludes the use of control flow (if statements, for-loops) the differentiated code.
2.2 Cross Country and Elimination Orders
The repeated application of the vertex elimination procedure computational graph (i.e. cross country elimination) until all intermediate vertices are eliminated will yield a graph where the input vertices and output vertices are directly connected by edges (no intermediate vertices left, see figure 2d.. This is called a bipartite graph and the edges of this graph contain the components of the long all intermediate vertices are eliminated, this Jacobian will always be exact up to machine precision [Griewank and Walther, 2008). There is no restriction the order in which the vertices are eliminated, but the choice will significantly influence computational cost and memory [Tadjouddine et al. 2006]. In the graph view, computational cost is straightforward to measure since every vertex elimination incurs known number of multiplications that depends on the shapes of the elemental Jacobians which can be used as a proxy for execution time. We ignore the cost of evaluating the partial derivatives since they have to be performed regardless of the elimination order. Thus, we use the number of multiplications as the optimization target for the remainder of this work. The two most common choices for elimination orders to either eliminate the vertices in the forward or reverse order. These two modes are called forderd-mode AD and reverse-mode AD (backpropagation), respectively.
Jacobian of the function f. In particular,
Forward-mode AD, where vertices are climinated in the same order as the computational graph is traversed, is particularly efficient for functions where the number of input variables is much smaller than the number of output variables m, i.c. n < m. In contrast, reverse-mode AD traverses the graph in the opposite direction and is particularly suited for the cases where > This is the case in machine learning and neural networks using scalar loss functions, which is why reverse-mode AD is the default choice in such workloads.
2.3 Minimal Markowitz Degree
A more advanced technique is to eliminate vertices with the lowest Markowitz degree first Griewank and Walther (2008). The Markowitz degree of vertex is defined as the number of incoming vertices times the number of outgoing vertices, i.e. Mark(j) = j < زاار < kl, where cardinality of the sets < j and j for fixed j. Thus the elimination order is constrained by finding the vertex with the lowest Markowitz. degree first, eliminating it and then finding the next vertex with minimal Markowitz degree the resulting graph. This elimination scheme is of the best known heuristics for finding efficient elimination orders and can incur savings of up to 20% over forward- and moverse-oode AD [Albrecht al.] 2003 Griewank and Walther] 2008 However, for computational graphs that have many inputs and few outputs, it is often outperformed by meverse-mode AD.
denotes the
2.4 Vector-valued Functions
In most applications, vector-valued functions are used as elemental building blocks of more complex functions. While in most cases, these vectorized operations could be broken down into scalar operations, this would be impractical since it would increase the size of the computational graph representation and action space by orders of magnitude. Thus, it is best to allow vertices of the computational graph to be vector-valued which results in the partial derivatives assigned to the edges becoming Jacobians in their own right. The multiplication operations during vertex elimination are then accordingly replaced with matrix multiplications or higher-order contractions of the elemental Jacobians. For many operations. the Jacobians themselves have a particular internal sparsity structure which can be exploited when performing the climinations. A simple example is the multiplication of a vector with matrix followed by the application of a non-linear function f(x, W) tanh (W x). The input vertices are given by the input x and weights W and the intermediate vertex is matrix multiplication a; = -E, Wirar with the partial derivatives
da,
0o: 05k
2
W ik.
(2)
aW
5

---

Ixpul Acken
W
Wsj
tanh
Graphax instead does fast elementwise multiplication
x
Output shape
Σ
Input shape
bybi
Sparsity type
x;0g
xz0p
Xy X2 2g
Smeciin
(a) Sparse vertex elimination.
(b) Computational graph representation.
Figure (a) Graphax implements sparse vertex elimination benefit from the advantages of cross country elimination. (b) Sketch of the diree-dimnnsional adjacency tensor that represents the computational graph. The colored surfaces represent the five different values encoded in the third dimension. The red and blue surfaces together contain the shape of the Jacobians while the green surface encodes their sparsity. The vertical dotted slices represent the input connectivity of single vertex. In this work, we compress and feed the vertical slices as tokens into the transformer backbone such that we build sequence running direction of the black arrow.
The output vertex represents the application of the activation function = tanh with the partial derivative
oy:
و. tanh? a.).
da,
According to the vertex elimination rule, upon elimination of the intermediate vertex the two Jacobians equation (2) are assigned the incoming edges are contracted together with the Jacobian from the outgoing edge in equation 33:
dy:
Σ (1 tanh? a,)6,; = 8,v tanh? az) az)x1,
dy Σ (1 - tanh? a.)o,, W = (1- 111-tan tanh? a. W.).
(5)
Oik
In both cases, instead of a matrix multiplication, one can perform simple element- wise multiplications as shown in figure 3al For vectorized cross-country elimination be efficient, it is paramount to exploit this property. Current state-of-the-arr AD frameworks typically lack the ability to perform cross-country elimination and subsequently can not deal with sparse Jacobians. The only exception the authors are aware of is EliAD, an AD interpreter in C++ which is fully capable of processing given elimination orders and create the derivative source code [Tadjouddine et al., 2002a However, we developed Graphax as a novel AD interpreter that builds on Google's 's JAX |Bradbury et al. 2018) in order to leverage it's defining features such as JIT compilation, automated batching, device parallelism and user fueeffryy Python front-end. Graphax is a fully fledged AD interpreter capable of performing cross-country elimination as described above and outperforms JAX' AD on the relevant tasks by several orders of magnitude (appendix (B). Graphax and AlphaGrad are available under and http://iww.oommmmmmoooofo/////ee
and https://github.com//mmam https//ghhabbbom//mmmooooffttthp
2.5 Computational Graph Representation and Network Architecture
We describe here how the computational graph is represented for optimization the RL algorithm, well the network architecture that is optimized with AlphaZero. In the scalar case, the computational graph can be represented by its adjacency matrix, meaning that for every pair of vertices (i.j) that share edge, the vectorized case, we define an extended adjacency tensor by extending the matrix into the third dimension. Along this third dimension, store values that describe the sparsity pattern and shape of the Jacobian associated with the respective edge. The first value is an integer between -10 and 10 which encodes the sparsity type of the Jacobian. Details about the supported sparsity types can
the i-th row and the j-th column of the matrix to 1. For

---

be found in Appendix C The next four values contain the shape of the Jacobian associated with the respective edge and thus imply that this representation can at most deal with Jacobians of the shape Brai where the first two values describe the shape of T&i and the other two values describe the shape of Viy. This can be expanded to arbitrary tensor sizes of x and y. but then also requires the definition of new sparsity types to account for the additional dimensions. Figure 3bj shows the representation of the entire computational graph and single selected edge with a Jacobian of shape (4, 2, 4, 2) with sparsity type 3. A horizontal or vertical slice of the extended adjacency tensor gives the input or output connectivity of particular vertex. These slices can be compressed and used as tokens to be fed into a transformer where they are processed simultaneously by the attention mechanism so that the model gets a full view of the graph's 's connectivity.
dus
In this work, we compress vertical slices into tokens using a convolutional layer with kernel size (3, 5) and use a linear projection to create a 64-dimensional embedding. We found it helpful to apply a positional encoding to the tokens Vaswani et al. [017]] The output of the transformer is then fed into a policy and value head. The policy head is a MLP mapped across every token separately, thus creating a probability distribution over the vertices to determine the next one to be eliminated.
Already eliminated vertices are masked. Similarly. the value head is also a MLP that predicts a score for every token. These scores are then summed to give the value prediction of the network.
3 Reinforcement Learning for Optimal Elimination Orders
Cross-country elimination is typically introduced as a means to reduce the computational cost of computing the Jacobian. We cast the problem of finding an efficient vertex elimination order as a single-player RL game called VertexGame. At every step of the game, the agent selects the next vertex to be eliminated by observing the current connectivity of the computational graph. Since it is difficult to directly optimize for execution time, it is common use the number of multiplications incurred by the climination order as a proxy value (Tadjouddine et al., 2006 2002b Albrecht et al. 2003]. Thus, we chose the negative number of multilications incurred by eliminating the selected vertex as reward. We use action masking to prevent the agent from eliminating the same vertex twice. This also ensures that the accumulated Jacobian is always exact and has clear terminal condition: when the extended computational graph is bipartite, the game ends.
Between elimination orders, the magnitude of the reward can range across multiple orders of magni- tude. To tackle this, we rescale the cumulative reward using a monotonous function. For the functions with scalar inputs as well as RoeFlux_36 and random function f, we found the method presented in Kapturowski al. 2019) performed well, i.e. we scaled with sr)) sgu(r) (VTl + 1) + er where e 10-3 For MLP and TrunsformerEncoder tasks, the best performance was achieved with logarithmic scaling s(r) = log r [Hafner et al., 2024). VertexGame is played by an AlphaZero agent, which successfully finds new AD algorithms. To reduce the computational cost of the AlphaZero agent (Silver et al. 2017]. we employed Gumbel action sampling Danihelka et al. 2022). Gumbel AlphaZero is policy improvement algorithm based on sampling actions without replacement which utilizes the Gumbel softmax trick and other augmentations. This algorithm is guaranteed to improve the policy while significantly reducing the number of necessary Monte-Carle Tree Search (MCTS) simulations. On most tasks, we found that 50 MCTS simulations were sufficient to reach satisfactory performance. Appendix contains more details about the training of the agent.
4 Experiments
To demonstrate the effectivvness of our approach, we devised a set of tasks sampled from different scientific domains where AD is used compute Jacobians. More details concerning the tasks are listed in appendix A
Deep Learning is prime example for the success of large-scale AD. We analyze two-layer MLP with layer norm as described in [Goodfellow et al. 2016) and a small-scale version of the transformer encoder (Dosovitskiy et al.) 2020).
Computational Fluid Dynamics relies on AD for computation of the flux Jacobian on the boundaries of the simulation grid cells. The RoeFlux is particularly relevant and has been studied extensively with vertex climination in the past [Roc] 1981 Tadjouddine et al., 2002b Zubair et al. 2023]. We test the ID and the 3D variants of this problem.
Differential Kinematics uses Jacobians to quantify the behavior of a robot or other mechanical system with respect to their controllable parameters (c.g. joints, actuators). There has been surge
7

---

Table Number of multiplications required by the best discovered elimination order for batch size of one. Results obtained from VertexGame played by the AlphaZero agent with 50 MCTS simulations and Gumbel noise scale of 1.0. marks the experiments where we employed a log-scaling of the cumulative reward instead of the default scaling. The values in parentheses were obtained for 250 MCTS simulations.
Task
Forward
Reverse
Markowitz
AlphaGrad
RoeFlux_1e 1d
620
364
407
320
RobotArm_6ln 6DOF
397
301
288
231
240
172
194
149
PropaneCombustion
151
90
111
88
Random function g
632
566
451
417
BlackScholes
545
572
350
312
RoeFlux_3a
1556
979
938
811
Random function f
17728
9333
12083
6374
2-layer MLP
10930
392
4796
398 (389)
Transformer'
135010
4688
51869
4831 (4656)
in interest of computing the Jacobian using AD Difitthhal et al., 2017]. We chose the forward kinematics of 6-DOF robot arm as representative problem and follow [Dikmenli 2022 for the implementation.
Non-Linear Equation Solving requires the computation of large Jacobians to apply state-of oftte-of-the- the- art solvers. The MINPACK problem collection provides a set of problems derived from real-life applications of non-lineat optimization and designed to be representative of commonly encountered problems. particular, we analyze the HumanHearzDiple and Propane Combustion tasks. for which vertex climination has also been analyzed thoroughly in [Forth et al. 2004b) Averick et al. 1992]. Computational Finance makes use of AD for fast computation of the called "greeks" which measure the sensitivities of the value of an option to the model parameters NNummann] 2010) Saaiine and Andreasen 2021]. Here, we compute the second-order greeks of the Black-Scholes equation using AD by computing the Hessian of the Black-Scholes equation through evaluation of the Jacobian of the Jacobian Black and Scholes [1973]. This way. this task serves two-fold purpose by also demonstrating how our approach is also useful for finding good AD algorithms for higher-order derivatives.
Random Functions are also commonly used to evaluate the performance of new AD algorithms AAlbrecht et al. 2003]. We generated two random funetions, and g with vector-valued and only scalar inputs respectively. The random code generator used to generate these arbitrary functions is included in the accompanying software package.
4.1 Finding Optimal Elimination Orders
Table I shows the number of multiplications required by the best elimination order found over 6 runs with different seeds. The model was trained from scratch on each task separately with a batch size of 1 to to keep the rewards small possible. The resulting AD algorithms nonetheless scalable to arbitrary batch sizes. We use forward-mode reverse-mos mode and the minimal Markowitz degree method baselines for comparison. The first six tasks are simple functions with only scalar inputs and simple operations and the cumulative reward stays within the same order of magnitude, making them easier to solve. For all tasks, approach improvements ranging from 2% to almost 20%. We found that even for only 5 MCTS simulations, the agent able to find better than state-of-the -ttthh-art art solutions for the scalar tasks.
able find new elimination orders with
On the opposite spectrum, our experiments with 250 MCTS simulations yielded no significant im- provement over the results presented in table [ The four remaining tasks arguably more difficult since the vector-valie valued inputs and large variance within possible rewards provide an additional chal- lenge. The RoeFlux_5 3d and random function were solved successfully with 50 MCTS simulations and yielded improvements of up to 33%. This is in stark contrast to prior work such as [Naumann] 1999]. where algorithms such simulated annealing beat common heuristics such as minimal Markowitz or reverse-mode AD. With a budget of only 50 MCTS simulations, AlphaGrad failed to find improvements for both deep learning tasks.
dynamic programming struggled to even

---

Table Median runtimes for the results obtained in table Results were measured with Graphax for batch size 512 on an AMD EPYC 9684X 2x96-Core processor. Uncertainties are given as 2.5- and 97.e-peres over 1000 trials. Execution time is given in milliseconds and default XLA compilation flags were used for all experiments. The size of the networks, i.c. the number of neurons were increased for the MLP and Transformer Encoder by a factor of 16 to create a more realistic sample. GPU experiments were run on a NVIDIA RTX 4090 with JIT compilation.
Task
Forward
Reverse
Markowitz
AlphaGrad
RoeFlux_Ia ld
3.03 0.11
3.08
2.87 H.
2.19
RobotArm_6DoO
8.85 0.17
+0:21
8.48+0.32
0,200
0.25
8.55 +0.38
0.36
0.66
6.05 +033
-0.35
HumantHearDDipole
16.97
16.87
16.42+
15.94+1.
36.91 +2.89
PropaneCombustion
36.47*
36.99+1.35
36.45+2
Random function g
82.41
2.88
81.64 81.64+2.82
82.97 1.17
80.56 2,30
2.07
3.56
BlackScholes
5.04 0.33
0.23
5.02 0.39
5.03
4.77 -0.29
0.29
RoeFlux _3d
72.26
83.98 +.669
91.27
63.92
6.02
6.68
5.76
Random function f
12.42
11.54
20.20
9.12+0.08
1:33
2-layer MLP
760.63*
29.67
317.65 19.34
28.73 3,73
53.30
2-layer (LP((GPP)
11.04+
0.30+0.02
1.01+0.02
0.29
0.13
0.03
Transformer
990.09
39.07 14.80
498.94 +2.34
40.38
31,11
3.62
Transformer (GPU)
21.38 0.06
0.22
0.29
16.17 16.17-0.09
15.04
0.28+0.02
0.03
0..0
The authors conjecture that this is not only due difficulties presented above, but because reverse- mode AD (backpropagation) is already a very well-suited algorithm for computing Jacobians of "funnel-like" computational graphs with many inputs and single, scalar output. Despite this, with an increase to 250 MCTS simulations, the agent marginally outperformed backpropagation for both deep learning models. Appendix E contains more information about the experiments, including reward curves, the actual elimination orders and more details about their implementation. Note that the results in table were obtained by separately training each single function/graph.
We also include joint training runs where the agent was trained on all tasks at once. While the results were inferior the separate training mode, the agent found new, improved elimination orders for almost all tasks except the MLP, "ransformerEEncoder, RoeFlux_54 and CropaneCombustionn tasks. For the random function f and BlackScholes Jacobban task, the multi-task training outperformed the results in table [] with new best results of 5884 and 307 respectively, thereby showing that the algorithm search might benefit from training diverse tasks simultaneously. This also hints the possibility of building a more general statistical model of AD applicable to workloads from many different domains. We also experimented with PPO as alternative (appendix F)..
4.2 Runtime Improvements and the Graphax library
The results in table aree mainly of theoretical value. Here, we investigate how these translate into actual runtime improvements. For this purpose, we implemented Graphax, to our knowledge the first wython-based AD interpreter able to leverage cross-country elimination. Graphax builds second program that computes the Jacobian by leveraging the climination orders found by AlphaGrad and using the source code of the function JAX' own representation of the computational graph of the function in question.
template by analyzing its Jaxpression. The Jaxpression is
2 runtime improvements for the elimination orders found in ta batch size of 512 with varying levels of improvement. This is due to the fact that the number of multiplications alone was only proxy to capture the complexity of the entire program. It ignores other relevant quantities such as memory accesses and operation fusion during compilation. Nonetheless, particularly impressive gain over the ofate-of-the-ee art methods can be observed for the RoeFlux tasks and the RobotArm _DDF task. Remarkably, we also observe a minor improvement for both deep learning tasks when executed on GPUs. Note that the TransformerEncooer task evaluated on batch size of because VertexGame only supports two-dimensional input tensors. Figure 4 shows how some of AlphaGrad's algorithms scale with growing batch size. Appendix B provides an in-depth comparison of our work and JAX' own AD modes. In general, the combination of AlphaGrad and Graphax is able to outperform the JAX AD modes in most cases, sometimes by orders of magnitude. While

---

[Forhh et al. 2004b and [Tadjouddine et al. [006] present state-of- teefthh-art results for some of the investigated tasks, we were unable to reproduce the experiments given their implementation details.
Roe_ _oe_3d 3d evaluation times for different mocies and batch sizes
MLP evaluation times different modes and batch
10a
16'
JAK reverse mode
Eveluation time [ms]
Evaluatton time [masl
Graphax AlphaGrad
10a
10h
100
JAX forward mode JAX reverse
Graphax AphaGred
10-1
16
256
512
1024
2048
166
M
Batchsize
Scale Factor
(a) RoeFlux 3d
(b) Multi-Layer Perceptron (GPU)
Figure 4: Runtime measurements over 1000 trials for the vectorized RoeFlux 3d and MLP tasks with different batch sizes using the same setup as in table[2 The MLP network sizes were scaled up with growing batch size by constant factor. The exact procedure of scaling explained in appendix B Error bars are the 2.5- and 97. -eeeeeeesssssees of the runtimes.
5 Conclusion
In this work we successfully demonstrated that AlphaGrad discovers new AD algorithms that out- perform the ofate-of-thh-arr We demonstrated that these theoretical gains translate into measurable runtime improvements with Graphax, Python-based interpreter we developed that leverages the AD algorithms discovered by AlphaGrad. However, AlphaGrad currently only optimizes for multiplica- tions which cannot capture the entire complexity of the AD algorithm. Future work could explore other optimization targets such as execution time, memory accesses, quantization and different hardware backends using hardware model for efficient simulation. Another promising research avenue would be the implementation of much more general framework for vertex elimination. Inspiration could be drawn from Enzyme, which operates on the intermediate representations (IR) using LLVM and is therefore less bound by the choice of programming language. Leveraging the LLVM approach would also open up new directions regarding AD-specific compiler optimizations similar to what was presented in LAGrad.
Two of the main shortcomings of our work are the lack of support for dynamic control flow and dynamically changing functions as well as the need for retraining of the algorithm for every single computational graph. To cireumvent the first issue, it would be possible to implement a version of vertex elimination that can deal with dynamic control flow although this would require some significant changes Graphax. However, demonstrated by the wide range of benchmark tasks, several applications are already possible without this feature.
The second issue, addressed partially in our work, is the training of the agent multiple graphs once. While the best results were still achieved with single-graph training, the multi-graph training still outperformed the other existing methods most benchmarks. This hints the possibility to train our agent on a large set of computational graphs at once, thereby effectively building a statistical model of AD. Finally, VertexGame offers novel way evaluate existing RL algorithms real-world problem that poses diverse challenges such as rewards across multiple scales and large action spaces. Thus, it could complement existing benchmarks such OpenAl gym and MuJoCo Towers et al. 2023
Todorov et al. 2012
Acknowledgments
This work sponsored by the Federal Ministry of Education. Germany (projects NEUROTEC-I grant no. 16ME0398K and 16ME0399 as well as GreenEdge-Fute funding no. 16ME0521). The authors also gratefully acknowledge the Gauss Centre for Supercomputing e.V. (www.gaust) gaauss-centre.eu for funding this project by providing computing time on the GCS Supercomputer JUWELS at Jülich Supercomputing Centre (JSC). Furthermore, thank Mark Schöne, Christian Pehle, Uwe Naumann and Matthew Johnson for the helpful discussions regarding the project.
10

---

References
Andreas Albrecht, Peter Gottschling. and Uwe Naumann. Markowitz-type heuristics for computing jacobian matrices efficiently. In Peter M. A. Sloot, David Abramson, Alexander V. Bogdanov, Yuriy E. Gorbachev, Jack J. Dongarra, and Albert Y. Zomaya, editors, Computational Science ICCS 2003, pages 575-584, Berlin, Heidelberg, 2003. Springer Berlin Heidelberg. ISBN
978-3-540-44862-4.
B. M. Averick, R. G. Carter, Guo- Guo-Liang Xue, and J. J. More. The minpack-2 test problem collection. Technical report, University of Illinois, June 1992. (ttLLLhtt:///iiii
library. unt edu/
ark: /tk:/675331 //////631//maac735552222
Accessed May 11, 2024.
Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation machine learning: survey. Journal of Machine Learning Research, 18(153):1 -3, 2018. URL http://jmlr wttpppapps//1///7--eeeo
html
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https: / /www. wandb com/ Software available from wandb.com.
Fischer Black and Myron Scholes. The pricing of options and corporate liabilities. The Journal of Political Economy, 81(3):337-64, 1973. bttp:://wwwwx jstorr org stttsaaaaaaa10//2
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and
Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
bettp://githn com/ google/jax
Luca Capriotti and Michael B. Giles. Algorithmic differentiation: Adjoint greeks made easy, Available at SSRN, 2011. doi: 10.2139/ssm, 1801522. URL https:// com/ abstract= April 2.
Jieqiu Chen, Paul Hovland, Todd Munson, and Jean Utke. An integer programming approach to optimal derivative accumulation. Lecture Notes in Computational Science and Engineering, 87, 01 2012. doi: 10.1007/978- 602-64422333332 20.
Ivo Danihelka, Arthur Guez, Julian Schrittwieser, and David Silver. Policy improvement by planning with gumbel. International Conference on Learning Representations, 2022. URL https
Sespenreview.net/fooree bd-EaaaNdooebb
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju. Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Clau- dio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King. Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik. Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojeiech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL httt:://ittub/ com/ cwtp:/oglooooooeeeeoo/ooooooooee deepmind
Serap Dikmenli. Forward & inverse kinematics solution 6-dof robots those have offset & spherical wrists. Eurasian Journal of Science Engineering and Technology, 3, 05 2022. doi: 10.55696/ejset. 1082648.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16 words: Transformers for image recognition at scale.
CoRR, abs/2010.1 1929, 2020. TRI.httts:: / //ap:v/...vvvvv wttp:iiw..rr/aa/////ooa// 11929
Alhussein Fawzi, Matej Balog, Andrew Huang, al. Discovering faster matrix multiplication algo- rithms with reinforcement learning. Nature, 610:47-53, 2022. doi: 10. 10.10088441588--02-051122... URL https: 1wtp:///011..100/111111110000110
Shaun A. Forth, John D. Pryce, Mohamed Tadjouddine, and John K. Reid. Jacobian code generated by source transformation and vertex climination can be as efficient as hand-coding. ACM Transactions on Mathematical Software (TOMS), 30(3):266-299, 2004a.

---

Shaun A. Forth, Mohamed Tadjouddine, John D. Pryce, and John K. Reid. Jacobian code generated by source transformation and vertex elimination can be as efficient as hand-coding, ACM Trans. Math. Softw, 30(3):266- 299, sep 2004b. ISSN 0098-3500. doi: 10. 10.1145/10240044102247766 URL
https://doi org/10 1145//02200441.... 1024076
Markus Giftthaler, Michael Neunert, Markus Stäuble, Marco Frigerio, Claudio Semini, and Jonas Buchli. Automatic differentiation of rigid body dynamics for optimal control and estimation. org/ htg//bb//709.. 03799
CoRR, abs/1709.03799, 2017. URL
Ian J. Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, Cambridge, org
MA, USA, 2016.
A. Griewank and A. Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Difserentiation, Second Edition. Other Titles Applied Mathematics. Society for Industrial and Applied Mathematics (SIAM, 3600 Market Street, Floor 6, Philadelphia, PA 19104), 2008. ISBN 9780898717761.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2024.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In AAAI Conference on Artificial Intelligence, 2017.
URL https://api. www.:/pi.sssaatacccccccsessseeee http:///p..aaatt..eooor.ooooorrr orgg/oorpus11:44747811111
Shengyi Huang, Rousslan Fernand Julien Dossa, Antonin Raffin, Anssi Kanervisto, and Weixun Wang. The 37 implementation details of proximal policy optimiza- tion. In ICLR Blog Truck, 2022. URL https://iclr-blog track.
03/ 23/25/ppo-oo www.:iieeeeettttoneettteaa implementation-detalls..
Yuu Jinnai, Arash Mehrjou. Kamil Ciosek, Anna Mitenkova, Alan Lawrence, Tom Ellis, Ryota Tomioka, Simon Peyton Jones, and Andrew Fitzgibbon. Knossos: Compiling ai with ai. 9 2019.
Steven Kapturowski, Georg Ostrovski, Will Dabney, John Quan, and Remi Munos. Recurrent experience replay in distributed reinforcement learning. In International Conference on Learning Representations, 2019, URL
Hiroharu Kato, Deniz Beker, Mihai Morariu, Takahiro Ando, Toru Matsuoka, Wadim Kehl, and Adrien Gaidon. Differentiable rendering: survey. CoRR, abs/2006. 2020. URL https
//arxiv. org/ 2006. 12057
Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CORR, abs/1412.69900, 2014. httpps//a///p http:/ppp..emmtiicccchlllllr org htt.aaaarrrrr//oopssDDd..6.-.
Seppo Linnainmaa. Taylor expansion of the accumulated rounding error. BTT Numerical Mathematics, 16:146-160, 1976. URL atteps::/app.ssseaaticcccooooo wtt.:oooaa..rooooooousooi.i.. 122357351
Daniel J. Mankowitz, Andrea Michi, Anton Zhernov, et Faster sorting algorithms discovered using deep reinforcement learning. Nature, 618:257-263, 2023. doi: 10. 10010088644118860003300000--9999 1033///41186-022-0000000
URL https://doi.
Charles C. Margossian. A review of automatic differentiation and its efficient implementation. CoRR,
Abs/181 1.050331, 2018. httl:attpp//arrdv.rr//aaasss
05031
William S. Moses and Valentin Churavy. Instead of rewriting foreign code for machine learn- ing, automatically synthesize fast gradients. In Proceedings of the 34th Conference on Neural Information Processing Systems (NeurIPS 2020), page Page Numbers, Vancouver, Canada, 2020. MIT CSAIL. URL https: httpp//weeddngg..ooo.ooooo/opppp 93330351oef444422223377222222223 Paper .df
12

---

Uwe Naumann. Save simulated annealing applied to the vertex elimination problem in com-- putational graphs. Technical Report RR-3660, INRIA, 1999. URL https://hal iitiaa.iiiiafii..
inria- 00073012
Uwe Naumann. Optimal jacobian accumulation is np-complete. Mathematical Programming. 112:427 441, 2008. doi: 10.1007/110107000-000422 URL 0011070006-002--2
Uwe Naumann. Exact first- and second-order greeks by algorithmic differentiation. 01 2010.
Uwe Naumann. Optimization generalized jacobian chain products without memory constraints, 2020.
Aditya Sanjay Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Regal: Transfer learning for fast optimization of computation graphs. ArXiv, 0oss/1905.02494 2019. URL httpsttps: http://app.semntiiiscccccacaaaae www.aaaroooooooCC.a....... 1468088444
Mai Jacob Peng and Christophe Dubach. Lagrad: Statically optimized differentiable programming in mlir. In Proceedings of the 32nd ACM SIGPLAN International Conference Compiler Construction, CC 2023. page 228-238, New York, NY, USA, 2023, Association for Computing Machinery. ISBN 9798400700880. doi: 10.1 10.114573578336.35582222 URL httpt:////ooii org/
10. 11.11/55577860 3580259
P.L Roe. Approximate riemann solvers, parameter vectors, and difference schemes. Journal of Computational Physics, 43(2):357-372, 1981. ISSN 0021-9991. doi: https://doi.org/10.100
0021-9991(881)0128--8 URL httph:tt////// htRL:ttte////w...cceeerrrrr...o wtwps/ciennedreet.coo//////////e 0021999181901285
A. Savine and J. Andreasen. Modern Computational Finance: Scripting for Derivatives and xVA. Wiley, 2021. ISBN 9781119540786. URL https: http://oook.ooglll..oe dww.de//bbookiiii
mBZDEAAAQBAJ
Jürgen Schmidhuber. Deep learning in neural networks: An overview. CoRR, abs/1404.7828, 2014.
URL http::/arxix:: org/ abs/ 7828
Patrick Schmidt, J. Born, Bommes, Campen, and Leif Kobbelt. Tinyad: Automatic differentia- tion in geometry processing made simple. Computer Graphics Forum, 41:113-124, 10 2022. doi: 10.111 10.1111/cgf.144 1/cgf.14607.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. CoRR, 0bs/1911.00266 2019. http::ttt///i...ooooo abs/aa/ss//// 1ttp:/991...22..
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
David Silver, Thomas Hubert, Julian Schritwieeser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, L. Sifre, Dharshan Kumaran, Thore Graepel, Timothy P. Lillicrap, Karen Simonyan, and Demis Hassabis. Mastering chess and shogi by self-play with general reinforcement learning algorithm. ArXiv, abx/1712.01815, 2017. URL hhtps:::pss : www.stiiscoooeee...o.m org/ CorpusID:33081033
M. Tadjouddine, Bodman, D. Pryce, and S. Forth. Improving the performance of the vertex climination algorithm for derivative calculation. In Martin Bücker, George Corliss, Uwe Naumann, Paul Hovland, and Boyana Norris, editors, Automatic Differentiation: Applications, Theory, and Implementations, pages 11--20, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg.
Mohamed Tadjouddine, Shaun A. Forth, John D. Pryce, and John K. Reid. Performance issues for vertex climination methods in computing jacobians using automatic differentiation. Peter M. A. Sloot, Alfons G. Hockstra, C. J. Kenneth Tan, and Jack J. Dongarra, editors, Computational Science - ICCS 2002, pages 1077-1086, Berlin, Heidelberg, 2002a. Springer Berlin Heidelberg. ISBN 978-3-540-46080-00
13

---

Mohamed Tadjouddine, Shaun A. Forth, John D. Pryce, and John K. Reid. Performance issues for vertex climination methods in computing jacobians using automatic differentiation. Peter M. A. Sloot, Alfons G. Hockstra, C. J. Kenneth Tan, and Jack J. Dongarra, editors, Computational Science ICCS 2002, pages 1077-1086, Berlin, Heidelberg, 2002b. Springer Berlin Heidelberg. ISBN 978-3-540-46080-0.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012. doi: 2.1109//RS..201..63888
Edan Toledo, Laurence Midgley, Donal Bymne Callum Rhys Tilbury, Matthew Macfarlane, Cy- prien Courtot, and Alexandre Laterre. Flashbax: Streamlining experience replay buffers for reinforcement learning with jax, 2023. Uttpttss:: github....... com http////staddpaa//lssbbbo.e
Mark Towers, Jordan K. Tenry, Ariel Kwiatkowski, John U. Balis, Gianluca de Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Arjun KG, Markus Krimmel, Rodrigo Perez- Vicente, Andrea Pierré, Sander Schulhoff, Jun Jet Tai, Andrew Tan Jin Shen, and Omar G. Younis. Gymnasium, March 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31 International Conference on Neural Information Processing Systems, NIPS' 17, page 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964
Yanqi Zhou, Sudip Roy, AmirAli Abdolrashidi, Daniel Wong, Peter C. Ma. Qiumin Xu, Hanxiao Liu, Mangpo Phitchaya Phothilimtha, Shen Wang, Anna Goldie, Azalia Mirhoseini, and James Laudon. Transferable graph optimizers for ml compilers. ArXiv, abs/2010.12438, 2020. URL https::/apt reeemantiischo.oo
org:/orrpusD::225023338
Muhammad Zubair, Dinesh Ranjan, Andrew Walden, Grigore Nastac, Eric Nielsen, Boris Diskin, Michael Paterno, Seung Jung, and James H. Davis. Efficient gpu implementation of automatic differentiation for computational fluid dynamics. IEEE 30th International Conference on High Performance Computing, Data, and Analytics (HiPC), page 11, 2023. doi: 10. 10.1109/HiPC58850.. 2023.00055.
14

---

Appendix
A
Task Descriptions
This section describes in-depth the mathematical formulation of the tasks that were evaluated in table and table 2 An implementation of the functions can be found in the accompanying software package. Note that if set of variables (21..... zn,) is referred to as {z,), they are treated separate scalar inputs to the function while x treats them as a single vectorized input.
A.1 RoeFlux_1d ld
For the implementation of the RoeFlux 1d task we thoroughly followed (Roe] 1981]. The pressure Pid. enthalpy H1d and flux term Fid of the one-dimensional Euler equations defined through:
uf
Pia (up u1 U2) = (7 - 1)(u2
2uo
U2 + P
Hia u2,p)
uo
uf u1 (p + uz)
Fia un u2. p)
U1.
ug
un
The Roe flux @noee (1o, un, U1, U13, 114, Uy0: Uy1, Ur2, Mr3, Ur4) between two adjacent cells u and ur is computed using the routine described below. First, we define some averaged quantities to simplify formulation:
Ug =4 - U40
Uir v 210
W1 v uio uro-
Then we define some state variables for the left cell through
un
U1
U10
Pl Pid (uto, un 42).
h Hia(uo, 422 Pi).
Similarly, we define the same variables for the right cell through
Pt Pia Uro, Uei U,2),
hr Hıa(uo, 4x2, Pr)-
We define some differences between the state variables of the two cells though
Ap pl Pr, Ar t Ur:
We proceed with the introduction of some auxiliary variables for further computation that
Vunoe + Vunuu,
W1 h iiyuu + uuohe
W1
a2 (7 1)
2 92),
where 92 u?, /az and uhoa. We define 4p to writing:
al.
|u] and 4, |u|| and proceed
Ap
Auo
1
02
Ap
Au
+
1p
Ap
Au
In
15

---

Next, we compute the fluxes between the cells through
Fio. Fi1, Fi2 = Fid (210, U1 U1, PI)
Pro, Frt, Fi2 Fia (Uro, Url, Ur2,
and then define P = Fi + FR and a = 2a
u1o- The flux differences are then given through
A co + (C1 ac2.
A 1 cou + ac (u + a) - a02((u aa,
AF2 Co42 aci (h ua) acz(h ua).
Then the output of the function is given through oo, Φ1 O2 with o; = 2 AF)).
A.2
RocFlux_36 3d
The implementation of the RoeFlux 3d task is similar to the RoeFlux 1d task. Again, we follow Roe, 1981] for the implementation and start by defining functions for the pressure Pad and enthalpy H34:
|u|2
Psa(uo, u, ua) =
( (7-11)
2no
UA +
H3a u4.p)
un
Note that now instead of a mme-dimenioonal state variable U1, we have a three-dimension state vector and the role of ug now taken over by u4. The flux in three dimensions is given through
Fsa uo u, Ug v, p) (11 P + ue, Un (p + u4))
where u1, u2, dg and en 02, Vs the three components that make up and respectively and = (p, T. We then again define the finite differences
Agg = 2p0
Au = U,
Au4 =44 24,
and furthermore set Vi ui and v,
u. We continue with defining some auxiliary variables
W1 Vuno Vue--
t /uiovi + Vucovi
W1
Then we define some state variables for the left cell through
pt = Psa (uoo un uja):
h Haa (u0, u4s pI).
Similarly, we define the same variables for the right cell through
Pad Psd ur Uz2), (u,o Ue4,
Then introduce
Vugh + v Vurhe
=
W1
16

---

and set = |t| and (7 1)(h (92)) and vag. Furthermore, we define the eigenvalues
of the Roe flux problem as 1p $ + a, 0 tı, In t1 a. Then coefficients a are given through
h1 k2
co
mm
Au2
Auo
C1
12
Au3
C2
Du1
tg
03
7
((h - Q2)
Auo +t. A.Au- -Aua))
02
k1
k2
C4
Au f1 Auo The definition of the fluxes of each cell
where we defined 41 Auo C3 and k2
a
is similar to the formulation in the one-dimensional case. The flux changes given through
A/o = + c3 + calp.
AFi coln c31 + cclp
A/2 = <nl2 + cı + 62l2 + + cal2,
AF3 cuts c2ts + csts cats,
AF co fıa) t + c2t
C392
+ ca(h tta).
The results of AF, AFt and A/'s are concatenated into the vector AF, The output of the function then (0o, ④; 24) as defined the one-dimensional case.
A.3 RobotArm RobotArm_6DOFF 6DOF
The RobotArm_6he _amO task models the forward differential kinematics of a f-degree-of-freedon (6-DOF) robot arm is often found in robotics labs and industrial manufacturing sites. For the implementation, we followed [Dikmenli 2022 and define G = cos I; and 8; = sin T4. Furthermore, define the functions
S(u, v) cos(u) sin(v) + sin(u) cos(e); C(u,v) cos(u) cos(v) sin (u) sin(v).
Then we also define 8;j 5(4 (j) and c.y C(4, (4) for the input variables {4: of the problem. Then we proceed with calculating the auxiliary intermediates
ag 8p (C102304 8184) ci 823%,
85 (810230 C184) + 81823(5,
a; 8236485 023
n. = Co(C2385 + 823C403) 8238486.
0, s6(c2335 + S23C405) $233106.
Then, we define the Tait Beit-Bryan angles through:
2 arctan ay
0gp
a
ÿ = arctan
ag
Os
arctan
Next, we calculate the positional parts of the kinematics by starting with the z-component:
æ 185( 55 (c10234 + 3131) + 8132)))) 2a = <i (175 + 890ca 50c23 + 1035s2a), Ps + 22
17

---

We continue with the y-component:
185(ss(C1225)) + 9184) + 51 Cs) y/2 51 5177 + 890c2 50c23 + 1035s23),
Py 1/1 + //2
Finally, the z-component given through:
Pz = 575 + 89082 + 50823 - 1035c23 + 185(8230485 c22c5)
The function then returns the six values (Px, Py, Pz. 2,
A.4
HumanHeartDipole
The HumanHeartDipole task derived from the experimental electrolytic determination of the resultant dipole moment in the human heart. For the inplementation, we followed [Averick et al. 1992] with { (I1, (T) such that
[1((;;)) x1 + Omx
fe((a;)) = Fy + I4 - Ony
fa((x;}) 25 T6 1733 Tg 0Λ
fa((x;}) = 2721 + 2g + T3 + T621 OB
[s((=;}) = (x - 22)) 2x 0577 + x2(46 [gg 2. ac
fo({x;}) = 23(2} 2 + 2x 2557 + 21 x3) + 2Iz 0D
fr((=;)) = < 3x . - 3x 12 3r3) + 1445 ( 3m
fo ((xx)) 22 3x7) - ", 3z3) + 2 3x 12 32 OF
with mmx. my, AA. OB.OC, OD. dF. oF being some arbitrary measured constants.
A.5 PropaneCombustion
The PropaneCombustion task arises in the determination of the chemical equilibrium of the com- bustion of propane in air. Each unknown is related product concentration given in mols formed during the combustion process. We implemented the PropaneCombustion task as defined [Averick]
et al. 1992 where with {x,} (#1, (21....,211. 211):
f1((x;)) + 3
fe({x;)) = 2 + x2 + 21 + 87 I8 +
Tg + 20
fa((x;)) 2x2 + 2xs +06 +
f4((x;}) 2r3 + 2g 4R
fo((x;)) Ksv 1125
fo((z;)) Ko K6V0r02 /x40T
p
T11
[r((x;)) 0102 V
p
T11
P
[s((x;)) = KATT I4J8
#11
fa((x;)) = KOT
Vzg
p
fu((fr))) K1027
-2.010
p
211
fu((x;}) = I11 I10 - F9
I8
26 -
255
24
33 322 31
Here, .s...., K are measured constants and = 10 is the relative amount of air and fuel, while p 40 is the pressure in atmospheres.
18

---

A.6 Black-Scholes Equation
The CrrnesSchols__acoobii task is derived from riskless portfolio management where the goal is the computation of so called second-order greeks which give insights about the price evolution of an option. The Black-Scholes partial differential equation is given through
0V + Øt
8V - 02
-V 0V
2
as
In this equation, models the volatility of the underlying geometric Brownian motion, while S describes the stock price of the underlying asset and r is the risk-free interest rate. is the price of the option at time f. For certain conditions described in [Black and Scholes 1973), this equation can be solved analytically. We define
-
<(z) d(z)=1+ - /
V /
f- de
(6)
which is the cumulative distribution function of the Gaussian distribution (also known as the error function). Next, we define
F
o²
di OVT
log
K
27
d2 di - -VTT
( ar o, c -- (FØ(dz) K K)()2))
where is the payoff. Then, the solution to the
equation is given through
f(S, o, (r, S). k,,
where Frr, = er We are interested in the second-order derivatives with respect to S, K, a, so we first calculate the Jacobian of equation (A.6) using reverse-mode AD with Graphax using gravhax.jacvelf. order="rev" The resulting computational graph is used to learn an optimal elimination order to compute the second-order derivatives with AD, i.e. the Hessian.
A.7 Multi-Layer Perceptron
The MLP task describes a simple 2-layer Perceptron with layer norm between the first and the second hidden layer. The input-size of the network is 4 and hidden layers have size 8 while the output layer has size of 4. The activation functions are tanh-functions and the output is processed using softmax cross-entropy loss. For the actual runtime experiments, we scaled the sizes of the inputs, outputs and hidden layers by a factor of 16 create a more realistic example of MIP.
A.8 Transformer Encoder
The TransformerEncoder task is inspired by recent advances natural language processing and image classification. The network consists of two attention blocks with single head only. The block itself consists of a softmax attention layer with residual connections followed by a layer norm and MLP with single hidden layer of size 4 and sigmoid linear unit activations. We stack two of these layers and process the output with softmax cross-entropy loss function. The input embeddings have size 4 with a sequence length of 4. For the actual runtime experiments, we scaled the sizes of the inputs, outputs and hidden layers by factor of to create realistic example of a transformer encoder.
A.9 Random Functions f and g
The random functions f and g randomly generated using custom JAX interpreter which consists of repository of elemental operations such as cos, log or + but also matrix multiplications and array reshape operations. The number of input and output variables can be specified as well as the number of intermediate vertices that will be eliminated by the vertex elimination algorithm. The random function generator then randomly samples elemental functions from the repository. The number of functions that are unary, binary or perform accumulation or reshape operations can be controlled by adjusting the respective sampling probabilities. At every step, the function is checked to guarantee that it is executable and well-defined. The function generator is part the accompanying AlphaGrad software package.
19

---

B Comparison to JAX
Figures 5 and 6 contain an in-depth analysis of the performance benefits of the combination of AlphaGrad and Graphax over default JAX forward-mose mode AD and reverse-mode AD. Unless stated otherwise, all runtimes were measured RoeFlux Id, RobotArm_ 2DOF.. random function g and BlackScholes Jacobian significantly outperforms both JAX AD modes for all batch sizes, sometimes by almost order of magnitude. For the HumanHeanDippole and PropaneCombusticn outperform both JAX modes significantly but 2048, For the RoeFlux_ 3d and random function f tasks, we found that the combination of AlphaGrad and Graphax consistently outperforms JAX' AD modes for large batch sizes. The deep learning tasks, i.c. the MLP and TransformerEncode differently. Both functions were vectorized only over their inputs and labels as is typically done in deep learning applications. Furthermore, we evaluated both networks for different scale factors where the entire network was scaled up by a constant factor and only compare to JAX reverse- mode AD since this is the default training mode for these kinds of networks. Also, both networks were evaluated on GPU as well this is the typical hardware backend on which they are executed. For the MLP task we found that the AlphaGrad and Graphax combination outperforms JAX reverse-mode AD for large scale factors, i.e. large networks with a large batch size. Since the gain through AlphaGrad is small, the authors conjecture that the gain in performance is mainly due to the sparse implementation the AD routines. Note that the batch size was also consistently scaled up with the other components of the network with a starting batch size of 8.
AMD EPYC 9684X 2x96-Core processor. On the
Jaccobbannn tasks, our work
tasks, AlphaGrad and Graphax observe crossover from batch size 1024 to
tasks, were analyzed slightly
Roe usstustonn different batch RobotArm different
Ene [mc]
101
100
JAX mascle
JAA forward mode JAX reveree mode Graprixx AlphaBrac
Jax reverss mode
Graphax AlpheGrad
1064 2015
1024
2018
(a) RoeFlux ld
(b) RobotArm 6DOF
evaluation times for different modes batch sizes
PC evalustion times for different modes and betch sizes
Fosbastion trme Emee
Exalcation Sma [mn]]
JAX forward mode JAX recersa made Grachas: AlphaGrad
100
JAX forward mode AXX reverse moce Graphax AlphsGrad
30
256
512
2024
64
512
1024
Satchstte
Batchsize
(c) HumanHeartDipol
(d) PropaneCombustion
evaluation times for different modes batch
US evalustion times for different modes and betch
10f
Eookation [mmq)
Evaluation time [mm]
10a
200
307
JAX forward mode
JAX forward mode JAX reverse mode Graphar Alshacirad
3.7
JAX reverse mode Graphax phaaGrrd
Betchsize
Datchsize
(e) Random function g
(1) BlackScholes_Jcconl
Figure 5: Runtime measurements over 100 trials for the scalar tasks. Error bars are the 2.5- and 97.5 97.5-percentiles of the runtimes.
20

---

For the TransformerEncoder task, we were only able to evaluate the runtime for a batch size of since the VertexGame implementation of AlphaGrad supports only inputs with maximum dimensionality of 2 while a typical batched transformer input has the shape (batch size, sequence length, embedding dimension). We found that the combination of AlphaGrad and Graphax performed on par with JAX' reverse-mode AD for all scale factors except 8 where we found a significant difference in performance.
evaluation times for
batch revallation times for different modes and batch sizes
10
209
JAX forward mode JAX reverse made Graphax + AlphaGred
JAX forward mode JAX recerse muce Graphax + AlphaGrad
10a
16
25
1924
2040
64
f12 Batcnsize
2049
Batthsize
(a) RocFlux_3d
(b) Random function f
MLP evalustion times different modes and betch sizes
avalustion times different and
AAX
Exsharden (n()
Erolutiinnmn [nn)]
127
Graphax
1.9
307
2001
32
Sealle Faclor
Statle Fat
(c) Multi-Layer Perceptron
(d) Transformer Encoder
MLP evaluation times for different modes and batch sizes transformer evelustion times for different modes and betch sizes
JAX racerse made
Evaladioon (mnd
Graphax AlphaGrad
Sesluation Emac)
ApaaGrod
204
10-3
16
32
64
36
32
64
Scale Ractor
Scale Factor
(c) Multi- Percoptron (GPU)
Transformer Encoder (GPU)
Figure 6: Runtime measurements over 100 trials for the vectorized tasks with different batch sizes. The MLP network sizes were scaled up with growing batch size. The transformer could only be evaluated on with batch size of due to limitations of VertexGame but the network size was scaled up as well. Both deep learning tasks were also evaluated on GPUs. Error bars are the 2.5- and 97.5 97.5-percentiles of the runtimes.
C Sparsity Types
The sparsity of the Jacobians associated with the edges in the computational graph representation is described with number ranging from -10 to 10. and an example tensor that is represented by the corresponding number. Our approach is directed only diagonal sparsity, meaning that we mainly consider tensors that can be decomposed into products of lower-dimensional tensors and Kronecker symbols diz. Since we might have operations like transposing or slicing in our computational graph. which just change the shape but not the value of the edge Jacobians when eliminated, we introduce a "copy gradient" sparsity type with value make a distinction between multiplication with
overview of the sparsity types
-1 to represent these operations. Furthermore, the unit tensor δαd and a constant multiple of the unit tensor cō.d since the latter incurs actual multiplications while the former can just be treated as a renaming of indices as demonstrated by the
21

---

Table 3: Different sparsity types that are required to represent all the possible tensor shapes that occur in tensor vertex elimination with vectors and matrices.
Sparsity Type
Example Tensor
-10
-9
-8
-7
-6
-5
Taosk
-4
-3
Tad
-2
-1
copy gradient operation
0
no edge
1
Tujal
2
Tasiock
3
Teykōji
4
5
Tadsk
6
7
Tyjdaōjk
8
9
10
following examples with tensors Sutmn dindin and Unimn cdkmdin
Σ
Tip 6kmorn = Timminnjan
(7)
k
Σ
Tipōndn cdkmôrn cieddimdin.
kd
In the first equation, we can just rename the indices of the tensor, but in the second equation we have to perform the product eTij which incurs |i] [i] multiplications. If two edges are multiplied with each other, we determine the resulting sparsity type of the new edge by looking up the combination in a large handcrafted table and then writing the result into the new graph representation. The new shape of the Jacobian determined according to the rules of tensor contraction. Finally we delete the old edges from the representation as required by the vertex elimination procedure. The sparse matrix multiplication table for the computational graph representation can be found in the accompanying source code. The number of multiplications incurred by a multiplication of two edge Jacobians is computed from the sparsity types involved and the Jacobian shapes. By multiplying all values of the two Jacobian shapes with each other and masking out certain values according to the sparsity types involved, we arrive at the correct number of multiplications. For examples consider two tensors Sigkt azyd+ with sizes (2, 3, 2, 3) with sparsity type 6 and Tiskt bςōwō,k with shape (2, 3, 2, 3) and sparsity type Thus, in the computational graph representation, we would have two non-zero entries (6, 2, 2, 3) and (9, 3, 2, 3). Then their contraction is given through
Σ Att
SignaTkma = Σ aarōddd bróandim = aig 6m
Then we form the product [i] -iil- [K]
2- 2- where |- gives the size of the
dimension associated with the index. Then use the sparsity types and mask out the relevant values. In this case, we mask out [k:, ||| and |m| such that we arrive at |i] [s] |rx| = 2- 3 = 18 multiplications which is exactly the number of multiplications incurred by the dense multiplication of aizhn-
22

---

D RL Algorithm and Network Architecture
We solved VertexGame with PPO (Schulman al. 2017] and AlphaZero [Schrittwieser al. 2019 and used a similar architecture backbone in both cases:
Convolutional embedding to compress the computational graph representation down using a 3x5 kernel mapped across the sequence dimension, i.e. the 3x5 filter was applied to all tokens to reduce the feature dimension. Since the computational graph representation is very sparse as each vertex typically only has few connections to other vertices.
Linear projection to an embedding size of 64.
Positional encoding as described in (Vaswani et al. 2017].
Multiple transformer layers with embedding size 64, noftmax-attention layer norm and two MLP layers with a hidden layer size of 256.
In the PPO case, we used two such backbones of with 5 and 4 transformer layers for separate policy and value networks. The policy head was an MLP of with hidden layers of sizes (256, 128) and output size 1 while the value head used an MLP with hidden layers of sizes (256, 128, 64) again with output size 1. The AlphaZero agent used the same heads on a single transformer backbone with 6 transformer layers. In all cases, the MLPs are mapped over the sequence dimension. For the policy head, this produces an unnormalized probability distribution over the actions while for the value head, first sum all outputs get an estimate of the input state value.
PPO Implementation Details We tightly followed [Huang et al. 2022 for the implementation of the PPO agent. Unless otherwise specified, all models were trained on 32 parallel environments with 4 minibatches and a clipping parameter of 0.2. We set the value and entropy weights to 1.0 and 0.01 respectively, while the learning rate was fixed to 2.5 10-4. However, found that the agent performed best if the rollout length was set to be equal to the number of intermediate vertices, although this is technically implementation error.
AlphaZero Implementation Details The implementation is loosely based on the implementation of Gumbel MuZero and uses the mcfx package provided by Google DeepMind (DeepMind et al. 2020). However, we modified the implementation by replacing the trainable model and reward functions with deterministic implementation of the VertexGame environment, thus effectively creating a Gumbel AlphaZero agent. Unless otherwise specified, all models were trained on 32 parallel environments with batch size 4096 and 50 MCTS simulations with 5 considered actions. We set the value and L2 weights to 10.0 and 0.01 respectively, while the learning rate was set to 1 10-3 with cosine learning rate annealing over 5000 episodes. For Gumbel MuZero work properly, it is necessary to rescale the rewards so that they lie in the interval 1). In the Gumbel MuZero paper, the authors designed specific transformation that normalizes the rewards and simultaneously completes the missing values using the Gumbel distribution. The parameters of this transformation have to be carefully tuned to enable learning. In our case, set the corresponding parameters to Cvisit 25 and Cacale 0.01 for all cases. We trained the agent using adaptive momentum gradient-bars based Icarning (Kingma and Ba 2014) with an initial learning rate of 10 and cosine learning rate scheduling over 5000 episodes on two to four NVIDIA RTX 4090 GPUs.
E AlphaZero Results
This section contains the reward curves and elimination orders for the results presented in tables T and Reward curves are shown for six different random seeds, namely 42, 123, 541, 1337, 1743 and 250197. The elimination orders can be directly used with the accompanying Graphax package using the graphax.jacve command that creates the Jacobian for a given function f and a given order through graphax.jacvel jacveff. order=order, argnums=argnums/ (xs). The orders given in tables 4 and can be directly used to compute the Jacobians.
In addition to the single- graph experiments where the agent was trained only single task, also experimented with training the agent on all tasks at once. For this, we randomly sampled from the 10 defined tasks to create 32 random environments. The agent trained with the same configuration as for the single task experiments. The results are shown in figures 9 and 10 and the best achieved numbers of multiplication are displayed in table
23

---

algorithm optim zation RoeFlux
AD algorithm optimization RobotArm, 6DOF
multppication:)
364
IM mullipicaatios)
400
288
450
500
550
Our work
Our work
A RoaHlux
SOTA RobotArm 6DO
samp es
1e6
samplees
106
(a) RocFlux_1d
(b) RobotArm _DOF
AD algorithm optimization HumantiearlDipole
AD algorithm optimization PropaneCombustion 87
150
return (# mutpiiiction))
return (a mumpitcationa
170
172.
97
180
100
190
Our
Our
SOTA DuppeiteertDipole
105
SOTA PropeneCombustion
0.5
2.0
samples
samples
166
(c) Homanne
(d) PropaneCombustion
algorithm optimization BlackScholes
AD algorithm optimization Random Function
300
multplications)
420-
multip icationsl
350
440
451
480-
400
300
450
520
540
Our work
SOTA for Random Furction g
BlackScholes
550
samples
leb
serppies
186
(c) Random function
(0) BlackScholes. Jacobian
Figure 7: Learning curves of the scalar tasks for the same six random seeds with the AlphaZero-basee agent. AlphaGrad manages to find better elimination orders for all tasks. The best result for each task displayed in table
24

---

algorithm optimization RoeFlux
algorithm optimization Random Function
1000
7500
(8 mullsplicatonn
retur 14 mutpp cattors)
3333
1250
10000
1500
12500
15000
17500
2000
Our wark
20000
Our work
2250
RooFlux
SOTA for Random Function
0.0
0.2
044 samples
0.6
0.8
1.0
samples
1e6
1s7
(a) RoeFlux_Be _3
(b) Random function f
AD algorithm optimization MLP
500
392
algorithm ostimization Trensformer
mullipicaelos))
4688
return (e muttilicationn
10000
1500
20000
43000
2000
50000
SOTA for MLP
SOTA for Tha seformer Frcoder
0.0
0.2
0.4
0.6 samples
0.8
1.0
1.2
samples
(c) Multi Layer Perceptron
Transformer Encoder
Figure 8: Learning curves of the vectorized tasks for the same six random seeds for the AlphaZero- based agent. AlphaGrad manages find better elimination orders for all tasks. The best result for each task is displayed in table
25

---

Table Elimination orders for the results obtained with AlphaGrad presented in table
Task
# multiplications
Elimination Order
RoeFlux_ld Id
320
[8, 82, 27, 66, 7, 78, 76, 13, 48
42, 68, 86, 95, 4, 59,
28, 77, 54, 1, 94, 5, 58, 72. 93. 75. 31, 33, 57, 90,
25, 89, 88, 84, 96, 72, 83,
81,
10, 85, 43, 73, 79,
18,
37, 38, 55, 49, 69, 65, 29, 64, 16, 15,
87, 80, 67, 24, 46, 63, 39, 6,, 3,, 26,
30, 98, 36, 12, 20, 100]
RobotArm_6DDo 6DOF
231
61, 4,
40,
82, 2,, 31, 30, 104, 103, 66, 4444 45, 49,4 47, 113, 12,11 28, 3.,. 109,27,
78, 51,
54, 69, 25, 17. 106, 1.11
96, 13. 12]
Humanle
148
[19, 85, 11, 19.
74, 3..
13,
82, 25, 26, 12,
PropaneCombustion
88
42,
36, 35, 34, 24, 26, 30, 24, 43, 58, 40,
55, 54, 38. 20, 21, 15, 7,2, 17, 53, 16,
14, 11,
Random function
417
43, 92, 89, 29, 16, 82, 86, 60, 24, 19, 79, [I, 99, 85, 90, 20, 20,
49, 64, 11,
78, 17, 68, 76, 76, 2, 52, 26,
10,59,
84, 27, 27,, 74]
BlackScholes
312
170, 16, 41,
101, 74, 71, 71,
54, 96, 64, 65, 119, 14, 118, 50, 76, 61, 15,
40, 59, 100, 49, 114, 83, 116, 113, 78,
25, 64,
133, 10, 13,
87, 4111101.99.4..3110009992.7444440 7555477755.2222240.49999970700.. 103, 82, 37, 94, 8, 1111, 106,
7,56, 12,, 67,
90, 192, 61, 34, 98, 35, 51, 80, 72, 115, 91, 33, 39,
88,
131, 123, 117, 73, 109, 26, 63, 128, 108, 40 4,,
3, 125, 130]
26

---

Table Elimination orders for the results obtained with AlphaGrad presented in table (cdd.).
Task
# multiplications
Elimination Order
RoeFlux_3d 3d
811
[124, 136, 56, 128, 78, 24, 1, 54, 101, 127, 121, 140,
47, 135, 67, 34, 111, 32, 100, 119, 99, 114, 125, 141,
122, 45, 65, 59, 117, 89, 116,
11,53,
36, 30, 108, 13,,55, 143.
132, 87,
18,
6,
57, 97, 16, 84, 143, 123, 90, 94, 137, 104, 69, 62, 130, 130, 49, 98, 134, 93, 138, 126, 115, 10,,
106,
52,
19, 20,
20, 86,
18 37,
13,38 60, 34, 61, 72, 81. 33
Random function f
6374
21..1110022231111130101121211111 53, 20, 46,
46, 17,39, 52, 30,
2-layer MLP
389
17, 25, 30, 19, 39, 13. 12, 36, 38, 37, 30, 23, 22, 18, 16.
25.
14.
Encoderf
4656
[60, 60..5........54444099944444.23.3
22,
71,
87, 19, 24, 24. 79, 67,
26, 89, 86, 18, 39, 30, 50,
73, 68, 62, 57, 13,
43, 45, 44, 42,
40, 33, 32, 21, 16, 15, 14, 14.44
27

---

Table Best number of multiplications achieved joint training mode with the AlphaZero-based agent.
RoeFlux_Ic
RobotArm_6DD 6DOF
HumanHeartDiiple
PropaneCombustion
331
248
158
n.a.
430
RoeFlux_3e 33
MLP
Encoder
BlackScholes, Macckcanles_Jacobian
907
n.a.
n.a.
307
5884
algorithm optimization
AD algorithm aptimization RobotArm 6DOF
250
return sntippiictaoons
400
multipiicelion)
300
268
500
350
600
400
700
450
800
Joint optimization
500
Joint optimization
SOTA RaeFlux, 3d
SOTA for Robot/vm, 6DOF
900
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0 semples les
1.b
2.0
samples
le7
le7
(a) RoeFlux_Ia ld
(b) RobotArm 6DOF
AD algorithm optimization HumanHeartDiiple
AD algorithm optimization PropaneCombustion
160
ndum 1# multo catono))
couum ( nuutppiccttoos)
170
172
100
180-
190
200
210
120
Joint optimization
Joint optimization
SOTA for Humerd Disptee
SOTA for PropaneCombusten
130
0.0
0.3
1U samples
1>
0.0
2.0 samples
1.
20
le7
le/
(c) HumanHeartDipole
(d) PropaneCombustion
algorithm optimization g
AD algorithm optimization Black5choles Jacobian
64 mulppccction)
450
recurn [P multppicanton)
451
400
500
500
Joint optimization SOTA for
Joins
600
SOTA for BlockScholes Jacobian
700
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
samples
le7
* casaples
1e7
(e) Random function
(f) BlackScholes_Jaceee
Figure 9: Learning curves of the scalar tasks joint training mode for same three random seeds with the AlphaZero-based agent. The best result for each task is displayed in table 6
28

---

algorithm optimlzation RoeFlux
algorithm optimization
1000
multpiizationa
10000
0333
rotur nuttpliration)
2000D
2000
.00000
2500
40000
joint ontimisation
joint optimization
RoeFlux 3d
50000
0.0
0.5
1.5
1.5
samples
le7
samples
1e7
(a) RoeFlux 3d
(b) Random function f
AD algorithm optimization Perceptron
AD algorithm optimlzation Transfomer
4668
retumm (# multilicattans
return te mu iiplicatiion)
2000
50000
4000
100000
8000
150000
oint optimizetion
200000
Joint optimization
8000
SOTA for Perceptron
SOTA Transformer
0.0
0.5
1.0 samplees
1.5
2.0
0.0
0.5
2.0 samples
1.5
2.0
1e7
(c) Multi-Layer Perceptron
Transformer Encoder
Figure 10: Learning curves of the vectorized tasks in joint training mode for the same three random seeds for the AlphaZero-base agent. The best result for each task displayed in table
PPO Results
We ran all experiments with the configuration described in and used the random seeds from the AlphaZero experiments. The PPO agent also manages to find better elimination orders that improve over the state-of-the-art but is outperformed by AlphaZero agent on all tasks, sometimes by significant margin as for example in the RoeFlux tasks and rundom function J or the MLP and TransformerEncoule where does not find better elimination order at all. This to expected since the AlphaZero agent can make use of the available model and thus select actions through planning. In other cases, the performance comes very close to the AlphaZero agent, for example in the BlackScholes_lacobian random function f tasks. Thus, the OPOO-bsed agent might still be a viable choice because it trained within minutes on single NVIDIA RTX 4090 GPU, even for large tasks such as the random function f and still find well well-peeforming elimination orders. Table shows the number of multiplications required by the best elimination order found by the PPO-agent. Figures and [2] contain the corresponding reward curves. We did not succeed in training a joint model using the PPO agent.
Table Best number of multiplications achieved by the PPO-based agent. HumanHeartDipple
RoeFiuux_e ld
RobotArm_6 -6OOF
PropaneCombustion
324
245
162
n.a.
422
RoeFlux 3d
MLP
Encoder
BlackScholes Jacobian
f
885
313
6497
29

---

algorithm optimization RoeFlux
AD algorithm optimization RobotArm 6DOF
multppiicatioss
muttpiicatoo
260
280
288
320
Our work
Our work SOTA RoborAm 6DOF
SOIA
340
1.5
2.0
0.0
1.0
samples
1e6
samples
(a) RoeFlux_ld
(b) RobotArm_óin _DOFF
AD algorithm optimization HumanHeertDippote
AD algorithm optimization PropaneCombustion
150
retur 1# muliiplicaaoo ons)
retum mumpliictiioss
160
172
180
100
102
190
work SOTA for HumanteertD pole
wark SOTA for PropareCambustion
105
2
0.0
0.5
1.0
1.5
2.0
1e
1a6
(c) HumanHcarDDipole
(d) PropancCombustion
AD algorithm optimization Random Function
algorithm aptimization BlackScholes
420
mutppiicaooos)
return mumplcattinn
350
350
451
480
400
500
450
520
Uur work
Our work
Function
BlackScholes
4
0.0
0.5
1.5 samples
2.0
3,0
# samples
1c6
1e6
(c) Random function
(f) MlackScholes_Jacoiie
Figure 11: Learning curves of the scalar tasks the same six random seeds with the PPO-based agent. It manages to find better elimination orders for many of the tasks. The best result for each task displayed in table
30

---

algorithm optim zation RoeFlux
agorrthhm optimization Random
938
8000
1000
retum (e mL ietiiiicaaterss
return (8 mu iipiliicttion
8000
1200
10000
1400
12000
1600
14000
1800
Our work
16000
work
SOTA Roeflux 3d
SOTA for Rardom Function
2000
0
2 samples
0
3
le6
samples
1e6
(a) RoeFlux__a 3d
(b) Random function f
AD algorithm optimization MLP
AD algorithm optimization Transfommer Encoder
2688
500
return carss)
10000
(uutplicaatonn)
20000
1000
1500
40000
50000
2000
Our work
Our work
MLP
60000
Transformer
0.0
1.0
samples
1e6
samples
1e6
(c) Multi-Layer Perceptron
(d) Transformer Encoder
Figure 12: Learning curves of the vectorized tasks for the same six random seeds for the PPO-based agent. It manages to find better elimination orders for all tasks. The best result for each task is displayed in table
31

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims contained in the abstract reflect exactly the content of the work. In the abstract, We make statements about theoretical and practical improvements in automatic differentiation. the experiments section we provide empirical proof for both.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
The paper clearly describes the limitations of the approach by including negative outcomes in the the experimental results table and provides more in-depth analysis in the discussion and appendix sections of the paper.
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
32

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
The entire codebase as well as comprehensive tutorial will be provided with the paper as part of GitHub repository. Furthermore. the appendix includes implementation details on the RL algorithms and neural networks as well as an in-depth description of the functions that were analyzed and all the best performing elimination orders. Furthermore, the reviewers will be provided with full access to the entire codebase on submission. Guidelines:
Justification:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility.
33

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: The entire codebase along with comprehensive introduction will be published on release of the work. Furthermore, the reviewers will be provided with full access to the entire codebase submission. Also, all novel elimination orders are provided in the appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
hetpo//////udddssdeeeebbbbn Sat////oeeSSSaaiiaaaeeeae issiiiissiii cassiiinnoccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: All necessary experimental parameters like learning rates, RL configuration etc. provided in the appendix and as yaml files in the source code of the model.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The key results of the paper runtime measurements that do not necessarily follow a Gaussian normal distribution. The authors therefore decided to display results using the median and 2.5- and 97.5-percentile quantify the uncertainty.
34

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details about the hardware requirements are included in the appendix of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///wnweiepeeeeepeeeiee/////
Answer: [Yes].
Justification: The authors have read the NeurIPS ethics code and are sure that the research conducted in this paper as well the results of the paper did not violate any of the points raised in the guidelines. The main result of the paper is way to find faster and more energy-efficieen automatic differentiation algorithms which to the authors best knowledge does not pose any harm to humanity and nature. Furthermore, all libraries used in this work are publicly available under very lenient licensing.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
35

---

10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The main result of the paper is way to find faster and more energy-efficient automatic differentiation algorithms which most likely has a positive impact on the environ- ment as it might reduce the energy consumption of many computational workloads. The authors are not aware of any possible negative influences.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA].
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
36

---

Justification:
All the external libraries used in this work have been properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Two new libraries will be released together with this work. Both libaries are well documented and the author of the paper is also the sole author of the libraries.
The source code is well documented and the package usage and installation guidelines are provided.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA].
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
37

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
38


Reviews:
[{"review_id": "sjPBHgsqh6", "reviewer": "Reviewer_PBG1", "summary": "This paper presents a novel method called AlphaGrad for optimizing computational graphs derived through automatic differentiation (AD) algorithms using deep reinforcement learning. The authors formulate the optimization of a computation graph as a single-player game where an AlphaZero agent aims to minimize the number of required multiplications. The authors demonstrate that AlphaGrad is able to improve over preexisting forward-mode AD methods in domains like computational fluid dynamics, robotics, optimization, and computational finance. Notably Deep Learning based tasks which leverage reverse-mode AD don't show significant improvements with AlphaGrad. Finally, the authors introduce Graphax, an AD interpreter for Jax that enables the optimization and execution of computational graphs discovered from elimination orders such as those proposed by AlphaGrad.", "questions": "- How challenging would it be to directly optimize for metrics like execution time or memory usage rather than the number of multiplications? This seems like an important direction for making the approach more practical.\n- Could you provide more details on how the random functions f and g were generated in the text?\n- Were hyperparameter sweeps performed per graph? Do optimal hyperparameters for the RL agent change depending on the structure of the graph? How might we overcome this limitation?\n- Why was Jax chosen instead of targeting more powerful IRs like MLIR?", "limitations": "The authors are forthcoming about the main limitations and highlight important areas for future work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Demonstrates a novel approach to optimizing AD by leveraging deep RL to discover optimized elimination orders.\n- The paper is technically sound and provides good descriptions of these various topics for an audience who may not be very familiar with AD.\n- Overall, the paper is well-written and the key ideas and results are clearly communicated.\n- The work has the potential for significant practical impact as even small improvements in AD efficiency can lead to large runtime and energy savings when applied at scale.", "weaknesses": "- Currently, the approach has a simple model based on optimizing the number of multiplications as a proxy for runtime. As the authors note this does not capture the full complexity and ignores factors like memory access patterns.\n- The baselines for this work seem somewhat limited, there are preexisting AD systems that operate on intermediate representations such as Enzyme [1] and LAGrad [2] that are able to perform much more sophisticated AD optimizations. This seems to be the future of AD and it would have been nice for AlphaGrad to operate on these IRs in a more general fashion. It would be nice if this was mentioned in the paper.\n- Results don't seem to be as strong for reverse-mode AD which makes sense. Deep learning is probably the largest user of AD and any gains here could have a significant real-world impact.\n- For every AD graph an RL algorithm must be run to solve for the elimination orderings. This can be quite computationally expensive and RL is known to be sensitive to hyperparameters and take quite a bit of tuning per environment.\n\n[1] Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients. William Moses and Valentin Churavy. Neural Information Processing Systems (NeurIPS), 2020.\n\n[2] LAGrad: Statically Optimized Differentiable Programming in MLIR. Mai Jacob Peng and Christophe Dubach.International Conference on Compiler Construction (CC), 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback, in particular for the additional references which we consider a useful contribution to shape the future of our project.\n\nWe agree with the reviewer's criticism regarding the use of the theoretical number of multiplications as a proxy for actual runtime improvements. Still, we want to emphasize that this somewhat crude approach already leads to several measurable improvements.\nThus, we are confident that when improving for actual runtime and memory access patterns, our results will only further improve. Achieving this with AlphaGrad would require an efficient hardware performance model, which we did not have available at the time of writing. We will leave this for future work, as it is promising research direction in its own right. \n\nTo answer **Q1**: Implementing new reward mechanism such as runtime and memory accesses will not be too hard, since it just requires to measure the quantity of interest either directly or using one of the aforementioned hardware models. The biggest challenge lies in doing it efficiently and in a statistically robust way, since runtime and memory consumption measurements are inherently noisy. We believe that by expanding our method to incorporate methods from distributional RL, we can alleviate this issue. Given the additional additional overhead of implementing the hardware models, reward measurement mechanisms and efficient implementation thereof, we believe that it would be out of the scope of this paper.\n\nWe partially agree with the statement that our baselines are limited with respect to state-of-the-art methods such as Enzyme an LAGrad. After a careful examination of both papers, we find that they operate on a different optimization level that is complementary to AlphaGrad. To prevent any misunderstandings upfront, we see this as a big advantage rather than a critique.\nLAGrad uses static optimizations to create and then optimize the differentiated code at the MLIR level. However, under the hood LAGrad is currently only able to leverage reverse-mode AD to create the derivative code. We conjecture that for an appropriate choice of function where reverse-mode AD performs worse than other AD algorithms (e.g. a function with few inputs but many outputs, which are often found in computer graphics problems such as differential rendering), the stochastic optimizations will in many cases not be able to bring reverse-mode to the same performance as other, better suited algorithms, e.g. forward-mode or minimal Markowitz. Thus, while being an amazing feat of engineering, the user is still required to select a specific AD algorithm for his use case from a very limited amount of choices. Our method instead aims to find the optimal AD algorithm, of which forward-mode, reverse-mode and minimal Markowitz are only certain instantiations. \nThus, AlphaGrad is a tool for algorithm search while LAGrad is a tool for algorithm optimization, even though depending on the level of granularity, such a distinction sometimes might be fuzzy. Thus, we do not entirely agree that including comparisons to LAGrad in our benchmarks are fair to either side, since they operate on different levels.\nNonetheless, we believe that combining both approaches offers a novel research direction to further extend their potential, thereby enabling even bigger gains. In particular, the handling of sparse operations seems quite intriguing.\nRegarding Enzyme, we completely agree with the reviewer that implementing AlphaGrad with a LLVM backend will surely be the more efficient choice. But due to our limited experience with LLVM and fitness to other concurrent projects in our research group, we selected JAX and its intermediate representation called a Jaxpression as an appropriate backend that enables fast prototyping (this should answer **Q4**).\nRegarding the reviewer's recommendation, we will add LAGrad and Enzyme to the introductory and discussion sections to properly appreciate their contributions and propose new research directions. In the future, we will seriously consider a reimplementation of our method using LLVM/Enzyme.\n\nTo answer **Q3** regarding the brittleness of RL, we agree that the performance of many RL algorithms can massively vary which is why we selected PPO and AlphaZero as our primary agents. They are known to generalize across many different tasks with the same choice of hyperparameters. In fact, all experiments in this paper were obtained with the same set of hyperparameters and we did not find any major improvement by performing a hyperparameter search. The only deviation are the best results for the MLP and transformer cases, where we used 250 Monte-Carlo simulations instead of the default of 50. Thus we are confident, that our training method works well for many other tasks as well, leveraging a plug-and-play style mechanism to improve AD algorithms.\n\nFinally, we want to address the reviewer's **Q2** on how the random functions $f$ and $g$ were generated in more detail.\nTo generate the random functions, we implemented another custom JAX interpreter which consists of a repository of elemental operations such as $\\cos$, $\\log$ but also matrix multiplications and array reshape operations. We then provide an interface to specify the number of input and output variables as well as the number of intermediate vertices that will be eliminated by the vertex elimination algorithm. The random function generator then randomly selects elemental functions from the repository. The number of functions that are unary, binary or perform accumulation or reshape operations can be controlled by adjusting the respective sampling probabilities.\nThe function generator is part of the *alphagrad* package and can be found in the source code under *src/alphagrad/vertexgame/codegeneration/random/random\\_codegenerator.py*.\nThe actual implementation of the random functions can be found in the examples directory of the *graphax* package.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lFrLmwVhdz", "reviewer": "Reviewer_kjpN", "summary": "The paper proposes to learn an efficient computation of the Jacobian of a program by reinforcement learning.\nThe authors base their learning procedure on cross-country elimination, a classical scheme that progressively transforms a computational graph to a bipartite graph representing the evaluation of the Jacobian. If all operations are simply scalar, a reward can simply be defined by the number multiplications accumulated in the final edges. The authors go a step further to accomodate for vector valued operations. The representation of actions and rewards can then feed to a deep reinforcement learning approach a la Alpha-zero. Experiments illsutrate the performance of the approach to find more efficient Jacobian evaluation algorithms than standard baselines on several computational graphs stemming from relevant applications. Comparisons in time are also provided to assess the practical benefits of the approach. \nThe authors developed a library enabling a direct use of the approach built on freely available package JAX.\nNumerous comparisons are also available in the supplementary material.", "questions": "- Are the performance displayed in Table 2 given after also some jit compilation?\n- To be sure I understand well: currently, to compute the best elimination procedure given some function, one needs to run a full training of a deep reinforcement learning agent and then use the final policy. In terms of computational time to find the right elimination strategy, this may be rather prohibitive? (Note that showing that better elimination strategies are possible is great.)\n\nDetail:\nline 288: typo \"two functions random functions\"", "limitations": "Limitations are discussed such as the poor generalization when performing joint training and the difficulties to handle the range of rewards.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The proposed approach demonstrates benefits across tasks.\n- The algorithm is developed in a package that will be released such that the method may quickly be adopted. \n- The problem can serve as an interesting benchmark for deep reinforcement learning.\n- The experimental details are throughly detailed to ensure preoducibility.", "weaknesses": "- As stated by the authors, the generalization performance of the approach may still be unclear, which may make the approach rather computationally prohibitive.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for his constructive feedback and the careful study of our manuscript. We will correct the typos immediately.\n\nRegarding the two questions posed by the reviewer:\n\nThe performances in table 2 are indeed given after JIT compilation of our AD algorithm. This was straight-forward to achieve because our package *graphax* is compatible with all other JAX transformations including JIT.\n\nThe reviewer's assessment is correct as the results presented in table 2 are obtained by training on a single computational graph. However, we also performed ablation studies on joint training where we optimized the AD algorithms of all benchmark tasks simultaneously. In this case, we still observed an improvement over the baseline in 7 of 10 cases with a major improvement of almost 10\\% for the random function $f$ over individual training. This makes us confident that with more data, i.e. computational graph examples and more compute we can train a model that can generalize out-of-distribution and predict optimal AD algorithms without requiring a training stage.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LOlJsSuezP", "reviewer": "Reviewer_ErgV", "summary": "In this paper, the authors propose 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. The author present the search for the optimal elimination order, aimed at minimizing the required number of multiplications, as a single-player game managed by an RL agent. The results show that this approach yields up to a 33% improvement over state-of-the-art methods across various relevant tasks from different domains.", "questions": "- How does the VertexGame ensures that the reduced computational graph gives correct derivatives? Is there a grad check procedure?\n- What is the application scope of the proposed method? Does it support to programs with complex control-flow (e.g., for-loop/while loop/if-condition and the combination of them)? \n- The mutable variables are notorious for a AD framework as the values of different version of these variables may needed to be stored for derivatives computations during the backward pass, how does the proposed method eliminate the vertex which is a mutable variable?\n- What level does the propose method works on? e.g., the source code or intermediate representation?\n- Does the proposed method more close to Source Code Transformation or Tracing/Operator Overloading in a context of AD?", "limitations": "yes.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Rephrasing the elimination order in AD as a deep reinforcement learning problem sounds somehow novel.\n- The paper is well-organized.", "weaknesses": "- The proposed method only shows a on obvious improvement on the RoeFlux_3d and Random function. While on other benchmarks especially on deep learning related MLP and Transformer cases, the improvements are very marginal (less than 3%).\n- It seems that the test cases provided are preliminary i.e., formula value evaluation instructions extracted from the whole graph, while not the the entire computational graph of a function/kernel with control-flow or mutable variables.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive feedback of our approach and want to take the time to address some of the feedback more in-depth as compared to the general rebuttal.\n\nFirstly, we noticed that our presentation of the obtained results is not optimal and invite the reviewer to have a quick look at figure 1 in the attached .pdf-document. This figure shows the actual runtime gains obtained with our algorithm for the MLP on different batch sizes and layer sizes. It clearly demonstrates that for larger MLPs, there is a significant runtime gain to be harvested. We will make our exposition more precise to properly appreciate this result.\nUnfortunately, the transformer benchmark does not reproduce the scaling properties of the MLP, but table 1 of the main body shows that there is potential for improvement.\nIn fact, we deliberately included this result in the hope of sparking future research and outlining the current limitations of our algorithm.\n\nSecondly, we partially agree with the sentiment that the results are preliminary.\nWhile the current benchmark tasks do not take into accounts actual kernels or control flow, many interesting practical applications do not require these features such as the MLP or the RoeFlux borrowed from computational fluid dynamics. However we appreciate the feedback and intend to include support for control flow in future iterations of the algorithm.\n\n**Q1**: How does the VertexGame ensure that the reduced computational graph gives correct derivatives?:\n\nThe vertex elimination algorithm which VertexGame is based on, always yields mathematically provable exact gradient up to machine precision. This is the base assumption of our work and we thank the reviewer for pointing out that we might not have made this assumption clear enough. We will improve our manuscript accordingly. For reference, we suggest the reviewer to check (Evaluating Derivatives, Griewank and Walter, 2008, Chapters 9+10, esp. Corollary 9.2). An in-depth explanation is available in the main rebuttal.\n\n**Q2:** What is the application scope of the proposed method? Does it support to programs with complex control-flow?\n\nOur work is targeted at every field that actively applies automatic differentiation for computing Jacobians. While our method currently does not support control flow primitives such as conditionals and loops, we are eager to implement support in future iterations of our methods. One however has to make a distinction between cases where the control flow depends on variables that are to be differentiated versus variables that are not differentiated. The latter case is already implicitly supported by our package while the former case will require additional work, since . There exists an extensive body of work that addresses the more complicated case which we will harvest for future improvements.\n\n**Q3**: The mutable variables are notorious for an AD framework as the values of different version of these variables may needed to be stored for derivatives computations during the backward pass, how does the proposed method eliminate the vertex which is a mutable variable?\n\nWe have considered the problems that mutable variables can cause but they are circumvented for many cases by the way JAX works. JAX leverages a functional programming paradigm which requires all functions to be pure functions with no side-effects, thereby not allowing for mutable variables that break the \"self-containedness\" of the pure function.\nRegarding mutable variables that do not break this, JAX guarantees referential transparency, i.e. at which point the variable has a given value by unrolling the respective computations and assigning all the intermediate values their own distinct static variables. This effectively creates a trace of the computations which can then be processed by the vertex elimination algorithm. In effect, we pay for being able to treat mutable variables with additional compute and memory consumption. It is important to note however, that in certain relevant cases such as for control flow and loops, there exist dedicated functions in JAX that avoid this costly rollout, e.g. the *jax.lax.scan* primitive. As already discussed, we intend to implement these cases in future iterations of the algorithm.\n\n**Q4**: What level does the propose method works on?\n\nOur framework works at the Jaxpression-level of JAX. Jaxpressions are a kind of intermediate representation introduced by JAX to handle many of JAX' features in a simple and elegant manner. In particular, JAX' own forward- and reverse-mode AD implementations are realized using Jaxpressions. If analyzed and transformed in the right way, the Jaxpression is actually the representation of the functions computational graph. We exploit this property in order implement vertex elimination which relies on the computational graph to work properly.\n\n**Q5**: Does the proposed method more close to Source Code Transformation or Tracing in a context of AD?\n\nOur proposed method is a mix of both, source code transformation and tracing since this is the modus operadi of JAX itself. All major features in JAX are implemented as function transformations and use a tracing feature to keep track of all operations performed. JAX introduces the Jaxpression intermediate representation to store the trace and perform the respective function transformations, including automatic differentiation. Our framework *graphax* is no exception to this. However, we want to emphasize here that to make vertex elimination work in JAX, we had to implement an entirely new interpreter from scratch. This interpreter first extracts the computational graph from the Jaxpression and then proceeds to apply the vertex elimination algorithm to obtain the optimized AD code. We consider our package *graphax* as a novel contribution in its own right because there exists no prior implementation of vertex elimination in any Python-based popular machine learning/numerical computing framework.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fe4hi8bPoU", "reviewer": "Reviewer_tWYs", "summary": "This paper considers the (combinatorially hard) optimization problem related to automatic differentiation algorithms represented via an \"elimination order\" based on prior work and proposes an RL formulation to solve it approximately. This follows in similar vein to recent celebrated results on matrix multiplication and sorting. The authors demonstrate that their approach is able to improve, conceptually on the first order metric of the number of multiplications as defined in the optimization search, and more practically that this can be sufficiently meaningful for runtime as they measure it on representative benchmarks , compared to the default forward/reverse mode implementations. The paper is well written and considers a reasonable variety of computational graphs in the evals. The main overhead in adopting the proposal in practice is primarily in the form of a specialized policy search for each computational graph, which factors in roughly speaking, at the compilation stage and is not reflected explicitly in the benchmarks.", "questions": "- When considering the number of multiplications, given the definition of a vertex, I assume this refers to the number of matmuls, but the costs of these vary by the dimensions of the jacobians. Do you take this into account anywhere? If not, what is the hurdle? Making the exposition a bit more explicit on the definition/implementation of the  reward is framed might be helpful. \n\nOther comments:\n- The labels for cross edges in Figure 2(d) appear to be swapped.\n- Caption for Figure 2 has intermediate variable definition of $v_2=cos(v_1)$, which needs to be $sin(v_1)$ instead?\n- L128 graph 2a --> Fig 2a\n- In Definition 1 L138: should this be $c_{ij}=c_{jk}=0$ rather than $c_{ij}=c_{ik}=0$?\n- L142, $\\phi_j \\circ \\phi_k$ needs to be $\\phi_k \\circ \\phi_j$ instead.\n- Equation (2) needs to be $W_{ik}$ rather than $W_{ki}$?\n- L324 refers to Table 2 being of theoretical value -- should this be Table 1?\n- L341 s/note/not\n- Figure 1 3rd box: s/executation/execution", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well motivated and proposes a policy search framework for a well studied combinatorial optimization problem (vertex elimination for AD).  Instead of targeting compilation (at a lower level) or dealing exclusively with optimizing individual matmuls as in prior work, the paper targets a novel alternative by focusing on vertex elimination in AD which is a ubiquitous part of all training compute graphs.\n- Any non-trivial improvements on large scale computational graphs could lead to significant compute savings, and this work proposes several novel small scale benchmarks to evaluate vertex elimination algorithms which maybe of interest in its own right.", "weaknesses": "- The primary weakness in terms of the impact appears to be that the policy must be trained from scratch separately for each graph. This unfortunately means that there is a whole training run necessary per graph unlike all of the other competing methods (including the minimum Markowitz degree heuristic, which needs no \"training\", and already appears to do better than fwd/reverse mode in many situations).\n \n- The overhead involved in running a separate optimization for every graph that is distinct from the compilation step (mostly under the hood) seems like it could makes the proposed method unwieldy for iterative usage inspite of saving some compute. Of course, any compute savings on large training graphs are valuable. However, the transformer baseline comparison indicates that the gains are more prominent in the smaller benchmarks and in the larger experiments, the reverse mode and AlphaGrad perform very close to each other.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the valuable feedback on our work and appreciate the time the reviewer took to double-check the equations. We will implement the corrections accordingly.\n\nWe partially agree with the reviewers assessment that our method requires retraining for every computational graph. While it is true that the results presented in tables 1 and 2 of the main paper were obtained with individual training, we also investigated the prospect of joint training and found that in 7 out of 10 cases our method still outperforms the baseline with a major improvement of around 10\\% for the random function $f$. These results make us confident that for a larger dataset and longer training time, our method will eventually generalize out-of-distribution and predict the optimal elimination order without requiring training.\nWe intend to implement this feature in future iterations of our algorithm.\n\nFurthermore, we would like to emphasize the point that this paper not only presents a novel RL-based method to optimize AD algorithms but also introduces, for the first time, a Python-based interpreter that actually allows to apply vertex elimination to a wide range of problems, including machine learning. In particular, the creation of the baselines using the minimal Markowitz heuristic are only possible due to our implementation in the first place. So while it is true that the minimal Markowitz heuristic can do better than forward-mode and reverse-mode in many cases already, it is only due to our implementation of the interpreter that this could be exploited in practice. We believe that this already makes a contribution in its own right. Furthermore, we were able to show that while the minimal Markowitz heuristic performs very well, it is outperformed by our method on all benchmark tasks, sometimes by a significant margin.\n\nRegarding the second major weakness that the reviewer pointed out, which is that our proposed method unwieldy for iterative usage, we are not entirely sure if we understand the remark correctly. Hence, we interpret it such that if the computational graph changes after the compilation stage, retraining is required.\n\nFirst of all, we want to clarify that the computational graph and the respective Jacobian function are independent of the input data, i.e. it is not necessary to retrain the graph every time the input of the Jacobian function changes (this might change when including control flow, which we have not considered yet).\n\nSecondly, for many practical use cases, the compute graph does not change after the compilation stage, i.e. the neural network architecture or the computational fluid dynamics equations are static. Therefore, the training phase is only necessary prior to the compilation stage once and the computed AD algorithm can be used iteratively after that for any input parameters.\nWe still agree with the reviewer on the fact that the performance of our method on transformers is not satisfactory yet, but we would like to point out that theoretically there is a benefit to be harnessed as shown by table 1 in the main body of the paper. Furthermore, we deliberately included the negative outcome on the transformer case in hope of sparking future research endeavours in this direction and point out current limitations of our approach.\n\nRegarding the questions about how the reward is actually calculated and whether the shape of the Jacobian is taken into account: \nWe thank the reviewer for pointing out that this part of the algorithm might not be described very well in the main body of the text and we will perform the appropriate changes to make this more explicit.\nWhile we gave a overview of the answer in the general rebuttal, we would like to take the time here to quickly explain in more detail how the reward is calculated:\nThe reward function takes into account the shape/dimension of the Jacobians, i.e. if we have a matrix-multiplication of two dense Jacobians of shape 3x3 and 3x2, the resulting reward would be -18.\nIt also takes into account the sparsity type of the Jacobian. i.e. some operations have a naturally diagonal Jacobian. In this case the matrix-multiplication becomes a pointwise operation so that for diagonal 3x3 Jacobian and a dense 3x2 Jacobian we have a reward of -6 instead. Furthermore, the *graphax* package provided together with this work is a novel interpreter that can take these sparse Jacobians into account, i.e. automatically detect these diagonal Jacobians and replace the matrix-multiplication or tensor-contraction by a pointwise multiply which is much faster in many cases. The native JAX AD package is not able to perform this kind of analysis.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "snkdqLrko6", "reviewer": "Reviewer_ZG2Q", "summary": "The authors study the problem of computing the Jacobian of arbitrary computation graphs. The typical approach in most autodiff libraries is to implement the classic backpropagation algorithm, aggregating gradients from the end to the beginning. This is because such approaches are well suited for cases where the number of inputs is significantly larger than the number of outputs, such as neural nets.\n\nThe paper first describe alternate approaches of computing the gradient from the literature, then describe how Jacobian computation can be treated as a sequential decision making problem. The state is the current computation graph, as represented by an adjacency matrix, and the actions are choosing vertices of the computation graph to eliminate. Each vertex elimination incurs a cost of updating the partial derivatives of its neighbors, and the terminal state is when all intermediate vertices are removed and only edges between the inputs + outputs remain. The goal is to minimize the cost, which this paper measures in multiplications.\n\nThe authors try both an AlphaZero style agent and a PPO agent, and demonstrate their agents can outperform standard methods (backward AD, forward AD, etc.) on a few example functions. The empirical runtime is also analyzed, showing that the AlphaZero agent can lead to faster computation, even though the reward function does not account for all runtime factors (i.e. memory access patterns)", "questions": "As part of the graph representation coding, there is a list of 21 different sparsity patterns to describe the intermediate edges of the computation graph. Is this list exhaustive to more general computation graphs, or are there other forms the partial derivative can take if the space of computation graphs is increased?", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is written well, and serves as a helpful introduction to how to optimize autodifferentiation. The paper acknowledges that it makes some simplifications to the classes of computation graphs it can be applied to (notably limits on input dimensionality make the Transformer experiment somewhat unrealistic), but demonstrate it is possible to achieve gains over typical methods using learning based methods. An analysis on real hardware is performed as well.", "weaknesses": "By my understanding, both the AlphaZero agent and PPO agent need to be instantiated separately for each class of functions. This is based off the note in Appendix F, that the authors did not succeed in training a joint model for the PPO agent. One of the papers cited as related work is the one for improving matrix multiplication with search algorithms, and one benefit of that work was that although the search was very specific to 4x4 matrices for example, it was possible to us the discovered results of that work elsewhere.\n\nIn comparison, I believe the current approach requires the time spent learning the PPO / AlphaZero agent to get made up by the time spent running the autodifferentiation calls. My suspicion is that the net time ends up being negative, assuming the AlphaZero setup takes longer than PPO, since the PPO setup is a few minutes of training to improve the MLP gradient computation by $$0.30 - 0.29 = 0.01$$ milliseconds  in the GPU case, which implies 18-24 million steps to hit a break-even point. (3-4 minutes / time-saved-per-step)\n\nStill, I overall think there are interesting signs of life here, even if the method itself is not entirely practical yet.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the helpful input to our work.\nIn particular, we agree with the reviewer that method is not entirely practical yet, but we intend to improve so that is becomes applicable to a wider range of problems.\nAutomatic differentiation is an ubiquitous tool that sees application in many scientific areas. Our experiments show that there are benefits to be reaped across fields on already the simplest problems.\nThus we are confident that the ramifications of scaling our algorithm and bringing it to practical use are huge with large potential benefits for users in all application areas.\n\nFurthermore, the reviewer's analysis made us realize that the presentation of our results does not entirely capture the actual runtime improvements achieved with our algorithm. We invite the reviewer to have a look at figure 1 in the accompanying .pdf-document where we demonstrate the scaling of our optimized algorithm for the MLP for increasing batch sizes and layer sizes. For larger MLPs and larger batch sizes, the actual runtime gains are significantly larger than what is presented in table 2 of the main text, thereby providing a stronger argument for the practical use of our method. In the words of the reviewer, with larger batch sizes, the break-even point shifts in favor of our method.\nWe will adjust the main body of our submission accordingly to accommodate this.\n\nFinally, to answer the question regarding the 21 different sparsity types and their exhaustiveness:\nThe list of sparsity types is exhaustive for Jacobian tensors of up to 4 dimensions i.e. $\\dfrac{\\mathrm{d}T_{ij}}{\\mathrm{d}x_{kl}}$. Note that dimension in this case refers to the size of the mathematical objects, i.e. a number is zero-dimensional, a vector is one-dimensional and a matrix is two-dimensional etc. All computation graphs that only rely on objects that are 2-dimensional (max. 4-dim Jacobian) can be analyzed and differentiated with this procedure without any further extensions. If we want support for higher-order Jacobians, e.g. a six dimensional Jacobian tensor like $\\dfrac{\\mathrm{d}T_{ijkl}}{\\mathrm{d}x_{mn}}$, we need to introduce new sparsity types. This is not too difficult, but for the sake of simplicity we limited ourselves to 4 dimensions as many interesting applications can already be treated in this case. In the future, we intend to extend the support for higher-dimensional Jacobian tensors to make our approach more generally applicable.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KtX7rkrnok", "reviewer": "Reviewer_N91J", "summary": "This paper studies automatic differentiation in a computational graph. The topic is classic with wide applications in scientific research, e.g., computing gradients of a neural network or Jacobians needed in numerical optimization. Classic numerical methods like standard forward- and reverse-mode differentiation (backpropagation) exist, and the novelty and contribution in this work lies in its proposal of training a reinforcement-learning (RL) agent to explore an optimized order in graph vertex elimination. The paper demonstrates the efficacy of its method by comparing the time cost of the found auto differentiation policy (measured by multiplication count and wall-clock time) with that of classic and heuristic methods.", "questions": "I don’t have any questions for now. The key idea described in this paper is quite straightforward, and the writing in the main paper is generally good.", "limitations": "Admittedly, the proposed RL method does not beat classic methods across the board in the experiments. I share the similar conjecture expressed in the paper: backpropagation is a quite strong baseline for large-input, small-output computational graphs, and it can be quite difficult to beat. This may also indicate that any small improvement on this strong baseline can have a significant impact, as backpropagation is so widely used in practice.\n\nOn a related note, I appreciate that this paper is upfront about its marginal improvement over baselines in certain cases. This gives readers a balanced and well-positioned view of the proposed method.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Training an RL agent to rethink classic numerical methods is an interesting research topic. Previous examples include DeepMind’s works on matrix multiplication and sorting algorithms. Because these numerical problems are basic and fundamental building blocks in today’s scientific research, improvements in its solutions can have high potential and impact. This paper follows this trend and applies RL to another specific example: automatic differentiation. I think the paper picks a good topic and makes a meaningful contribution to the community.\n\nThe technical method in this paper also looks good to me. Using RL to explore better numerical and computing methods typically requires a careful, RL-friendly (re-)formulation of the original numerical problem, including the design of the state space, the action space, and the reward function. I think choosing to tackle vertex elimination ordering is a smart and reasonable move, which leads to a well-formulated game that captures the key structure in automatic differentiation.", "weaknesses": "I don’t have major concerns with the paper, but I want to mention that there are quite a few typos in the equations in the main paper. For example, I suspect “i > j” should be replaced with “i < j”, and “sum_i” in line 192 should be “sum_j”. Please double-check all equations in the paper and make sure there are no typos.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for pointing out the typos and will definitely fix them for the final submission.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank reviewers for their thoughtful, detailed reviews and the additional literature suggestions. We first discuss two common themes, then address certain individual comments.

**Retraining is necessary for every function**

The best results for most of the graphs were indeed achieved with individual training. However, we also performed an ablation study where we jointly trained on all graphs at once.
The results of joint training still outperform the baseline for 7/10 tasks with even a major improvement for the random function $f$ over individual training by almost 10% (see table 1 in the attached .pdf-document). Thus, we are confident that with more computational graph examples, it is possible to train a model that eventually generalizes out-of-distribution and can predict high-performance, individually tailored elimination orders/AD algorithms for almost arbitrary functions.

**AlphaGrad does not outperform reverse-mode on ML workloads**

We acknowledge that the theoretical gain for the transformer architecture does not translate into practical runtime improvements. In fact, we chose to include this result to outline the current limitations of our approach. However, we feel that it is important to emphasize that, our algorithm, for the first time, has found a new automatic differentiation algorithm that is more efficient, i.e. requires fewer multiplications than reverse-mode/backpropagation. This shows that there is some potential to be harvested here which might spark further research since the transformer architecture is now an integral part of the ML community.
Furthermore, we invite the reviewers to have a quick look at the accompanying .pdf-document that exemplifies the scaling of the results obtained for the MLP for different workload sizes. 
The relevance of MLPs for small-scale ML applications is evident and we show that for reasonably sized MLPs, we can achieve significant runtime improvements. 

**Response to Reviewer #3**

**Q1**: How does the VertexGame ensure that the reduced computational graph gives correct derivatives? Is there a grad check procedure?

There exists a series of mathematical proofs that guarantee that vertex elimination always yields the correct derivatives up to machine precision. For more information, we suggest (Evaluating Derivatives, Griewank and Walter, 2008, Chapters 9+10, esp. Corollary 9.2).
Computing derivatives of a function can be seen as multiplying a long chain of Jacobians of elemental functions (such as matmuls and $
\cos$) with each other. No matter in what order they are multiplied, they will always yield the same exact Jacobian of the function. However, the order in which they are multiplied with each other severely impacts the computational performance, i.e. the number of operations that are needed to arrive at the final result. As an example take three matrices A,B,C of shapes 3x3, 3x2, 2x1. First multiplying A and B and then C requires $3\cdot 3\cdot 2 + 3 \cdot 2\cdot 1 = 24$ multiplications. Solving B and C first and then A requires $3\cdot 2\cdot 1 + 3\cdot 3\cdot 1 = 15$ multiplications. AlphaGrad tries to find the best possible order of multiplication for the chain of elemental Jacobians. The elemental Jacobians are hardcoded into the algorithm as for every other automatic differentiation tool such as Tensorflow or PyTorch. Thus, there are no approximations and vertex elimination always gives the exact gradient.

**Q2**: What is the proposed application scope? Does it support to programs with complex control-flow?

The application scope of this method includes major areas that actively utilizes automatic differentiation for gradient computation, including computational fluid dynamics, robotics and deep learning. While the current method is limited to functions without control flow, it is possible to extend the notion of vertex elimination to these cases. However, many interesting applications already arise without the need for control flow e.g. MLPs or differential kinematics of robots. Nonetheless, we agree that expanding the method to include control flow will significantly widen the scope of AlphaGrad.

**Q3**: The mutable variables are notorious for a AD framework as the values of different version of these variables may needed to be stored for derivatives computations during the backward pass, how does the proposed method eliminate the vertex which is a mutable variable?

The functional programming paradigm of JAX requires functions to be pure, i.e. to have no side effects outside the scope of the current function. This already precludes many potential problems arising from mutable variables. However, it is still possible to have mutable variables that do not violate the scope of the function. In these cases, while converting the function of interest into a Jaxpression, JAX rolls out these computations and assigns each new value of the mutable variable a new static variable (similar to static single assignment, SSA), thereby paying the price of additional compute and memory to allow for such constructs. The resulting computational graph can then be differentiated with vertex elimination.

**Q4**: What level does the propose method works on? 

The method works on the Jaxpression level, which is a type of intermediate representation introduced by JAX to handle function transformations.

**Q5**: Does the proposed method more close to Source Code Transformation or Tracing/Operator Overloading in a context of AD?

The proposed method is a mixture of both approaches, operator overloading and source code transformation, since JAX also utilizes this hybrid approach.

**Response to Reviewer #4**

**Q1**: The number of multiplications is dependent on the shape of the Jacobians?

Correct. We take this into account when computing the reward and will adjust the description to make this more explicit. A full-blown explanation of how the rewards are calculated for (potentially sparse) Jacobians can be found in Appendix C.

Meta-review: This paper is on optimizing the order of operations for Jacobian computation in a computational graph, focusing on reducing the number of multiplications required. Instead of using the standard practice of backward automatic differentiation (AD), this paper uses a reinforcement learning-based method for optimizing the automatic differentiation process.

There is a general agreement among the reviewers that the work is novel, potentially impactful, and practical. The reviewers also pointed out some weaknesses regarding computational expense for dynamic graphs and only having a modest performance improvement compared to reverse-mode AD for MLP and transformer tasks.

Overall, this is a strong paper with no major issues, and hence, I recommend accepting the paper, as do all reviewers. However, another weakness I would like the authors to address in the camera-ready paper is that the paper does not discuss a particular scope of the work. The work seems specific to static computational graphs and, hence, is more suitable for building on JAX than PyTorch. The authors mentioned in the rebuttal that “for many practical use cases, the compute graph does not change after the compilation stage, i.e. the neural network architecture or the computational fluid dynamics equations are static.” The authors should specify that and acknowledge that this work is going to be prohibitive for those using dynamically changing network architectures. The authors should also mention the implication of this issue for building on JAX vs. PyTorch in terms of implementation difficulty and overall computations.
Decision: Accept (spotlight)