Optimizing Automatic Differentiation with Deep Reinforcement Learning

Authors: jamie  lohoff
emre  neftci
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:19:21

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.

Main paper: 
 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, Figure 1 : Summary of the AlphaGrad pipeline. We trained a neural network to produce new Automatic Differentiation (AD) algorithms using Deep RL that can be used in JAX. The resulting algorithms significantly outperform the current state of the art.2008]. Inspired by recent advances in finding optimal matrix-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 [Naumann, , 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.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. 
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 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.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 ∂vi . 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 updatec ki += c kj c ji ∀(i, k) ∈ V × V where i ≺ j and j ≺ k (1)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.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.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.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 The output vertex represents the application of the activation function y i = tanh a i with the partial derivativeEQUATIONEQUATIONAccording 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):EQUATIONEQUATIONIn 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.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 ∂yij ∂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.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. 
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.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. 
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.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).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 [Forth et al., 2004b] and [Tadjouddine et al., 2006] present state-of-the-art results for some of the investigated tasks, we were unable to reproduce the experiments given their implementation details.  
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] .Next, we compute the fluxes between the cells throughF l0 , F l1 , F l2 = F 1d (u l0 , u l1 , u l2 , p l ) F r0 , F r1 , F r2 = F 1d (u r0 , u r1 , u r2 , p r )and then define F i = F li + F ri and α = 1 2a u lr0 . The flux differences are then given through∆F 0 = c 0 + αc 1 -αc 2 , ∆F 1 = c 0 u + αc 1 (u + a) -αc 2 (u -a), ∆F 2 = 1 2 c 0 q 2 + αc 1 (h + ua) -αc 2 (h -ua).Then the output of the function is given through ϕ 0 , ϕ 1 , ϕ 2 with ϕ i = 1 2 (F i -∆F i ).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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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. 
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. 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).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. 
 
 Appendix: 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 :p 3d (u 0 , u, u 4 ) = (γ -1) u 4 - |u| 2 2u 0 , H 3d (u 0 , u 4 , p) = u 4 + p u 0 .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 throughF 3d (u 0 , u, u 4 , v, p) = (u 1 , p + uv, v 1 (p + u 4 ))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∆u 0 =u l0 -u r0 , ∆u =u l -u r , ∆u 4 =u l4 -u r4and furthermore set v l = u l u l0 and v r = ur ur0 We continue with defining some auxiliary variablesw 1 = √ u l0 - √ u r0 . t = √ u l0 v l + √ u r0 v l w 1 .Then we define some state variables for the left cell throughp l = p 3d (u l0 , u l , u l4 ), h l = H 3d (u l0 , u l4 , p l ).Similarly, we define the same variables for the right cell throughp r = p 3d (u r0 , u r , u r2 ), h r = H 3d (u r0 , u r4 , p r ).Then we introduceh = √ u l0 h l + √ u r0 h r w 1 ,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 throughc 0 = k 1 -k 2 2 l m , c 1 = l ∆u 2 t 2 -∆u 0 , c 2 = l ∆u 3 t 3 -∆u 1 , c 3 = l γ -1 a 2 ((h -q 2 )∆u 0 + t • ∆u -∆u 4 ) , c 4 = k 1 + k 2 2 l p .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∆F 0 = c 0 + c 3 + c 4 l p , ∆F 1 = c 0 l n + c 3 t 1 + c 4 l p , ∆F 2 = c 0 t 2 + c 1 t 2 + c 2 t 2 + c 3 t 2 + c 4 t 2 , ∆F 3 = c 0 t 3 + c 2 t 3 + c 3 t 3 + c 4 t 3 , ∆F 4 = c 0 (h -t 1 a) + c 1 t 2 2 + c 2 t 2 3 + c 3 q 2 2 + c 4 (h + t 1 a).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.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 S(u, v) = cos(u) sin(v) + sin(u) cos(v), C(u, v) = cos(u) cos(v) -sin(u) sin(v).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 intermediatesa y = s 5 (c 1 c 23 c 4 + s 1 s 4 ) + c 1 s 23 c 5 , a y = s 5 (s 1 c 23 c 4 -c 1 s 4 ) + s 1 s 23 c 5 , a z = s 23 c 4 s 5 -c 23 c 5 , n z = c 6 (c 23 s 5 + s 23 c 4 c 5 ) -s 23 s 4 s 6 , o z = -s 6 (c 23 s 5 + s 23 c 4 c 5 ) -s 23 s 4 c 6 .Then, we define the Tait-Bryan angles through:z = arctan a y a x ŷ = arctan 1 -a 2 z a z ẑ = arctan - o zn z Next, we calculate the positional parts of the kinematics by starting with the x-component:x 1 = 185(s 5 (c 1 c 23 c 4 + s 1 s 4 ) + c 1 s 23 c 5 ), x 2 = c 1 (175 + 890c 2 + 50c 23 + 1035s 23 ),p x = x 1 + x 2We continue with the y-component:y 1 = 185(s 5 (c 1 c 23 c 4 + c 1 s 4 ) + s 1 s 23 c 5 ), y 2 = s 1 (175 + 890c 2 + 50c 23 + 1035s 23 ), p y = y 1 + y 2Finally, the z-component is given through:p z = 575 + 890s 2 + 50s 23 -1035c 23 + 185(s 23 c 4 s 5 -c 23 c 5 )The function then returns the six values (p x , p y , p z , z, ŷ, ẑ)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 thatf 1 ({x i }) = x 1 + x 2 -σ mx f 2 ({x i }) = x 3 + x 4 -σ my f 3 ({x i }) = x 5 x 1 + x 6 x 2 -x 7 x 3 -x 8 x 4 -σ A f 4 ({x i }) = x 7 x 1 + x 8 x 2 + x 5 x 3 + x 6 x 4 -σ B f 5 ({x i }) = x 1 (x 2 5 -x 2 7 ) -2x 1 x 5 x 7 + x 2 (x 2 6 -x 2 8 ) -2x 4 x 6 x 8 -σ C f 6 ({x i }) = x 3 (x 2 5 -x 2 7 ) + 2x 1 x 5 x 7 + x 4 (x 2 6 -x 2 8 ) + 2x 2 x 6 x 8 -σ D f 7 ({x i }) = x 1 x 5 (x 2 5 -3x 2 7 ) + x 3 x 7 (x 2 7 -3x 2 5 ) + x 2 x 6 (x 2 6 -3x 2 8 ) + x 4 x 8 (x 2 8 -3x 2 6 ) -σ E f 8 ({x i }) = x 3 x 5 (x 2 5 -3x 2 7 ) -x 1 x 7 (x 2 7 -3x 2 5 ) + x 4 x 6 (x 2 6 -3x 2 8 ) -x 2 x 8 (x 2 8 -3x 2 6 ) -σ F with σ mx , σ my , σ A , σ B ,σ C , σ D , σ E , σ Fbeing some arbitrary measured constants.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 ):f 1 ({x i }) = x 1 + x 4 -3 f 2 ({x i }) = 2x 1 + x 2 + x 4 + x 7 + x 8 + x 9 + 2x 10 -R f 3 ({x i }) = 2x 2 + 2x 5 + x 6 + x 7 -8 f 4 ({x i }) = 2x 3 + x 9 -4R f 5 ({x i }) = K 5 √ x 2 x 4 + x 1 x 5 f 6 ({x i }) = K 6 √ x 1 x 2 - √ x 4 x 7 p x 11 f 7 ({x i }) = K 7 √ x 1 x 2 - √ x 4 x 7 p x 11 f 8 ({x i }) = K 8 x 1 -x 4 x 8 p x 11 f 9 ({x i }) = K 9 x 1 √ x 3 -x 4 x 9 p x 11 f 10 ({x i }) = K 10 x 2 1 -x 2 4 x 10 p x 11 f 11 ({x i }) = x 11 -x 10 -x 9 -x 8 -x 7 -x 6 -x 5 -x 4 -x 3 -x 2 -x 1Here, 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.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∂V ∂t + 1 2 σ 2 S 2 ∂ 2 V ∂S 2 = rV -rS ∂V ∂S .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ϕ(x) = 1 + 1 √ π x -∞ e -1 2 z 2 dz (6)which is the cumulative distribution function of the Gaussian distribution (also known as the error function). Next, we defined 1 = 1 σ √ T log F K + σ 2 2T d 2 = d 1 -σ √ T Φ(F, K, r, σ, T ) = e -rT (F ϕ(d 1 ) -Kϕ(d 2 ))where K is the payoff. Then, the solution to the Black-Scholes equation is given through f (S, K, r, σ, T ) = Φ(F (r, T, S), K, r, σ, T )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.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.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 gThe 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.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. 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. The MLP network sizes were scaled up with growing batch size. The transformer could only be evaluated on with a batch size of 1 due to limitations of VertexGame but the network size was scaled up as well. Both deep learning tasks were also evaluated on GPUs. Error bars are the 2.5-and 97.5-percentiles of the runtimes.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 T j δ il δ jk -8T j δ ik δ jl -7 δ il δ jk -6 δ ik δ jl -5 T il δ jk -4 T jk δ il -3 T ik δ jl -2 T jl δ ik -1 copy gradient operation 0 no edge 1 T ijkl 2 T ijl δ ik 3 T ijk δ jl 4 T jk δ il 5 T il δ jk 6 T ij δ ik δ jl 7 T ij δ il δ jk 8 T i δ ik δ jl 9 T i δ il δ jk 10 cδ ik δ jlfollowing examples with tensors S klmn = δ km δ ln and U klmn = cδ km δ ln :klEQUATIONklEQUATIONIn 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 throughEQUATIONThen we form the product |i| • |j| • |k| • |l| • |m| • |n| = 2 • 3 • 2 • 3 • 2 • 3 where | • |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.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.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. 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. 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.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


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

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

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

