JAWS: Enhancing Long-term Rollout of Neural Operators via Spatially-Adaptive Jacobian Regularization

• URL: http://arxiv.org/abs/2603.05538v1
• Date: Wed, 04 Mar 2026 06:15:09 GMT
• Title: JAWS: Enhancing Long-term Rollout of Neural Operators via Spatially-Adaptive Jacobian Regularization
• Abstract summary: Jacobian-Adaptive Weighting for Stability (JAWS) is a probabilistic regularization strategy designed to mitigate these limitations.<n>By framing operator learning as Maximum A Posteriori (MAP) estimation with spatially heteroscedastic uncertainty, JAWS dynamically modulates the regularization strength based on local physical complexity.<n> Experiments demonstrate that this spatially-adaptive prior serves as an effective spectral pre-conditioner, which reduces the base operator's burden of handling high-frequency instabilities.
• Score: 0.0
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Data-driven surrogate models improve the efficiency of simulating continuous dynamical systems, yet their autoregressive rollouts are often limited by instability and spectral blow-up. While global regularization techniques can enforce contractive dynamics, they uniformly damp high-frequency features, introducing a contraction-dissipation dilemma. Furthermore, long-horizon trajectory optimization methods that explicitly correct drift are bottlenecked by memory constraints. In this work, we propose Jacobian-Adaptive Weighting for Stability (JAWS), a probabilistic regularization strategy designed to mitigate these limitations. By framing operator learning as Maximum A Posteriori (MAP) estimation with spatially heteroscedastic uncertainty, JAWS dynamically modulates the regularization strength based on local physical complexity. This allows the model to enforce contraction in smooth regions to suppress noise, while relaxing constraints near singular features to preserve gradients, effectively realizing a behavior similar to numerical shock-capturing schemes. Experiments demonstrate that this spatially-adaptive prior serves as an effective spectral pre-conditioner, which reduces the base operator's burden of handling high-frequency instabilities. This reduction enables memory-efficient, short-horizon trajectory optimization to match or exceed the long-term accuracy of long-horizon baselines. Evaluated on the 1D viscous Burgers' equation, our hybrid approach improves long-term stability, shock fidelity, and out-of-distribution generalization while reducing training computational costs.

Related papers

• When Does Scale-Invariant Optimization Become Unstable? An Exact Schedule Law with Weight Decay [6.625270757109195]
  We show that a single quantity captures all schedule and decay forcing, while norm growth induces an opposing geometric self-quenching effect.<n>This yields a sharp boundary that cleanly separates contraction- and expansion-dominated effective learning rate regimes.
  arXiv  Detail & Related papers  (2026-09-08T17:46:17Z)
• Motion Planning with Model-Based Diffusion via Constraint Optimization and Adaptive Scheduling [7.761673975696802]
  Single-Robot Motion Planning (SRMP) is challenging under complex constraints and computational limits.<n>Recent Model-Based Diffusion (MBD) recasts the SRMP to guide samples from a posterior diffusion trajectories.<n>Our method achieves higher safety & success rates, faster convergence, and lower final costs than baseline planners.
  arXiv  Detail & Related papers  (2026-07-16T01:04:32Z)
• Clipping Bottleneck: Stabilizing RLVR via Stochastic Recovery of Near-Boundary Signals [83.0127582612634]
  Near-boundary Rescue (NSR) is a minimal, plug-and-play modification that retains out-of-bound tokens to recover lost signals.<n>NSR substantially improves training stability and delivers consistent gains over strong baselines such as DAPO and GSPO.
  arXiv  Detail & Related papers  (2026-05-21T16:45:31Z)
• Taming the Instability: A Robust Second-Order Optimizer for Federated Learning over Non-IID Data [77.19117648492104]
  We present Federated Robust Curvature Optimization (FedRCO), a novel second-order optimization framework designed to improve convergence speed and reduce communication cost.<n>FedRCO addresses these challenges by integrating an efficient approximate curvature with a provable stability mechanism.<n>We show that FedRCO achieves superior robustness against diverse non-IID scenarios while achieving higher accuracy and faster convergence than both state-of-the-art first-order and second-order methods.
  arXiv  Detail & Related papers  (2026-03-30T11:37:46Z)
• LEVDA: Latent Ensemble Variational Data Assimilation via Differentiable Dynamics [6.953554594702111]
  We propose Latent Ensemble Data Assimilation (LEVDA), an ensemble-space variational smoother.<n>It assimilates states and unknown parameters without the need for adjoint code or auxiliary observation-to-latent encoders.<n>It substantially improved assimilation accuracy and computational efficiency compared to full-state 4DEnVar.
  arXiv  Detail & Related papers  (2026-02-23T00:54:59Z)
• An Adaptive Differentially Private Federated Learning Framework with Bi-level Optimization [10.218291445871435]
  Federated learning enables collaborative model training across distributed clients while preserving data privacy.<n>In practical deployments, device heterogeneity, non-independent, and identically distributed (Non-IID) data often lead to highly unstable and biased gradient updates.<n>We propose an adaptive differentially private federated learning framework that explicitly targets model efficiency under heterogeneous and privacy-constrained settings.
  arXiv  Detail & Related papers  (2026-02-06T16:27:33Z)
