A research framework for writing differentiable PDE discretizations in JAX

• URL: http://arxiv.org/abs/2111.05218v1
• Date: Tue, 9 Nov 2021 15:58:44 GMT
• Title: A research framework for writing differentiable PDE discretizations in JAX
• Authors: Antonio Stanziola, Simon R. Arridge, Ben T. Cox, Bradley E. Treeby
• Abstract summary: Differentiable simulators are an emerging concept with applications in several fields, from reinforcement learning to optimal control. We propose a library of differentiable operators and discretizations, by representing operators as mappings between families of continuous functions, parametrized by finite vectors. We demonstrate the approach on an acoustic optimization problem, where the Helmholtz equation is discretized using Fourier spectral methods, and differentiability is demonstrated using gradient descent to optimize the speed of sound of an acoustic lens.
• Score: 3.4389358108344257
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: Differentiable simulators are an emerging concept with applications in several fields, from reinforcement learning to optimal control. Their distinguishing feature is the ability to calculate analytic gradients with respect to the input parameters. Like neural networks, which are constructed by composing several building blocks called layers, a simulation often requires computing the output of an operator that can itself be decomposed into elementary units chained together. While each layer of a neural network represents a specific discrete operation, the same operator can have multiple representations, depending on the discretization employed and the research question that needs to be addressed. Here, we propose a simple design pattern to construct a library of differentiable operators and discretizations, by representing operators as mappings between families of continuous functions, parametrized by finite vectors. We demonstrate the approach on an acoustic optimization problem, where the Helmholtz equation is discretized using Fourier spectral methods, and differentiability is demonstrated using gradient descent to optimize the speed of sound of an acoustic lens. The proposed framework is open-sourced and available at \url{https://github.com/ucl-bug/jaxdf}

Related papers

• DimOL: Dimensional Awareness as A New 'Dimension' in Operator Learning [63.5925701087252]
  We introduce DimOL (Dimension-aware Operator Learning), drawing insights from dimensional analysis. To implement DimOL, we propose the ProdLayer, which can be seamlessly integrated into FNO-based and Transformer-based PDE solvers. Empirically, DimOL models achieve up to 48% performance gain within the PDE datasets.
  arXiv  Detail & Related papers  (2024-10-08T10:48:50Z)
• Neural Control Variates with Automatic Integration [49.91408797261987]
  This paper proposes a novel approach to construct learnable parametric control variates functions from arbitrary neural network architectures. We use the network to approximate the anti-derivative of the integrand. We apply our method to solve partial differential equations using the Walk-on-sphere algorithm.
  arXiv  Detail & Related papers  (2024-09-23T06:04:28Z)
• Dilated convolution neural operator for multiscale partial differential equations [11.093527996062058]
  We propose the Dilated Convolutional Neural Operator (DCNO) for multiscale partial differential equations. The DCNO architecture effectively captures both high-frequency and low-frequency features while maintaining a low computational cost. We show that DCNO strikes an optimal balance between accuracy and computational cost and offers a promising solution for multiscale operator learning.
  arXiv  Detail & Related papers  (2024-07-16T08:17:02Z)
• Linearization Turns Neural Operators into Function-Valued Gaussian Processes [23.85470417458593]
  We introduce a new framework for approximate Bayesian uncertainty quantification in neural operators. Our approach can be interpreted as a probabilistic analogue of the concept of currying from functional programming. We showcase the efficacy of our approach through applications to different types of partial differential equations.
  arXiv  Detail & Related papers  (2024-06-07T16:43:54Z)
• Score-based Diffusion Models in Function Space [140.792362459734]
  Diffusion models have recently emerged as a powerful framework for generative modeling. We introduce a mathematically rigorous framework called Denoising Diffusion Operators (DDOs) for training diffusion models in function space. We show that the corresponding discretized algorithm generates accurate samples at a fixed cost independent of the data resolution.
  arXiv  Detail & Related papers  (2023-02-14T23:50:53Z)
• Equivariance with Learned Canonicalization Functions [77.32483958400282]
  We show that learning a small neural network to perform canonicalization is better than using predefineds. Our experiments show that learning the canonicalization function is competitive with existing techniques for learning equivariant functions across many tasks.
  arXiv  Detail & Related papers  (2022-11-11T21:58:15Z)
• Transformer for Partial Differential Equations' Operator Learning [0.0]
  We present an attention-based framework for data-driven operator learning, which we term Operator Transformer (OFormer) Our framework is built upon self-attention, cross-attention, and a set of point-wise multilayer perceptrons (MLPs)
  arXiv  Detail & Related papers  (2022-05-26T23:17:53Z)
• A Unified Framework for Implicit Sinkhorn Differentiation [58.56866763433335]
  We propose an algorithm that obtains analytical gradients of a Sinkhorn layer via implicit differentiation. We show that it is computationally more efficient, particularly when resources like GPU memory are scarce.
  arXiv  Detail & Related papers  (2022-05-13T14:45:31Z)
• Feature Engineering with Regularity Structures [4.082216579462797]
  We investigate the use of models from the theory of regularity structures as features in machine learning tasks. We provide a flexible definition of a model feature vector associated to a space-time signal, along with two algorithms which illustrate ways in which these features can be combined with linear regression. We apply these algorithms in several numerical experiments designed to learn solutions to PDEs with a given forcing and boundary data.
  arXiv  Detail & Related papers  (2021-08-12T17:53:47Z)
• Multipole Graph Neural Operator for Parametric Partial Differential Equations [57.90284928158383]
  One of the main challenges in using deep learning-based methods for simulating physical systems is formulating physics-based data. We propose a novel multi-level graph neural network framework that captures interaction at all ranges with only linear complexity. Experiments confirm our multi-graph network learns discretization-invariant solution operators to PDEs and can be evaluated in linear time.
  arXiv  Detail & Related papers  (2020-06-16T21:56:22Z)
• Deep neural networks for inverse problems with pseudodifferential operators: an application to limited-angle tomography [0.4110409960377149]
  We propose a novel convolutional neural network (CNN) designed for learning pseudodifferential operators ($Psi$DOs) in the context of linear inverse problems. We show that, under rather general assumptions on the forward operator, the unfolded iterations of ISTA can be interpreted as the successive layers of a CNN. In particular, we prove that, in the case of LA-CT, the operations of upscaling, downscaling and convolution, can be exactly determined by combining the convolutional nature of the limited angle X-ray transform and basic properties defining a wavelet system.
  arXiv  Detail & Related papers  (2020-06-02T14:03:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.