Reverse Mode Automatic Differentiation Research Articles

Inverse analysis of granular flows using differentiable graph neural network simulator

Inverse problems in granular flows, such as landslides and debris flows, involve estimating material parameters or boundary conditions based on a target runout profile. Traditional high-fidelity simulators for these inverse problems are computationally expensive, restricting the number of possible simulations. These simulators are also non-differentiable, making efficient gradient-based optimization methods for high-dimensional spaces inapplicable. Machine learning-based surrogate models offer computational efficiency and differentiability. However, they often struggle to generalize beyond their training data due to relying on low-dimensional input–output mappings that fail to capture the complete physics of granular flows. We propose a novel differentiable graph neural network simulator (GNS) that combines reverse mode automatic differentiation of graph neural networks with gradient-based optimization for solving inverse problems. GNS learns the dynamics of granular flow by representing the system as a graph and predicts the evolution of the graph at the next timestep, given the current state. The differentiable GNS demonstrates optimization capabilities beyond the training data. We demonstrate the effectiveness of our method for inverse estimation across single and multi-parameter optimization problems, including evaluating material properties and boundary conditions for a target runout distance and designing baffle locations to limit a landslide runout. Our proposed differentiable GNS framework solves inverse problems with orders of magnitude faster convergence than the conventional gradient-based optimization approach using finite difference.

Geometrically exact beam theory for gradient-based optimization

Decades of research have progressed geometrically exact beam theory to the point where it is now an invaluable resource for analyzing and modeling highly flexible slender structures. Large-scale optimization using geometrically exact beam theory remains nontrivial, however, due to the inability of gradient-free optimizers to handle large numbers of design variables in a computationally efficient manner and the difficulties associated with obtaining smooth, accurate, and efficiently calculated design sensitivities for gradient-based optimization. To overcome these challenges, this paper presents a finite-element implementation of geometrically exact beam theory which has been developed specifically for gradient-based optimization. A key feature of this implementation of geometrically exact beam theory is its compatibility with forward and reverse-mode automatic differentiation. Another key feature is its support for both continuous and discrete adjoint sensitivity analysis. Other features are also presented which build upon previous implementations of geometrically exact beam theory, including a singularity-free rotation parameterization based on Wiener-Milenković parameters, an implementation of stiffness-proportional structural damping using a discretized form of the compatibility equations, and a reformulation of the equations of motion for geometrically exact beam theory as a semi-explicit system. Several examples are presented which verify the utility and validity of each of these features.

Zero coordinate shift: Whetted automatic differentiation for physics-informed operator learning

Automatic differentiation (AD) is a critical step in physics-informed machine learning, required for computing the high-order derivatives of network output w.r.t. coordinates of collocation points. In this paper, we present a novel and lightweight algorithm to conduct AD for physics-informed operator learning, which we call the trick of Zero Coordinate Shift (ZCS). Instead of making all sampled coordinates as leaf variables, ZCS introduces only one scalar-valued leaf variable for each spatial or temporal dimension, simplifying the wanted derivatives from “many-roots-many-leaves” to “one-root-many-leaves” whereby reverse-mode AD becomes directly utilisable. It has led to an outstanding performance leap by avoiding the duplication of the computational graph along the dimension of functions (physical parameters). ZCS is easy to implement with current deep learning libraries; our own implementation is achieved by extending the DeepXDE package. We carry out a comprehensive benchmark analysis and several case studies, training physics-informed DeepONets to solve partial differential equations (PDEs) without data. The results show that ZCS has persistently reduced GPU memory consumption and wall time for training by an order of magnitude, and such reduction factor scales with the number of functions. As a low-level optimisation technique, ZCS imposes no restrictions on data, physics (PDE) or network architecture and does not compromise training results from any aspect.

Qualitative and quantitative enhancement of parameter estimation for model-based diagnostics using automatic differentiation with an application to inertial fusion

