Can AI agents build a modular differentiable Earth system model in JAX?

Agents wrote code; humans defined correctness

The question behind legoESM is not whether a language model can produce plausible Fortran, but whether agent-written code can be trusted as a physical model. The paper's answer rests on a division of labour. Humans specified the scientific contract — the governing logic, conserved quantities and evaluation criteria — and the agents recast decades of community parameterizations and numerical methods into one unified JAX framework. Trust came from verification, not from the agents themselves: automated continuous integration enforced the physical contracts, and community benchmark cases supplied unambiguous targets that the agents iterated against.

The scale of that effort is part of the evidence. The infrastructure, developed over 9 months, comprised about 620,000 lines of model code in the packages, about 824,000 lines of tests, and about 555,000 lines of simulation and related code. The authors are explicit that adversarial agent review and persistent iteration loops mattered most where physical intuition is hard to encode, such as grid-imprint artifacts at cubed-sphere edges.

This is also where the paper locates its own dependency. Agentic development worked because the community had already distilled its understanding into open, well-posed benchmark cases, from dynamical-core tests to intercomparison projects. The authors note that such benchmarks remain sparse for the land surface, sea ice, biogeochemistry and the coupled system — precisely the components where agent verification has the least to check against.

Differentiability turns a bias into a gradient

Conventional physics-based ESMs generally do not provide an adjoint, meaning the gradient of model outputs with respect to uncertain parameters is unavailable; calibration therefore relies on derivative-free Monte Carlo or ensemble methods whose cost grows rapidly with the number of parameters. legoESM is differentiable end to end, so a single reverse pass yields the gradient of a scalar diagnostic with respect to every internal parameter and boundary input, at a cost essentially independent of parameter count.

The land-surface calibration is the clearest demonstration. Per-plant-functional-type surface parameters were tuned by gradient descent through the differentiable land model against ERA5 mean-annual fields, with physical bounds imposed by differentiable sigmoid transforms. Starting from CLM5 default parameters, the global mean warm bias in skin temperature fell from +1.09 K to +0.30 K, with the largest gains over arid and semi-arid land. Surface albedo moved only from +0.002 to +0.004, which the authors read as evidence that the temperature fit was not bought at the cost of the radiative surface state.

The same adjoint capability underpins variational data assimilation without a separately coded adjoint model. In a single-observation 4D-Var experiment with the MPAS dynamical core, a synthetic 500 hPa temperature innovation at t0 + 6 h was propagated back to t0, producing an upstream increment that the westerlies carry to the observation location; 3D-Var, which treats the observation as valid at analysis time, instead placed the increment centre east of the observation. The tangent-linear propagation generated automatically by JAX agreed with the nonlinear forecast to within 1%.

Exact gradients are not always useful gradients

The paper is careful about what differentiability buys. Gradients are exact for the discretized model, but for chaotic, autoregressive dynamics, gradients through long integrations can vanish or explode; implicit solvers, iterative fixers and threshold-like physics can make them noisy or ill-conditioned. There is also a trade-off between forward accuracy and reverse-pass stability — the authors note that NeuralGCM, for instance, builds on a comparatively simple spectral dynamical core.

legoESM does not remove these limits. What it offers instead is choice: users can check gradients against finite differences and select, per application, the grid, dynamical core, rollout length and loss — for example time-averaged or ensemble statistics — that best balance forward fidelity against gradient quality.

A second caveat concerns what calibration can fix. Gradient-based tuning can reduce bias and errors, but it does not remove structural error in physics-based parameterizations, so a well-tuned simulation can still carry substantial biases from its inherent structure. The authors' proposed response is to use the model's swappability to replace default structures with pure machine-learning components and test whether those address the bias.

One code base, from eddies to global climate

The multiscale claim is validated component by component rather than asserted. As a large-eddy simulation, the model reproduces expected mean profiles and turbulence structure against reference cases and against the MicroHH reference code. Enlarging the domain and adding moist physics turns it into a cloud-resolving model tested in radiative–convective equilibrium against the SAM model, where convection spontaneously self-aggregates; warm-rain microphysics closely matches a particle-based super-droplet reference on the Golovin analytic benchmark. Microphysics schemes were additionally checked against their original Fortran or C implementations, used as oracles for code translation.