Reviews: review: paper_summary: paper_summary: This paper presents a novel method called AlphaGrad for optimizing computational graphs derived through automatic differentiation (AD) algorithms using deep reinforcement learning. The authors formulate the optimization of a computation graph as a single-player game where an AlphaZero agent aims to minimize the number of required multiplications. The authors demonstrate that AlphaGrad is able to improve over preexisting forward-mode AD methods in domains like computational fluid dynamics, robotics, optimization, and computational finance. Notably Deep Learning based tasks which leverage reverse-mode AD don't show significant improvements with AlphaGrad. Finally, the authors introduce Graphax, an AD interpreter for Jax that enables the optimization and execution of computational graphs discovered from elimination orders such as those proposed by AlphaGrad.
strength_weakness: strength_weakness: - 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.
- 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.
- 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.
- 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.

[1] Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients. William Moses and Valentin Churavy. Neural Information Processing Systems (NeurIPS), 2020.

[2] LAGrad: Statically Optimized Differentiable Programming in MLIR. Mai Jacob Peng and Christophe Dubach.International Conference on Compiler Construction (CC), 2023.
questions: questions: - How challenging would it be to directly optimize for metrics like execution time or memory usage rather than the number of multiplications? This seems like an important direction for making the approach more practical.
- Could you provide more details on how the random functions f and g were generated in the text?
- 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?
- Why was Jax chosen instead of targeting more powerful IRs like MLIR?
limitations: limitations: The authors are forthcoming about the main limitations and highlight important areas for future work.
paper_summary: paper_summary: The paper proposes to learn an efficient computation of the Jacobian of a program by reinforcement learning.
The 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. 
The authors developed a library enabling a direct use of the approach built on freely available package JAX.
Numerous comparisons are also available in the supplementary material.
strength_weakness: strength_weakness: - As stated by the authors, the generalization performance of the approach may still be unclear, which may make the approach rather computationally prohibitive.
questions: questions: - Are the performance displayed in Table 2 given after also some jit compilation?
- 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.)

