LLM4ED

A large language model proposes and revises candidate equations round by round, filtered by a sparse reward.

algorithm

llm4ed

score

LLM4ED sparse reward max

cost

medium

derivatives

finite_diff

equation form

EVOLUTION

data layout

grid

Usage

import kd
from kd import Llm4edConfig

dataset = kd.load_llm4ed_heat()
model = kd.Model(
    algorithm="llm4ed",
    generations=12,
    config=Llm4edConfig(base_url="https://api.openai.com/v1"),
)
model.fit(dataset)

Main parameters

Parameter	Default	Resume	Description
temperature	0.8	resume-safe	LLM sampling temperature.
max_tokens	1024	resume-safe	Maximum decoded tokens per request.
stop_threshold	0.995	resume-safe	Reward threshold for completion.
reward_limit	0.5	resume-safe	Minimum admitted sparse reward.
pool_size	5	resume-safe	Elite-pool capacity.

All fields of Llm4edConfig

Field	Type	Default
temperature	float	0.8
max_tokens	int	1024
stop_threshold	float	0.995
reward_limit	float	0.5
pool_size	int	5
init_num	int	20
samples_per_epoch	int	8
max_llm_calls_per_propose	int	50
max_llm_calls_per_run	int	200
seed	int	0
model	str	gpt-4o-mini
base_url	str \| None	None
tape_record_path	str \| None	None
extra_body	dict[str, Any] \| None	None

Search space

LLM4ED asks a language model for coefficient-free equations over a fixed symbol library and fits their coefficients by sparse regression. It operates on one scalar field with exactly one spatial axis and one time axis. The symbols are the field u, derivatives u_x, u_xx, u_xxx, and spatial coordinate x. The prompts specify +, -, *, /, squares and cubes; the parser also accepts integer powers through five. Explicit spatial dependence is expressible. Extra fields, other-axis derivatives, mixed space-time derivatives, higher spatial derivatives, trigonometric functions, exponentials, square roots and standalone constants are outside this route's symbol library. Extra fields and other spatial axes are refused before proposal generation.

Operand construction needs at least four spatial and three time points. Spacing is taken from x[2] - x[1] and t[1] - t[0], so use uniform axes. The method itself does not smooth or train a surrogate; the platform's savgol_window applies before differentiation. Higher derivatives therefore depend on measurement noise and resolution.

Method

The language model receives the symbol library and scored proposal history, not the observed arrays. The first round requests a numbered batch of coefficient-free equations. Later rounds alternate an optimization prompt, which lists earlier equations in ascending score order and asks for improved ones, with an evolution prompt that requests selection, crossover and mutation of term sets.

Each proposal is parsed with SymPy, expanded into additive terms, and executed as design-matrix columns over second-order finite-difference operands. Division allows reciprocals, but a zero denominator that makes a column non-finite rejects that candidate. STRidge sweeps a tolerance ladder to fit coefficients; fitted coefficient magnitudes outside [1e-4, 1e4] also reject a candidate.

A round resamples until enough distinct-scoring survivors arrive or its call budget is exhausted. Scores at or below reward_limit and scores already kept in the round are discarded. The best samples_per_epoch survivors are offered to an elite pool of pool_size; after filling, it admits only a strict improvement over its minimum with an unseen score. The next prompt receives that batch and the pool as it stood before the update. One round is one KD iteration, controlled by generations. Per-round and per-run call limits bound endpoint requests. Once the run limit is exhausted, subsequent rounds can return no proposals; a stop threshold above the attainable reward cannot end the search early.

Result interpretation

The sparse reward is (1 - 0.01 * k) / (1 + RMSE / std(target)), where k counts fitted nonzero coefficients and is capped at five. It is rounded to four decimals before ranking and duplicate filtering. A one-term perfect fit reaches 0.99, three terms reach 0.97 and four terms reach 0.96. The same reward therefore implies different errors for different support sizes. Before rounding, the inverse is nmse = ((1 - 0.01 * k) / reward - 1)**2.

