Biological World Models for Drug Discovery

Connecting simulation-based inference to world models, and what that framing buys us for adaptive drug discovery.

Encoding biology in a joint predictive model

Many hallmarks of a world model are the same as a probabilistic scientific simulators such as those trained using Simulation-based inference (SBI) techniques. For example, in biology a scientific simulator of the Bone Morphogenetic Protein (BMP) signaling pathway is modeled as a joint probabilistic model

\[p_\phi(y, \theta, \xi),\]

which is parameterized by model weights $\phi$ and models the joint distribution of what phenotypic response $y$ you would see given that you perturbed with an action $\xi$ and assuming latent variables $\theta$ that describe the underlying biology. You pre-train on tuples of simulations $(y, \theta, \xi)$ that you think are true a priori — we’ll get to where this assumption falls apart later. (Note: SBI models don’t always encode an action variable $\xi$.)

World models: a brief primer

A world model, generally, is a predictive model that captures how an environment evolves over time and how that evolution depends on an agent’s previous actions. Let $o_t$ be the observation at time $t$ given an action $a_t$; this results in a reward $r_t$ that depends on a utility function. A world model is a parameterized predictive system with model parameters $\psi$ that approximates the environment’s dynamics

\[p_\psi(s_{t+1}, o_{t+1}, r_t \mid s_t, a_t),\]

where the true environment state $s_t$ may be fully or partially observed, or hidden. In partially observed settings, the model needs to infer a latent belief-like state from an interaction history tuple $h_t = (o_{\leq t}, a_{\leq t})$, which can optionally be done with the help of an encoder $s_t \sim q_\eta (s_t \mid h_t)$.

Relating SBI models to world models

If you squint long enough you can see a similarity between an action-oriented SBI model and world models. Let’s convert from world model terminology to that of SBI models. Observations can be redefined as $y=o$, actions are $\xi=a$, and state variables are the latent parameters that SBI models seek to infer $\theta=s$. Let’s rewrite a world model in SBI terminology

\[p_\psi(\theta_{t+1}, y_{t+1}, r_t \mid \theta_t, \xi_t),\]

where SBI models can encode a history of action-observation pairs into a belief over latent parameters via an encoder similar to a classic world model, such that $\theta \sim q_\eta(\theta \mid h_t)$, where $h_t$ now consists of $y$ and $\xi$ values. Alternatively, SBI often leverage Bayes’ theorem to model a belief update over latent (biological) parameters:

\[\begin{aligned} p_\psi(\theta \mid h_t, \xi_t, y_{t+1}) &\propto p_\psi(\theta, y_{t+1} \mid h_t, \xi_t) \\ &= p_\psi(y_{t+1} \mid \theta, \xi_t)\, p_\psi(\theta \mid h_t). \end{aligned}\]

So the update is still grounded in a joint model, but now the “next state” is also a posterior distribution over latent biology after observing the latest experimental response.

This is also where SBI-style biological world models can have an advantage. Many learned world models compress histories into latent states that are optimized for prediction. That may be useful, but biology often asks a more specific question. Given what we observed in one biological context, what would have happened if we had made a different intervention?

For that question to be meaningful, the latent variables should line up with mechanisms we can reason about and perturb, such as in our biological example of receptor abundance, binding affinity, or catalytic efficiency. Mechanistic SBI models already start from named scientific variables and simulator-defined dependencies. In Simformer-style models, those dependencies can be encoded through the model architecture, for example with attention masks that reflect an expert graph over parameters, experiments, and observations. This does not make the graph true, but it gives the world model an explicit causal hypothesis over named simulator variables that can be conditioned on, criticized, and refined as new experiments reveal where the model succeeds or fails.

From prediction to utility

The remaining world-model ingredient is the reward. As in many world-model settings, this reward is best thought of as a utility used to evaluate imagined outcomes, rather than as the core dynamics model itself. In SBI, one natural choice is expected information gain (EIG), which rewards experiments that maximally shrink uncertainty about the latent parameters $\theta$

\[\mathrm{EIG}(\xi_t) := \mathbb{E}_{\theta \sim q_\eta(\theta \mid h_t),\, y \sim p_\text{eval}(y \mid \theta, \xi_t)} \left[ \log \frac{ p_\psi(\theta \mid h_t, \xi_t, y) }{ q_\eta(\theta \mid h_t) } \right].\]

In words, this asks how much we expect the next experiment $\xi_t$ to change our beliefs about the latent biology given what we believe the latent parameters are before the experiment.

