CORTEXA
← Browse
arxivstat.MLcs.AIcs.LG2026-07-06

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Markus Heinonen, Yair Shenfeld, Ricardo Baptista, Daniel Waxman, Dmitry Batenkov, Tim Cooijmans, Eli Bingham

Reconstructing population dynamics is a central problem in the physical and data sciences. Often, the dynamics are modeled as a Wasserstein gradient flow (WGF): a curve of distributions driven by an energy functional. Though there are multiple mathematical characterizations of a WGF, the dominant algorithmic approach relies on the Jordan--Kinderlehrer--Otto (JKO) scheme. JKO-based methods are inflexible to time discretisation and require solving costly optimal transport problems. We take a residual approach, enforcing the continuity equations via a non-negative loss function whose minimum is the WGF. Combined with a data-fitting divergence, this gives a single global objective. This perspective unifies several existing methods and leads to a new particle-based method, stitching, that is simulation-free and robust to large gaps between observations. We demonstrate that the stitching method achieves state-of-the-art performance across trajectory inference benchmarks. For code see github.com/BasisResearch/wasserstein-residuals.

View free PDFSource page

Related papers

arxivcs.LGcs.AIstat.ML2026-07-12

Learning from Local Walks on Dynamic Graphs with Bandit Feedback

Sourav Chakraborty, Amit Kiran Rege, Claire Monteleoni, Lijun Chen

We study stochastic multi-armed bandits on dynamic graphs, where arms correspond to the vertices of a network with time-varying edges. In this setting, the learner is restricted to local movement, selecting only its current node or an immediate neighbor at each round. This constr…

View free PDFSource page
arxivstat.MLcs.AIcs.LGeess.SYmath.OC2026-07-07

EHR-MPC: Inference-Time Control for Sepsis Treatment with Generative Patient Digital Twins

Joshua Pickard, Wei Qi, Na Li, Ann Woolley, Lisa Cosimi, Roy Kishony, et al.

Sepsis is a leading cause of mortality, yet optimal treatment policies remain contested. Existing reinforcement learning (RL) approaches learn fixed strategies for sepsis treatment, limiting adaptability to changing clinical objectives during inference. We propose EHRMPC, a frame…

View free PDFSource page
arxivcs.CYcs.AIcs.LGstat.ML2026-07-02

Simulating Eutopia: Revisiting Long-term Fairness with Outcomes, Performativity, and Dynamics

Vedant Palit, Udvas Das, Brahim Driss, Debabrota Basu

As AI-driven Decision Makers (ADMs) influence our socioeconomic reality, their roles in both enhancing efficiency and amplifying the social biases have drawn attention. In this paper, we revisit the nuances of long-term `fairness' achievable by an ADM, specifically in the context…

View free PDFSource page
arxivstat.MLcs.AIcs.LG2026-07-06

To Retain or to Adapt? Generalizing Continual Learning

Giulia Lanzillotta, Mandana Samiei, Doina Precup, Razvan Pascanu, Claire Vernade

The Continual Learning (CL) literature has long been driven by the goal of mitigating catastrophic forgetting. This objective rests on a pervasive, often unstated assumption: that a lifelong learner should approximate the Joint-Task Learning (JTL) solution and retain all previous…

View free PDFSource page
arxivcs.LGcs.AIstat.ML2026-06-29

What Drives the Inlier-Memorization Effect? A Theory of Outlier Detection via Early Training Dynamics

Kunwoong Kim, Dongha Kim

Outlier detection (OD) aims to identify anomalous instances by learning the underlying structure of normal data (inliers), and is particularly challenging in fully unsupervised settings where no information about anomalies is available during training. Recent advances have levera…

View free PDFSource page