Detail:
line 288: typo "two functions random functions"
limitations: limitations: Limitations are discussed such as the poor generalization when performing joint training and the difficulties to handle the range of rewards.
paper_summary: paper_summary: In this paper, the authors propose a novel method to optimize the number of necessary multiplications for Jacobian computation by leveraging deep reinforcement learning (RL) and a concept called cross-country elimination while still computing the exact Jacobian. The author present the search for the optimal elimination order, aimed at minimizing the required number of multiplications, as a single-player game managed by an RL agent. The results show that this approach yields up to a 33% improvement over state-of-the-art methods across various relevant tasks from different domains.
strength_weakness: strength_weakness: - 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%).
- 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.
questions: questions: - How does the VertexGame ensures that the reduced computational graph gives correct derivatives? Is there a grad check procedure?
- 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)? 
- 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?
- What level does the propose method works on? e.g., the source code or intermediate representation?
- Does the proposed method more close to Source Code Transformation or Tracing/Operator Overloading in a context of AD?
limitations: limitations: yes.
paper_summary: paper_summary: This paper considers the (combinatorially hard) optimization problem related to automatic differentiation algorithms represented via an "elimination order" based on prior work and proposes an RL formulation to solve it approximately. This follows in similar vein to recent celebrated results on matrix multiplication and sorting. The authors demonstrate that their approach is able to improve, conceptually on the first order metric of the number of multiplications as defined in the optimization search, and more practically that this can be sufficiently meaningful for runtime as they measure it on representative benchmarks , compared to the default forward/reverse mode implementations. The paper is well written and considers a reasonable variety of computational graphs in the evals. The main overhead in adopting the proposal in practice is primarily in the form of a specialized policy search for each computational graph, which factors in roughly speaking, at the compilation stage and is not reflected explicitly in the benchmarks.
strength_weakness: strength_weakness: - 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).
 
