[Paper Review] Acceleration of Non-Linear Minimisation with PyTorch

Bojan Nikolic|arXiv (Cornell University)|May 18, 2018

Model Reduction and Neural Networks9 references3 citations

TL;DR

This paper demonstrates that PyTorch can accelerate non-linear function minimisation in scientific computing by leveraging automatic differentiation and GPU acceleration, achieving up to a 100-fold speedup over traditional NumPy implementations with minimal code changes. The approach uses PyTorch's automatic differentiation and GPU offloading to efficiently compute gradients and evaluate complex functions, enabling faster convergence in optimization problems like maximum-likelihood phase retrieval.

ABSTRACT

I show that a software framework intended primarily for training of neural networks, PyTorch, is easily applied to a general function minimisation problem in science. The qualities of PyTorch of ease-of-use and very high efficiency are found to be applicable in this domain and lead to two orders of magnitude improvement in time-to-solution with very small software engineering effort.

Motivation & Objective

• To demonstrate that PyTorch, originally designed for neural network training, can be effectively applied to general scientific function minimization problems.
• To reduce the time-to-solution for non-linear minimization by leveraging PyTorch’s automatic differentiation and GPU acceleration.
• To show that minimal software engineering effort is required to achieve substantial performance gains compared to traditional NumPy-based implementations.
• To maintain code readability and correctness while significantly improving computational efficiency in scientific optimization tasks.

Proposed method

• Use of PyTorch’s automatic reverse-mode differentiation to compute gradients of complex functions without manual derivation or numerical approximation.
• Implementation of scientific functions using PyTorch tensors, enabling seamless CPU and GPU execution with minimal code changes from NumPy.
• Offloading tensor operations to NVIDIA Tesla K20c GPUs via PyTorch’s CUDA support to accelerate data-parallel computations.
• Adoption of standard optimization algorithms like BFGS with the jac=True flag, enabled by PyTorch’s automatic gradient computation.
• Use of high-level tensor operations to express complex scientific models (e.g., Zernike polynomial-based phase retrieval) in a concise and readable form.
• Benchmarking of PyTorch and NumPy implementations across varying grid sizes and model parameter counts to measure performance gains.

Experimental results

Research questions

• RQ1 Can PyTorch’s automatic differentiation and GPU acceleration be effectively repurposed for general scientific function minimization beyond neural networks?
• RQ2 To what extent can PyTorch reduce time-to-solution in non-linear minimization problems compared to traditional NumPy-based approaches?
• RQ3 Does the performance gain from PyTorch scale with increasing model complexity and data size in scientific optimization tasks?
• RQ4 Can PyTorch maintain code readability and maintainability while delivering high computational performance in scientific computing?

Key findings

• PyTorch achieved up to a 100-fold improvement in time-to-solution compared to CPU-only NumPy implementations for phase retrieval problems.
• The GPU-offloaded PyTorch implementation was approximately 10 times faster than the CPU-only PyTorch version for grid sizes N ≥ 256.
• The CPU-only PyTorch implementation was about an order of magnitude faster than the NumPy implementation for intermediate and large grid sizes.
• The performance improvement was observed to be largely independent of the number of model parameters, as measured by Zernike polynomial order up to n=8.
• All implementations converged to the same solution within algorithmic tolerances, confirming correctness despite performance gains.
• The PyTorch implementation maintained code readability and maintainability comparable to NumPy, reducing the risk of subtle errors in scientific models.