Searcharxiv⌕ Search
Home/Archive/arXiv

arXiv · 2609.38916

Warm-starting PDE solvers with any-dimensional machine learning

Wilson G. Gregory·George A. Kevrekidis·Ben Blum-Smith·Soledad Villar
View original source ↗Read source document ↗Export metadata ↓

Abstract

Any-dimensional machine learning models, such as graph neural networks (GNNs), can be naturally trained and evaluated on inputs of different sizes and dimensions. Inspired by the GNN transferability literature, we show mathematical conditions under which a partial differential equation (PDE) learning-based solver can be trained in small dimensions and directly applied to solve a higher dimensional PDE in a zero-shot fashion. These conditions are based on symmetries in both the partial differential equation and the initial data. When the equations satisfy the symmetries but the data does not, which is the case for many PDEs arising from physics, we show that our theory gives a principled way of warm-starting low-dimensional PDE solvers for higher dimensional PDEs. We apply this method on the heat equation, Burgers' equation, and the compressible Navier--Stokes equations, improving the performance in both zero-shot and typical training regimes on high dimensional data. For example, we train a surrogate model on 2D Navier--Stokes data and achieve better results on 3D test data than a baseline surrogate model trained on 3D data, while only using 12$\%$ of the flops and 20$\%$ of the total data size.

Explore related subjects

stat.MLcs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Wilson G. Gregory, George A. Kevrekidis, Ben Blum-Smith, Soledad Villar. 2026-09-30. Warm-starting PDE solvers with any-dimensional machine learning. https://arxiv.org/abs/2609.38916

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Beyond Uncertainty Sets: Leveraging Optimal Transport to Extend Conformal Predictive Distributions to Multivariate Settings

Conformal prediction (CP) constructs uncertainty sets for model outputs with finite-sample coverage guarantees. Yet ranking scores is straightforward only when they are scalar-valued, limiting CP to real-valued scores or ad-hoc one-dimensional reductions. Vector-valued scores arise naturally in multi-output regression and model aggregation, where each predictor in an ensemble provides its own score. Optimal transport (OT) defines vector ranks and center-outward multivariate quantile regions, though generally with asymptotic coverage guarantees. Applying a fixed transport map learned from calibration data to a new point introduces an uncontrolled approximation error. We restore finite-sample, distribution-free coverage by conformalizing vector-valued OT quantile regions. Each candidate's rank is defined by transporting the calibration scores augmented with that candidate's score, preserving the symmetry needed for validity. This appears to require a continuum of OT problems. However, we prove that the optimal assignment is piecewise constant across a fixed polyhedral partition of score space. This lets us characterize the entire prediction set in $O(n^3)$ time, matching the cost of a single assignment solve. It also addresses a limitation of prediction sets: they indicate which outcomes are plausible, but not their relative likelihood. In one dimension, conformal predictive distributions (CPDs) fill this gap by producing a predictive distribution with finite-sample calibration. Extending CPDs beyond one dimension remained an open problem. We construct, to our knowledge, the first multivariate CPDs with finite-sample calibration: a center-outward predictive distribution whose derived uncertainty regions have conformal coverage. We present both conservative and exact randomized versions; the latter generalizes the classical Dempster-Hill procedure.

stat.ML↗
arXiv2026

Closing the Approximation Gap in Simulation-free Latent SDEs

Recovering dynamical systems from noisy observations is a recurring challenge across scientific domains, including neuroscience and physics. Latent stochastic differential equations (SDEs) address this by modeling the system as an unobserved state that evolves according to a learnable SDE and generates the observations. Variational inference (VI) provides a tractable objective for fitting latent SDEs. Traditional VI algorithms evaluate this objective by numerical simulation over a time discretization, trading fidelity for computational cost. A recent class of algorithms, simulation-free VI, sidesteps this tradeoff by parameterizing the posterior through its instantaneous marginals rather than its drift. In this work, we show that the efficiency of existing simulation-free VI algorithms comes at a price: their parameterizations restrict the approximate posterior to a subset of the SDEs available to simulation-based methods, degrading posterior inference and parameter learning. We propose Helmholtz-SDE, a simulation-free VI algorithm that closes this gap by optimizing over path laws compatible with a prescribed collection of marginals. Helmholtz-SDE recovers dynamics more faithfully than prior simulation-free methods, with the largest gains under high posterior uncertainty. It further matches the performance of simulation-based VI at a fraction of the runtime. Code is available at https://github.com/lindermanlab/helmholtz-sde .

stat.ML↗
arXiv2026

Sharp Statistical Rates for Asynchronous TD Learning with Markovian Data

We study the last iterate of standard tabular temporal-difference (TD) learning from a single trajectory of a finite Markov reward process. For discount factor $γ$, write $H=(1-γ)^{-1}$, and let $μ_{\min}$ and $t_{\operatorname{mix}}$ denote the minimum stationary probability and total-variation mixing time. We prove that last-iterate TD achieves sup-norm error at most $\varepsilon$ with high probability using$\widetilde O\left( \frac{H^3}{μ_{\min}\varepsilon^2} +\frac{t_{\operatorname{mix}}}{μ_{\min}} \right)$ transitions, for $0<\varepsilon\leq1$. This rate holds both for a constant step size selected for the target accuracy and for a decreasing schedule independent of the target accuracy and terminal time. The latter gives a simultaneous guarantee over all times beyond an explicit transient threshold. The statistical term retains the cubic effective-horizon dependence of synchronous TD, and the additive mixing transient has no extra horizon factor. The result allows non-reversible chains, arbitrary initial state distributions, and bounded rewards that may depend on the next state. The proof uses an anchored local Poisson equation in reverse time to control stochastic fluctuations without a mixing-time factor, and a hitting-time compensation identity to bound initialization error. The latter also yields a finer transient in terms of the worst expected reverse hitting time. A bound on the expected cumulative propagation mass extends this argument to decreasing step sizes. A three-state construction with known deterministic rewards gives matching minimax lower bounds for the statistical and mixing terms, up to logarithms, over specified model classes in a slow-mixing parameter regime.

stat.ML↗