From a LinkedIn reply to a NeurIPS workshop paper

It started with a LinkedIn thread in February 2025. Igor Shuryak, a researcher at Columbia University, asked for help with a problem in cancer research. It was close to my causal AI research at Google and my PSPS work, so I replied in the thread and pointed him to the architecture and algorithms I had published.

That reply led to a conversation and then to a collaboration. The paper is now public: Setting a Dose versus Observing One: Continuous Dose-Response Estimation under Right-Censoring, accepted at the NeurIPS 2026 Workshop on World Models for High-Stakes Health.

What the paper does

Many treatments come in amounts rather than as a yes or no. In radiotherapy, our example, we want to know what would happen at a different dose, not just whether a patient was treated. Outcomes in this setting are survival times, and many are right-censored, meaning we only know the patient survived up to some point. Confounding is strong, since doctors choose doses for reasons that also affect outcomes.

The model, PSPS-CS, is a causal model and deep learning architecture (a deep mixture-of-experts) that estimates whole dose-response curves under those conditions. In simulations it was competitive with five other estimators. On a real head and neck cancer cohort, the data support no dose contrast and no change in current clinical practice. We see it as a tool for generating hypotheses, and will try it next on more cohorts and treatments, and on any setting where someone or something sets a continuous amount and we need to know what follows, from drug doses to robots choosing how hard to push.

Where this goes

There is already other work on deep learning for causal AI, so this is not the first of its kind. What’s new is handling continuous treatments with survival outcomes in one end-to-end model with gradient based learning. The question the paper asks, what happens if we set a value instead of just observing one, is the same one an agent faces when it picks an action and needs to know what follows. Think of a robot arm choosing how hard to grip an object: grip force is a continuous setting, and we want to know how long the gripper lasts before it fails if we change it. That is a survival outcome, and it is right-censored whenever the machine is still running when we stop watching. That is what a world model has to answer, and I think this architecture can work there. More work is coming.

The paper is on the Research page and the code is on GitHub.