At the global scale, the atmosphere runs shallow-water cases, 3D radiative–convective equilibrium coupled to ocean or land, Held–Suarez and AMIP-style prescribed-ocean simulations on any grid, and is evaluated with ClimateEval, which builds on ESMValTool. Forced three-dimensional ocean integrations were compared with NEMO, including a twin experiment in the idealized DINO basin and additional salinity and mixed-layer-depth diagnostics. The ocean benchmark suite passes on every grid on which each case is defined, across lat–lon, MPAS/Voronoi, FESOM-type unstructured and tripolar meshes.

The authors state that some established frameworks already span large-eddy to global scales, but that spanning a metre-scale eddy to a global simulation within one differentiable code base is, to their knowledge, not available in existing frameworks. That is a claim about the combination of properties, not about any single simulation being superior.

Where the simulated climate still misses

The validation results are mixed in ways that matter for judging readiness. Evaluation criteria were met within only a few iterations for most land-surface targets, with the exception of soil evaporation, where high biases desiccated soil and a modernized soil-evaporation parameterization had to be implemented.

The carbon side shows a clearer residual bias. Comparing initial carbon pools with observational estimates revealed too little soil carbon in northern high-latitude soils, and the agents iterated on parameters and structure until the pools were realistic. But gross primary productivity is overestimated: legoESM produces a global total of 199 PgC yr⁻¹, against 157 PgC yr⁻¹ for CLM5-SP and 120–146 PgC yr⁻¹ for the data-driven CEDAR and FLUXCOM products. The latitudinal structure — a tropical maximum and a northern mid-latitude peak — is reproduced, and the authors flag the global total as a bias that remains to be addressed.

These are component-level evaluations, not a verdict on the coupled system. The authors anticipate the same hard coupled-system problems that challenge established models, such as spurious deep-water formation in the Labrador Sea, drift in the Atlantic meridional overturning circulation, and tropical biases, and present faster calibration as a way to test candidate fixes more rapidly rather than as a solution already demonstrated.

GPU scaling depends on grid structure

Hardware performance is reported as strong scaling on the DKRZ Levante system with A100-80GB GPUs and AMD EPYC 7763 Milan CPU ranks, at single and double precision, with dashed lines showing ideal scaling anchored at the first point. The atmosphere was tested on a lat–lon grid at 4096×8192 (about 5 km) and 2048×4096 with 26 levels, and on MPAS icosahedral meshes at subdivision levels 9 and 10 with 26 levels. The ocean was tested on an ORCA12-class tripolar grid (3072×4352, 75 levels, GPU only), an MPAS Voronoi mesh with 2.6 million cells and 40 levels, and a FESOM2 FORCA20 mesh with 2.1 million nodes and 70 levels in double precision.

The pattern is consistent: structured grids follow the ideal line to large device counts, whereas unstructured meshes roll off once per-device work is too small to amortize halo exchange. That distinction is the practically relevant one for anyone planning kilometre-scale runs, because it means the cost of going to more devices depends on which mesh the configuration uses, not only on the number of GPUs available.

A template, not a finished ESM

The authors describe the paper as a demonstration of a new architecture rather than a finished Earth system model, and state that the model as of now is not as complex as established operational models. The novelty they claim is the set of capabilities — end-to-end sensitivity and attribution, online training of embedded neural components through the model's own evolution, and clean cross-scale comparison within one differentiable code — rather than any single simulation.

That framing carries a governance question the paper raises directly. A federated framework can be curated by a core team that tests and integrates contributions into a reference version, or forked by users who adapt it with less long-term maintenance but make improvements harder to integrate and evaluate across the framework. The authors favour a combination in which contributions passing shared benchmark and contract tests are merged into a reference version.

The broader claim is that in an agentic workflow, source code becomes comparatively inexpensive to regenerate, so tests, benchmarks and reference datasets increasingly encode what the community knows and must be curated most carefully. For readers weighing legoESM, the decisive comparison is therefore not against an operational ESM's complexity but between two development models: one where physical understanding lives in hand-maintained code, and one where it lives in verifiable benchmarks that agent-written code must pass.

Sources used in this answer