What in the World is a World Model?

It's clear that LLMs+RL+<Insert ML/AI Approach> is leading to some pretty impressive AI tools that can solve problems¹ in novel and unique ways. Additionally, LLMs have shown that when the task or thinking is best defined through language, code, or mathematics, these AI tools can clearly outperform humans at some of these tasks. Software engineering benchmarks [1] clearly show rapidly advancing capabilities in autonomous issue resolution, though the resolved fractions on real-world repositories remain far below typical human developers. However, the internal mechanisms or state representations, which are clearly rich and capturing something we humans would call "intelligence", these AI tools do not, as far as I know², out of the box, maintain a reliable internal "simulator" of state and dynamics. They can roll out scenarios in language/code/math (chain-of-thought, tree search, and similar), but those rollouts are often inconsistent or ungrounded when you need faithful prediction over many steps of inputs, information, constraints, goals, and actions [2]. This seems to me to be a clear limitation of the current SOTA, but keep in mind I'm not a researcher in this space, 🤷‍♂️.

Note

There are a lot of references for this post and I did not read most of them all in great detail, but rather "ran" through them and then asked questions to clarify my understanding using Gemini and other LLMs.

If I take a moment here to think about how a human thinks internally it seems we excel at abstraction and planning in a model space of the world we've experienced. Take the scenario of when you or I think about a ball rolling down a hill, we have this mental model for what the ball will do based on our knowledge of the environment around us that we have observed. Now true, a LLM does have knowledge of the governing physics, possible initial/boundary conditions, and thus can numerically simulate the physics exactly but that is not what we do, we often do this by approximate forward simulation in the head, like an "intuitive physics engine" [3]. This knowledge can be shaped early by maturation and perception³, and it is directly reinforced from our continued observations of the same interactions and patterns. We humans have a world model [4] that we use to make decisions and predictions about the world around us; the same idea was formalized in AI as a learned simulator of environment dynamics [5]. Our brains likely evolved for this because it allows us to make informed decisions towards our actual actions and goals ... not a reseacher here just thinking out loud.

Again just stop and think about it, when you are preparing for say a trip or presentation how you mentally play things out the scenarios that may occur. Or think about when you are doing any lab or simulation work, how you mentally scope various scenarios and hypothesize about the outcomes. As you do these things you narrow down on the possible set of actions and targeted goals you want to achieve based on how your mental perspective of these form.

I still don't know if I truly gauge or grasp what a world model is, but I think what it tries to represent is the abstract planning of actions and goals based on the latent (i.e. largely unconscious) knowledge we have acquired. The better the predictive model, the better our real world actions and goals are often defined and executed ... though "world model" is an analogy here it seems Craik's internal models, episodic simulation of the future, and motor internal models in the brain are clearly related ideas, not one single thing [6].

Take riding a bike. Everyone who has learned to ride a bike will tell you that if they rode a bike for 5 years and then stopped for 10 years and then tried again to ride the bike, they often recover riding ability within 10 minutes or so. This is the procedural memory for skilled actions that is stored separately from everyday declarative memory [7]. Truth though is it seems to be that the persistence is usually explained by your implicit motor circuits (basal ganglia, cerebellum, and related systems), not by consciously rehearsing the physics in your mind's eye before you mount, although though mental imagery does help. Still, I argue that there is a motor "model" of riding a bike such that it encodes balance, pedaling, and steering in a given context. If we did the same thing but instead of riding the bike on a road we told you to ride the bike on a tightrope or snow, your learned skills for riding on a road would still be invoked (i.e. you would think about how it might feel to ride a bike on a tightrope) but because that latent representation was never tuned or updated with actual feedback from those environments (i.e., you never actually rode a bike on a tightrope or snow), you would likely struggle [8]. Interestingly though, you could imagine that some might do better than others because their prior experience and internal forward models are richer, so mind-planning in the tightrope action space is more robust.

LeCun's World Model

Yann LeCun has been fairly vocal about world models as part of a broader cognitive architecture for embodied AI. He sees it as a orthogonal direction to LLMs. My understanding on his position is that systems that only predict the next token (or pixel) in some observation space will struggle to reach the kind of sample-efficient reasoning and planning that animals show. The hypothesis is because they waste capacity on unpredictable surface detail instead of learning the proper abstract state and dynamics. The recipe he and collaborators have worked on is a configurable predictive world model: encode what is observed, then roll forward in a compact latent space (optionally conditioned on actions), and use those "imagined" trajectories to choose goals and controls [9].

The way, to my understanding, that this idea departs from classic supervised vision is exactly what I was gesturing in the cat example. Don't train a network to output the label "cat" from pixels and don't ask the model to reconstruct every RGB value either. Instead have encoders that map observations into a high-dimensional latent representation, and a predictor that forecasts how that representation should evolve. For just static images the "action" conditioning the predictor can be absent (e.g., I-JEPA⁴ [10]). From this we get a self-supervised prediction in representation space and avoids forcing two augmented views of a cat to share one embedding via contrastive learning. In ref. [9] it seems he actually argues against leaning on contrastive losses and favors regularized, predictive objectives instead, but not sure I understand, per usual.

So what does JEPA (Joint-Embedding Predictive Architecture) look like? Rather than me go through the details, for which my visuals may not be as clear, The interview with LeCun on Welch Labs had some really good visuals and explanations:

There are many variations of JEPA, each built on the same overall idea ofmapping observed data into latent representations with a predictor and training via prediction errors in latent space rather than on raw outputs [9]. There is I-JEPA for images, V-JEPA for video, action-conditioned versions for control tasks, and hierarchical JEPA for planning at multiple time scales [11],[9].

The LeWorldModel [12] is a recent implementation of an action-conditioned small encoder-predictor trained and includes a SIGReg in the loss that keeps latents from collapsing. There is no reconstructiong the pixels for forward frames; everything is done in latent space. At test time it encodes start and goal images, rolls the predictor forward under candidate actions, and picks the sequence whose final latent is closest to the goal. The paper's claim to fame is that the planning (i.e. action conditioned predictor) in latent space is cheap comapred to using pixels. An interesting thing here is that nothing in the self-supervised loss contains any biasing towards known physics (i.e., how 3D classical mechanics should work).