Parameter estimation using observables is a fundamental concept in the experimental sciences. Mathematical models that represent the physical processes can enable reconstructions of the experimental observables and greatly assist in parameter estimation by turning it into an optimization problem which can be solved by gradient-free or gradient-based methods. In this work, the recent rise in flexible frameworks for developing differentiable scientific computing programs is leveraged in order to dramatically accelerate data analysis of a common experimental diagnostic relevant to laser–plasma and inertial fusion experiments, Thomson scattering. A differentiable Thomson-scattering data analysis tool is developed that uses reverse-mode automatic differentiation (AD) to calculate gradients. By switching from finite differencing to reverse-mode AD, three distinct outcomes are achieved. First, gradient descent is accelerated dramatically to the extent that it enables near real-time usage in laser–plasma experiments. Second, qualitatively novel quantities which require O(103) parameters can now be included in the analysis of data which enables unprecedented measurements of small-scale laser–plasma phenomena. Third, uncertainty estimation approaches that leverage the value of the Hessian become accurate and efficient because reverse-mode AD can be used for calculating the Hessian.

Efficient CHAD

We show how the basic Combinatory Homomorphic Automatic Differentiation (CHAD) algorithm can be optimised, using well-known methods, to yield a simple, composable, and generally applicable reverse-mode automatic differentiation (AD) technique that has the correct computational complexity that we would expect of reverse-mode AD. Specifically, we show that the standard optimisations of sparse vectors and state-passing style code (as well as defunctionalisation/closure conversion, for higher-order languages) give us a purely functional algorithm that is most of the way to the correct complexity, with (functional) mutable updates taking care of the final log-factors. We provide an Agda formalisation of our complexity proof. Finally, we discuss how the techniques apply to differentiating parallel functional array programs: the key observations are 1) that all required mutability is (commutative, associative) accumulation, which lets us preserve task-parallelism and 2) that we can write down data-parallel derivatives for most data-parallel array primitives.

Forward- or reverse-mode automatic differentiation: What's the difference?

Forward- or reverse-mode automatic differentiation: What's the difference?

Efficient Automatic Subdifferentiation for Programs with Linear Branches

Computing an element of the Clarke subdifferential of a function represented by a program is an important problem in modern non-smooth optimization. Existing algorithms either are computationally inefficient in the sense that the computational cost depends on the input dimension or can only cover simple programs such as polynomial functions with branches. In this work, we show that a generalization of the latter algorithm can efficiently compute an element of the Clarke subdifferential for programs consisting of analytic functions and linear branches, which can represent various non-smooth functions such as max, absolute values, and piecewise analytic functions with linear boundaries, as well as any program consisting of these functions such as neural networks with non-smooth activation functions. Our algorithm first finds a sequence of branches used for computing the function value at a random perturbation of the input; then, it returns an element of the Clarke subdifferential by running the backward pass of the reverse-mode automatic differentiation following those branches. The computational cost of our algorithm is at most that of the function evaluation multiplied by some constant independent of the input dimension n, if a program consists of piecewise analytic functions defined by linear branches, whose arities and maximum depths of branches are independent of n.

Modeling Mixtures with PCP-SAFT: Insights from Large-Scale Parametrization and Group-Contribution Method for Binary Interaction Parameters

Vapor/liquid equilibria of mixtures are of utmost importance in the design of chemical processes. Because of the combinatorial complexity of mixtures, the available experimental data are small when considering the large molecular space. To fill this knowledge gap, molecular equations of state like PCP-SAFT show promise due to their explicit consideration of intermolecular interactions that can be transferred to mixtures. In this work, we comprehensively assess and exploit PCP-SAFT for modeling phase equilibria of mixtures. First, we provide binary interaction parameters for 7861 binary systems for which pure-component parameters and experimental data are available. Bubble and dew point pressures are described with a median deviation of 2.3 %. Secondly, we adjust a matrix of binary group/group interaction parameters for the homosegmented and heterosegmented group-contribution (GC) methods for PCP-SAFT. Among 1389 mixtures that can be described with the GC methods, the median deviation in bubble and dew point pressures are 6.4 % for the homosegmented approach and 5.1 % for the heterosegmented approach. The detailed analysis shows the importance of hydrogen bonds in mixtures of non-self-associating components with self-associating components. The parametrization is only possible by introducing a fast numerical method to calculate the derivative of bubble and dew point pressures with respect to arbitrarily many model parameters. The approach leverages reverse mode automatic differentiation (backpropagation), the same method used in machine learning to regress millions of model parameters to large datasets.

