Title: Optimizing Automatic Differentiation with Deep Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

Automatic Differentiation

AlphaZero

Deep Learning
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

Abstract

Computing Jacobians with automatic differentiation is ubiquitous in many scientific domains such as machine learning, computational fluid dynamics, robotics, and finance. Even small savings in the number of computations or memory usage in Jacobian computations can already incur massive savings in energy consumption and runtime. While there exist many methods that allow for such savings, they generally trade computational efficiency for approximations of the exact Jacobian. In this paper, we present a novel method to optimize the number of necessary multiplications for Jacobian computation by leveraging deep reinforcement learning (RL) and a concept called cross-country elimination while still computing the exact Jacobian. Cross-country elimination is a framework for automatic differentiation that phrases Jacobian accumulation as ordered elimination of all vertices on the computational graph where every elimination incurs a certain computational cost. We formulate the search for the optimal elimination order that minimizes the number of necessary multiplications as a single player game which is played by an RL agent. We demonstrate that this method achieves up to 33% improvements over state-of-the-art methods on several relevant tasks taken from diverse domains. Furthermore, we show that these theoretical gains translate into actual runtime improvements by providing a cross-country elimination interpreter in JAX that can efficiently execute the obtained elimination orders.

1 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 et al., 2018, Margossian, 2018, Forth et al., 2004a, Tadjouddine et al., 2002b, Giftthaler et al., 2017, Kato et al., 2020, Schmidt et al., 2022, Capriotti and Giles, 2011, Savine and Andreasen, 2021]. To many researchers in the machine learning community, AD is synonymous with the backpropagation algorithm [Linnainmaa, 1976, Schmidhuber, 2014]. However, backpropagation is just one particular way of algorithmically computing the Jacobian that is very efficient in terms of computations for 'funnel-like' functions, i.e. with many inputs and a single scalar output such as in neural networks. In many other domains, we may find functions that 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 advantages and drawbacks regarding computational cost and memory consumption depending on 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-County Elimination frames AD as an ordered vertex elimination 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,

ComputationalGraph

ReinforcementLearning

OptimalAutmaticDifferentiation

State

eliminationorder=actionsequence

Function inPython +JAX

Feedeliminationorderand

def f(x，y):

Python codeintoGraphax

v1=x*y

v2 = sin(vl)

GraphRepr.

Policy

jax.jit(jacve(f， order=order, argnums=(0,1)))(1.0，2.θ)

return log(v2)，v1-v2

1

AlphaZero

1

Optimal AD algorithm as"jaxpr"

Computationalgraphand

1

forexecutioninJAX

adjacencymatrixrepresentation

1

1

{lanbda ;a:f32[]b:f32[].let

c:f32[] =nul a b

d:f32[] = stn c

e:f32[]=cos c

_:f32[] =log d

f:f32[] =dtv 1.0 d

1

action=eliminatevertex

_:f32[]

h:f32[]

= sub c d

=nul f e

(=application of chain rule)

t:f32[] = nul g e

j:f32[] = add 1 1.0

k:f32[] = nul j b

reward=-computational cost

yn=[]

l:f32[]=nul

n:f32[]= nu]

h b

)a

tn (n, n, k, 1) )

2008]. Inspired by recent advances in finding optimal matrix-multiplication 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 1). 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 an 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 elimination techniques through heuristics, simulated annealing or dynamic programming and minimizing related quantities such as fill-in [Naumann, 1999, 2020]. However, none of these works was successful in optimizing with respect to relevant quantities such as number of multiplications or memory consumption. We set up our optimization problem by formulating crosscountry elimination as a single player RL game called VertexGame . At each step of VertexGame, the 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 is played by an AlphaZero-based agent [Silver et al., 2017, Schrittwieser et al., 2019, Danihelka et al., 2022] with policy and value functions modeled with a transformer architecture that processes the graph representation and predicts the next vertex to eliminate, thereby incrementally building the AD algorithm.

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 of 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 a deep RL agent finds new, tailored AD algorithms that improve the state-ofthe-art on several relevant tasks,

We investigate how the discovered novel elimination procedures translate into actual runtime improvements by implementing Graphax, a cross-country elimination interpreter in JAX allowing the efficient execution of newly found elimination orders.

1.1 Related Work

RL for Algorithm Research: AlphaTensor and AlphaDev successfully demonstrated that modelbased deep RL finds new and improved matrix-multiplication and sorting algorithms [Fawzi et al., 2022, Mankowitz et al., 2023]. In particular, AlphaTensor used an extension of the AlphaZero agent to search for new matrix-multiplications algorithms that require fewer multiplication operations by directly using this quantity as a reward. The key insight is that different matrix-multiplication

algorithms have a common, simple representation through the three-dimensional matrix-multiplication 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-multiplication algorithms for 4x4 matrices by beating Strassen's algorithm, the current state-of-the-art, with an improvement from 49 to 47 multiplications.

In a 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 as an optimization target, the 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 .

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 to 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 REINFORCE-based genetic algorithm to 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 on 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 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.

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 a 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 AD which can be suboptimal for certain tasks. In 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 RL.

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-art even when it 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 to accumulate the Jacobian on real-world problems. Another approach described in [Chen et al., 2012] utilized integer linear programming to find optimal elimination 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-Country Elimination

AD is a systematic approach to computing the derivatives of dependent variables y = f ( x ) ∈ R m with respect to the independent variables x ∈ R n 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

(a)

Computational graph of

f

(

x

1

, x

2

)

.

(b)

Adding partial derivatives.

(c)

Elimination of vertex 2.

(d)

Graph after eliminating vertex 1.

x1

log

x2

x1

1og)

1

cOSU1

U2

X

-1·coSU1

x2

x1

10g)

02

COSU1

×

x2

x1

x2021

COSU1

log)

x2-x2COSU1