How does a world model connect to physics? Probably not a drop-in replacement for a classical PDE solver on a macroscopic grid or an atomistic MLIP, but the idea of a JEPA-style scientific world model makes sense to me because it matches what I've seen from learned simulators and neural operators, such as Fourier Neural Operators [13] to learn a compact state, predicted evolution under physics biases, rolled forward without reconstructing every emergent degree of freedom at each step.

Footnotes

---

1. Frontier LLM have clearly shown in mathematics and perhaps elsewhere solutions that had yet to been-thought-of by humans. The clear demonstration for this is Erdos problems where solutions that experts state are novel have been produced by frotnier models from OpenAI and Anthropic [14–15]. ↩

2. Keep in mind, I am not a frontier researcher in AI research in any shape or form but more of a domain specialist applying AI tooling. ↩

3. I use "bootstrapping" to mean that without any prescribed inductive biases, the correct mental model exist and can be used to make the correct prediction or action. I don't know if this is the correct wording. A good example is the visual cliff experiment [16]: crawling infants on a glass floor where one half looks solid and safe while the other looks like a cliff. Most refuse to cross to their mother when she calls from the cliff side, consistent with a bootstrapped visual model of depth. ↩

4. Take the visible patches of an image as context and train a predictor to match the latent embeddings of other masked regions in the same image, where the target embeddings come from a slow-moving target encoder (exponential moving average of the context encoder) so the problem does not collapse to a trivial constant latent representation. ↩

References

[1] C.E. Jimenez, J. Yang, A. Wettig, S. Yao, K. Pei, O. Press, K.R. Narasimhan, {SWE}-bench: Can Language Models Resolve Real-World GitHub Issues?, in: International Conference on Learning Representations, 2024. https://doi.org/10.48550/arXiv.2310.06770.

[2] S. Yao, D. Yu, J. Zhao, I. Shafran, T.L. Griffiths, Y. Cao, K. Narasimhan, Reasoning with Language Model is Planning with World Model, in: Empirical Methods in Natural Language Processing, 2023. https://doi.org/10.18653/v1/2023.emnlp-main.507.

[3] P.W. Battaglia, J.B. Hamrick, J.B. Tenenbaum, Simulation as an Engine of Physical Scene Understanding, Proceedings of the National Academy of Sciences. 110 (2013) 18327--18332. https://doi.org/10.1073/pnas.1306572110.

…

Show remaining references

[4] K.J.W. Craik, The Nature of Explanation, Cambridge University Press, Cambridge, 1943. https://books.google.com/books?id=wT04AAAAIAAJ.

[5] J. Schmidhuber, Making the World Differentiable: On Using Self-Supervised Fully Recurrent Neural Networks for Dynamic Reinforcement Learning and Planning in Non-Stationary Environments, (1990). https://www.idsia.ch/~juergen/FKI-126-90.pdf.

[6] D. McNamee, D.M. Wolpert, Internal Models in Biological Control, Annual Review of Control, Robotics, and Autonomous Systems. 2 (2019) 339--364. https://doi.org/10.1146/annurev-control-060117-105206.

[7] N.J. Cohen, L.R. Squire, Preserved Learning and Retention of Pattern-Analyzing Skill in Amnesia: Dissociation of Knowing How and Knowing That, Science. 210 (1980) 207--210. https://doi.org/10.1126/science.7414331.

[8] K.E. Adolph, Specificity of Learning: Why Infants Fall Over a Veritable Cliff, Psychological Science. 11 (2000) 290--295. https://doi.org/10.1111/1467-9280.00258.

[9] Y. LeCun, A Path Towards Autonomous Machine Intelligence, (2022). https://openreview.net/forum?id=BZ5a1r-kVsf.

[10] M. Assran, Q. Duval, I. Misra, P. Bojanowski, P. Vincent, M. Rabbat, Y. LeCun, N. Ballas, Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture, in: IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023. https://doi.org/10.48550/arXiv.2301.08243.

[11] M. Assran, A. Bardes, D. Fan, Q. Garrido, R. Howes, M. Komeili, M. Muckley, A. Rizvi, C. Roberts, K. Sinha, A. Zholus, S. Arnaud, A. Gejji, A. Martin, F.R. Hogan, D. Dugas, P. Bojanowski, V. Khalidov, P. Labatut, F. Massa, M. Szafraniec, K. Krishnakumar, Y. Li, X. Ma, S. Chandar, F. Meier, Y. LeCun, M. Rabbat, N. Ballas, {{V-JEPA} 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning}, arXiv Preprint arXiv:2506.09985. (2025). https://doi.org/10.48550/arXiv.2506.09985.

[12] L. Maes, Q.L. Lidec, D. Scieur, Y. LeCun, R. Balestriero, {{LeWorldModel}: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels}, arXiv Preprint arXiv:2603.19312. (2026). https://doi.org/10.48550/arXiv.2603.19312.

[13] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, A. Anandkumar, Fourier Neural Operator for Parametric Partial Differential Equations, in: International Conference on Learning Representations, 2021. https://doi.org/10.48550/arXiv.2010.08895.

[14] G. DeepMind, Advancing Mathematics Research with {AI}-Driven Formal Proof Search, (2026). https://doi.org/10.48550/arXiv.2605.22763.

[15] {OpenAI}, An OpenAI model has disproved a central conjecture in discrete geometry, (2026). https://openai.com/index/model-disproves-discrete-geometry-conjecture/.

[16] E.J. Gibson, R.D. Walk, The Visual Cliff, Scientific American. 202 (1960) 64--71. https://doi.org/10.1038/scientificamerican0460-64.

Reuse and Attribution

Conformal Prediction

An area that has gained my specific attention lately is the idea of trying to quantify how confident a model is in its prediction. There are various techniques for doing this, the familiar Bayesian methods are some of the most popular, while others are still new to me, the one that I've been reading on this weekend is the idea of conformal prediction [1]. So what is Conformal Prediction? A Conformal prediction is a method for uncertainty quantification where instead of giving a single prediction we provide a set (e.g., ${y_1, y_2, ..., y_n}$) of most predicted values with a guarantee about how frequently the true value will be in that set. To put it even more simply, the model spits out a bunch of possible answers given several inputs and we just ask, “Is the real answer somewhere in this pile? If so are you 95% sure?”.