- 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.
questions: questions: - When considering the number of multiplications, given the definition of a vertex, I assume this refers to the number of matmuls, but the costs of these vary by the dimensions of the jacobians. Do you take this into account anywhere? If not, what is the hurdle? Making the exposition a bit more explicit on the definition/implementation of the  reward is framed might be helpful. 

Other comments:
- The labels for cross edges in Figure 2(d) appear to be swapped.
- Caption for Figure 2 has intermediate variable definition of $v_2=cos(v_1)$, which needs to be $sin(v_1)$ instead?
- L128 graph 2a --> Fig 2a
- In Definition 1 L138: should this be $c_{ij}=c_{jk}=0$ rather than $c_{ij}=c_{ik}=0$?
- L142, $\phi_j \circ \phi_k$ needs to be $\phi_k \circ \phi_j$ instead.
- Equation (2) needs to be $W_{ik}$ rather than $W_{ki}$?
- L324 refers to Table 2 being of theoretical value -- should this be Table 1?
- L341 s/note/not
- Figure 1 3rd box: s/executation/execution
limitations: limitations: Yes
paper_summary: paper_summary: The authors study the problem of computing the Jacobian of arbitrary computation graphs. The typical approach in most autodiff libraries is to implement the classic backpropagation algorithm, aggregating gradients from the end to the beginning. This is because such approaches are well suited for cases where the number of inputs is significantly larger than the number of outputs, such as neural nets.

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

The 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)
strength_weakness: strength_weakness: By my understanding, both the AlphaZero agent and PPO agent need to be instantiated separately for each class of functions. This is based off the note in Appendix F, that the authors did not succeed in training a joint model for the PPO agent. One of the papers cited as related work is the one for improving matrix multiplication with search algorithms, and one benefit of that work was that although the search was very specific to 4x4 matrices for example, it was possible to us the discovered results of that work elsewhere.

In 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)

Still, I overall think there are interesting signs of life here, even if the method itself is not entirely practical yet.
questions: questions: As part of the graph representation coding, there is a list of 21 different sparsity patterns to describe the intermediate edges of the computation graph. Is this list exhaustive to more general computation graphs, or are there other forms the partial derivative can take if the space of computation graphs is increased?
limitations: limitations: Yes
paper_summary: paper_summary: This paper studies automatic differentiation in a computational graph. The topic is classic with wide applications in scientific research, e.g., computing gradients of a neural network or Jacobians needed in numerical optimization. Classic numerical methods like standard forward- and reverse-mode differentiation (backpropagation) exist, and the novelty and contribution in this work lies in its proposal of training a reinforcement-learning (RL) agent to explore an optimized order in graph vertex elimination. The paper demonstrates the efficacy of its method by comparing the time cost of the found auto differentiation policy (measured by multiplication count and wall-clock time) with that of classic and heuristic methods.
strength_weakness: strength_weakness: I don’t have major concerns with the paper, but I want to mention that there are quite a few typos in the equations in the main paper. For example, I suspect “i > j” should be replaced with “i < j”, and “sum_i” in line 192 should be “sum_j”. Please double-check all equations in the paper and make sure there are no typos.
questions: questions: I don’t have any questions for now. The key idea described in this paper is quite straightforward, and the writing in the main paper is generally good.
limitations: limitations: Admittedly, the proposed RL method does not beat classic methods across the board in the experiments. I share the similar conjecture expressed in the paper: backpropagation is a quite strong baseline for large-input, small-output computational graphs, and it can be quite difficult to beat. This may also indicate that any small improvement on this strong baseline can have a significant impact, as backpropagation is so widely used in practice.