The utility function can also be tied to a downstream task. For example, suppose $y$ is a positive signal related to a disease-relevant phenotype, such as fluorescence intensity from a pathway reporter, and our goal is to inhibit that signal. A simple utility is to reward perturbations that make the response fall below a desired threshold $\tau$:

\[U_{\mathrm{inhibit}}(\xi_t) = \mathbb{E}_{\theta \sim q_\eta(\theta \mid h_t)} \left[ \Pr(y \leq \tau \mid h_t, \theta, \xi_t) \right].\]

Here $h_t$ plays two roles. It informs our current belief about the latent biology through $q_\eta(\theta \mid h_t)$, and it can also contain task-relevant experimental history, such as which perturbations or drugs have already been tried. In other words, what we think will inhibit the signal depends both on what we currently believe the biology is and on what the previous experiments have taught us about how perturbations behave in that context.

This is close in spirit to a phenotypic screen, where assay quality is often judged by how cleanly positive and negative controls separate. The Z-prime factor compares the distance between the control means with their variability, asking whether the assay has enough dynamic range to detect meaningful perturbations. Here, the world-model version asks a slightly different question. A phenotypic screen first asks whether the assay can reliably distinguish effect from no effect. The world model asks which perturbation is most likely to push the predicted phenotype into the desired range, given our current uncertainty about the biology.

Side information in drug discovery

Once we define a utility, the next question is what information the model needs in order to make better decisions under that utility. This is where the world-model framing becomes particularly useful. We can ask which additional variables should be brought into the joint model, whether they actually improve decisions, and how uncertainty propagates through the downstream decision.

In drug discovery, this question shows up in the debate between ligand-based and structure-based modeling. A ligand-only model may predict activity from chemical similarity, while a structure-aware model tries to use information about the target, its binding site, and protein-ligand interactions. Let $\zeta$ represent this external structural or biophysical information. In the BMP setting, this could include ligand-receptor binding information, which overlaps with my previous work on BMP ligand-receptor structure and with Huber et al.’s use of parameter-level Bayesian updates to guide biological hypothesis formation. Let $\beta$ represent a small-molecule perturbation, distinct from the BMP ligand perturbation $\xi$. We can then write a conditional model

\[p_\phi (y \mid \theta, \xi, \zeta, \beta).\]

We can also think of this as a joint distribution over the variables involved:

\[p_\phi(y, \theta, \xi, \zeta, \beta).\]

The joint formulation provides an intuitive formualtion of what were the necessary ingredients to achieve an outcome of interest. For example, given a small molecule $\beta$, an experimental context $\xi$, structural information $\zeta$, and latent biology $\theta$, what downstream responses $y$ are we expecting? Ideally, adding more information helps us become more confident in the outcomes of interest.

The key word is ideally. Adding variables to a joint model is not the same thing as adding useful information. A structural variable $\zeta$ helps only if it is relevant to the outcome $y$, calibrated enough to trust, and connected to the decision we actually care about. Otherwise, it can make the model more confident for the wrong reason, which can be a bad thing if it fails to find new drugs, or, a good thing if it reflects the performance of some successful hedge fund managers–correct for the wrong reasons (but still correct).

The world-model framing adds one more twist. When side information fails, that failure is itself informative. If a binding affinity model predicts that a molecule should perturb the pathway, but the observed response $y$ does not move, the model can help separate different explanations. The affinity prediction may be wrong, the structural representation may be misleading, or the molecule may bind correctly while remaining irrelevant to the downstream phenotype. These are different failure modes, and they imply different next experiments. In this sense, the biological world model gives us a place to propagate uncertainty from structure prediction, affinity prediction, pathway simulation, and experimental measurement into the decision we actually care about.

This leaves a few layers of uncertainty:

There are more layers of uncertainty, but these may be the most salient for small-molecule drug discovery and convey the challenge of reducing uncertainty in a world model of biology.

From automation to adaptive discovery

The upshot is not that a biological world model automates drug discovery by itself, but that it gives us a way to make discovery adaptive. Different diseases, targets, assays, and therapeutic goals come with different prior information and different failure modes. A useful model should therefore help decide what to do next given the current goal, the current uncertainty, and the reliability of the tools being used.

This matters because drug discovery is not a steady march from target to molecule to clinic. A new perturbation may update our belief about pathway biology, reveal an unreliable binding mode, or show that a structural hypothesis was correct but irrelevant for the downstream phenotype. In that sense, the model is useful not only when it predicts correctly, but also when its failures point to which assumption should be tested next.

