LLM-Based Scientific Equation Discovery via Physics-Informed Token-Regularized Policy Optimization

AI summary

Overview

Research area: Machine learning for scientific discovery — specifically symbolic regression (distilling mathematical equations from data) using large language models.

Technical level: Advanced. The paper assumes familiarity with reinforcement learning (policy gradients, GRPO), transformer fine-tuning (LoRA), symbolic regression benchmarks, and turbulence modeling concepts.

Scope: A single framework, PiT-PO, that turns an LLM from a passive equation proposer into a self-improving equation discoverer by combining physics constraints with token-level penalties during reinforcement-learning updates.

What This Paper Is About

Symbolic regression tries to recover the exact mathematical equation behind a set of measurements, but the search space of possible formulas is astronomically large and formally NP-hard. Recent work uses LLMs to guess candidate equations, but almost all of it keeps the model frozen and steers it only through prompts, so the model never actually learns from what the search finds. This paper asks whether the LLM itself can be trained during the search — using physics awareness and fine-grained structural penalties — so that it generates better and simpler equations over time.

Key Contributions

1. In-search LLM evolution. PiT-PO updates the LLM's weights (via LoRA) while the evolutionary equation search is running, converting prompt-level verbal feedback into numerical gradient feedback that the model internalizes.

2. A dual-constraint reward system. Candidate equations are scored under two orthogonal pressures: hierarchical physical constraints (dimensional homogeneity, differentiability, plus domain-specific laws) that ensure scientific validity, and theorem-guided mathematical constraints that detect and penalize redundant terms.

3. The Support Exclusion Theorem with token-level penalties. The authors prove a bound (involving "internal" and "external interference" between basis functions) that identifies when a fitted term is almost certainly not part of the true equation. This is translated into a logarithmic token penalty applied only to the tokens that produced the redundant term.

4. A token-aware advantage estimator. Unlike standard GRPO, which assigns the same scalar advantage to every token in a generated answer, PiT-PO combines a globally standardized reward with a per-token pruning penalty, so credit is assigned at the granularity of individual terms.

Main Findings

• State-of-the-art on standard benchmarks. PiT-PO leads on the LLM-SR Suite across all four tasks (Oscillation 1 and 2, E. coli growth, Stress-Strain) and on all categories of LLM-SRBench (Chemistry, Biology, Physics, Material Science), recording the highest symbolic accuracy and lowest NMSE in nearly every cell of the comparison tables. It is the only method to recover the exact ground-truth equation for Oscillation 1.

• Small models beat large closed-source ones. Llama-3.2-1B and 3B backbones trained with PiT-PO match or outperform LLM-SR running on Mixtral 8x7B and GPT-4o-mini. On Oscillation 1, PiT-PO with a 3B model reaches 100% accuracy while LLM-SR with 4o-mini reaches 99.92% and with Mixtral 100.00% — competitive without proprietary APIs.

• The constraints break search stagnation. Training curves show that once error drops into a low regime, plain LLM-SR plateaus while PiT-PO repeatedly escapes the plateau with step-wise, orders-of-magnitude NMSE reductions. The authors attribute this to the moment when MSE alone stops being informative and the physics and redundancy signals take over.

• Physics-aware discovery in turbulence modeling. On flow over periodic hills at Re_b = 5600, PiT-PO learns an anisotropic Reynolds stress correction that is closer to Direct Numerical Simulation reference data than the standard k-ω SST model, with fewer non-physical extremes and better flow-field predictions.

• Every component is necessary. Ablations removing the physics penalty (w/o Phy) or the token regularization (w/o TokenReg) both degrade NMSE substantially and widen the in-distribution/out-of-distribution generalization gap.

• Computational cost is offset by convergence speed. The in-search fine-tuning adds overhead, but under an equal wall-clock budget PiT-PO retains a clear performance advantage over LLM-SR.

Methodology in Plain English

The researchers start with the standard LLM symbolic regression setup: the model writes candidate equations with unknown constants, a numerical optimizer fits those constants to the data, and the best equations are fed back as examples in later prompts. Their change is to also let the model learn from that feedback loop.

