Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

From Uniform to Adaptive: General Skip-Block Mechanisms for Efficient PDE Neural Operators

Lei Liu, Zhongyi Yu, Hong Wang, Huanshuo Dong, Haiyang Xin, Hongwei Zhao, Bin Li

TL;DR

Transformer-based neural operators for PDEs suffer from a uniform computation bottleneck across heterogeneous physical domains. The authors propose Skip-Block Routing (SBR), which decouples complexity assessment via a Global Router from an Adaptive Processing Backbone that progressively concentrates computation on the most important tokens, guided by a sparsity schedule. Across NS2D, Pipe, Airfoil, and Heat2D, SBR achieves roughly 50% FLOPs reduction and up to 2x end-to-end speedups with accuracy largely preserved, and in some cases improved, highlighting its practical impact for large-scale PDE solving. This physics-informed adaptive computation enables scalable, efficient neural operators suitable for industrial-scale PDE tasks.

Abstract

In recent years, Neural Operators(NO) have gradually emerged as a popular approach for solving Partial Differential Equations (PDEs). However, their application to large-scale engineering tasks suffers from significant computational overhead. And the fact that current models impose a uniform computational cost while physical fields exhibit vastly different complexities constitutes a fundamental mismatch, which is the root of this inefficiency. For instance, in turbulence flows, intricate vortex regions require deeper network processing compared to stable flows. To address this, we introduce a framework: Skip-Block Routing (SBR), a general framework designed for Transformer-based neural operators, capable of being integrated into their multi-layer architectures. First, SBR uses a routing mechanism to learn the complexity and ranking of tokens, which is then applied during inference. Then, in later layers, it decides how many tokens are passed forward based on this ranking. This way, the model focuses more processing capacity on the tokens that are more complex. Experiments demonstrate that SBR is a general framework that seamlessly integrates into various neural operators. Our method reduces computational cost by approximately 50% in terms of Floating Point Operations (FLOPs), while still delivering up to 2x faster inference without sacrificing accuracy.

From Uniform to Adaptive: General Skip-Block Mechanisms for Efficient PDE Neural Operators

TL;DR

Transformer-based neural operators for PDEs suffer from a uniform computation bottleneck across heterogeneous physical domains. The authors propose Skip-Block Routing (SBR), which decouples complexity assessment via a Global Router from an Adaptive Processing Backbone that progressively concentrates computation on the most important tokens, guided by a sparsity schedule. Across NS2D, Pipe, Airfoil, and Heat2D, SBR achieves roughly 50% FLOPs reduction and up to 2x end-to-end speedups with accuracy largely preserved, and in some cases improved, highlighting its practical impact for large-scale PDE solving. This physics-informed adaptive computation enables scalable, efficient neural operators suitable for industrial-scale PDE tasks.

Abstract

In recent years, Neural Operators(NO) have gradually emerged as a popular approach for solving Partial Differential Equations (PDEs). However, their application to large-scale engineering tasks suffers from significant computational overhead. And the fact that current models impose a uniform computational cost while physical fields exhibit vastly different complexities constitutes a fundamental mismatch, which is the root of this inefficiency. For instance, in turbulence flows, intricate vortex regions require deeper network processing compared to stable flows. To address this, we introduce a framework: Skip-Block Routing (SBR), a general framework designed for Transformer-based neural operators, capable of being integrated into their multi-layer architectures. First, SBR uses a routing mechanism to learn the complexity and ranking of tokens, which is then applied during inference. Then, in later layers, it decides how many tokens are passed forward based on this ranking. This way, the model focuses more processing capacity on the tokens that are more complex. Experiments demonstrate that SBR is a general framework that seamlessly integrates into various neural operators. Our method reduces computational cost by approximately 50% in terms of Floating Point Operations (FLOPs), while still delivering up to 2x faster inference without sacrificing accuracy.

Paper Structure

This paper contains 44 sections, 5 equations, 7 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Neural Operators for PDE Solving
  4. Adaptive Computation for Efficient Deep Learning
  5. Motivation
  6. Non-Uniform Complexity in PDE Tasks
  7. Concentrated Activations in Model Representations
  8. Methods
  9. Problem Formulation
  10. The Global Router Module
  11. The Adaptive Processing Backbone
  12. Sparsity Schedule.
  13. Structured Computational Flow.
  14. Experiments
  15. Experimental Setup and Evaluation Protocol
  16. ...and 29 more sections

Figures (7)

  • Figure 1: The heatmap on the right represents the intensity of change at this location across all time steps within the NS2D data; the two rows of images on the left respectively depict the temporal evolution of the most intense and most stable regions.
  • Figure 2: Transolver and GNOT model feature activation norms on the Pipe dataset, quantified as the L2 norm across all feature channels.
  • Figure 3: Flowchart of the Proposed SBR Method.
  • Figure 4: Load fluctuations of SBR and MoR under equal total capacity (Pipe)
  • Figure 5: Load fluctuations of SBR and MoR under equal total capacity (Airfoil)
  • ...and 2 more figures