Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

• URL: http://arxiv.org/abs/2412.00088v2
• Date: Mon, 13 Jan 2025 01:43:15 GMT
• Title: Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators
• Authors: Zekun Shi, Zheyuan Hu, Min Lin, Kenji Kawaguchi,
• Abstract summary: We show how to efficiently perform arbitrary contraction of the derivative tensor of arbitrary order for multivariate functions. When applied to Physics-Informed Neural Networks (PINNs), our method provides >1000$times$ speed-up and. 30$times$ memory reduction over randomization with first-order AD.
• Score: 29.063441432499776
• License:
• Abstract: Optimizing neural networks with loss that contain high-dimensional and high-order differential operators is expensive to evaluate with back-propagation due to $\mathcal{O}(d^{k})$ scaling of the derivative tensor size and the $\mathcal{O}(2^{k-1}L)$ scaling in the computation graph, where $d$ is the dimension of the domain, $L$ is the number of ops in the forward computation graph, and $k$ is the derivative order. In previous works, the polynomial scaling in $d$ was addressed by amortizing the computation over the optimization process via randomization. Separately, the exponential scaling in $k$ for univariate functions ($d=1$) was addressed with high-order auto-differentiation (AD). In this work, we show how to efficiently perform arbitrary contraction of the derivative tensor of arbitrary order for multivariate functions, by properly constructing the input tangents to univariate high-order AD, which can be used to efficiently randomize any differential operator. When applied to Physics-Informed Neural Networks (PINNs), our method provides >1000$\times$ speed-up and >30$\times$ memory reduction over randomization with first-order AD, and we can now solve \emph{1-million-dimensional PDEs in 8 minutes on a single NVIDIA A100 GPU}. This work opens the possibility of using high-order differential operators in large-scale problems.

Related papers

• A Quasilinear Algorithm for Computing Higher-Order Derivatives of Deep Feed-Forward Neural Networks [0.0]
  $n$-TangentProp computes the exact derivative $dn/dxn f(x)$ in quasilinear, instead of exponential time. We demonstrate that our method is particularly beneficial in the context of physics-informed neural networks.
  arXiv  Detail & Related papers  (2024-12-12T22:57:28Z)
• Fast and scalable Wasserstein-1 neural optimal transport solver for single-cell perturbation prediction [55.89763969583124]
  Optimal transport theory provides a principled framework for constructing such mappings. We propose a novel optimal transport solver based on Wasserstein-1. Our experiments demonstrate that the proposed solver can mimic the $W$ OT solvers in finding a unique and monotonic" map on 2D datasets.
  arXiv  Detail & Related papers  (2024-11-01T14:23:19Z)
• On Convergence of Incremental Gradient for Non-Convex Smooth Functions [63.51187646914962]
  In machine learning and network optimization, algorithms like shuffle SGD are popular due to minimizing the number of misses and good cache. This paper delves into the convergence properties SGD algorithms with arbitrary data ordering.
  arXiv  Detail & Related papers  (2023-05-30T17:47:27Z)
• Convergence Rates for Non-Log-Concave Sampling and Log-Partition Estimation [0.0]
  It is known that for $m$times differentiable functions in $d$, the optimal rate for algorithms with $n$(m/d) is to be $n(m/d). We show that similar rates for sampling and computation are possible, and whether they can be realized in time with an independent rate of $d$.
  arXiv  Detail & Related papers  (2023-03-06T15:53:44Z)
• Scalable First-Order Bayesian Optimization via Structured Automatic Differentiation [4.061135251278187]
  We show that a wide range of kernels gives rise to structured matrices, enabling an exact $mathcalO(n2d)$ matrix-vector multiply for gradient observations and $mathcalO(n2d2)$ for Hessian observations. Our methods apply to virtually all canonical kernels and automatically extend to complex kernels, like the neural network, radial basis function network, and spectral mixture kernels.
  arXiv  Detail & Related papers  (2022-06-16T17:59:48Z)
• Efficient and robust high-dimensional sparse logistic regression via nonlinear primal-dual hybrid gradient algorithms [0.0]
  We propose an iterative algorithm that provably computes a solution to a logistic regression problem regularized by an elastic net penalty. This result improves on the known complexity bound of $O(min(m2n,mn2)log (1/epsilon))$ for first-order optimization methods.
  arXiv  Detail & Related papers  (2021-11-30T14:16:48Z)
• Scaling Gaussian Processes with Derivative Information Using Variational Inference [17.746842802181256]
  We introduce methods to achieve fully scalable Gaussian process regression with derivatives using variational inference. We demonstrate the full scalability of our approach on a variety of tasks, ranging from a high dimensional stellarator fusion regression task to training graph convolutional neural networks on Pubmed.
  arXiv  Detail & Related papers  (2021-07-08T18:23:59Z)
• Reducing the Variance of Gaussian Process Hyperparameter Optimization with Preconditioning [54.01682318834995]
  Preconditioning is a highly effective step for any iterative method involving matrix-vector multiplication. We prove that preconditioning has an additional benefit that has been previously unexplored. It simultaneously can reduce variance at essentially negligible cost.
  arXiv  Detail & Related papers  (2021-07-01T06:43:11Z)
• Asynchronous Stochastic Optimization Robust to Arbitrary Delays [54.61797739710608]
  We consider optimization with delayed gradients where, at each time stept$, the algorithm makes an update using a stale computation - d_t$ for arbitrary delay $d_t gradient. Our experiments demonstrate the efficacy and robustness of our algorithm in cases where the delay distribution is skewed or heavy-tailed.
  arXiv  Detail & Related papers  (2021-06-22T15:50:45Z)
• Finding Global Minima via Kernel Approximations [90.42048080064849]
  We consider the global minimization of smooth functions based solely on function evaluations. In this paper, we consider an approach that jointly models the function to approximate and finds a global minimum.
  arXiv  Detail & Related papers  (2020-12-22T12:59:30Z)
• Linear Time Sinkhorn Divergences using Positive Features [51.50788603386766]
  Solving optimal transport with an entropic regularization requires computing a $ntimes n$ kernel matrix that is repeatedly applied to a vector. We propose to use instead ground costs of the form $c(x,y)=-logdotpvarphi(x)varphi(y)$ where $varphi$ is a map from the ground space onto the positive orthant $RRr_+$, with $rll n$.
  arXiv  Detail & Related papers  (2020-06-12T10:21:40Z)

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.