This idea is more aligned with the frequentist thinking, where we don't concern ourselves with a prior or the full posterior distribution, but only seeking how well the prediction is in the set.

The prediction interval when doing regression is then some range where we specify the values of the predictions should fall in. So say its the bandgap we are predicting the we specify say , 3.0 to 3.5 eV and we attached some confidence level that the true bandgap value will be in that range. If we are focused on assigning category labels, the it would be something like the crystal structure is either rocksalt, zincblende, or perovskite and again the confidence that the true label will be in that set is given some value (e.g. 90%).

Important Distinction

These are prediction intervals, not confidence intervals. Prediction intervals quantify uncertainty about individual predictions, while confidence intervals quantify uncertainty about distribution parameters like the mean.

Another aspect to keep in mind is that conformal prediction uses calibration residuals from a held-out set to size the prediction interval. This differs from Bayesian or bootstrap methods, which model the data-generating process (e.g., via likelihoods) to derive uncertainty.

Lets put the conformal predictions words into math. First we define $P(\cdot)$ as the probability of some value occurring ($P(E_b=3.1\text{eV})$), then we define the true value as $E_b^{n+1}$, where the $n+1$ is the new input. Now we define the prediction interval as $\Gamma(X_{n+1})$, where $X_{n+1}$ is the new input, so again this just tells us given an input evaluate a set-valued function that returns a set of possible values. For regression, this is typically a continuous interval like $[3.1, 3.5]$ eV for the bandgap prediction. Then we have to define our confidence level, i.e., quantify how well we want to or do know the prediction is in the set. To do this we use $1-\alpha$, where $\alpha$ is the confidence level. If $\alpha = 0$ we are saying that with 100% confidence the true value will be in the prediction interval. If $\alpha = 0.05$ we are saying that with 95% confidence the true value will be in the prediction interval, and so on. Now we can write the probability statement as:

$$ \begin{equation} P(E_b^{n+1} \in \Gamma(X_{n+1})) \geq 1 - \alpha. \label{eq:conformal_prediction} \end{equation} $$

Reading through this, we have the probability of the true value $E_b^{n+1}$ (i.e., the true bandgap value at next input $X_{n+1}$) being in the prediction interval $\Gamma(X_{n+1})$ set with a confidence level of $1 - \alpha$.

The thing to note is this is more of a guarantee statement; it is saying that if we define the confidence level and interval set then we can make a statement about how likely the true value is in the interval. We never need to directly compute the probability of the true value being in the interval in $\eqref{eq:conformal_prediction}$, which is why this method is called "distribution-free".

We also need a way to measure how "bad" the model's guess was compared to the truth, that is compute some loss function, in other words a nonconformity score. There are several different ways to define the nonconformity score depending on the type of prediction we are making. I'll just show the L1 norm as an example for regression, which would be

$$ \begin{equation} s_i = \vert E_{b}^{i} - \hat{E}_{b}^{i} \vert \nonumber \end{equation} $$

where $\hat{E}_b^{i}$ is the model's prediction for the bandgap value given input $X_{i}$, and the true value is $E_b^{i}$; $s_i$ is the nonconformity score for the $i$-th prediction. For example, a nonconformity score of 0.5 eV quantifies how much the model's prediction deviates from the true value. This is standard in regression. Importantly, the nonconformity scores are computed on a calibration set, that is, for $i \in {1, 2, \ldots, m}$, where this set is held out and not used for training.

So how do we use this in practice?

Switching from bandgap to forces, consider a GNN trained to predict atomic forces from atomic structures. We hold out a calibration set and compute nonconformity scores $s_i = |F_i^{true} - \hat{F}_i|$ for each force component, where $F_i^{true}$ is the true force and $\hat{F}_i$ is the model prediction.

We now build the prediction interval. Choose a desired coverage level, say $1-\alpha = 0.9$. Sort the $m$ calibration scores $s_{(1)} \le \cdots \le s_{(m)}$ and take $(r = \lceil (m+1)(1-\alpha) \rceil)$. The threshold is the $r$-th smallest score, $\hat q = s_{(r)}$. This is the standard split-conformal finite-sample rule.

Now we take a new atomic structure $X_{n+1}$ and compute the updated prediction interval.

$$ \begin{equation} \Gamma(F^{n+1}) = [\hat{F}^{n+1} - \hat q, \hat{F}^{n+1} + \hat q] \label{eq:conformal_prediction_updated} \end{equation} $$

With the updated prediction interval (with $\hat q$ chosen as above) we can then make a statement about how likely the true value is in the interval in $\eqref{eq:conformal_prediction_updated}$.

Again this is pretty basic concepts that are to follow and its more that the theory guarantees the prediction interval/set will contain the true value with the desired confidence level. This makes it simpler than Bayesian methods, where you have to "try" and compute the full posterior distribution through likelihood and prior probability distributions, all which require reasonably good choices.

To further illustrate let me work through this simple example of computing the conformal prediction for the forces using the orb-v3 pre-trained models on structures from the load-atoms GST dataset. The basic steps are to define the calibration and test sets, compute the nonconformity scores, build the prediction interval/set, and then make a statement about how likely the true value is in the interval. Below is the code snippet that shows the conformal prediction for the forces in addition to the model predictions.

Code

  # /// script
# requires-python = ">=3.8"
# dependencies = [
#   "numpy",
#   "matplotlib",
#   "orb-models",
#   "load-atoms",
#   "seaborn",
# ]
# ///
import math
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from matplotlib.patches import Patch
from sklearn.neural_network import MLPRegressor
from load_atoms import load_dataset
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

data = load_dataset("GST-GAP-22")
calibration, tests = data.random_split([500, 100], seed=42)

orb = pretrained.orb_v3_direct_20_omat(device="cpu")
calc = ORBCalculator(orb, device="cpu")


def pred_forces(atoms):
    atoms.calc = calc
    return atoms.get_forces()