Efficient integration of deep neural networks in sequential multiscale simulations

AbstractMultiscale computations involving finite elements are often unfeasible due to their substantial computational costs arising from numerous microstructure evaluations. This necessitates the utilization of suitable surrogate models, which can be rapidly evaluated. In this work, we apply a purely data‐based deep neural network as a surrogate model for the microstructure evaluation. More precisely, the surrogate model predicts the homogenized stresses, which are typically obtained by the solution of (initial) boundary‐value problems of a microstructure and subsequent homogenization. Furthermore, the required consistent tangent matrix is computed by leveraging reverse mode automatic differentiation. To improve data efficiency and ensure high prediction quality, the well‐known Sobolev training is chosen for creating the surrogate model. This surrogate model is seamlessly integrated into a Fortran‐based finite element code using an open‐source library. As a result, this integration, combined with just‐in‐time compilation, leads to a speed‐up of more than 6000× , as demonstrated in the example of a plate with a hole examined in this work. Furthermore, the surrogate model proves to be applicable even in load‐step size controlled simulations, where it overcomes certain load‐step size limitations associated with the microstructure computations.

CHAD for expressive total languages

AbstractWe show how to apply forward and reverse mode Combinatory Homomorphic Automatic Differentiation (CHAD) (Vákár (2021). ESOP, 607–634; Vákár and Smeding (2022). ACM Transactions on Programming Languages and Systems44 (3) 20:1–20:49.) to total functional programming languages with expressive type systems featuring the combination of •tuple types;•sum types;•inductive types;•coinductive types;•function types.We achieve this by analyzing the categorical semantics of such types in $\Sigma$ -types (Grothendieck constructions) of suitable categories. Using a novel categorical logical relations technique for such expressive type systems, we give a correctness proof of CHAD in this setting by showing that it computes the usual mathematical derivative of the function that the original program implements. The result is a principled, purely functional and provably correct method for performing forward- and reverse-mode automatic differentiation (AD) on total functional programming languages with expressive type systems.

You Only Linearize Once: Tangents Transpose to Gradients

Automatic differentiation (AD) is conventionally understood as a family of distinct algorithms, rooted in two “modes”—forward and reverse—which are typically presented (and implemented) separately. Can there be only one? Following up on the AD systems developed in the JAX and Dex projects, we formalize a decomposition of reverse-mode AD into (i) forward-mode AD followed by (ii) unzipping the linear and non-linear parts and then (iii) transposition of the linear part. To that end, we define a (substructurally) linear type system that can prove a class of functions are (algebraically) linear. Our main results are that forward-mode AD produces such linear functions, and that we can unzip and transpose any such linear function, conserving cost, size, and linearity. Composing these three transformations recovers reverse-mode AD. This decomposition also sheds light on checkpointing, which emerges naturally from a free choice in unzipping let expressions. As a corollary, checkpointing techniques are applicable to general-purpose partial evaluation, not just AD. We hope that our formalization will lead to a deeper understanding of automatic differentiation and that it will simplify implementations, by separating the concerns of differentiation proper from the concerns of gaining efficiency (namely, separating the derivative computation from the act of running it backward).

Efficient Dual-Numbers Reverse AD via Well-Known Program Transformations

Where dual-numbers forward-mode automatic differentiation (AD) pairs each scalar value with its tangent value, dual-numbers reverse-mode AD attempts to achieve reverse AD using a similarly simple idea: by pairing each scalar value with a backpropagator function. Its correctness and efficiency on higher-order input languages have been analysed by Brunel, Mazza and Pagani, but this analysis used a custom operational semantics for which it is unclear whether it can be implemented efficiently. We take inspiration from their use of linear factoring to optimise dual-numbers reverse-mode AD to an algorithm that has the correct complexity and enjoys an efficient implementation in a standard functional language with support for mutable arrays, such as Haskell. Aside from the linear factoring ingredient, our optimisation steps consist of well-known ideas from the functional programming community. We demonstrate the use of our technique by providing a practical implementation that differentiates most of Haskell98.

