An adaptive wavelet-based PINN for problems with localized high-magnitude source
| Authors | Himanshu Pandey & Ratikanta Behera |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2604.28180 |
| Download | |
| Categories | cs.LG |
Abstract
In recent years, physics-informed neural networks (PINNs) have gained significant attention for solving differential equations, although they suffer from two fundamental limitations, namely, spectral bias inherent in neural networks and loss imbalance arising from multiscale phenomena. This paper proposes an adaptive wavelet-based PINN (AW-PINN) to address the extreme loss imbalance characteristic of problems with localized high-magnitude source terms. Such problems frequently arise in various physical applications, such as thermal processing, electro-magnetics, impact mechanics, and fluid dynamics involving localized forcing. The proposed framework dynamically adjusts the wavelet basis function based on residual and supervised loss. This adaptive nature makes AW-PINN handle problems with high-scale features effectively without being memory-intensive. Additionally, AW-PINN does not rely on automatic differentiation to obtain derivatives involved in the loss function, which accelerates the training process. The method operates in two stages, an initial short pre-training phase with fixed bases to select physically relevant wavelet families, followed by an adaptive refinement that adapts scales and translations without populating high-resolution bases across entire domains. Theoretically, we show that under certain assumptions, AW-PINN admits a Gaussian process limit and derive its associated NTK structure. We evaluate AW-PINN on several challenging PDEs featuring localized high-magnitude source terms with extreme loss imbalances having ratios up to 10^{10}:1. Across these PDEs, including transient heat conduction, highly localized Poisson problems, oscillatory flow equations, and Maxwell equations with a point charge source, AW-PINN consistently outperforms existing methods in its class.
Engineering Breakdown
Plain English
This paper addresses a fundamental problem in physics-informed neural networks (PINNs): they struggle when solving differential equations with localized high-magnitude source terms, common in thermal processing, electromagnetics, and fluid dynamics. The authors propose Adaptive Wavelet-based PINNs (AW-PINN), which dynamically adjusts wavelet basis functions during training by analyzing residual and supervised loss to handle extreme loss imbalance across different scales of the problem. The key innovation is replacing fixed neural network activations with adaptive wavelets that can concentrate representational power where the physics is most complex, particularly in regions with sharp gradients or localized forcing.
Core Technical Contribution
The core novelty is introducing dynamic wavelet basis adaptation into the PINN framework to directly combat two well-known PINN limitations: spectral bias (neural networks struggle to learn high-frequency components) and loss imbalance (multiscale phenomena cause some loss terms to dominate training). Rather than using static ReLU or sine activations, AW-PINN learns which wavelet basis functions best represent different regions of the solution domain by monitoring both the PDE residual loss and supervised data loss. This adaptive mechanism automatically identifies and focuses computational capacity on localized regions where source terms create sharp transitions, enabling the network to simultaneously satisfy physics constraints across vastly different scales without manual loss weighting or curriculum learning tricks.
How It Works
The AW-PINN architecture takes input variables (spatial coordinates, time, or parameter values) and passes them through a learnable wavelet transform layer instead of standard linear-activation pairs. During each training iteration, the framework computes two loss components: (1) the PDE residual loss measuring how well the network satisfies the differential equation at collocation points, and (2) supervised loss from known boundary conditions or labeled data. The network evaluates which regions have high residual error and which wavelet basis functions best reduce loss in those regions. Based on this feedback, the adaptive mechanism updates the wavelet selection—essentially choosing which wavelets to emphasize and de-emphasize—allowing the model to concentrate its representational capacity exactly where localized source terms create the most challenging learning dynamics. The output is a solution field that accurately represents both the smooth background solution and sharp localized features.
Production Impact
For engineers building simulation systems, AW-PINN directly reduces the manual tuning burden in PINN deployment. Instead of hand-crafting loss weights or using domain-specific preprocessing to handle multiscale problems, this approach automatically learns what features matter where, making PINNs more practical for engineers unfamiliar with each specific physics domain. In production pipelines, you could replace fixed PINN architectures with AW-PINN when solving PDEs with localized forcing (e.g., thermal spikes in semiconductor processing, concentrated loads in structural analysis, fuel injection plumes in combustion). The trade-off is modest: you incur extra compute cost during training (wavelet selection and adaptation), and the approach assumes localized phenomena are the dominant challenge rather than, say, extremely long-range interactions. Memory overhead is minimal since wavelets are typically stored compactly, and inference speed remains similar to standard PINNs. Integration complexity is moderate—this requires careful initialization of wavelet parameters and monitoring of loss components, not just plugging in a pretrained model.
Limitations and When Not to Use This
The paper assumes that loss imbalance primarily arises from localized high-magnitude source terms; it may be less effective for problems where imbalance comes from fundamentally different physics scales (e.g., vastly different speeds of wave propagation in coupled systems). The approach requires careful tuning of how aggressively to adapt the wavelet basis during training—too aggressive and training destabilizes, too conservative and you get no benefit. The paper does not appear to address spectral bias as thoroughly as the abstract claims; wavelets help with localization but do not fundamentally solve the high-frequency learning problem that affects all neural networks. Additional follow-up work is needed on: (1) how to select the wavelet family (Morlet, Daubechies, etc.) for different PDE types, (2) theoretical guarantees on convergence with adaptive basis changes, and (3) scalability to very high-dimensional problems (>10 dimensions) where the wavelet advantage may diminish.
Research Context
This work builds on the substantial PINN literature initiated by Raissi et al. (2019) and addresses concrete pain points identified in subsequent work on PINN limitations. It relates to parallel efforts in adaptive neural network design (e.g., adaptive activation functions, dynamic networks) but specifically targets the PDE solver domain where loss balancing is notorious. The contribution sits between pure numerical methods (which handle multiscale problems well but are slower) and vanilla PINNs (which are fast but fragile on multiscale problems), filling a practical gap. This opens research directions on other adaptive basis families beyond wavelets, integration with domain decomposition for large-scale problems, and hybrid schemes combining AW-PINN with traditional solvers in critical regions.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::