# Locally adaptive split-conformal prediction using normalized residuals
# Note: These are prediction intervals, not confidence intervals
calibration_true = np.concatenate([a.arrays["forces"].ravel() for a in calibration])
calibration_pred = np.concatenate([pred_forces(a).ravel() for a in calibration])
scores = np.abs(calibration_true - calibration_pred)  # L1 nonconformity scores
z_calibration = np.abs(calibration_pred)  # feature: force magnitude


def fit_cqr_model(X, y, alpha=0.001):
    """CQR model for adaptive conformal prediction - learns local uncertainty scales.

    This model learns the relationship between force magnitude and prediction
    uncertainty to create adaptive prediction intervals.
    """
    model = MLPRegressor(
        hidden_layer_sizes=(100, 50, 25),
        activation="relu",
        solver="adam",
        alpha=alpha,
        learning_rate="adaptive",
        learning_rate_init=0.001,
        max_iter=1000,
        early_stopping=True,
        validation_fraction=0.1,
        n_iter_no_change=20,
        random_state=42,
    )
    model.fit(X, y)
    return model


# Fit adaptive scale model and compute normalized scores
cqr_scale = fit_cqr_model(z_calibration.reshape(-1, 1), scores, alpha=0.01)
sigma_calibration = np.clip(cqr_scale.predict(z_calibration.reshape(-1, 1)), 1e-8, None)
scores = np.sort(scores / sigma_calibration)

test_true = np.concatenate([a.arrays["forces"].flatten() for a in tests])
test_pred = np.concatenate([pred_forces(a).flatten() for a in tests])

z_test = np.abs(test_pred)
sigma_test = np.clip(cqr_scale.predict(z_test.reshape(-1, 1)), 1e-8, None)

levels = {
    "2σ": math.erf(2 / np.sqrt(2)),
    "3σ": math.erf(3 / np.sqrt(2)),
}

# ---- Plotting ----
lo, hi = float(min(test_true.min(), test_pred.min())), float(
    max(test_true.max(), test_pred.max())
)
ax_lim = max(abs(lo), abs(hi))
lo, hi = -ax_lim, ax_lim
idx = np.argsort(test_true)
xs, ys = test_true[idx], test_pred[idx]

fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True)
hexbin_for_colorbar = None
for ax, (tag, p) in zip(axes, levels.items()):
    delta_n = scores[
        min(max(int(np.ceil((len(scores) + 1) * p)) - 1, 0), len(scores) - 1)
    ]
    halfw = delta_n * sigma_test
    hw_s = halfw[idx]

    inside_band = np.abs(test_pred - test_true) <= halfw
    ax.plot([lo, hi], [lo, hi], "k-", lw=0.2)
    inside_band_sorted = np.abs(ys - xs) <= hw_s

    main_mask = (xs >= -1.5) & (xs <= 1.5) & (ys >= -1.5) & (ys <= 1.5)
    main_xs = xs[main_mask]
    main_ys = ys[main_mask]
    main_hw_s = hw_s[main_mask]
    main_inside = inside_band_sorted[main_mask]

    # Plot adaptive prediction interval boundaries
    ax.plot(main_xs, main_xs - main_hw_s, "b-", lw=0.2, alpha=0.6, zorder=2)
    ax.plot(main_xs, main_xs + main_hw_s, "b-", lw=0.2, alpha=0.6, zorder=2)

    confident_xs = main_xs[main_inside]
    confident_ys = main_ys[main_inside]
    if len(confident_xs) > 0:
        hb = ax.hexbin(
            confident_xs, confident_ys, gridsize=25, cmap="Greens", alpha=0.7, mincnt=1
        )

    uncertain_xs = main_xs[~main_inside]
    uncertain_ys = main_ys[~main_inside]
    if len(uncertain_xs) > 0:
        ax.scatter(
            uncertain_xs,
            uncertain_ys,
            s=1.5,
            alpha=0.5,
            color="red",
            marker="x",
            zorder=1,
        )

    coverage = np.mean(inside_band) * 100
    ax.set_title(
        f"{tag}: Test Coverage = {coverage:.1f}%", fontsize=14, fontweight="bold"
    )
    ax.set_xlabel("True Forces (eV/Å)", fontsize=12)
    if ax is axes[0]:
        ax.set_ylabel("Predicted Forces (eV/Å)", fontsize=12)
    ax.tick_params(labelsize=11)
    ax.set_xlim(-1.5, 1.5)
    ax.set_ylim(-1.5, 1.5)
    ax.set_aspect("equal")

    inset = inset_axes(ax, width="35%", height="35%", loc="upper left")
    inset.plot([lo, hi], [lo, hi], "k-", lw=0.2)
    inset.plot(xs, xs - hw_s, "b-", lw=0.15, alpha=0.6, zorder=2)
    inset.plot(xs, xs + hw_s, "b-", lw=0.15, alpha=0.6, zorder=2)

    confident_xs_full = xs[inside_band_sorted]
    confident_ys_full = ys[inside_band_sorted]
    if len(confident_xs_full) > 0:
        hb_inset = inset.hexbin(
            confident_xs_full,
            confident_ys_full,
            gridsize=20,
            cmap="Greens",
            alpha=0.7,
            mincnt=1,
        )
        hexbin_for_colorbar = hb_inset

    uncertain_xs_full = xs[~inside_band_sorted]
    uncertain_ys_full = ys[~inside_band_sorted]
    if len(uncertain_xs_full) > 0:
        inset.scatter(
            uncertain_xs_full,
            uncertain_ys_full,
            s=1.0,
            alpha=0.5,
            color="red",
            marker="x",
            zorder=1,
        )

    inset.set_xlim(lo, hi)
    inset.set_ylim(lo, hi)
    inset.set_aspect("equal")
    inset.tick_params(labelsize=8)

plt.suptitle(
    "Adaptive Conformal Prediction for Force Prediction Intervals",
    fontsize=16,
    fontweight="bold",
)

legend_elements = [
    Patch(facecolor="red", label="Model Uncertain"),
    Patch(facecolor="blue", label="Model Adaptive PI"),
]

fig.legend(
    handles=legend_elements,
    loc="center",
    bbox_to_anchor=(0.35, 0.15),
    ncol=1,
    fontsize=10,
    framealpha=0.9,
)