Automatic Differentiation for the Direct Minimization Approach to the Hartree-Fock Method.

Automatic differentiation (AD) has become an important tool for optimization problems in computational science, and it has been applied to the Hartree-Fock method. Although the reverse-mode AD is more efficient than the forward-mode, eigenvalue calculation in the self-consistent field (SCF) method has impeded the use of the reverse-mode AD. Here, we propose a method to directly minimize Hartree-Fock energy under the orthonormality constraint of the molecular orbitals using reverse-mode AD by avoiding eigenvalue calculation. According to our validation, the proposed method was more stable than the conventional SCF method and achieved comparable accuracy.

TORCWA: GPU-accelerated Fourier modal method and gradient-based optimization for metasurface design

TORCWA is an electromagnetic wave simulation and optimization tool utilizing rigorous coupled-wave analysis. One of the advantages of TORCWA is that it provides GPU-accelerated simulation. It shows a greatly accelerated simulation speed compared to when the same simulation is performed on a CPU-based. Although it has accelerated speed, the simulation results are almost identical to the commercialized electromagnetic wave simulations. The second advantage is that it provides GPU-accelerated gradient calculation for the simulation results with reverse-mode automatic differentiation of PyTorch version 1.10.1. In particular, the instability of gradient calculation of eigendecomposition is also improved. With this property, TORCWA can be utilized for the optimization of various nanophotonic devices. Here, we first introduce the formulation used in TORCWA, compare it with other commercial simulations, and show the computational performance in multiple environments. Then, the gradient calculation and optimization examples are shown. Thanks to accelerated computational performance and gradient calculation, TORCWA is a worthy program for designing and optimizing various nanophotonic devices. Program summaryProgram title: TORCWACPC Library link to program files:https://doi.org/10.17632/2dybvpk42g.1Developer's repository link:https://github.com/kch3782/torcwaLicensing provisions: LGPL v3Programming language: Python 3.8Nature of Problem: Time-harmonic electromagnetic fields.Solution method: Fourier modal methods and extended scattering matrix method.Additional comments including restrictions and unusual features: TORCWA utilizes PyTorch version 1.10.1 [1].

CHAD: Combinatory Homomorphic Automatic Differentiation

We introduce Combinatory Homomorphic Automatic Differentiation (CHAD), a principled, pure, provably correct define-then-run method for performing forward and reverse mode automatic differentiation (AD) on programming languages with expressive features. It implements AD as a compositional, type-respecting source-code transformation that generates purely functional code. This code transformation is principled in the sense that it is the unique homomorphic (structure preserving) extension to expressive languages of Elliott’s well-known and unambiguous definitions of AD for a first-order functional language. Correctness of the method follows by a (compositional) logical relations argument that shows that the semantics of the syntactic derivative is the usual calculus derivative of the semantics of the original program.In their most elegant formulation, the transformations generate code with linear types. However, the code transformations can be implemented in a standard functional language lacking linear types: While the correctness proof requires tracking of linearity, the actual transformations do not. In fact, even in a standard functional language, we can get all of the type-safety that linear types give us: We can implement all linear types used to type the transformations as abstract types by using a basic module system.In this article, we detail the method when applied to a simple higher-order language for manipulating statically sized arrays. However, we explain how the methodology applies, more generally, to functional languages with other expressive features. Finally, we discuss how the scope of CHAD extends beyond applications in AD to other dynamic program analyses that accumulate data in a commutative monoid.

A δ