On 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.

Comments: title: Reply
Response
Reply
Response
Response to the reply
Response
Reviewer response
Acknowledging rebuttal
Reply
comment: I thank the authors for answering my questions and agreeing to update the manuscript to discuss future work on optimizing over other IRs. I want to update my score to reflect this, I think this work is important and a good first step towards applying RL for optimizing AD.
Thank you for the clarification. I will keep my score as-is.
Thank you for the clarification, this was very helpful.
We apologize for misunderstanding the question and try to answer the questions to the best of our understanding: 

In the following everything is described for a single example, say RoeFlux_1d. We create an initial computational graph representation in form of the adjacency matrix. This adjacency matrix serves as an initial state for the agent similar to how for AlphaZero the initial position of the chess or Go board is fixed. In this sense, VertexGame starts from a fixed initial input for every example task, i.e every row.

However, every time a vertex is eliminated, new edges are created and others are deleted and the current state of the graph and as such the adjacency matrix changes. Thus the computational graph representation evolves with every action and this evolving computational graph representation (the adjacency matrix) is fed into the agent's neural network model at every step to decide the next action. The selection of the next action is stochastic but becomes more deterministic over time since the policy, i.e. the probability distribution over actions becomes more deterministic. Thus we sample the space of possible computational graphs that can arise from the initial graph through vertex elimination with a focus on exploration early on.
Throughout the training on a single task and its initial computational graph, the agent is thus confronted with a distribution over different graph configurations and learns the optimal action to take, similar to how in chess for different board positions the agent learns to make the next optimal move.

AlphaZero learns two quantities, the policy and the value function which are realized as two distinct heads forked from the transformer backbone. The policy learns to select the next action and the value function learns to predict the value of the state, which in our case is the number of multiplications remaining until all remaining vertices are eliminated if following the current policy.
A pure search without learning (given that we understand learning as an improvement of policy and value functions) means that those quantities would be random. The performance with random policy and value and their subsequent improvement is show in the reward curves presented in Appendix E, pages 24 and 25. We clearly see that at initialization, when the policy and value functions are random, the performance of the tree search is usually very bad and far worse than the baseline. However the performance improves over time, as evidenced by the evolving reward curves that become increasingly better until they beat the baselines. This is in our opinion clear evidence that the policy and value functions are learned over time and improve the tree search by focusing the search budget on the relevant, most promising parts of the tree. The distribution that the agent is learning on is the distribution over all possible graph configurations and their optimal next actions that can arise from performing vertex elimination on the initial computational graph of a given task example, e.g. RoeFlux_1d. Thus for each row in table 1, we start with a fixed initial input, i.e. the state at the root node in the tree is fixed. But when expanding the tree by sampling and taking actions, we create state-action-reward triplets that are sampled from the the distribution over possible graphs that can arise from the initial input graph due to vertex elimination. These samples are then used to train the policy and value functions to improve them, i.e. learn them.

If we did not answer the reviewer's question, we are happy to try to clarify further!
Thanks for the detailed response.

My comment was specific to the current paper's evaluation in e.g. Table 1 -- it is unclear whether each row is a fixed input instance or an evaluation over a distribution of inputs. If it is the former, then the performance being evaluated is one of a pure search procedure with no explicit learning involved. Clearly, the above critique does not apply to training across a range of inputs and/or graphs.
We thank the reviewer for their reconsideration and appreciate the additional questions since they are indeed quite relevant for our work. 

> However, even while considering the core results for a fixed graph, I am curious what defines distribution of the inputs that makes this evaluation distinct from that of a one-off search procedure looking for a single best elimination order not conditioned on any other inputs.
	
The idea behind generalization across graphs or even within one graph is that our agent recognizes reoccurring connectivity patterns and learns how these should be optimally eliminated, conditioned on the global structure of the graph.
The distribution of graphs is in that sense the distribution of all possible ways to connect $N$ vertices in a meaningful way such that they correspond to proper, mathematically well-defined and observable functions. A single computational graph and its different stages during vertex elimination do no sample this space entirely (for a single graph with 100 vertices, which is a typical size for the problems in this work, we have $100!=10^157$ possible elimination orders), thereby making it hard to find enough patterns so that the algorithm can generalize out-of-distribution without requiring further training. Also for a single graph, the patterns tend to be highly correlated, which is detrimental for generalization.
	
