% IMPORTANT: The following is UTF-8 encoded.  This means that in the presence
% of non-ASCII characters, it will not work with BibTeX 0.99 or older.
% Instead, you should use an up-to-date BibTeX implementation like “bibtex8” or
% “biber”.

@INPROCEEDINGS{Lohoff:1038109,
      author       = {Lohoff, Jamie and Neftci, Emre},
      title        = {{O}ptimizing {A}utomatic {D}ifferentiation with {D}eep
                      {R}einforcement {L}earning},
      volume       = {38},
      issn         = {1049-5258},
      reportid     = {FZJ-2025-01156},
      series       = {Advances in neural information processing systems},
      pages        = {n/a},
      year         = {2024},
      note         = {Accepted as a spotlight paper.},
      abstract     = {Computing Jacobians with automatic differentiation is
                      ubiquitous in many scientific domains such as machine
                      learning, computational fluid dynamics, robotics, and
                      finance. Even small savings in the number of computations or
                      memory usage in Jacobian computations can already incur
                      massive savings in energy consumption and runtime. While
                      there exist many methods that allow for such savings, they
                      generally trade computational efficiency for approximations
                      of the exact Jacobian. In this paper, we present a novel
                      method to optimize the number of necessary multiplications
                      for Jacobian computation by leveraging deep reinforcement
                      learning (RL) and a concept called cross-country elimination
                      while still computing the exact Jacobian. Cross-country
                      elimination is a framework for automatic differentiation
                      that phrases Jacobian accumulation as ordered elimination of
                      all vertices on the computational graph where every
                      elimination incurs a certain computational cost. We
                      formulate the search for the optimal elimination order that
                      minimizes the number of necessary multiplications as a
                      single player game which is played by an RL agent. We
                      demonstrate that this method achieves up to $33\%$
                      improvements over state-of-the-art methods on several
                      relevant tasks taken from diverse domains. Furthermore, we
                      show that these theoretical gains translate into actual
                      runtime improvements by providing a cross-country
                      elimination interpreter in JAX that can efficiently execute
                      the obtained elimination orders.},
      month         = {Dec},
      date          = {2024-12-09},
      organization  = {38th Conference on Neural Information
                       Processing Systems, Vancouver (Canada),
                       9 Dec 2024 - 16 Dec 2024},
      cin          = {PGI-15},
      ddc          = {500},
      cid          = {I:(DE-Juel1)PGI-15-20210701},
      pnm          = {5234 - Emerging NC Architectures (POF4-523) / BMBF 16ME0400
                      - Verbundprojekt: Neuro-inspirierte Technologien der
                      künstlichen Intelligenz für die Elektronik der Zukunft -
                      NEUROTEC II - (16ME0400) / GREENEDGE - Taming the
                      environmental impact of mobile networks through GREEN EDGE
                      computing platforms (953775)},
      pid          = {G:(DE-HGF)POF4-5234 / G:(BMBF)16ME0400 /
                      G:(EU-Grant)953775},
      typ          = {PUB:(DE-HGF)8 / PUB:(DE-HGF)7},
      doi          = {10.34734/FZJ-2025-01156},
      url          = {https://juser.fz-juelich.de/record/1038109},
}