For each candidate equation the system computes a total reward that combines three things: how well it fits the data, a penalty for being structurally complex (measured by counting nodes in the equation's syntax tree), and a physics penalty. The physics penalty is deliberately gated — it only switches on once a candidate has already fit reasonably well, so early exploration is not choked off. The physical checks are layered: general ones like unit consistency and smoothness apply everywhere, while domain checks (realizability, wall boundary conditions, near-wall scaling, energy consistency) are added for the turbulence task.

The second ingredient is a mathematical test. Because the candidate terms are basis functions, the authors define a projection matrix between them and derive a bound on fitted coefficients. A term whose fitted coefficient falls below that bound is provably redundant — it is fitting noise or mimicking other terms rather than representing real structure. Each such term gets a penalty proportional to the negative log of its coefficient size, and that penalty is applied only to the tokens that generated the offending term.

The training loop runs in three repeating phases: parallel "islands" of search diversity each generate candidate equations; accumulated trajectories update the LLM through a modified GRPO objective using the token-aware advantage; and a selection step keeps the best candidates while resetting weak islands with high-performing seeds. The whole thing runs on a single RTX 3090 with 4-bit quantized Llama models.

Why This Matters

This work shifts the role of an LLM in scientific discovery from tool to trainee: instead of hoping a general-purpose model already knows the right physics, it teaches the model the physics of the specific problem it is solving. Because the method works well on small, locally runnable models, it also lowers the practical and financial barrier to automated equation discovery — an important equity consideration for research groups without access to frontier APIs.

Real-world applications:

• Turbulence closure modeling in CFD. Learning symbolic corrections to RANS stress tensors can improve aerodynamic simulation accuracy without the cost of full DNS, relevant to aircraft, wind turbine, and automotive design.
• Accelerated scientific modeling in chemistry and biology. Fast recovery of interpretable kinetics or population-dynamics equations supports drug development and bioprocess engineering where mechanistic understanding matters more than prediction alone.
• Materials science and engineering. Discovering compact, physically consistent constitutive laws from experimental stress-strain data helps characterize new alloys and composites.
• Climate and geophysical modeling. Symbolic expressions for subgrid processes are easier to audit and trust than black-box neural surrogates, which matters for regulatory and policy-facing models.

Industry relevance: Any organization that runs physics simulations — aerospace, energy, automotive, semiconductor manufacturing — has an interest in replacing expensive or inaccurate empirical sub-models with compact, physically valid symbolic alternatives. The ability to run this on consumer hardware makes pilot adoption feasible without new infrastructure.

Future Directions

• Generalizing the Support Exclusion Theorem beyond linear-in-basis-function formulations. The current bound assumes a fixed support set of basis functions with bounded true coefficients; extending it to nested or non-linear parameterizations is open.

• Automatic construction of domain-specific constraint sets. The turbulence constraints were hand-specified by experts. Whether an LLM can propose and validate domain priors on its own would broaden the method's applicability to fields lacking established invariants.

• Scaling analysis and stability. The paper evaluates 1B, 3B, and 8B backbones; how the approach behaves on much larger models, and whether the LoRA-based in-search updates remain stably beneficial at scale, is untested.

• Broader physical validation beyond periodic hills and synthetic benchmarks. Real experimental data with noise, sparse sampling, and unmodeled physics would test whether the physics gate and redundancy penalty hold up outside curated settings.

Target Audience

Researchers in machine learning for science, symbolic regression, and physics-informed ML; computational scientists in fluid dynamics or materials modeling looking for interpretable surrogate models; and RL practitioners interested in fine-grained token-level credit assignment. Readers without a background in reinforcement learning or symbolic regression will find the theoretical sections (especially the Support Exclusion Theorem and GRPO formulation) difficult, though the benchmark results and case study are accessible to a broader scientific audience.

Authors’ abstract

Symbolic regression aims to distill mathematical equations from observational data. Recent approaches have successfully leveraged Large Language Models (LLMs) to generate equation hypotheses, capitalizing on their vast pre-trained scientific priors. However, existing frameworks predominantly treat the LLM as a static generator, relying on prompt-level guidance to steer exploration. This paradigm fails to update the model's internal representations based on search feedback, often yielding physically inconsistent or mathematically redundant expressions. In this work, we propose PiT-PO (Physics-informed Token-regularized Policy Optimization), a unified framework that evolves the LLM into an adaptive generator via reinforcement learning. Central to PiT-PO is a dual-constraint mechanism that rigorously enforces hierarchical physical validity while simultaneously applying fine-grained, token-level penalties to suppress redundant structures. Consequently, PiT-PO aligns LLM to produce equations that are both scientifically consistent and structurally parsimonious. Empirically, PiT-PO achieves state-of-the-art performance on standard benchmarks and successfully discovers novel turbulence models for challenging fluid dynamics problems. We also demonstrate that PiT-PO empowers small-scale models to outperform closed-source giants, democratizing access to high-performance scientific discovery.