The returned law contains only surviving structural terms and their native sparse-fit coefficients. Terms whose coefficients were zeroed are removed from both lists. model.result_.equation carries the aligned values, and NMSE and MSE use the finite-difference domain that supplied the candidate columns.

The per-round recorder includes pool_best, pool_median, pool_worst, n_valid, n_invalid and n_llm_calls. The admitted count n_valid is measured after truncation and saturates at samples_per_epoch. n_invalid counts dropped scoring events, broken down into eight series: n_invalid_undefined_operands, n_invalid_undefined_operators, n_invalid_inexpressible, n_invalid_parse, n_invalid_non_finite, n_invalid_lstsq, n_invalid_abnormal_coef and n_invalid_other. They sum to the total in each round; unused causes record zero, and other collects labels outside those named categories. Pool progress and admissions help distinguish proposal failures from valid but poorly fitting equations.

Python setup

Endpoint configuration

Set an OpenAI-compatible base_url and OPENAI_API_KEY for the default client; the OpenAI SDK is installed with KD. This is the algorithm that calls an external language-model service during search. Alternatively, pass a provider= implementing kd.llm.LLMProvider: a prompt goes in and text comes out. A custom client can set timeouts and headers, use BudgetedProvider for metering, or use TapeRecordingProvider to save a JSONL tape for offline replay. The authors' mirrored datasets download on first load.

References

Du et al. (2024). "Large language models for automatic equation discovery of nonlinear dynamics". Phys. Fluids 36, 097121. Paper · arXiv:2405.07761

Code: menggedu/EDL

Implementation notes. KD ports the reference's finite-difference templates, STRidge tolerance sweep, coefficient gate, reward formula and rounding. The NumPy STRidge port matches its reference on identical design matrices. Prompt term permutation uses a plugin-owned seeded generator, including on resume. Columns use KD's executor instead of NumPy lambdify, so floating-point differences can affect candidates near the coefficient gate. Single-character candidates are scored rather than skipped. Constructs outside the prompt vocabulary, such as trigonometric functions and non-integer exponents, are rejected explicitly. The call budgets are KD additions. The paper used GPT-3.5-turbo; the KD client accepts OpenAI-compatible endpoints. The method's contribution is the prompting loop that generates equations and evaluates them by sparse regression.

Worked example

The run below uses the heat-equation dataset from the LLM4ED authors, mirrored on Hugging Face Hub: a 200 × 201 grid over x in [0, 9.95] with periodic boundaries and t in [0, 10]. The initial profile is a localized oscillation centred near x = 5, a positive lobe at x = 4.55 against a negative one at x = 5.45 with weaker lobes to either side, and diffusion spreads and flattens it: the amplitude falls from 0.89 at t = 0 to 0.049 at t = 10, half of it gone by t = 1.45.

The observed field over its physical spatial and time coordinates.

The recorded requests and responses were replayed offline for twelve rounds. The original endpoint served nvidia/llama-3.3-nemotron-super-49b-v1.5, with a limit of six calls per round and seed 42. The equation returned by the run:

	Equation
Discovered	\(u_{t} = 0.05014 u_{xx}\)
Reference	\(u_t = 0.05\,u_{xx}\)

The term and the coefficient match. The model wrote a two-term right-hand side, u_xx + u u_x; the sparse regression fitted the second coefficient to zero, so one term survives and the reward counts one term.

The fitted derivative target, its equation prediction, and predicted minus measured values. The first two panels share a color scale; the residual has its own scale, with any clipping disclosed on the figure.

Example diagnostics

This run has LLM4ED sparse reward 0.9889 and NMSE 1.198e-06.

The search is recorded as well. kd.VizEngine renders the per-round record; alongside the universal convergence curve, LLM4ED contributes three diagnostic panels of its own (the full list is under Visualization):

The best native score is shown above the algorithm's population diagnostic, when one is recorded. Each panel retains its own score scale.

The algorithm's diagnostic panels, using the same data and plot functions as the full report.

viz = kd.VizEngine(output_dir="out/heat")
viz.render_all(model.result_, algorithm=model.algorithm_, dataset=dataset)

Full report

Full HTML report from the same run, with all figures, captions and raw result data.