if hexbin_for_colorbar is not None:
    cbar = fig.colorbar(
        hexbin_for_colorbar,
        ax=axes,
        orientation="vertical",
        shrink=0.6,
        aspect=10,
        pad=0.15,
    )
    cbar.set_label("Point Density", fontsize=9, labelpad=-60)
    cbar.ax.tick_params(labelsize=8, labelleft=False, labelright=True)
    cbar.ax.set_position([0.52, 0.15, 0.02, 0.7])

    ticks = cbar.get_ticks()
    formatted_labels = []
    for tick in ticks:
        if tick >= 1000:
            formatted_labels.append(f"{tick/1000:.0f}K")
        elif tick >= 100:
            formatted_labels.append(f"{tick/1000:.1f}K")
        else:
            formatted_labels.append(f"{tick:.0f}")
    cbar.set_ticklabels(formatted_labels)

plt.tight_layout()
plt.subplots_adjust(bottom=0.1, wspace=0.30, top=0.85)
plt.savefig("force_parity_adaptive.png", dpi=600, bbox_inches="tight")

Figure 1. Force Parity plot for test data including adaptive conformal prediction bands.

The results for the prediction interval levels of 95.3% (2$\sigma$) and 99.6% (3$\sigma$) are shown in the figure above for 500 calibration structures and 100 test structures. I use a locally adaptive split-conformal method: fit a function $\sigma(|\hat F|)$ on the calibration set to model residual scale, compute normalized scores $s_i/\sigma_i$, select the threshold $\hat q$ on these normalized scores as above, and then set the half-width $h(X)=\hat q\,\sigma(|\hat F|)$. This provides heteroscedastic, input-dependent intervals.

The adaptive conformal prediction bands reveal details about model performance. The achieved coverage (95.3% for 2$\sigma$, 99.6% for 3$\sigma$) closely matches the theoretical expectations (95.45% and 99.73% respectively), so we can say the model is well-calibrated. Points outside the prediction bands represent cases where the model's uncertainty estimate was insufficient, i.e., these are the "hard" cases where the model struggles to make the correct predictions. While 3$\sigma$ bands provide higher confidence (99.6% vs 95.3%), they are wider and less informative. The ideal model would provide narrow bands with high confidence, indicating both precision and reliability.

Notice in the inset that the bands have large uncertainty in the large force regime. This points to the fact that the model may be unable to properly capture regimes where the potential energy surface exhibits large gradients.

Beyond Distribution-Free Guarantees

Exact split-conformal validity holds when the score function (and any normalization) is fixed before seeing the calibration labels. Here the scale function $\sigma(\cdot)$ is fit on the calibration set itself, which can break exchangeability of the scores and thus the finite-sample guarantee. The trade-off often yields tighter, more efficient intervals, but coverage should be validated empirically. Using a separate proper-training set to learn $\sigma(\cdot)$ and reserving calibration solely for ranking scores restores the classical guarantee.

There is an inherent tension in uncertainty quantification between confidence and precision. Higher sigma values (3$\sigma$) provide greater confidence that the true value lies within the band, but at the cost of wider, less informative intervals. Conversely, lower sigma values (2$\sigma$) provide narrower, more precise intervals but with lower confidence guarantees.

The other aspect is we can asses the honesty of the model by looking and the test coverage, i.e., how often the true value is in the prediction interval. In Figure 1 we see the test coverage is close to the confidence level for 2$\sigma$ and 3$\sigma$.

I think this was an interesting approach to model uncertainty that is simple in a lot of ways. So what are the downsides? The main one is that it requires a lot of calibration and test data to compute reliable prediction intervals. If you have a small dataset, you may not be able to compute the prediction interval/set well. Another issue is that conformal prediction gives you marginal coverage, not conditional coverage, meaning the guarantees hold on average over your whole dataset but not uniformly for every input, so you might have a model that's well-calibrated overall but terrible for specific material classes or rare chemistries. The method also assumes your data is i.i.d. or exchangeable, but if you have time dependence, selection bias, or covariate shift, the guarantees can break down completely. The efficiency depends heavily on your choice of nonconformity score, so if you pick a poor one you'll get conservative, overly wide prediction intervals. It also doesn't give you per-point probabilities, just sets with frequency guarantees, so if you need likelihoods. Distribution shift can inflate your residual quantiles, giving you uninformative bands, and for structured targets like energy+forces+stress, joint coverage can be quite conservative.

Where does one go from here?

The prediction bands provide a natural criterion for active learning, i.e., points falling outside bands represent high-uncertainty cases that would benefit from additional training data. The idea then is to use conformal predictions to create uncertainty quantification for ML models and therefore in an active learning framework think about how to query the most informative samples to make improvements. With conformal predictions you would propose a bunch of unlabeled test inputs, pass through the model, and then see if the true value is in the prediction interval to establish the degree of confidence (i.e., uncertainty) in the prediction. Then based on your sampling strategy you would query a oracle to label the data point(s). Then one would retrain/fine-tune the model with new data, update the calibration set, and repeat the process.

Uncertainty quantification isn't about being right, it's about being right with the right amount of confidence. A model that's always right but with huge prediction bands is less useful than one that's right 95% of the time with narrow, informative bands. Our results demonstrate a well-calibrated model that appropriately increases uncertainty for challenging cases, which is exactly what you want in a reliable uncertainty quantification system.

---

References

[1] A.N. Angelopoulos, S. Bates, A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification, (2022). https://doi.org/10.48550/arXiv.2107.07511.

---

Reuse and Attribution

Synthetic Datasets & Model Distillation

Say your given an all-knowing oracle machine that you can query to get the ground truth data but the thing is this oracle system is fairly expensive¹ to query. Furthermore say that your not really interested in all ground truths the oracle has to offer, but rather you just want it for a specific set of inquiries. What would you do? the most direct thing would be to carefully curate a set of inquiries that are the most valuable and then spend your resources to get the answers and use those to teach a sage system. The idea is not new and is a form of model compression/distillation [1] and teacher-student paradigm [2].

Oracle 🧙‍♂️ > Sage 🧑‍🏫