Over the last decade, automatic differentiation (AD) has profoundly impacted graphics and vision applications --- both broadly via deep learning and specifically for inverse rendering. Traditional AD methods ignore gradients at discontinuities, instead treating functions as continuous. Rendering algorithms intrinsically rely on discontinuities, crucial at object silhouettes and in general for any branching operation. Researchers have proposed fully- automatic differentiation approaches for handling discontinuities by restricting to affine functions, or semi- automatic processes restricted either to invertible functions or to specialized applications like vector graphics. This paper describes a compiler-based approach to extend reverse mode AD so as to accept arbitrary programs involving discontinuities. Our novel gradient rules generalize differentiation to work correctly, assuming there is a single discontinuity in a local neighborhood, by approximating the prefiltered gradient over a box kernel oriented along a 1D sampling axis. We describe when such approximation rules are first-order correct, and show that this correctness criterion applies to a relatively broad class of functions. Moreover, we show that the method is effective in practice for arbitrary programs, including features for which we cannot prove correctness. We evaluate this approach on procedural shader programs, where the task is to optimize unknown parameters in order to match a target image, and our method outperforms baselines in terms of both convergence and efficiency. Our compiler outputs gradient programs in TensorFlow, PyTorch (for quick prototypes) and Halide with an optional auto-scheduler (for efficiency). The compiler also outputs GLSL that renders the target image, allowing users to interactively modify and animate the shader, which would otherwise be cumbersome in other representations such as triangle meshes or vector art.

Combinatory Adjoints and Differentiation

We develop a compositional approach for automatic and symbolic differentiation based on categorical constructions in functional analysis where derivatives are linear functions on abstract vectors rather than being limited to scalars, vectors, matrices or tensors represented as multi-dimensional arrays. We show that both symbolic and automatic differentiation can be performed using a differential calculus for generating linear functions representing Fr\'echet derivatives based on rules for primitive, constant, linear and bilinear functions as well as their sequential and parallel composition. Linear functions are represented in a combinatory domain-specific language. Finally, we provide a calculus for symbolically computing the adjoint of a derivative without using matrices, which are too inefficient to use on high-dimensional spaces. The resulting symbolic representation of a derivative retains the data-parallel operations from the input program. The combination of combinatory differentiation and computing formal adjoints turns out to be behaviorally equivalent to reverse-mode automatic differentiation. In particular, it provides opportunities for optimizations where matrices are too inefficient to represent linear functions.

Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation

In this paper, we give a simple and efficient implementation of reverse-mode automatic differentiation, which both extends easily to higher-order functions, and has run time and memory consumption linear in the run time of the original program. In addition to a formal description of the translation, we also describe an implementation of this algorithm, and prove its correctness by means of a logical relations argument.

Physics constrained learning for data-driven inverse modeling from sparse observations

Deep neural networks (DNN) can model nonlinear relations between physical quantities. Those DNNs are embedded in physical systems described by partial differential equations (PDE) and trained by minimizing a loss function that measures the discrepancy between predictions and observations in some chosen norm. This loss function often includes the PDE constraints as a penalty term when only sparse observations are available. As a result, the PDE is only satisfied approximately by the solution. However, the penalty term typically slows down the convergence of the optimizer for stiff problems. We present a new approach that trains the embedded DNNs while numerically satisfying the PDE constraints. We developed an algorithm that enables differentiating both explicit and implicit numerical solvers in reverse-mode automatic differentiation. Our method allows the gradients of the DNNs and the PDE solvers to be computed in a unified framework. We demonstrate that our approach enjoys faster convergence and better stability in relatively stiff problems compared to the penalty method. Our approach could solve and accelerate a wide range of data-driven inverse modeling, where the physical constraints are described by PDEs and need to be satisfied accurately.

AI education matters: Model AI assignment

Many of the most significant breakthroughs in artificial intelligence over the past decade have been based on progress in deep neural networks. That progress has been facilitated by deep-learning libraries like Theano (Al-Rfou et al., 2016), TensorFlow (Abadi et al., 2015) and PyTorch (Paszke et al., 2019) that allow rapid prototyping and efficient execution. The key algorithm at the heart of all of these libraries is reverse-mode automatic differentiation. This column introduces the Model AI Assignment ScalarFlow: Implementing Reverse Mode Automatic Differentiation. This assignment gives students the opportunity to gain a deeper understanding of modern deeplearning frameworks by building their own automatic differentiation engine and using it to experiment with some important concepts in deep learning. In this column we will review some basic background on training neural networks, provide a brief overview of the reverse-mode automatic differentiation algorithm, describe the model assignment and provide some pointers to additional resources.