There has been a lot of prior work in the AD community that aims to identify these connectivity patterns. (Griewank and Walter, 2008) contains some examples in chapters 9, 10 and 11. They call this *local preaccumulation*. In particular, they show that in many cases, an efficient solution is to first find the best elimination order locally for one or more subgraphs. These motifs can be generalized across tasks. However, in practice these recipes do not work very well because it is hard for a human to identify and generalize these for already reasonably small graphs. Given enough training samples and a reasonable search budget, we believe that our agent learns to identify these locally optimal elimination procedures and generalize them to novel computational graphs.
Thus, the difference between the one-off search and the search with the model trained on a multitude of graphs is that the agent has learned more diverse wiring patterns and the corresponding elimination orders it could potentially apply.
This is compounded by the fact that RL algorithms are prone to be stuck in local minima, i.e. unaware of better solutions\different connectivity patterns due to a lack of exploration. This could be partly alleviated with more compute/training time, but sometimes the one-off training on a single graph is not sufficient and the agent actually benefits from the knowledge gained from training on other graphs. This is evident when looking at the random function $f$, which is particularly hard to solve because there are no regularities and a diverse set of solutions exist. In this difficult domain, the joint training increased performance by almost 10%. We believe this is due to the agent being aware of more possible wiring patterns and their optimal elimination orders which it learned from the other tasks.
	
> Given the small sizes of the graphs involved, it would be helpful to understand the extent of learning and generalization demonstrated even if within the same graph. 
	
This is a good recommendation for a future direction of research. To learn more about learning and generalization in our setting, it might be worth taking a look at the actual elimination orders (see Appendix C) and correlate them with the functions that were executed in the graph. This could provide some insights into some common themes that algorithm uses to optimize the AD algorithm. For example, we found that the algorithm quite often performs elimination of vertices with small Markowitz degree first, except for those that are involved in accumulation operations like summations. Those are eliminated last. More insight could possibly be gained by looking at the actual attention maps of the transformer which could give more information about which parts of the graph actually matter for optimizing the elimination orders.
However, a more in-depth understanding of our method from this perspective will require additional tools which we first need to implement.
Thanks for the detailed reply to my concerns and questions. I can see the performance improvements comparing to default Jax reverse mode in Figure 1 of the attached PDF. Though it does not work the same on transformer benchmark, I agree that it can be a future work as the author mention. I don't have further questions and would like to improve my score.
We thank the reviewer for their response and will try our best to include further results in the final version.
I remain positive about this work and will maintain my score.
I thank the authors to answer my questions and comments.
I maintain my score. 
- Applying RL methods for such a problem is an impressive feat, it even serves as a new benchmark for such approaches. 
- Even if the results do not "significantly outperform" baselines, the fact that they can for example find better strategies than reverse-mode is quite surprising. 
- The authors also clearly answered reviewer ErgV concerns.

As the authors acknowledge themselves, it would be great if, for a final version, additional generalization results could be reported.
Thanks for the rebuttal clarifications. 

> ...These results make us confident that for a larger dataset and longer training time, our method will eventually generalize out-of-distribution and predict the optimal elimination order without requiring training. We intend to implement this feature in future iterations of our algorithm.

Amortizing the search costs across different graphs by learning to generalize between them would be nice, and the authors point out they have already seen evidence of this. However, even while considering the core results for a fixed graph, I am curious what defines distribution of the inputs that makes this evaluation distinct from that of a one-off search procedure looking for a single best elimination order not conditioned on any other inputs. Given the small sizes of the graphs involved, it would be helpful to understand the extent of learning and generalization demonstrated even if within the same graph. 

>...Furthermore, we would like to emphasize the point that this paper not only presents a novel RL-based method to optimize AD algorithms but also introduces, for the first time, a Python-based interpreter that actually allows to apply vertex elimination to a wide range of problems, including machine learning....

I agree this is a significant and valuable contribution of the work (and distinct from the policy learning/generalization aspects). I will update my final score to take this into account.