Jax is well known as a popular framework for large-scale AI model development, but it has also gained rapid adoption in a wider set of scientific domains. I am particularly excited to see an increase in use in computationally intensive fields such as physics-based machine learning. Jax Support Configurable conversions, A set of higher order functions. For example, GRAD takes a function as input and returns another function that calculates the gradient.Importantly, you can freely nest (create) these transformations. This design makes Jax particularly elegant Higher-order Derivatives Other complex transformations.
Recently I was able to talk to Zekun Shi and Min Lin, researchers at the National University of Singapore and the SEA AI Lab. Their experience clearly shows how JAX can address fundamental challenges in scientific research, particularly around the computational cliffs that they face when solving complex partial differential equations (PDEs). Their journey from tackling the limitations of their traditional framework to leveraging Jax's unique Taylor mode auto-differentiation is a story that resonates with many researchers.
A new approach to solving PDE: in researchers' own words
Our work focuses on challenging areas of scientific computing. Resolve high-order PDEs using neural networks. Neural networks are universal function approximators and provide a promising alternative to traditional methods such as finite elements. However, the main hurdle in solving PDEs in neural networks is that it is necessary to evaluate their higher-order derivatives, sometimes up to 4th or even higher, including mixed partial derivatives.
Standard deep learning frameworks optimized primarily for training models via backpropagation are not suitable for this task as it is very expensive to calculate high-order derivatives. The cost of repeatedly applying backpropagation (backward mode AD) to higher-order derivative scales is exponentially polynomial in derivative order (k) and polynomial in domain dimension (D). This “dimensional curse” and exponential scaling in order of origin makes it virtually impossible to tackle large, complex real-world problems.
Calculate graph scale exponentially in derivative order k
There are other popular libraries for deep learning, but our research required more basic competencies. Taylor Mode Auto Differentiation (AD). Jacks was a game changer for us.
Jax's key architecture distinction is its powerful Function representation and transformation mechanismis implemented by tracing Python code and compiled for high performance. The system is designed with generality enough to enable versatile applications, from just-in-time compilation to standard derivative computing. This one Underlying flexibility This allows for advanced operations that other frameworks cannot easily achieve. For us, the key application was Taylor Mode AD support. Learned that this is a direct and powerful result of this unique architecture, it is best equipped with Jax for scientific work. Taylor Mode AD advances the Taylor series of functions, allowing efficient calculation of higher-order derivatives by efficiently calculating higher-order derivatives in a single pass, rather than through repeated, expensive backpropagation. This allowed us to develop an algorithm, the Stochastic Taylor Derivative Estimator (STDE), to efficiently randomize and estimate differential operators.
Taylor-Mode-no exponential scaling for secondary derivatives.
The recent paper, “Stochastic Taylor Derivative Timator: Efficient Amortization of Arbitraryary Dhirverial Operator,” which won the Best Paper Award at Neurips 2024, demonstrated how this approach could be used. By using Jax's Taylor mode, we have shown that algorithms can be created to efficiently extract these higher order partial derivatives. The core idea was to utilize Taylor Mode Ad to efficiently calculate the shrinkage of higher-order derived tensors displayed in PDES. By constructing a special random tangent vector (or “jets”), you can get a fair estimate of arbitrarily complex differential operators in a single efficient forward pass.
The results were dramatic. I achieved a using the STDE method in JAX > 1000X speed up and > 30x memory reduction Compare with the baseline method. This efficiency gain allowed us to solve a 1 million dimension PDE in just 8 minutes with a single nvidia A100 GPUpreviously unmanageable tasks.
This was not possible with a framework that only targets standard machine learning workloads. Other frameworks are highly optimized for backpropagation, but do not focus much on end-to-end computational graph representations than JAX. This helps JAX shine in operations such as transposed functions and implementing higher-order Taylor mode differentiation.
Beyond Taylor Mode, Jax's modular design and support for general data types and functional transformations is important for our research. Another task, “Automatic Functional Differentiation of JAX,” even dealt with infinite dimension vectors (functions of Hilbert space) by generalizing JAX, describing them as custom arrays, and registering them with JAX. This allows you to reuse existing machines to calculate the variation derivatives of functions and operators, which are functions that are completely out of reach of other frameworks.
For these reasons, we adopted JAX for a wide range of research not only in this project but also in areas such as quantum chemistry. The basic design as a general, scalable and symbolically powerful system makes it an ideal option for driving the frontier of scientific calculations. We believe it is important for the scientific community to know about these abilities.
Explore the Jax Scientific Computing ecosystem
Zekun and Min's experience demonstrates the power and flexibility of Jax. The STDE method, developed using JAX, is a major contributor to the field of physics-based machine learning, allowing you to tackle classes of previously unruly questions. We recommend reading the award-winning papers, dig deeper into technical details and exploring their open source STDE libraries on Github.
Such stories highlight growth trends. JAX is more than a tool for deep learning. This is a fundamental library of differentiable programming that enhances the scientific discovery of a new generation. The Jax team at Google is committed to supporting and growing this vibrant ecosystem, and it starts with hearing directly from you.
- Share your story: Do you use JAX to tackle challenging scientific questions? We want to learn how Jax is accelerating your research and potentially showcase your work.
- Help us guide our roadmap: Are there any new features or features that will unlock the next breakthrough? Feature requests are essential to guide the evolution of JAX in the scientific community.
We look forward to partnering with our customers to build the next generation of scientific calculators. Contact our team to share your work or discuss what you need from Jax.
We would like to express our sincere gratitude to Zekun and Min for sharing their insightful journey with us.
reference
Shi, Z., Hu, Z., Lin, M. , & Kawaguchi, K. (2025). Probability Taylor Derivative Estimator: Efficient amortization of any differential operator. Advances in neural information processing systems, 37.
Rin, M. (2023). Automatic functional differentiation of JAX. 12th International Conference on Learning Expressions.