1cosU1

x102

x1-x1COS71

x2

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 trigonometric functions. Partial derivatives of these elemental operations are coded into the AD software and the Jacobian is accumulated by recursively applying the chain rule to 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 V being the elemental operations ϕ j and directed edges E that describe the data dependencies between the operations (figure 2a). The relation i ≺ j states that vertex i has an edge connecting it with vertex j , meaning that the output v i of ϕ i is an input of ϕ j . The partial derivatives of the elemental operations with respect to their dependents are assigned to the connecting edges (figure 2b). We can then identify the edges of the graph with their respective partial derivatives c ji = ∂ϕ j ∂v i . The cross-country elimination algorithm computes the Jacobian by a procedure called vertex elimination .

Definition 1 [Griewank and Walther, 2008] For a computational graph G = ( V, E ) with partial derivatives c ij , vertex elimination of vertex j is defined as the update

and then setting c ji = c kj = 0 for all involved vertices. The += operator creates a new edge if there is no edge c ki and otherwise adds the new value to the existing value.

Intuitively, vertex elimination can be understood as the local application of the chain rule to a single vertex in the graph since the multiplication c ij c jk is exactly the result of applying the chain rule to ( ϕ k ◦ ϕ j )( v i ) . If a 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 are added to the graph, we delete all edges connected to the eliminated vertex since all the derivative information is now contained in the new edges (figure 2c). Note that the new graph resulting from a vertex elimination no longer directly represents the data dependencies of the function since the eliminated vertex is now disconnected. Furthermore, the application of vertex elimination as described above

requires the computational graph to be static and precludes the use of control flow ( if -statements, for -loops) in the differentiated code.

2.2 Cross Country and Elimination Orders

The repeated application of the vertex elimination procedure to a 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 Jacobian of the function f . In particular, as long as all intermediate vertices are eliminated, this Jacobian will always be exact up to machine precision [Griewank and Walther, 2008]. There is no restriction on 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 a known number of multiplications that depends on the shapes of the elemental Jacobians which can be used as a proxy for execution time. We can 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 are to either eliminate the vertices in the forward or reverse order. These two modes are called forward-mode AD and reverse-mode AD (backpropagation), respectively.

Forward-mode AD, where vertices are eliminated in the same order as the computational graph is traversed, is particularly efficient for functions where the number of input variables n is much smaller than the number of output variables m , i.e. n ≪ m . In contrast, reverse-mode AD traverses the graph in the opposite direction and is particularly suited for the cases where n ≫ m . 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 a vertex is defined as the number of incoming vertices times the number of outgoing vertices, i.e. Mark( j ) = | i ≺ j || j ≺ k | , where | · | denotes the cardinality of the sets i ≺ j and j ≺ k 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 on the resulting graph. This elimination scheme is one of the best known heuristics for finding efficient elimination orders and can incur savings of up to 20% over forward- and reverse-mode AD [Albrecht et al., 2003, Griewank and Walther, 2008]. However, for computational graphs that have many inputs and few outputs, it is often outperformed by reverse-mode AD.

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 eliminations. A simple example is the multiplication of a vector with a 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 i = ∑ j W ij x j with the partial derivatives

(a) Sparse vertex elimination.

(b) Computational graph representation.

W

Wki

bidii

@

tanh

Graphax

insteaddoes

x10jk

fastelementwise

x

multiplication

3

1

b;dij

x18jk

x20jk

X38jk

xbixbixbi

Input token

Output shape

Inputshape

Sparsitytype

Sequence dimension

The output vertex represents the application of the activation function y i = tanh a i with the partial derivative

According to the vertex elimination rule, upon elimination of the intermediate vertex the two Jacobians in equation (2) are assigned to the incoming edges are contracted together with the Jacobian from the outgoing edge in equation (3):

In both cases, instead of a matrix multiplication, one can perform simple element-wise multiplications as shown in figure 3a. For vectorized cross-country elimination to be efficient, it is paramount to exploit this property. Current state-of-the-art 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 JAX [Bradbury et al., 2018] in order to leverage it's defining features such as JIT compilation, automated batching, device parallelism and a user-friendly 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 https://github.com/jamielohoff/graphax and https://github.com/jamielohoff/alphagrad.

2.5 Computational Graph Representation and Network Architecture

We describe here how the computational graph is represented for optimization in the RL algorithm, as well as 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 an edge, we set the i -th row and the j -th column of the matrix to 1 . For the vectorized case, we define an extended adjacency tensor by extending the matrix into the third dimension. Along this third dimension, we store 5 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

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 ∂y ij ∂x kl where the first two values describe the shape of x kl and the other two values describe the shape of y ij . 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 3b shows the representation of the entire computational graph and a 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 a 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 connectivity.

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., 2017]. The output of the transformer is then fed into a policy and a 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 to use the number of multiplications incurred by the elimination order as a proxy value [Tadjouddine et al., 2006, 2002b, Albrecht et al., 2003]. Thus, we chose the negative number of multiplications 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 a 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 magnitude. To tackle this, we rescale the cumulative reward using a monotonous function. For the functions with scalar inputs as well as RoeFlux_3d and random function f , we found the method presented in [Kapturowski et al., 2019] performed well, i.e. we scaled with s ( r ) = sgn( r )( √ | r | +1 -1) + ϵr where ϵ = 10 -3 . For MLP and TransformerEncoder 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 a 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-Carlo Tree Search (MCTS) simulations. On most tasks, we found that 50 MCTS simulations were sufficient to reach satisfactory performance. Appendix D contains more details about the training of the agent.

4 Experiments

To demonstrate the effectiveness of our approach, we devised a set of tasks sampled from different scientific domains where AD is used to compute Jacobians. More details concerning the tasks are listed in appendix A.