To make this concrete, imagine we have three structural hypotheses for a target–say three candidate binding site conformations $\zeta^1, \zeta^2, \zeta^3$. We want the world model to help determine how “good” each one is. One way to answer this is through the downstream prediction: conditioning on each structure hypothesis, does the model’s predicted response $y$ better match what we actually observe? A structure hypothesis that leads to more accurate predictions provides evidence that it captures something real about the biology. This still requires cross-validation against held-out experimental data, but each comparison provides an informative data point that updates our belief about which structural representation is most useful. This is closely related to the use of parameter-level Bayesian updates for structural hypothesis testing in BMP signaling. However, care must be taken as this can still fall prey to the hedge fund fallacy.

A compounding effect comes from how the model relates one context to another. A drug that works against one kinase may update our beliefs about a binding pocket, a scaffold, a target family, an assay regime, or a downstream pathway. If another disease target shares some of those variables, the model can partially transfer what was learned. The same is true for failures. If a structure-based affinity model fails in one pocket or chemical series, that failure should update how much we trust similar predictions elsewhere.

This is where the world model can help localize failure. Was the binding prediction wrong, or was the predicted binding real but irrelevant to the downstream phenotype? Those are different failure modes and they suggest different subsequent experiments. Over time, each result updates the world model’s shared variables and assumptions, including which binding pockets behave similarly, which chemical scaffolds transfer across targets, which assays are comparable, and which signals are useful in a given biological context.

As an example, suppose we screen three candidate molecules $\beta^1, \beta^2, \beta^3$ against a target and the world model predicts how each will affect a downstream phenotypic response. An affinity model may predict that $\beta^1$ and $\beta^2$ should both strongly inhibit the pathway, while $\beta^3$ has a weaker predicted effect. After running the experiment, $\beta^1$ inhibits as expected, $\beta^3$ shows the predicted weak effect, but $\beta^2$ does nothing despite a confident binding prediction. That discrepancy is the interesting case. The world model does not just flag $\beta^2$ as a miss; it provides a framework for asking why: was the affinity prediction wrong, or did the molecule bind but fail to perturb the relevant downstream biology? This gets more intersting for small moelcule discovery as we go one level up and ask what are the similarities between working and failing moleucles in a given biological context.

Combining these two ideas, we can ask a joint question: given a structural hypothesis and a candidate molecule, does the world model predict that the downstream response falls within a therapeutically useful range? A correct structure paired with the right molecule should yield a tight, well-placed prediction. An approximate structure widens the uncertainty, and a wrong structure shifts the prediction away from the therapeutic window entirely. This is where the structural hypothesis testing and the molecule screening feed into the same decision: which combination of structure and molecule is worth pursuing next?

This also surfaces a deeper question that biologists care about. For a given disease, is this even the right target to be drugging in the first place? If the best available structure paired with the most promising molecule still places the predicted response outside the therapeutic window, that is evidence that the target itself may not be relevant to the disease phenotype. The world model can help separate “wrong molecule” from “wrong structure” from “wrong target,” and each of those failure modes points to a fundamentally different next step.

Takeaways

The core idea is that probabilistic scientific simulators trained with SBI techniques already contain most of the ingredients of a world model, encoding a joint distribution over latent biology, experimental context, and downstream responses that updates as new data arrives. Framing these simulators as world models lets us connect them to utilities, side information, and adaptive decision-making in a way that is natural for drug discovery. In particular, a biological world model can help with several concrete tasks across small molecule and target discovery:

What the world-model framing adds is a way to evaluate that side information through a utility. Adding a variable to the joint model is only useful if it improves decisions under the objective we care about. Structure helps when it sharpens predictions within a therapeutic window. It hurts when it makes the model more confident for the wrong reasons. And when a prediction fails, the joint model gives us a place to ask whether the failure came from the molecule, the structure, or the target itself.

Agents can sit on top of this loop, but they are just one actor in a larger system. A language model can propose experiments, molecules, or explanations, but it still needs tools that represent biological uncertainty and competing objectives. A biological world model can act as the decision layer that connects structure prediction, binding affinity models, pathway simulators, experimental data, and downstream disease utilities. The goal then is to ask how each component changes the decision, where uncertainty enters, and which source of uncertainty matters the most for the task at hand.

The useful world model is not the one that predicts everything. It is the one that helps us choose the next action under uncertainty in biology.