让神经算子按需计算，复杂区域多处理，节省一半算力。
From Uniform to Adaptive: General Skip-Block Mechanisms for Efficient PDE Neural Operators
- 根据输入复杂度动态决定每层跳过多少令牌，实现自适应计算。
- 在多个模型上降低约50%浮点运算量，推理速度提升2倍。
- 适合需要高效求解复杂偏微分方程的工程应用，如湍流模拟。
近年来，神经算子（NO）逐渐成为求解偏微分方程（PDEs）的热门方法。然而，其在大规模工程任务中的应用面临显著的计算开销。当前模型对所有区域施加均匀计算成本，而物理场实际复杂度差异巨大，这是效率低下的根本原因。例如，在湍流中，涡旋区域比稳定区更复杂，需要更深的网络处理。为此，我们提出跳块路由（SBR）框架，适用于基于Transformer的神经算子，可无缝嵌入其多层结构。SBR通过路由机制学习令牌的复杂度与排序，并在推理时决定后续层传递的令牌数量。这样，模型将更多计算资源集中于复杂区域。实验表明，该方法能广泛适配各类神经算子，在保持精度的前提下，将浮点运算量（FLOPs）降低约50%，推理速度提升至2倍。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。