• Breaking the Stochasticity Barrier: An Adaptive Variance-Reduced Method for Variational Inequalities [0.0]
  We propose VR-A-A (VarianceReduced-Ascent with Armijo) as a novel algorithm for non-con optimization tasks.<n>We show that our method effectively suppresses limit cycles and accelerates convergence with reduced dependence on manual learning scheduling.
  arXiv  Detail & Related papers  (2026-01-30T14:43:07Z)
• The Curious Case of In-Training Compression of State Space Models [49.819321766705514]
  State Space Models (SSMs) tackle long sequence modeling tasks efficiently, offer both parallelizable training and fast inference.<n>Key design challenge is striking the right balance between maximizing expressivity and limiting this computational burden.<n>Our approach, textscCompreSSM, applies to Linear Time-Invariant SSMs such as Linear Recurrent Units, but is also extendable to selective models.
  arXiv  Detail & Related papers  (2025-10-03T09:02:33Z)
• SUMO: Subspace-Aware Moment-Orthogonalization for Accelerating Memory-Efficient LLM Training [25.244065166421517]
  Low-rank gradient-based optimization methods have significantly improved memory efficiency during the training of large language models (LLMs)<n>These methods primarily emphasize memory savings, often overlooking potential acceleration in convergence.<n>In this paper, we propose SUMO (Subspace-Aware Moment-Orthogonalization), an norm that employs exact singular value decomposition.<n>We show that SUMO accelerates convergence, enhances stability, improves performance, and reduces memory requirements by up to 20% compared to state-of-the-art methods.
  arXiv  Detail & Related papers  (2025-05-30T16:08:40Z)
• Adaptive Deadline and Batch Layered Synchronized Federated Learning [66.93447103966439]
  Federated learning (FL) enables collaborative model training across distributed edge devices while preserving data privacy, and typically operates in a round-based synchronous manner.<n>We propose ADEL-FL, a novel framework that jointly optimize per-round deadlines and user-specific batch sizes for layer-wise aggregation.
  arXiv  Detail & Related papers  (2025-05-29T19:59:18Z)
• Generative emulation of chaotic dynamics with coherent prior [0.129182926254119]
  We present an efficient generative framework for dynamics emulation, unifying principles of turbulence with diffusion-based modeling: Cohesion.<n>Specifically, our method estimates large-scale coherent structure of the underlying dynamics as guidance during the denoising process.<n>Cohesion superior long-range forecasting skill can efficiently generate physically-consistent simulations, even in the presence of partially-observed guidance.
  arXiv  Detail & Related papers  (2025-04-19T11:14:40Z)
• Efficient Transformed Gaussian Process State-Space Models for Non-Stationary High-Dimensional Dynamical Systems [49.819436680336786]
  We propose an efficient transformed Gaussian process state-space model (ETGPSSM) for scalable and flexible modeling of high-dimensional, non-stationary dynamical systems.<n>Specifically, our ETGPSSM integrates a single shared GP with input-dependent normalizing flows, yielding an expressive implicit process prior that captures complex, non-stationary transition dynamics.<n>Our ETGPSSM outperforms existing GPSSMs and neural network-based SSMs in terms of computational efficiency and accuracy.
  arXiv  Detail & Related papers  (2025-03-24T03:19:45Z)
• Linear-Time User-Level DP-SCO via Robust Statistics [55.350093142673316]
  User-level differentially private convex optimization (DP-SCO) has garnered significant attention due to the importance of safeguarding user privacy in machine learning applications.<n>Current methods, such as those based on differentially private gradient descent (DP-SGD), often struggle with high noise accumulation and suboptimal utility.<n>We introduce a novel linear-time algorithm that leverages robust statistics, specifically the median and trimmed mean, to overcome these challenges.
  arXiv  Detail & Related papers  (2025-02-13T02:05:45Z)
• Towards Continual Learning Desiderata via HSIC-Bottleneck Orthogonalization and Equiangular Embedding [55.107555305760954]
  We propose a conceptually simple yet effective method that attributes forgetting to layer-wise parameter overwriting and the resulting decision boundary distortion. Our method achieves competitive accuracy performance, even with absolute superiority of zero exemplar buffer and 1.02x the base model.
  arXiv  Detail & Related papers  (2024-01-17T09:01:29Z)
• Achieving Constraints in Neural Networks: A Stochastic Augmented Lagrangian Approach [49.1574468325115]
  Regularizing Deep Neural Networks (DNNs) is essential for improving generalizability and preventing overfitting. We propose a novel approach to DNN regularization by framing the training process as a constrained optimization problem. We employ the Augmented Lagrangian (SAL) method to achieve a more flexible and efficient regularization mechanism.
  arXiv  Detail & Related papers  (2023-10-25T13:55:35Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.