If our oracle is a all-knowing system then a sage system is one that has equivalent deep wisdom and knoweldge on specific area(s) as the oracle but is inferior in all other areas. Our oracle is all-knowing but slow (think of an old professor 🧙‍♂️) but our sage is more youthful, but less expierienced (i.e., assistant professor 🧑‍🏫)

For the curation of the inquiries, one could use a Bayesian optimization approach that would iteratively select the most valuable inquiries to query the oracle based on the aquisition function the user chooses. This could be a exploration or exploitation function, that is do you want to find a broad set of valuable inquiries or do you want to find a specific set of valuable inquiries. Then there are more traditional statistical sampling methods. I won't touch on this since its not the focus of this post, but keep that in mind.

This is the idea behind synthetic dataset generation, that is the creation of inquiries and answers (i.e., labels) to create a dataset that is dervived from the oracle/teacher but is not a direct copy from the corpus of knowledge/truths. This is a form of data augmentation and data compression.

Once you have the set of inquiries and the answers, you can faithful teach the a sage system. This leads to what is called model distillation. Note this isn't model compression since we didn't take the underlying knowledge-structure of the oracle and modify it, we just extract stochastic answers from the oracle. Typically the model is more efficient and computationallyheaper to both train and infer with.

In the wild

The recent paper by Gardner et al. [3] provides a nice demonstration of synthetic dataset generation and model distillation for creating fast and accurate, chemical system specific, interatomic ML potentials. Their work offers a general, architecture-agnostic protocol for distilling knowledge from large, expensive "foundation models" (our Oracle 🧙‍♂️) into smaller, more efficient "student" models (our Sage 🧑‍🏫).

The process is essentially a three-step workflow:

%%{init: {'theme':'base', 'themeVariables': {'primaryColor':'#4f46e5','primaryTextColor':'#ffffff','primaryBorderColor':'#3730a3','lineColor':'#6b7280','sectionBkgColor':'#f8fafc','altSectionBkgColor':'#e2e8f0','gridColor':'#e5e7eb','tertiaryColor':'#f1f5f9'}}}%% graph LR; subgraph S1 ["Step 1: Fine-tuning"] A[/"Small DFT
dataset"/] --> C; B["Foundation
Model"] --> C{"Fine-tune"}; C --> D(("Oracle")); end subgraph S2 ["Step 2: Synthetic Data"] D --> E{"Generate &
Label"}; E --> F[/"Large Synthetic
Dataset"/]; end subgraph S3 ["Step 3: Distillation"] F --> G{"Train
Student"}; H["Fast
Architecture"] --> G; G --> I(("Sage")); end

1. Fine-tuning the Teacher: They start with a pre-trained foundation model (like MACE-MP-0). These models are powerful in that they support chemical systems of any composition including atomic numbers 1-89 (or so). The downside is the architecture is slow to evaluate and there fore not suitble for large-scale atomic simulations (e.g. 100,000's of atoms). As part of Gardner et al. protocol they fine-tune the teacher/oracle² model on a very small set of high-quality, domain-specific calculations (e.g, ~25 DFT calculations for liquid water). This turns³ the general model into a specialized "teacher" for that specific system, without the high cost of creating a large DFT dataset from scratch and still remaining general to all other chemical systems.

2. Generating Synthetic Data: For the liquid water system, the fine-tuned teacher is used to generate a large dataset of ~1,000 atomic configurations and their corresponding energies and forces are labeled with the fine-tuned teacher. For sampling, they employ a fairly efficient "rattle-and-relax" scheme that constructs a "family tree" of configurations. At each step, a "parent" structure is selected from the tree, a new "child" structure is generated from it, and this child is added back to the tree (see Figure 1). This is much faster than running molecular dynamics simulations to explore the potential energy surface. The scheme works by taking a starting structure and iteratively:

   • Rattling: Displace the atomic positions \( \mathbf{R} \) and unit cell \( C_0 \) by applying random perturbations:
     • \( \mathbf{R}' \leftarrow [(\mathbf{A} + \mathbf{I}) \times \mathbf{R}] + \mathbf{B} \)
     • \( \mathbf{C}' \leftarrow (\mathbf{A} + \mathbf{I}) \times \mathbf{C}_0 \)

   • Relaxing: Nudging the atoms in the direction of the forces predicted by the teacher model to find a new, stable configuration. At each relaxation step ( x ), the atomic positions ( \mathbf{R} ) are updated according to the equation:

     \[ \mathbf{R}' \leftarrow \mathbf{R} + \frac{\sigma_B}{x} \cdot \frac{\mathbf{F}}{||\mathbf{F}||} \]

     where \( \mathbf{F} \) represents the forces predicted by the teacher model, and \( \sigma_B \) is a hyperparameter controlling the rattle intensity that is also used to scale the relaxation step size. This iterative process — the Robbins-Monro algorithm — allows for an efficient exploration of the local energy landscape to generate new structures.

3. Distilling the Student (Sage): Finally, they train a much smaller, computationally cheaper model, the "student" (our Sage 🧑‍🏫), on this large synthetic dataset. Because the student model architecture (e.g., PaiNN or ACE) is simpler and has fewer parameters, both training and inference are much faster.

Figure 1. Synthetic data generation process (Fig. 6 from Gardner et al.).

As a proof-of-concept for liquid water, thier results suggest that you can get massive speed-ups with only a minor hit to accuracy.

Results from distilling a fine-tuned MACE-MP-0 foundation model for liquid water.

Model	Type	Relative Speed	Force MAE (meV/Å)
MACE-MP-0b3	Teacher (Fine-tuned, Oracle 🧙‍♂️)	1x	32
TensorNet	Student (Sage 🧑‍🏫)	> 10x	37
PaiNN	Student (Sage 🧑‍🏫)	> 30x	39
ACE	Student (Sage 🧑‍🏫)	> 100x	51

The student/sage models: PaiNN, TensorNet, and ACE are more computationally efficient for a couple of key reasons (see Figure 2). First being less complexity in the model, i.e., less parameters and layers. The next is that these models have computational cost scales dramatically with the interaction cut-off radius (\(r\)), often as \( \mathcal{O}(r^3) \). The big foundation models need a large radius (e.g., 6 Å) to be general to capture many-body interactions. However, sage/student models, however, can get away with a smaller radius (e.g., 4.5 Å for PaiNN) without losing much accuracy for the specific system they're trained on. This hyperparameter adjustment can lead to large reductions in computational cost, which is how the authors get >10x to >100x speedups. They are still "good" because they have learned the essential physics for that one system from the huge synthetic dataset provided by the all-knowing (but slow) teacher.

Figure 2. Computational efficiency of distilled models (Fig. 2 from Gardner et al.).

Outlook

The synthetic dataset generation and distillation model isn't just a numerical trick; the resulting sage models can produce physically meaningful simulations. For the liquid water, Gardner et al. report that the distilled models reproduce the structural properties of the much more expensive teacher model well. The distilled PaiNN model, for example, almost perfectly reproduces the radial distribution function (RDF).

Other structural and dynamical properties such as ring-size distributions and tetrahedral order parameters, which describe the medium-range ordering and local geometry, show the PaiNN student model behaves similarly to the teacher. The ACE model, which is very fast, leads to a slightly more ordered, crystal-like water structure, but still captures the key chemical/material physics.

Broader Applicability

The Garnder et al. paper [3] appears to be a versatility distillation method. Demonstrations on some systems included are:

• Metallic Hydrogen: Distilled models seem to reproduced DFT EOS for hydrogen at the extreme temperatures and pressures.
• Porous Amorphous Silica: The distilled models correctly captured the structure factor for porous amorphous silica.
• Hybrid Perovskites: The distilled models replicate the rotational dynamics of the molecular cations inside frameworks.
• Organic Reactions: Distilled models successfully reproduced the intended SN2 reaction mechanism. However, long-timescale simulations led to unphysical products.

Prospects

This fairly broad range of applications showcases that the distillation protocol is a powerful and general tool for creating more computationally efficient, specialized potentials for higher-throughput research. I'm bullish on this approach and think this is one-direction atomistic modeling and simulation is headed if there is no need for electronic effects are needed; although who knows these models might get good enough to predict charge transfer and polarization effects.

Footnotes

---

1. Expensive here can mean time to get the answer or cost in terms of resources. ↩

2. The all-knowing oracle in this fine-tuning process would be the DFT calculations, although this doesn't truly fit the notion of a black-box since we know the physics that is being solved, in contrast with GNN potentials we kind of know what embeddings it is learning. So in the protocol of Gardner et al., what we have is the all-knowing DFT oracle and a semi-oracle (fine-tuned model) that then teaches the sage. ↩

3. Fine-tuning is a technique in machine learning where a pre-trained model is adapted to a specific task by updating its parameters on a smaller dataset. This is often used to improve the performance of a model. While this helps the student model significantly, I feel the better approach is for the community to create larger more accurate foundation models and just use those as a true oracle. ↩

References

[1] G. Hinton, O. Vinyals, J. Dean, Distilling the Knowledge in a Neural Network, (2015). DOI.

[2] C. Buciluǎ, R. Caruana, A. Niculescu-Mizil, Model compression, in: Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, ACM, Philadelphia PA USA, 2006: pp. 535–541. DOI.

[3] J.L.A. Gardner, et al., Distillation of atomistic foundation models across architectures and chemical domains, (2025). DOI.

---

Reuse and Attribution

DFT is gettings its share of AI

I made a post on LinkedIn the other day sharing a recent preprint[1] by Microsoft AI that introduces a significant advancement in exchange-correlation functionals (XC) for DFT. This, to me, seemed like an important piece of news given that much of the data used for training the recent explosion in foundational interatomic potentials for atomistics uses massive DFT datasets.

The biggest achievement seems to be the ability for Microsoft AI's DFT XC functional, called Skala, to achieve near chemical accuracy (i.e., 1 kcal/mol) for molecular systems while doing it with reasonable computational cost. What this means is that we may get a new rung added to the chemistry Jacob's ladder of accuracy for DFT. I wanted to illusrate this so I made a diagram of the chemistry Jacob's ladder with the new Skala XC added as shown in Figure 1. I got the idea to annotate it from Keiran Rowell's blog[2].

The idea of Jacobs Ladder in chemistry, I think, was first introduced by J. Perdew in 2013(?) [3], I believe. The primary idea is that like the in the biblical story, Jacob dreams of a ladder reaching from earth to heaven and God standing above it promising Jacob divine blessing and protection. In the spirit of this, we the scientist dream of reaching chemical accuracy using our ladder of DFT approximations. Why is this important? Well, chemical accuracy with DFT would enable accurate enough predictions in reasonable computational time to improve discovery, design, and understanding of chemicals and materials. In particular understanding reaction mechanisms and kinetics would be much more tractable if chemical accuracy is reached.

Figure 1. Computational chemistry Jacob's ladder for DFT with the Microsoft AI Skala XC added.

DFT XC Overview

First, let me briefly summarize some basics of density functional theory (DFT) to explain why the exchange-correlation (XC) functional is so important. The Nobel-worthy Hohenberg–Kohn theorems and the formulation of the Kohn–Sham equations established that the total ground-state (i.e., 0 K) energy of a many-body electronic system can be written as:

$ \begin{equation} E[\rho] = T_s[\rho] + V_{ext}[\rho] + J[\rho] + E_{xc}[\rho] \end{equation} \label{eq:dft_energy} $

The first three terms are kinetic energy, external potential energy, and Coulombic repulsion energy. The last term is the XC energy which is the only term we don't know exactly and is crucial because it encapsulates the many-body quantum mechanical effects of the electrons whereas the others are single-particle terms.

Throughout the successful history of DFT, different approximations have been used for the XC term. The first was local density approximation (LDA), with the idea being that the XC energy is simply a function of the local density of the electrons. For all its simplicity, LDA did work reasonably for some systems but shortcomings were quickly discovered.

Then the generalized gradient approximation (GGA) was introduced, where the XC energy is now not only a function of the local density but also the gradient of the density. This introduced the idea of non-local XC functionals (i.e., the energy now depends on how the density changes with position, rather than just at that position).

After GGA, what followed was a "flavor-of-the-month" approach toward improving XC functions. There were hybrid functionals that mixed DFT XC with the Hartree-Fock exchange. Then meta-GGA was introduced with the idea of adding higher-order derivatives to the density. Add to that expensive XC functionals that include perturbation quantum chemistry methods. All have their improvements and shortfalls. You can see all the major players in the rungs of the ladder in Figure 1 showing how different XC functionals improve the range of accuracy as you climb up the ladder.

These are just the formalisms; there are also ad-hoc tweaks and specializations that get added based on your domain use and expertise. The limitations of such approaches are:

1. Handcrafted Features: Most functionals rely on fixed analytic forms and known constraints.
2. Non-locality Deficiency: Electron correlation is inherently non-local; local/semi-local functionals cannot fully capture this.
3. Empirical Tuning vs. Physical Justification: Empirical functionals may generalize poorly or violate constraints.
4. Slow Progress at Higher Rungs: Despite decades of work, hybrid/double-hybrid methods still fall short of universal chemical accuracy (error < 1 kcal/mol).

Skala Neural XC

The efforts by Microsoft AI represent a significant leap in XC functional design. Unlike traditional functionals that have relied on hand-crafted mathematical forms, Skala incorporates deep learning (i.e., Neural Networks) to achieve near-chemical accuracy while maintaining computational efficiency comparable to meta-GGA functionals[1]. Skala is particularly impressive because it navigates the trade-off between accuracy and computational cost. The Microsoft AI team has designed it to bridge the gap between semi-local functionals (fast but less accurate) and hybrid/double-hybrid functionals (accurate but computationally expensive).

Architecture and Design Philosophy

As mentioned, Skala XC is a neural network architecture trained to learn the non-local electron density interactions without requiring the full computational burden of exact exchange calculations. Like meta-GGA, it starts with seven semi-local density-derived features, but the design then employs what the authors call a "coarse-fine grid structure" that captures long-range density correlations consistent with multipole-like behavior.

For the exchange, the authors use LDA but incorporate a neural enhancement factor $f_\theta$:

$ \begin{equation} E_{xc}^\theta[\rho] = -\frac{3}{4}\left(\frac{6}{\pi}\right)^{1/3} \int \left(\rho_\uparrow(r)^{4/3} + \rho_\downarrow(r)^{4/3}\right) f_\theta[x[\rho]](r) \, dr \label{eq:skala_xc} \end{equation} $

Not knowing too much about XC design, this does seem clever though because $f_\theta$ operates on a feature vector $x[\rho]$ that encodes both local and non-local density information while maintaining computational efficiency (would like to understand this but above my head for now). This is reminiscent of delta-learning in atomistic ML, though Skala performs a direct functional approximation rather than a residual correction atop an existing functional.

What data was used?

The Microsoft AI team used a quantum chemistry dataset that seems only they could have curated, something like 150,000 data points spanning thermochemistry, conformational energies, noncovalent interactions, and ionization potentials. This was all done with gold-standard wavefunction methods like CCSD(T) to generate the data, 🤯. Their MSR-ACC/TAE dataset alone includes about 80,000 total atomization energies with errors below 1 kcal/mol, which is just staggering when you consider the computational expense of generating such reference data. The datasets alone might be useful for others.

Peformance Metrics

I'm not too familiar with benchmarks in this space, but they mention the challenging W4-17 benchmark and Skala achieves a mean absolute error (MAE) of around 1.0 kcal/mol, outperforming more established functionals. On the GMTKN55 benchmark, it scores a WTMAD-2 of 3.89 kcal/mol, which puts it in competition with the best hybrid functionals while requiring significantly less computational resources.

Another thing that is super interesting is that for large atomic numbers that are out of distribution for the trained Scala XC, They are showing that it maintains very good accuracy. This is usually a failure point for data-driven ML models where out of distribution data is not well handled particularly well [4]. The authors also method that the self-consistent field convergence is stable, crucial for practical application in DFT --- since, unlike post-SCF models, Skala is trained and used self-consistently.

Redefining Jacob's Ladder?

So going back to the Jacob's ladder of DFT functionals in Figure 1. Whether Skala constitutes a sixth rung or a bypass of the traditional ladder remains obviously open to interpretation, but it clearly represents a shift. I'm not a veteran in this space so hard for me to be definitive, but it probably represents some departure to non-locality from semi-local features to effectively bypass the traditional way of doing things.

It will be interesting to see from the "Kings of DFT" to see how they view Skala XC. Based on the Microsoft AI preprint this functional(s) is showing systematic improvement through use of high-quality data and training procedures, and the authors claim it can encode known physics through appropriate constraints and design. On the other hand, they lack the theoretical transparency and interpretability that traditional functional forms offer will probably be a critiqued, although existing functionals are not perfect either.

I'm just interested to see how this gets used for the other down stream applications in solid-state physics and materials science. Will we get DFT training sets for Materials Project structures that range in the 1 million of structures and all have chemical accuracy? If so that could in turn make these foundation models even better for materials discovery and classical atomistic modeling.

---

References

[1] G. Luise, C.-W. Huang, T. Vogels, D.P. Kooi, S. Ehlert, S. Lanius, K.J.H. Giesbertz, A. Karton, D. Gunceler, M. Stanley, W.P. Bruinsma, L. Huang, X. Wei, J.G. Torres, A. Katbashev, B. Máté, S.-O. Kaba, R. Sordillo, Y. Chen, D.B. Williams-Young, C.M. Bishop, J. Hermann, R. van den Berg, P. Gori-Giorgi, Accurate and scalable exchange-correlation with deep learning, (2025). DOI.

[2] K. Rowell, An Ersatz Ansatz, Blog (2023). https://keiran-rowell.github.io/guide/2023-04-12-compchem-methods-basics (accessed June 21, 2025).

[3] J.P. Perdew, Climbing the ladder of density functional approximations, MRS Bull. 38 (2013) 743–750. DOI.

[4] K. Li, A.N. Rubungo, X. Lei, D. Persaud, K. Choudhary, B. DeCost, A.B. Dieng, J. Hattrick-Simpers, Probing out-of-distribution generalization in machine learning for materials, Commun Mater 6 (2025) 1–10. DOI.

---

Reuse and Attribution