用可训练的波动方程替代注意力，实现高效全局建模
Wave-PDE Nets: Trainable Wave-Equation Layers as an Alternative to Attention
- 用波动方程作为神经层，通过可学习的传播速度和阻尼模拟信息传播
- 在语言与视觉任务上性能媲美甚至超越Transformer，提速30%，内存降25%
- 具有物理可解释性，适合追求高效与可解释性的模型设计者
我们提出Wave-PDE Nets，一种以二阶波动方程的可微分仿真为基本操作的神经架构。每一层将隐藏状态作为连续场，在可训练的空间速度c(x)和阻尼γ(x)定义的介质中传播。基于FFT的辛谱求解器实现该传播，时间复杂度为O(n log n)。这种振荡且全局的机制为注意力和一阶状态空间模型提供了有力替代方案。我们证明单个Wave-PDE层是通用近似器。在语言与视觉基准测试中，Wave-PDE Nets达到或超过Transformer性能，同时展现出更优的实际效率：壁时减少最高达30%，峰值内存降低25%。消融实验确认辛积分和谱拉普拉斯对稳定性和性能的关键作用。对学习到的物理参数可视化显示，模型学会了直观的信息传播策略。这些结果表明，Wave-PDE Nets是一种计算高效、鲁棒性强且具有强物理先验偏见的架构。
原文摘要 · Abstract (English)
We introduce Wave-PDE Nets, a neural architecture whose elementary operation is a differentiable simulation of the second-order wave equation. Each layer propagates its hidden state as a continuous field through a medium with trainable spatial velocity c(x) and damping γ(x). A symplectic spectral solver based on FFTs realises this propagation in O(nlog n) time. This oscillatory, global mechanism provides a powerful alternative to attention and first-order state-space models. We prove that a single Wave-PDE layer is a universal approximator. On language and vision benchmarks, Wave-PDE Nets match or exceed Transformer performance while demonstrating superior practical efficiency, reducing wall-clock time by up to 30% and peak memory by 25%. Ablation studies confirm the critical role of symplectic integration and a spectral Laplacian for stability and performance. Visualizations of the learned physical parameters reveal that the model learns intuitive strategies for information propagation. These results position Wave-PDE Nets as a computationally efficient and robust architecture with a strong physical inductive bias.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。