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

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.

• The paper introduces Skip-Block Routing, a mechanism that adaptively prunes tokens based on their computed importance and physical complexity.
• The paper leverages a global router and an adaptive processing backbone with a tailored sparsity schedule to optimize both efficiency and accuracy.
• The paper demonstrates up to 50% reduction in FLOPs with maintained or improved predictive accuracy, highlighting practical gains for large-scale PDE modeling.

Skip-Block Routing: Adaptive Computation for Transformer-based PDE Neural Operators

Motivation and Problem Statement

Transformer-based neural operators (NOs) have become a primary approach for data-driven modeling of Partial Differential Equations (PDEs), offering high predictive performance across complex scientific domains. However, these models impose uniform, full-depth computation across all spatial and temporal regions, misaligning resource allocation with the heterogeneous complexity characteristic of physical fields. In large-scale engineering contexts, such inefficiency is prohibitive due to the excessive computational overhead, especially where repeated model inference is required. The paper identifies and formalizes this "uniform computation" bottleneck as the root cause of the mismatch between model computation and physical task complexity.

Two phenomena further motivate the need for adaptivity: (1) physical complexity is typically concentrated in sparse regions with high temporal and spatial gradients, whereas most of the domain is relatively simple; (2) transformer-based neural operators reveal highly non-uniform activation energy distributions, with deep processing consistently required only for a minority of critical regions. These insights directly motivate designing mechanisms to focus computation adaptively according to region-specific complexity.

Skip-Block Routing (SBR) Framework

The proposed SBR framework introduces efficient, adaptive computation into transformer-based NOs through two primary innovations: a learned static global router and an adaptive backbone controlled via a user-defined sparsity schedule.

Global Router Module

The router is a lightweight, differentiable module operating on the complete input field (tokens), designed as a single feed-forward layer with sigmoid normalization. It computes per-token importance scores reflecting their anticipated physical or computational complexity. The router emits a static, task-dependent ranking of all tokens, forming a persistent computational priority schedule for the entire inference.

Adaptive Processing Backbone

The backbone leverages the router's ranking to enact structured, depth-variant sparsification. At each layer, only the top-$k$ most important tokens (as determined via a sparsity schedule) are selected for expensive operations (self-attention, MLP), while a residual connection ensures inactive token information retention. The schedule typically retains full participation in shallow layers for contextual aggregation and prunes aggressively in deeper layers, focusing computation on complex regions. Crucially, all operations are fully differentiable, allowing end-to-end training under the standard PDE loss without auxiliary constraints or reinforcement mechanisms.

Empirical Evaluation and Analysis

Extensive empirical studies on standard PDE benchmarks (NS2D, Pipe, Airfoil, Heat2d) and state-of-the-art operators (GNOT, OFormer, Transolver, IPOT) validate the efficacy and generality of SBR.

Key findings:

• Computational Cost: SBR reduces backbone FLOPs by approximately 50% on average, yielding end-to-end speedups of up to 2x (and as high as 4.46x in deep models).
• Predictive Accuracy: SBR-enhanced models maintain or slightly improve accuracy versus dense baselines (e.g., SBR-GNOT and SBR-OFormer achieve lower errors on the Pipe dataset). Degradations, where observed (NS2D/Transolver), are marginal and linked to the task/model suitability for token pruning.
• Regularization: Focusing depth on high-complexity tokens can act as a regularizer, enhancing generalization.
• Router Effectiveness: Ablations replacing the learned router with random token selection result in significant accuracy degradation (up to 20% in some models), underscoring the criticality of informed, physics-aligned routing.

Theoretical analysis and controlled experiments further confirm the hypothesis that deeper processing is most beneficial for physically complex regions, with the largest gains observed in tokens exhibiting high ground-truth gradient magnitudes.

Comparative and Ablation Studies

SBR is contrasted with Mixture-of-Recursions (MoR)-style routing. While MoR assigns per-token depths, it yields unstable per-layer token counts, higher accuracy loss, and less predictable computational loads. SBR's deterministic, schedule-guided adaptivity guarantees stable throughput and minimal variance, a strong advantage for deployment in real-world platforms with hard resource budgets.

The structure of the sparsity schedule is also critical. Schedules concentrating computational effort in middle layers yield superior accuracy, aligning with prior findings on functional specialization in deep transformers. This reflects the necessity of maintaining non-local interactions at intermediate depths for PDE modeling.

Practical and Theoretical Implications

SBR advances transformer-based NOs towards practical, scalable deployment in computational science and engineering. The hardware-friendly, static routing supports easy implementation and integration into existing modeling pipelines. SBR’s focus on static, importance-based plans eliminates the variance and instability of purely dynamic schemes, aligning well with high-throughput and latency-sensitive inferential applications.

Theoretically, the work underscores that optimal computation for PDEs should reflect the local physical state, aligning with the spirit of adaptive mesh refinement in numerical solvers but implemented directly in the neural operator architecture. The SBR principle—decoupling region complexity assessment from the main computation, and then adaptively assigning network depth—could motivate further research into physics-informed routing mechanisms, integrating domain invariances, or hybridizing static and lightweight dynamic adaptation.

Conclusion

Skip-Block Routing presents a significant step toward efficiency-aligned transformer-based PDE solvers by enabling structured, model-agnostic adaptive computation within neural operators. It delivers major cost savings, strong accuracy preservation, and predictable computation, facilitating deployment in resource-constrained scientific and engineering workflows. Future research may explore dynamic or physics-integrated routers, as well as extend these principles to more general operator learning and physical simulation tasks.

For comprehensive technical details and full methodology, refer to "From Uniform to Adaptive: General Skip-Block Mechanisms for Efficient PDE Neural Operators" (2511.00032).