Deep Learning is a prime example for the success of large-scale AD. We analyze a 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 elimination in the past [Roe, 1981, Tadjouddine et al., 2002b, Zubair et al., 2023]. We test on the 1D 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 (e.g. joints, actuators). There has been a surge

in interest of computing the Jacobian using AD[Giftthaler et al., 2017]. We chose the forward kinematics of a 6-DOF robot arm as a representative problem and follow [Dikmenli, 2022] for the implementation.

Non-Linear Equation Solving requires the computation of large Jacobians to apply state-of-theart solvers. The MINPACK problem collection provides a set of problems derived from real-life applications of non-linear optimization and designed to be representative of commonly encountered problems. In particular, we analyze the HumanHeartDiple and PropaneCombustion tasks, for which vertex elimination 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 so called 'greeks' which measure the sensitivities of the value of an option to the model parameters[Naumann, 2010, Savine 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 a 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 [Albrecht et al., 2003]. We generated two random functions f 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 1 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 as small as possible. The resulting AD algorithms are nonetheless scalable to arbitrary batch sizes. We use forward-mode, reverse-mode and the minimal Markowitz degree method as 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, our approach was able find new elimination orders with improvements ranging from 2% to almost 20%. We found that even for only 5 MCTS simulations, the agent was able to find better than state-of-the-art solutions for the scalar tasks.

On the opposite spectrum, our experiments with 250 MCTS simulations yielded no significant improvement over the results presented in table 1. The four remaining tasks are arguably more difficult since the vector-valued inputs and large variance within possible rewards provide an additional challenge. The RoeFlux_3d and random function f 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 as simulated annealing or dynamic programming struggled to even 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.

The authors conjecture that this is not only due to difficulties presented above, but because reversemode AD (backpropagation) is already a very well-suited algorithm for computing Jacobians of 'funnel-like' computational graphs with many inputs and a 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 1 were obtained by separately training on 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 to the separate training mode, the agent found new, improved elimination orders for almost all tasks except the MLP , TransformerEncoder , RoeFlux_3d and PropaneCombustion tasks. For the random function f and BlackScholes_Jacobian task, the multi-task training outperformed the results in table 1 with new best results of 5884 and 307 respectively, thereby showing that the algorithm search might benefit from training on diverse tasks simultaneously. This also hints at the possibility of building a more general statistical model of AD applicable to workloads from many different domains. We also experimented with PPO as an alternative (appendix F).

4.2 Runtime Improvements and the Graphax library

The results in table 1 are 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 Python-based AD interpreter able to leverage cross-country elimination. Graphax builds a second program that computes the Jacobian by leveraging the elimination orders found by AlphaGrad and using the source code of the function as a template by analyzing its Jaxpression . The Jaxpression is JAX' own representation of the computational graph of the function in question.

Table 2 shows runtime improvements for the elimination orders found in table 1 for a batch size of 512 with varying levels of improvement. This is due to the fact that the number of multiplications alone was only a proxy to capture the complexity of the entire program. It ignores other relevant quantities such as memory accesses and operation fusion during compilation. Nonetheless, a particularly impressive gain over the state-of-the-art methods can be observed for the RoeFlux tasks and the RobotArm_6DOF task. Remarkably, we also observe a minor improvement for both deep learning tasks when executed on GPUs. Note that the TransformerEncoder task was evaluated on a batch size of 1 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

(a) RoeFlux_3d

(b) Multi-Layer Perceptron (GPU)

Roe_3devaluationtimesfordifferentmodesandbatchsizes MLPevaluationtimesfordifferentmodesandbatchsizes

103

101

JAXreverse mode

[sw]u!

Evaluation time in [ms]

Graphax+AlphaGrad

102

Evaluation timei

100

101

100

JAxforwardmode

JAXreversemode

10-1

Graphax+AlphaGrad

10-1

f

工

16

64

256

512

1024

2048

2

4

8

16

32

64

Batchsize

ScaleFactor

5 Conclusion

In this work we successfully demonstrated that AlphaGrad discovers new AD algorithms that outperform the state-of-the-art. We demonstrated that these theoretical gains translate into measurable runtime improvements with Graphax, a Python-based interpreter we developed that leverages the AD algorithms discovered by AlphaGrad. However, AlphaGrad currently only optimizes for multiplications 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 a hardware model for efficient simulation. Another promising research avenue would be the implementation of a 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 circumvent 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 to Graphax. However, as 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 on multiple graphs at once. While the best results were still achieved with single-graph training, the multi-graph training still outperformed the other existing methods on most benchmarks. This hints at 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 a novel way to evaluate existing RL algorithms on a real-world problem that poses diverse challenges such as rewards across multiple scales and large action spaces. Thus, it could complement existing benchmarks such as OpenAI gym and MuJoCo [Todorov et al., 2012, Towers et al., 2023].

Acknowledgments

This work was sponsored by the Federal Ministry of Education, Germany (projects NEUROTEC-II grant no. 16ME0398K and 16ME0399 as well as GreenEdge-FuE, funding no. 16ME0521). The authors also gratefully acknowledge the Gauss Centre for Supercomputing e.V. (www.gauss-centre.eu) for funding this project by providing computing time on the GCS Supercomputer JUWELS at Jülich Supercomputing Centre (JSC). Furthermore, we thank Mark Schöne, Christian Pehle, Uwe Naumann and Matthew Johnson for the helpful discussions regarding the project.

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.

M. Averick, R. G. Carter, Guo-Liang Xue, and J. J. More. The minpack-2 test problem collection. Technical report, University of Illinois, June 1992. URL https://digital.library.unt.edu/ ark:/67531/metadc734562/ . Accessed May 11, 2024.

Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of Machine Learning Research , 18(153):1-43, 2018. URL http://jmlr.org/papers/v18/17-468.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):637-654, 1973. URL http://www.jstor.org/stable/1831029 .

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 http://github.com/google/jax .

Luca Capriotti and Michael B. Giles. Algorithmic differentiation: Adjoint greeks made easy. Available at SSRN , 2011. doi: 10.2139/ssrn.1801522. URL https://ssrn.com/abstract=1801522 . 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-3-642-30023-3_20.

Ivo Danihelka, Arthur Guez, Julian Schrittwieser, and David Silver. Policy improvement by planning with gumbel. In International Conference on Learning Representations , 2022. URL https: //openreview.net/forum?id=bERaNdoegnO .

DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio 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´ c, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http://github.com/google-deepmind .

Serap Dikmenli. Forward & inverse kinematics solution of 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 16x16 words: Transformers for image recognition at scale. CoRR , abs/2010.11929, 2020. URL https://arxiv.org/abs/2010.11929 .

Alhussein Fawzi, Matej Balog, Andrew Huang, et al. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature , 610:47-53, 2022. doi: 10.1038/s41586-022-05172-4. URL https://doi.org/10.1038/s41586-022-05172-4 .

Shaun A. Forth, John D. Pryce, Mohamed Tadjouddine, and John K. Reid. Jacobian code generated by source transformation and vertex elimination 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.1145/1024074.1024076. URL https://doi.org/10.1145/1024074.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. CoRR , abs/1709.03799, 2017. URL http://arxiv.org/abs/1709.03799 .

Ian J. Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning . MIT Press, Cambridge, MA, USA, 2016. http://www.deeplearningbook.org .

Griewank and A. Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, Second Edition . Other Titles in Applied Mathematics. Society for Industrial and Applied Mathematics (SIAM, 3600 Market Street, Floor 6, Philadelphia, PA 19104), 2008. ISBN 9780898717761. URL https://books.google.de/books?id=xoiiLaRxcbEC .

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.semanticscholar.org/CorpusID:4674781 .

Shengyi Huang, Rousslan Fernand Julien Dossa, Antonin Raffin, Anssi Kanervisto, and Weixun Wang. The 37 implementation details of proximal policy optimization. In ICLR Blog Track , 2022. URL https://iclr-blog-track.github.io/2022/ 03/25/ppo-implementation-details/ . https://iclr-blog-track.github.io/2022/03/25/ppoimplementation-details/.

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 https://openreview.net/forum?id=r1lyTjAqYX .

Hiroharu Kato, Deniz Beker, Mihai Morariu, Takahiro Ando, Toru Matsuoka, Wadim Kehl, and Adrien Gaidon. Differentiable rendering: A survey. CoRR , abs/2006.12057, 2020. URL https: //arxiv.org/abs/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 P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR , abs/1412.6980, 2014. URL https://api.semanticscholar.org/CorpusID:6628106 .

Seppo Linnainmaa. Taylor expansion of the accumulated rounding error. BIT Numerical Mathematics , 16:146-160, 1976. URL https://api.semanticscholar.org/CorpusID:122357351 .

Daniel J. Mankowitz, Andrea Michi, Anton Zhernov, et al. Faster sorting algorithms discovered using deep reinforcement learning. Nature , 618:257-263, 2023. doi: 10.1038/s41586-023-06004-9. URL https://doi.org/10.1038/s41586-023-06004-9 .

Charles C. Margossian. A review of automatic differentiation and its efficient implementation. CoRR , abs/1811.05031, 2018. URL http://arxiv.org/abs/1811.05031 .

William S. Moses and Valentin Churavy. Instead of rewriting foreign code for machine learning, 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://proceedings.neurips.cc/paper/2020/file/ 9332c513ef44b682e9347822c2e457ac-Paper.pdf .

Uwe Naumann. Save - simulated annealing applied to the vertex elimination problem in computational graphs. Technical Report RR-3660, INRIA, 1999. URL https://hal.inria.fr/ inria-00073012 .

Uwe Naumann. Optimal jacobian accumulation is np-complete. Mathematical Programming , 112:427-441, 2008. doi: 10.1007/s10107-006-0042-z. URL https://doi.org/10.1007/ s10107-006-0042-z .

Uwe Naumann. Exact first- and second-order greeks by algorithmic differentiation. 01 2010.

Uwe Naumann. Optimization of 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 , abs/1905.02494, 2019. URL https://api.semanticscholar.org/CorpusID:146808144 .

Mai Jacob Peng and Christophe Dubach. Lagrad: Statically optimized differentiable programming in mlir. In Proceedings of the 32nd ACM SIGPLAN International Conference on Compiler Construction , CC 2023, page 228-238, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400700880. doi: 10.1145/3578360.3580259. URL https://doi.org/ 10.1145/3578360.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.1016/ 0021-9991(81)90128-5. URL https://www.sciencedirect.com/science/article/pii/ 0021999181901285 .

Savine and J. Andreasen. Modern Computational Finance: Scripting for Derivatives and xVA . Wiley, 2021. ISBN 9781119540786. URL https://books.google.de/books?id= mBZDEAAAQBAJ .

Jürgen Schmidhuber. Deep learning in neural networks: An overview. CoRR , abs/1404.7828, 2014. URL http://arxiv.org/abs/1404.7828 .

Patrick Schmidt, J. Born, D. Bommes, M. Campen, and Leif Kobbelt. Tinyad: Automatic differentiation in geometry processing made simple. Computer Graphics Forum , 41:113-124, 10 2022. doi: 10.1111/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 , abs/1911.08265, 2019. URL http://arxiv.org/abs/1911.08265 .

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.

David Silver, Thomas Hubert, Julian Schrittwieser, 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 a general reinforcement learning algorithm. ArXiv , abs/1712.01815, 2017. URL https://api.semanticscholar.org/ CorpusID:33081038 .

Tadjouddine, F. Bodman, J. D. Pryce, and S. A. Forth. Improving the performance of the vertex elimination 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 111-120, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg.

Mohamed Tadjouddine, Shaun A. Forth, John D. Pryce, and John K. Reid. Performance issues for vertex elimination methods in computing jacobians using automatic differentiation. In Peter M. A. Sloot, Alfons G. Hoekstra, 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-0.

Mohamed Tadjouddine, Shaun A. Forth, John D. Pryce, and John K. Reid. Performance issues for vertex elimination methods in computing jacobians using automatic differentiation. In Peter M. A. Sloot, Alfons G. Hoekstra, 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: 10.1109/IROS.2012.6386109.

Edan Toledo, Laurence Midgley, Donal Byrne, Callum Rhys Tilbury, Matthew Macfarlane, Cyprien Courtot, and Alexandre Laterre. Flashbax: Streamlining experience replay buffers for reinforcement learning with jax, 2023. URL https://github.com/instadeepai/flashbax/ .

Mark Towers, Jordan K. Terry, 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. URL https://zenodo.org/record/8127025 .

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31st 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://api.semanticscholar.org/CorpusID:225062396 .

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. In IEEE 30th International Conference on High Performance Computing, Data, and Analytics (HiPC) , page 11, 2023. doi: 10.1109/HiPC58850. 2023.00055.

Appendix

A Task Descriptions

This section describes in-depth the mathematical formulation of the tasks that were evaluated in table 1 and table 2. An implementation of the functions can be found in the accompanying software package. Note that if a set of variables ( x 1 , . . . , x n ) is referred to as { x i } , they are treated as a separate scalar inputs to the function while x treats them as a single vectorized input.

A.1 RoeFlux_1d

For the implementation of the RoeFlux_1d task we thoroughly followed [Roe, 1981]. The pressure p 1d , enthalpy H 1d and flux term F 1d of the one-dimensional Euler equations are defined through:

The Roe flux ϕ Roe ( u l0 , u l1 , u l2 , u l3 , u l4 , u r0 , u r1 , u r2 , u r3 , u r4 ) between two adjacent cells u l and u r is computed using the routine described below. First, we define some averaged quantities to simplify formulation:

Then we define some state variables for the left cell through

Similarly, we define the same variables for the right cell through

We define some differences between the state variables of the two cells though

We proceed with the introduction of some auxiliary variables for further computation so that

where q 2 = u 2 , a = √ a 2 and n = u lr0 a . We define l p = | u + a | , l = | u | and l n = | u | and proceed to writing:

Next, we compute the fluxes between the cells through

and then define F i = F l i + F r i and α = 1 2 a u lr0 . The flux differences are then given through

Then the output of the function is given through ϕ 0 , ϕ 1 , ϕ 2 with ϕ i = 1 2 ( F i -∆ F i ) .

A.2 RoeFlux_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 p 3d and enthalpy H 3d :

Note that now instead of a one-dimensional state variable u 1 , we have a three-dimensional state vector u and the role of u 2 is now taken over by u 4 . The flux in three dimensions is given through

where u 1 , u 2 , u 3 and v 1 , v 2 , v 3 are the three components that make up u and v respectively and p = ( p, 0 , 0) ⊺ . We then again define the finite differences

and furthermore set v l = u l u l0 and v r = u r u r0 We continue with defining some auxiliary variables

Then we define some state variables for the left cell through

Similarly, we define the same variables for the right cell through

Then we introduce

and set q 2 = | t | 2 and a 2 = ( γ -1)( h -1 2 q 2 ) and a = √ a 2 . Furthermore, we define the eigenvalues of the Roe flux problem as l p = t 1 + a , l = t 1 , l n = t 1 -a . Then coefficients c i are given through

where we defined k 1 = ∆ u 0 -c 3 and k 2 = ∆ u 1 -t 1 ∆ u 0 a . The definition of the fluxes of each cell is similar to the formulation in the one-dimensional case. The flux changes are given through

The results of ∆ F 1 , ∆ F 2 and ∆ F 3 are concatenated into the vector ∆ F . The output of the function is then ( ϕ 0 , ϕ, ϕ 4 ) as defined in the one-dimensional case.

A.3 RobotArm_6DOF

The RobotArm_6DOF task models the forward differential kinematics of a 6-degree-of-freedom (6-DOF) robot arm as is often found in robotics labs and industrial manufacturing sites. For the implementation, we followed [Dikmenli, 2022] and define c i = cos x i and s i = sin x i . Furthermore, we define the functions

Then we also define s ij = S ( t i , t j ) and c ij = C ( t i , t j ) for the input variables { t i } of the problem. Then we proceed with calculating the auxiliary intermediates

Then, we define the Tait-Bryan angles through:

Next, we calculate the positional parts of the kinematics by starting with the x -component:

We continue with the y -component:

Finally, the z -component is given through:

The function then returns the six values ( p x , p y , p z , z, ˆ y, ˆ z )

A.4 HumanHeartDipole

The HumanHeartDipole task is derived from the experimental electrolytic determination of the resultant dipole moment in the human heart. For the implementation, we followed [Averick et al., 1992] with { x i } = ( x 1 , . . . , x 8 ) such that

with σ mx , σ my , σ A , σ B , σ C , σ D , σ E , σ F being some arbitrary measured constants.

A.5 PropaneCombustion

The PropaneCombustion task arises in the determination of the chemical equilibrium of the combustion of propane in air. Each unknown is related to a product concentration given in mols formed during the combustion process. We implemented the PropaneCombustion task as defined in [Averick et al., 1992] where with { x i } := ( x 1 , . . . , x 11 ) :

Here, K 5 , . . . , K 10 are measured constants and R = 10 is the relative amount of air and fuel, while p = 40 is the pressure in atmospheres.

A.6 Black-Scholes Equation

The BlackScholes_Jacobian 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

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. V is the price of the option at time t . For certain conditions described in [Black and Scholes, 1973], this equation can be solved analytically. We define

which is the cumulative distribution function of the Gaussian distribution (also known as the error function). Next, we define

where K is the payoff. Then, the solution to the Black-Scholes equation is given through

where F ( r, T, S ) = e rT S . We are interested in the second-order derivatives with respect to S, K, σ, r , so we first calculate the Jacobian of equation (A.6) using reverse-mode AD with Graphax using graphax.jacve(f, 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 a 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 a size of 4. The activation functions are tanh -functions and the output is processed using a 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 to create a more realistic example of a MLP.

A.8 Transformer Encoder

The TransformerEncoder task is inspired by recent advances in natural language processing and image classification. The network consists of two attention blocks with a single head only. The block itself consists of a softmax attention layer with residual connections followed by a layer norm and a MLP with a single hidden layer of size 4 and sigmoid linear unit activations. We stack two of these layers and process the output with a 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 a factor of 16 to create a realistic example of a transformer encoder.

A.9 Random Functions f and g

The random functions f and g are randomly generated using a custom JAX interpreter which consists of a 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 of the accompanying AlphaGrad software package.

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-mode AD and reverse-mode AD. Unless stated otherwise, all runtimes were measured on an AMD EPYC 9684X 2x96-Core processor. On the RoeFlux_1d , RobotArm_6DOF , random function g and BlackScholes_Jacobian tasks, our work significantly outperforms both JAX AD modes for all batch sizes, sometimes by almost an order of magnitude. For the HumanHeartDipole and PropaneCombustion tasks, AlphaGrad and Graphax outperform both JAX modes significantly but we can observe a crossover from batch size 1024 to 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.e. the MLP and TransformerEncoder tasks, were analyzed slightly 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 as 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 of 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.

(a) RoeFlux_1d

(b) RobotArm_6DOF

(c) HumanHeartDipole

(d) PropaneCombustion

(e) Random function

g

(f) BlackScholes_Jacobian

Roe_1devaluationtimesfordifferentmodesandbatchsizesRobotArmevaluationtimesfordifferentmodesandbatchsizes

103

Evaluation time in [ms]

102

Evaluation time in [ms]

102

101

101

10°

100

JAX forward mode

JAxforwardmode

10-1

JAX reverse mode

10-1

JAX reverse mode

十

Graphax + AlphaGrad

Graphax + AlphaGrad

16

64

256

512

1024

2048

16

64

256

512

1024

2048

Batchsize

Batchsize

HHDevaluationtimesfordifferentmodesandbatchsizesPCevaluationtimesfordifferentmodesandbatchsizes

103

[s] u aw uoeniens

102

Evaluation time in [ms]

102

101

101

100

JAX forward mode

100

JAXforward mode

JAX reverse mode

JAX reverse mode

10-1

Graphax + AlphaGrad

Graphax + AlphaGrad

10-1

16

64

256

512

1024

2048

16

64

256

512

1024

2048

Batchsize

Batchsize

gevaluationtimesfordifferentmodesandbatchsizesBSjevaluationtimesfordifferentmodesandbatchsizes

104

Evaluation time in [ms]

Evaluation time in [ms]

102

101

102

101

100

JAX forward mode

JAXforward mode

100

JAX reverse mode

10-1

JAX reverse mode

Graphax + AlphaGrad

十

Graphax + AlphaGrad

16

64

256

512

1024

2048

16

64

256

512

1024

2048

Batchsize

Batchsize

For the TransformerEncoder task, we were only able to evaluate the runtime for a batch size of 1 since the VertexGame implementation of AlphaGrad supports only inputs with a 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.

(a) RoeFlux_3d

(b) Random function

f

(c) Multi-Layer Perceptron

(d) Transformer Encoder

(e) Multi-Layer Perceptron (GPU)

(f) Transformer Encoder (GPU)

Roe_3d evaluationtimesfordifferentmodesandbatchsizesfevaluationtimesfordifferentmodesandbatchsizes

103

[w] ui a onienn

[s] u aw uonins

102

101

100

JAX forward mode

100

JAXforward mode

十

JAX reverse mode

JAX reverse mode

十

Graphax+AlphaGrad

十

Graphax + AlphaGrad

10-1

16

64

256

512

1024

2048

16

64

256

512

1024

2048

Batchsize

Batchsize

MLPevaluationtimesfordifferentmodes andbatchsizesTransformerevaluationtimesfordifferentmodes andbatchsizes

103

1

JAX reverse mode

JAXreversemode

Evaluation time in [ms]

102

Graphax+AlphaGrad

Evaluation time in [ms]

102

Graphax+AlphaGrad

101

101

100

100

10-1

2

4

8

16

32

64

2

4

8

16

32

64

ScaleFactor

Scale Factor

MLPevaluationtimesfor different modes and batch sizesTransformer evaluationtimesfor different modesand batch sizes

101

JAXreversemode

JAX reverse mode

Evaluation time in [ms]

Graphax+ AlphaGrad

Evaluation time in [ms]

102

十

Graphax+ AlphaGrad

101

100

10

10-1

10-1

2

4

8

16

32

64

4

8

16

32

64

Scale Factor

Scale Factor

C Sparsity Types

The sparsity of the Jacobians associated with the edges in the computational graph representation is described with a number ranging from -10 to 10. Table 3 contains an overview of the sparsity types and an example tensor that is represented by the corresponding number. Our approach is directed only at diagonal sparsity, meaning that we mainly consider tensors that can be decomposed into products of lower-dimensional tensors and Kronecker symbols δ ij . 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 -1 to represent these operations. Furthermore, we make a distinction between multiplication with the unit tensor δ il δ jk and a constant multiple of the unit tensor cδ il δ jk since the latter incurs actual multiplications while the former can just be treated as a renaming of indices as demonstrated by the

following examples with tensors S klmn = δ km δ ln and U klmn = cδ km δ ln :

In the first equation, we can just rename the indices of the tensor, but in the second equation we have to perform the product cT ij , which incurs | i | · | j | 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 is 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 S ijkl = a ij δ ik δ jl with sizes (2 , 3 , 2 , 3) with sparsity type 6 and T ijkl = b i δ il δ jk with shape (2 , 3 , 2 , 3) and sparsity type 9. Thus, in the computational graph representation, we would have two non-zero entries (6 , 2 , 3 , 2 , 3) and (9 , 2 , 3 , 2 , 3) . Then their contraction is given through

Then we form the product | i | · | j | · | k | · | l | · | m | · | n | = 2 · 3 · 2 · 3 · 2 · 3 where | · | gives the size of the dimension associated with the index. Then we use the sparsity types 7 and 9 to mask out the relevant values. In this case, we mask out | k | , | l | and | m | such that we arrive at | i | · | j | · | n | = 2 · 3 · 3 = 18 multiplications which is exactly the number of multiplications incurred by the dense multiplication of a ij b n .

D RL Algorithm and Network Architecture

We solved VertexGame with PPO [Schulman et al., 2017] and AlphaZero [Schrittwieser et 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 a 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, softmax-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, we first sum all outputs to 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, we 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 an implementation error.

AlphaZero Implementation Details The implementation is loosely based on the implementation of Gumbel MuZero and uses the mctx package provided by Google DeepMind [DeepMind et al., 2020]. However, we modified the implementation by replacing the trainable model and reward functions with our 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 a cosine learning rate annealing over 5000 episodes. For Gumbel MuZero to work properly, it is necessary to rescale the rewards so that they lie in the interval [0,1). In the Gumbel MuZero paper, the authors designed a 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, we set the corresponding parameters to c visit = 25 and c scale = 0 . 01 for all cases. We trained the agent using adaptive momentum gradient-based learning [Kingma and Ba, 2014] with an initial learning rate of 10 -3 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 1 and 2. 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.jacve(f, order=order, argnums=argnums)(*xs) . The orders given in tables 4 and 5 can be directly used to compute the Jacobians.

In addition to the single-graph experiments where the agent was trained only on a single task, we 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 was 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 6.

(a) RoeFlux_1d

(b) RobotArm_6DOF

(c) HumanHeartDipole

(d) PropaneCombustion

(e) Random function

g

(f) BlackScholes_Jacobian

AD algorithm optimizationRoeFlux_1d

ADalgorithmoptimizationRobotArm_6DoF

multiplications)

350

364

250

400

288

300-

450

#

500

return (#

350

return

550

Ourwork

400

Ourwork

SOTA for RoeFlux_1d

0

i

2

3

4

0

2

4

6

#samples

1e6

#samples

1e6

ADalgorithmoptimization HumanHeartDipole

ADalgorithmoptimizationPropaneCombustion

87

150

90

90

160

92

95

170

172.

L6

180

100

102

190

Ourwork

Ourwork

SOTAforHumanHeartDipole

105

SOTAforPropaneCombustion

0

1

2

3

4

0.0

0.5

1.0

1.5

2.0

2.5

#samples

1e6

#samples

1e6

ADalgorithmoptimizationBlackScholes

ADalgorithmoptimizationRandomFunctiong

300

440-

420

350

350

451

460-

400

480

500

450

520

500

540

Ourwork

Our work

SOTA for Random Function g

SOTA for BlackScholes

560

550

0

i

2

3

4

0

2

4

6

8

#samples

1e6

#samples

1e6

(a) RoeFlux_3d

(b) Random function

f

(c) Multi-Layer Perceptron

(d) Transformer Encoder

ADalgorithmoptimizationRoeFlux_3d

750

ADalgorithmoptimizationRandomFunctionf

938

multiplications)

1000

7500

multiplications)

9333

1250-

10000

12500

1500

15000

#

1750

#

return

return(

17500

2000-

2250

SOTAforRoeFlux_3d

Ourwork

20000

Our work

0

1

2

3

4

0.0

0.2

0.4

0.6

0.8

1.0

#samples

1e6

#samples

1e7

ADalgorithmoptimizationMLP

500392---

ADalgorithmoptimizationTransformerEncoder

multiplications)

4688

1000

multiplications)

