# Towards Reasoning for PDE Foundation Models: A Reward-Model-Driven Inference-Time-Scaling Algorithm

arXiv: https://arxiv.org/abs/2509.02846 · HTML: https://mlsanity.com/paper/2509.02846/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion reasoning and test-time scaling, one of 17 directions in LLM reasoning.
- Builds on 17 of its 24 references that are in the map.
- First posted on arXiv 2 Sep 2025.

## Placed in

- LLM reasoning / [Diffusion reasoning and test-time scaling](https://mlsanity.com/lr/diffusion-reasoning-and-test-time-scaling/)
- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 17 of 17)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Improve Mathematical Reasoning in Language Models by Automated Process Supervision](https://mlsanity.com/paper/2406.06592/) (2024)
- [ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search](https://mlsanity.com/paper/2406.03816/) (2024)
- [PDE-Net: Learning PDEs from Data](https://mlsanity.com/paper/1710.09668/) (2017)
- [PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers](https://mlsanity.com/paper/2308.05732/) (2023)
- [PDE-Net 2.0: Learning PDEs from Data with A Numeric-Symbolic Hybrid Deep Network](https://mlsanity.com/paper/1812.04426/) (2018)
- [Hidden Physics Models: Machine Learning of Nonlinear Partial Differential Equations](https://mlsanity.com/paper/1708.00588/) (2017)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [Self-Evaluation Guided Beam Search for Reasoning](https://mlsanity.com/paper/2305.00633/) (2023)
- [DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training](https://mlsanity.com/paper/2403.03542/) (2024)
- [Deductive Beam Search: Decoding Deducible Rationale for Chain-of-Thought Reasoning](https://mlsanity.com/paper/2401.17686/) (2024)
- [Active Inference AI Systems for Scientific Discovery](https://mlsanity.com/paper/2506.21329/) (2025)

## Built on by (most cited first, 4 of 4)

- [Tadpole: Autoencoders as Foundation Models for 3D PDEs with Online Learning](https://mlsanity.com/paper/2605.15284/) (2026)
- [PDE foundation model-accelerated inverse estimation of system parameters in inertial confinement fusion](https://mlsanity.com/paper/2603.04606/) (2026)
- [Enhancing Transformer Representations of Symbolic ODE Expressions](https://mlsanity.com/paper/2609.24746/) (2026)
- [Cross-Domain Pretraining for Steady-State Neural CFD Surrogates](https://mlsanity.com/paper/2610.10398/) (2026)

## Abstract (arXiv)

Partial Differential Equations (PDEs) are the bedrock for modern computational sciences and engineering, and inherently computationally expensive. While PDE foundation models have shown much promise for simulating such complex spatio-temporal phenomena, existing models remain constrained by the pretraining datasets and struggle with auto-regressive rollout performance, especially in out-of-distribution (OOD) cases. Furthermore, they have significant compute and training data requirements which hamper their use in many critical applications. Inspired by recent advances in ``thinking" strategies used in large language models (LLMs), we introduce the first test-time computing (TTC) strategy for PDEs that utilizes computational resources during inference to achieve more accurate predictions with fewer training samples and smaller models. We accomplish this with two types of reward models that evaluate predictions of a stochastic based model for spatio-temporal consistency. We demonstrate this method on compressible Euler-equation simulations from the PDEGym benchmark and show that TTC captures improved predictions relative to standard non-adaptive auto-regressive inference. This TTC framework marks a foundational step towards more advanced reasoning algorithms or PDE modeling, inluding building reinforcement-learning-based approaches, potentially transforming computational workflows in physics and engineering.