EqGPT

A pretrained transformer reads the field and proposes equations, then fine-tunes on reward.

algorithm

eqgpt

score

EqGPT reward max

cost

heavy

derivatives

finite_diff

equation form

EVOLUTION HOMOGENEOUS

data layout

grid scattered

Usage

import kd
from kd import EqGPTConfig

dataset = kd.generate_burgers_data(nx=256, nt=101, nu=0.1, seed=42)
model = kd.Model(algorithm="eqgpt", generations=15, config=EqGPTConfig.burgers_preset())
model.fit(dataset)

EqGPT proposes from a pretrained transformer, so a run needs the published checkpoint. Point KD_EQGPT_ASSET_DIR at the directory holding it, or pass EqGPTConfig(weights_path=...). The checkpoint is distributed by the authors of the paper below and is used here with their permission.

config= is required rather than optional: sparsity_alpha weights the penalty on the number of terms in the reward, and its working value depends on the size of the coefficients being recovered, so there is no default that fits every equation. EqGPTConfig.burgers_preset() carries the value used below, 0.02.

Main parameters

Parameter	Default	Resume	Description
`samples_per_epoch`	`400`	resume-safe	GPT samples per Runner iteration.
`top_k`	`10`	init-only	Elite-pool and fine-tune slice size.
`sparsity_alpha`	required	init-only	Problem-specific sparsity weight.
`finetune_lr`	`1e-05`	resume-safe	GPT fine-tuning learning rate.
`exploration_rate`	`0.2`	resume-safe	Sampling exploration rate.

All fields of EqGPTConfig

Field	Type	Default
`sparsity_alpha`	`float`	required
`seed`	`int`	`0`
`samples_per_epoch`	`int`	`400`
`top_k`	`int`	`10`
`finetune_lr`	`float`	`1e-05`
`finetune_steps`	`int`	`5`
`exploration_rate`	`float`	`0.2`
`max_length`	`int`	`49`
`variables`	`tuple[str, ...] \\| None`	`None`
`start_words`	`tuple[str, ...]`	`('S', 'ut', '+')`
`masked_tokens`	`frozenset[int]`	`frozenset()`
`weights_path`	`Path \\| None`	`None`
`asset_dir`	`Path \\| None`	`None`
`case_filter`	`str \\| None`	`None`
`wave_pkl_path`	`Path \\| None`	`None`
`v1_asset_dir`	`Path \\| None`	`None`
`reward_points_per_window`	`int`	`50`
`coeff_points_per_window`	`int`	`100`
`primary_case`	`str \\| None`	`None`
`steady`	`bool`	`False`
`steady_activation`	`Literal['sin', 'rational'] \\| None`	`None`
`steady_boundary_delete_num`	`int \\| None`	`None`
`steady_polar_eval`	`bool`	`False`
`steady_constant_column`	`bool`	`False`
`steady_train_points`	`int`	`10000`
`steady_validate_points`	`int`	`1000`
`steady_train_iters`	`int`	`50000`
`steady_surrogate_seed`	`int`	`525`

Worked example

The run below uses synthetic Burgers data: a 256 × 101 grid over x in [-1, 1] with periodic boundaries and t in [0, 1], started from u(x, 0) = -sin(pi x) and integrated at a viscosity of 0.1. The nonlinear term carries the profile toward the origin and steepens it there (the largest |u_x| of the run, 4.1, falls at t = 0.3), while diffusion flattens the whole field, from an amplitude of 1.0 at t = 0 to 0.32 at t = 1.

Heatmap of the Burgers field u(x, t) — The observed field u over x and t.

The equation returned by the run:

	Equation
Discovered	\(u_t = -0.9986\,u\,u_x +0.09995\,u_{xx}\)
Reference	\(u_t = -u\,u_x +0.1\,u_{xx}\)

Both terms and both coefficients match. The transformer proposes a whole equation at a time as a sequence of tokens, so no term was assembled from a candidate library. The terms come from the distribution the model learned over token sequences, and the coefficients from the least-squares fit inside the reward.

Measured u_t, predicted u_t, and the residual between them — The discovered equation evaluated against the data. The first two panels share a color scale and are indistinguishable by eye; the third is their difference, at most 7.0e-02 against a 2.25 peak of u_t, and confined to the first time step.

Result interpretation

model.best_score_ is EqGPT's own reward, (1 - sparsity_alpha * log10(k)) * R² for an equation of k columns: higher is better. The R² is measured on the time derivative the equation predicts, and the factor in front discounts it by the number of terms, so a longer equation has to fit better to score the same. This run ends at 0.9904, with an nmse of 4.4e-05.

The coefficients are the algorithm's own, not a platform refit: the same least-squares solve that produces the R² produces them, and model.result_.equation carries them term by term.

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

Best EqGPT reward and the pool's worst reward per epoch — The best reward is 0.9904 at the first epoch and unchanged over the remaining 14: the pretrained proposer writes the answer in its first batch. The elite pool behind it continues to improve, with its worst member starting at 0.8803 and reaching 0.9879 by the third epoch.

Pool reward spread and fine-tuning loss panels — Left: the reward band of the elite pool, best against median against worst. The band closes as the fine-tune concentrates the proposals, and the median settles at 0.9879. Right: the cross-entropy loss of each epoch's fine-tune on that pool, from 1.267 to 0.7587, not monotonically: it rises again around epoch 9 and is still falling at the last epoch.

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

Method

EqGPT is a transformer trained to write equations. It was pretrained on a corpus of equations collected from mathematical handbooks, so it starts a run already able to produce well-formed physical laws, and the search is a matter of directing it rather than of assembling terms.

One epoch samples a batch of sentences from the model, token by token, keeping only sequences the grammar accepts as an equation. Each is converted to right-hand-side terms, and the platform builds the term columns from the data. The reward is then a least-squares fit of the time derivative onto those columns, scored by centered R² and discounted by sparsity_alpha times the log of the column count. The best candidates so far are kept as an elite pool, and each epoch ends by fine-tuning the model on that pool for a few Adam steps, so the next epoch's samples resemble the ones that scored well. One epoch is one KD iteration, and generations sets how many are run.

References

Xu et al. (2025). "Generative discovery of partial differential equations by learning from math handbooks". Nat. Commun. 16, 10255. Paper

Code: woshixuhao/EqGPT

Implementation note: KD drives the published checkpoint through its own sampler and reward, and reproduces the paper's single-case pipeline. The same plugin also runs the paper's multi-case wave-breaking setup, where one structure is scored against 12 experiments at once; that path is shown under Examples.