Computing Jacobians with automatic differentiation is ubiquitous in many scientific domains such as machine learning, computational fluid dynamics, robotics and finance. Even small savings in the number of computations or memory usage in Jacobian computations can already incur massive savings in energy consumption and runtime. While there exist many methods that allow for such savings, they generally trade computational efficiency for approximations of the exact Jacobian. In this paper, we present a novel method to optimize the number of necessary multiplications for Jacobian computation by leveraging deep reinforcement learning (RL) and a concept called cross-country elimination while still computing the exact Jacobian. Cross-country elimination is a framework for automatic differentiation that phrases Jacobian accumulation as ordered elimination of all vertices on the computational graph where every elimination incurs a certain computational cost. We formulate the search for the optimal elimination order that minimizes the number of necessary multiplications as a single player game which is played by an RL agent. We demonstrate that this method achieves up to 33% improvements over state-of-the-art methods on several relevant tasks taken from diverse domains. Furthermore, we show that these theoretical gains translate into actual runtime improvements by providing a cross-country elimination interpreter in JAX that can efficiently execute the obtained elimination orders.
Paper
Similar papers
Peer review
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.
Strengths
Training an RL agent to rethink classic numerical methods is an interesting research topic. Previous examples include DeepMind’s works on matrix multiplication and sorting algorithms. Because these numerical problems are basic and fundamental building blocks in today’s scientific research, improvements in its solutions can have high potential and impact. This paper follows this trend and applies RL to another specific example: automatic differentiation. I think the paper picks a good topic and makes a meaningful contribution to the community. The technical method in this paper also looks good to me. Using RL to explore better numerical and computing methods typically requires a careful, RL-friendly (re-)formulation of the original numerical problem, including the design of the state space, the action space, and the reward function. I think choosing to tackle vertex elimination ordering is a smart and reasonable move, which leads to a well-formulated game that captures the key structure in automatic differentiation.
Weaknesses
I don’t have major concerns with the paper, but I want to mention that there are quite a few typos in the equations in the main paper. For example, I suspect “i > j” should be replaced with “i < j”, and “sum_i” in line 192 should be “sum_j”. Please double-check all equations in the paper and make sure there are no typos.
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.
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
3
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.
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)
Strengths
The paper is written well, and serves as a helpful introduction to how to optimize autodifferentiation. The paper acknowledges that it makes some simplifications to the classes of computation graphs it can be applied to (notably limits on input dimensionality make the Transformer experiment somewhat unrealistic), but demonstrate it is possible to achieve gains over typical methods using learning based methods. An analysis on real hardware is performed as well.
Weaknesses
By my understanding, both the AlphaZero agent and PPO agent need to be instantiated separately for each class of functions. This is based off the note in Appendix F, that the authors did not succeed in training a joint model for the PPO agent. One of the papers cited as related work is the one for improving matrix multiplication with search algorithms, and one benefit of that work was that although the search was very specific to 4x4 matrices for example, it was possible to us the discovered results of that work elsewhere. 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
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?
Rating
7
Confidence
3
Soundness
4
Presentation
3
Contribution
3
Limitations
Yes
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.
Strengths
- The paper is well motivated and proposes a policy search framework for a well studied combinatorial optimization problem (vertex elimination for AD). Instead of targeting compilation (at a lower level) or dealing exclusively with optimizing individual matmuls as in prior work, the paper targets a novel alternative by focusing on vertex elimination in AD which is a ubiquitous part of all training compute graphs. - Any non-trivial improvements on large scale computational graphs could lead to significant compute savings, and this work proposes several novel small scale benchmarks to evaluate vertex elimination algorithms which maybe of interest in its own right.
Weaknesses
- The primary weakness in terms of the impact appears to be that the policy must be trained from scratch separately for each graph. This unfortunately means that there is a whole training run necessary per graph unlike all of the other competing methods (including the minimum Markowitz degree heuristic, which needs no "training", and already appears to do better than fwd/reverse mode in many situations). - 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
- 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
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
Yes
Response
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!
Reply
Thank you for the clarification, this was very helpful.
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.
Strengths
- Rephrasing the elimination order in AD as a deep reinforcement learning problem sounds somehow novel. - The paper is well-organized.
Weaknesses
- The proposed method only shows a on obvious improvement on the RoeFlux_3d and Random function. While on other benchmarks especially on deep learning related MLP and Transformer cases, the improvements are very marginal (less than 3%). - 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
- 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?
Rating
5
Confidence
3
Soundness
2
Presentation
3
Contribution
2
Limitations
yes.
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.
Strengths
- The proposed approach demonstrates benefits across tasks. - The algorithm is developed in a package that will be released such that the method may quickly be adopted. - The problem can serve as an interesting benchmark for deep reinforcement learning. - The experimental details are throughly detailed to ensure preoducibility.
Weaknesses
- As stated by the authors, the generalization performance of the approach may still be unclear, which may make the approach rather computationally prohibitive.
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"
Rating
8
Confidence
4
Soundness
4
Presentation
4
Contribution
4
Limitations
Limitations are discussed such as the poor generalization when performing joint training and the difficulties to handle the range of rewards.
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.
Strengths
- Demonstrates a novel approach to optimizing AD by leveraging deep RL to discover optimized elimination orders. - The paper is technically sound and provides good descriptions of these various topics for an audience who may not be very familiar with AD. - Overall, the paper is well-written and the key ideas and results are clearly communicated. - The work has the potential for significant practical impact as even small improvements in AD efficiency can lead to large runtime and energy savings when applied at scale.
Weaknesses
- Currently, the approach has a simple model based on optimizing the number of multiplications as a proxy for runtime. As the authors note this does not capture the full complexity and ignores factors like memory access patterns. - 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
- 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?
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
The authors are forthcoming about the main limitations and highlight important areas for future work.
Reply
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.
Response
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.
Reply
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.
Acknowledging rebuttal
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.
Response
We thank the reviewer for their response and will try our best to include further results in the final version.
Response to the reply
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.
Thank you for the clarification. I will keep my score as-is.
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.
Decision
Accept (spotlight)