10000

20000

1500

30000

return (#r

40000

2000

#

return

50000

Our work

2500

SOTA for MLP

60000

SOTAforTransformer Encoder

Our work

0

1

2

3

4

5

0.0

0.2

0.4

0.6

80

1.0

1.2

#samples

1e6

#samples

1e7

(a) RoeFlux_1d

(b) RobotArm_6DOF

(c) HumanHeartDipole

(d) PropaneCombustion

(e) Random function

g

(f) BlackScholes_Jacobian

ADalgorithmoptimizationRoeFlux_1d

ADalgorithmoptimizationRobotArm_6DOF

250-

400

364

300

288

500

350

600

400

700

450

800

Joint optimization

Joint optimization

SOTA for RoeFlux_1d

SOTA for RobotArm_6DOF

006

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

#samples

1e7

#samples

1e7

ADalgorithmoptimizationHumanHeartDipole

ADalgorithmoptimizationPropaneCombustion

160 -

90

90

180-

172

100

190

110

200

210-

120

Joint optimization

Joint optimization

220

SOTAforHumanHeartDipole

SOTAforPropaneCombustion

130

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

#samples

1e7

#samples

1e7

ADalgorithmoptimizationg

ADalgorithmoptimizationBlackScholes_Jacobian

451

350

400

500

500

550

600

Joint optimization

Jointoptimization

600

SOTA for g

700

----

SOTAforBlackScholesJacobian

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

#samples

1e7

#samples

1e7

(suoesdinu #) uin

(a) RoeFlux_3d

(b) Random function

f

(c) Multi-Layer Perceptron

(d) Transformer Encoder

ADalgorithmoptimizationf

ADalgorithmoptimizationRoeFlux_3d

1000938

100009333

1500

20000-

Tl

2000

30000

return(#

2500

40000

Joint optimization

Joint optimization

3000

SOTA for RoeFlux_3d

50000

SOTA for f

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

#samples

1e7

#samples

1e7

ADalgorithmoptimizationPerceptron

ADalgorithmoptimizationTransformer

0

392

4688

multiplications)

2000

50000

4000

100000

6000

#150000

return

8000-

Joint optimization

200000

Jointoptimization

SOTAforPerceptron

SOTAforTransformer

0.0

0.5

1.0

1.5

2.0

0.0

0.5

1.0

1.5

2.0

#samples

1e7

#samples

1e7

F PPO Results

We ran all experiments with the configuration described in D 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 the AlphaZero agent on all tasks, sometimes by a significant margin as for example in the RoeFlux tasks and random function f or the MLP and TransformerEncoder where it does not find a better elimination order at all. This is to be 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_Jacobian or random function f tasks. Thus, the PPO-based agent might still be a viable choice because it is trained within minutes on a single NVIDIA RTX 4090 GPU, even for large tasks such as the random function f and still find well-performing elimination orders. Table 7 shows the number of multiplications required by the best elimination order found by the PPO-agent. Figures 11 and 12 contain the corresponding reward curves. We did not succeed in training a joint model using the PPO agent.

return(#

(a) RoeFlux_1d

(b) RobotArm_6DOF

(c) HumanHeartDipole

(d) PropaneCombustion

(e) Random function

g

(f) BlackScholes_Jacobian

ADalgorithmoptimizationRoeFlux_1d

ADalgorithmoptimizationRobotArm_6DOF

240

350 364.

multiplications)

260

400-

280

450

288

300

return(#

500

return(#

550

320

Ourwork

Ourwork

600

SOTA for RoeFlux_1d

340

0.0

0.5

1.0

1.5

2.0

2.5

3.0

0.0

0.5

1.0

1.5

2.0

2.5

# samples

1e6

# samples

1e6

ADalgorithmoptimizationHumanHeartDipole

ADalgorithmoptimizationPropaneCombustion

87

150

90.

multiplications)

06

160

92

95

170

172

97

180

return(#

100

102

190

Ourwork

Ourwork

SOTAforHumanHeartDipole

105

SOTAforPropaneCombustion

0

i

2

3

4

0.0

0.5

1.0

1.5

2.0

2.5

#samples

1e6

#samples

1e6

AD algorithm optimizationRandomFunctiong

ADalgorithmoptimizationBlackScholes

420

300

440

multiplications)

350

350

451

460

400

480-

return(#

500

#

450

return

520

Our work

500

Our work

540

SOTAforRandomFunctiong

SOTA for BlackScholes

0

1

2

3

4

0.0

0.5

1.0

1.5

2.0

2.5

3.0

#samples

1e6

#samples

1e6

(a) RoeFlux_3d

(b) Random function

f

(c) Multi-Layer Perceptron

(d) Transformer Encoder

AD algorithmoptimizationRoeFlux_3d

ADalgorithmoptimizationRandomFunctionf

6000

938

multiplications)

1000

8000

1200-

10000

1400

12000

return(#

1600

#

return

14000

1800-

Ourwork

16000

Ourwork

SOTAforRoeFlux_3d

SOTAforRandomFunctionf

2000-

0

1

2

3

4

0

1

2

3

#samples

1e6

#samples

1e6

ADalgorithm optimization MLP

ADalgorithmoptimizationTransformerEncoder

500 392--

4688

10000

multiplications)

20000

1000

30000

1500

return(#

40000

50000

2000

Our work

Our work

SOTA for MLP

60000

SOTAforTransformer Encoder

0.0

0.5

1.0

1.5

0

i

2

3

4

#samples

1e6

#samples

1e6

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. In 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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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]

Justification: The paper clearly describes the limitations of the approach by including negative outcomes in the the experimental results table and provides a more in-depth analysis in the discussion and appendix sections of the paper.

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 a 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 well-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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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

Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The entire codebase as well as a comprehensive tutorial will be provided with the paper as part of a 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:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a 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 is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The entire codebase along with a comprehensive introduction will be published on release of the work. Furthermore, the reviewers will be provided with full access to the entire codebase on 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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) is 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 (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All necessary experimental parameters like learning rates, RL configuration etc. are 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 to 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 are 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-percentiles to quantify the uncertainty.

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, confidence intervals, or statistical significance tests, at 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.

It is 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 or 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 as well as 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 it 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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes] .

Justification: The authors have read the NeurIPS ethics code and are sure that the research conducted in this paper as well as the results of the paper did not violate any of the points raised in the guidelines. The main result of the paper is a way to find faster and more energy-efficient 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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

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 a way to find faster and more energy-efficient automatic differentiation algorithms which most likely has a positive impact on the environment 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 or 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA] .

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or 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]

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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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 are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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 as 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 to 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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

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 crowdsourcing nor research with human subjects.

Depending on 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nTraining 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nAdmittedly, 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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)\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nBy 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.\n\nQUESTIONS:\nAs 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- Rephrasing the elimination order in AD as a deep reinforcement learning problem sounds somehow novel.\n- The paper is well-organized.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- As stated by the authors, the generalization performance of the approach may still be unclear, which may make the approach rather computationally prohibitive.\n\nQUESTIONS:\n- 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\"\n\nLIMITATIONS:\nLimitations are discussed such as the poor generalization when performing joint training and the difficulties to handle the range of rewards.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors are forthcoming about the main limitations and highlight important areas for future work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}