Publications

PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling

Yashi Zhang

Hongyu Guo

Building Virtual Cells that can accurately simulate cellular responses to perturbations is a long-standing goal in systems biology. A fundam… (see more)ental challenge is that high-throughput single-cell sequencing is destructive: the same cell cannot be observed both before and after a perturbation. Thus, perturbation prediction requires mapping unpaired control and perturbed populations. Existing models address this by learning maps between distributions, but typically assume a single fixed response distribution when conditioned on observed cellular context (e.g., cell type) and the perturbation type. In reality, responses vary systematically due to unobservable latent factors such as microenvironmental fluctuations and complex batch effects, forming a manifold of possible distributions for the same observed conditions. To account for this variability, we introduce PerturbDiff, which shifts modeling from individual cells to entire distributions. By embedding distributions as points in a Hilbert space, we define a diffusion-based generative process operating directly over probability distributions. This allows PerturbDiff to capture population-level response shifts across hidden factors. Benchmarks on established datasets show that PerturbDiff achieves state-of-the-art performance in single-cell response prediction and generalizes substantially better to unseen perturbations. See our project page (https://katarinayuan.github.io/PerturbDiff-ProjectPage/), where code and data will be made publicly available (https://github.com/DeepGraphLearning/PerturbDiff).

2026-04-29

International Conference on Machine Learning (accepted)

doi.org

openreview.net

Position: Theory-Level Autoformalization, From Isolated Statements to Unified Formal Knowledge Bases

Marcus J. Min

Mike He

Zhaoyu Li

Zixuan Yi

Sharad Malik

Aarti Gupta

Xujie Si

Osbert Bastani

Autoformalization translates informal natural language into formal, machine-verifiable languages. While most work focuses on individual stat… (see more)ements, real formalization efforts are inherently theory-level: they require an entire web of axioms, definitions, and lemmas before target theorems can even be stated. In this position paper, we argue for theory-level autoformalization: formalizing complete theories, including all their inter-dependencies, as structured libraries. We examine the significance of this shift, address alternative views, identify open challenges, and propose three promising paths forward.

2026-04-29

Position Paper Track @ International Conference on Machine Learning (spotlight)

openreview.net

Predictive Spatio-Temporal Scene Graphs for Semi-Static Scenes

Miguel Saavedra-Ruiz

Steven Parkison

We have seen tremendous recent progress in our ability to build "spatio-semantic" representations that enable robots to perform complex reas… (see more)oning across geometry and semantics. However, the vast majority of these methods lack any ability to perform reasoning across time. This is a desirable property in situations where a robot repeatedly observes an environment where instances may change in between observations, but in a structured way. Consider as an example a home environment where the location of a mug typically moves from the cupboard to a countertop to the sink and then back to the cupboard on a daily basis. We should be able to learn this cyclic behavior and use it to predict the state of the mug in the future. In this work, we propose a method that is able to perform this type of tempo-spatio-semantic reasoning. Underpinning the method is a filter, Perpetua

2026-04-29

arXiv (preprint)

doi.org

arxiv.org

Privileged Information Distillation for Language Models

Emiliano Penaloza

Dheeraj Vattikonda

Nicolas Gontier

Alexandre Lacoste

Laurent Charlin

Massimo Caccia

Training-time privileged information (PI) can enable language models to succeed on tasks they would otherwise fail, making it a powerful too… (see more)l for reinforcement learning in hard, long-horizon settings. However, transferring capabilities learned with PI to policies that must act without it at inference time remains a fundamental challenge. We study this problem in the context of distilling frontier models for multi-turn agentic environments, which typically hide their internal reasoning and expose only action trajectories. This breaks standard distillation pipelines, since successful behavior is observable, but the reasoning process is not. For this, we introduce {\pi}-Distill, a joint teacher-student objective that trains a PI-conditioned teacher and an unconditioned student simultaneously using the same model. Additionally, we also introduce On-Policy Self-Distillation (OPSD), an alternative approach that trains using Reinforcement Learning (RL) with a reverse KL-penalty between the student and the PI-conditioned teacher. We show that both of these algorithms effectively distill frontier agents using action-only PI. Specifically, we find that {\pi}-Distill and, in some cases, OPSD, outperform industry standard practices (Supervised finetuning followed by RL) that assume access to full Chain-of-Thought supervision across multiple agentic benchmarks, models, and forms of PI. We complement our results with extensive analysis that characterizes the factors enabling effective learning with PI, focusing primarily on {\pi}-Distill and characterizing when OPSD is competitive.

2026-04-29

International Conference on Machine Learning (accepted)

doi.org

openreview.net

Reinforcement Learning with Pairwise Preferences in Long-Term Decision Problems

Jonathan Colaço Carr

Prakash Panangaden

Doina Precup

Benjamin Van Roy

Reinforcement learning problems typically define the goal as maximizing the expected value of a scalar reward function. But, pairwise prefer… (see more)ences are often easier to specify than scalar rewards, and they express certain goals that scalar rewards cannot. Methods for reinforcement learning with pairwise preferences have thus received growing interest. Unfortunately, these methods are inefficient in problems with long time horizons, and they lack guarantees on the performance of Markov policies relative to history-dependent policies, which bridge the theory and practice of reinforcement learning. We therefore propose the \textit{Markov decision contest} as a new problem model for reinforcement learning with pairwise preferences. We prove that stationary Markov policies are optimal among all history-dependent policies, that solving a Markov decision contest exactly is in P, and that a simple iterative algorithm converges to an optimal policy at a sublinear rate. Lastly, in a set of high-dimensional decision problems with long time horizons, we show that our approximate algorithm is significantly more learning-efficient than prior work.

2026-04-29

International Conference on Machine Learning (accepted)

doi.org

openreview.net

Reward Redistribution for CVaR MDPs using a Bellman Operator on L-infinity

Tail-end risk measures such as static conditional value-at-risk (CVaR) are used in safety-critical applications to prevent rare, yet catastr… (see more)ophic events. Unlike risk-neutral objectives, the static CVaR of the return depends on entire trajectories without admitting a recursive Bellman decomposition in the underlying Markov decision process. A classical resolution relies on state augmentation with a continuous variable. However, unless restricted to a specialized class of admissible value functions, this formulation induces sparse rewards and degenerate fixed points. In this work, we propose a novel formulation of the static CVaR objective based on augmentation. Our alternative approach leads to a Bellman operator with: (1) dense per-step rewards; (2) contracting properties on the full space of bounded value functions. Building on this theoretical foundation, we develop risk-averse value iteration and model-free Q-learning algorithms that rely on discretized augmented states. We further provide convergence guarantees and approximation error bounds due to discretization. Empirical results demonstrate that our algorithms successfully learn CVaR-sensitive policies and achieve effective performance-safety trade-offs.

2026-04-29

International Conference on Machine Learning (spotlight)

doi.org

openreview.net

SAD-Flower: Flow Matching for Safe, Admissible, and Dynamically Consistent Planning

Tzu-Yuan Huang

Armin Lederer

Dai-Jie Wu

Xiaobing Dai

Sihua Zhang

Hsiu-Chin Lin

Shao-Hua Sun

Stefan Sosnowski

Sandra Hirche

Flow matching (FM) has shown promising results in data-driven planning. However, it inherently lacks formal guarantees for ensuring state an… (see more)d action constraints, whose satisfaction is a fundamental and crucial requirement for the safety and admissibility of planned trajectories on various systems. Moreover, existing FM planners do not ensure the dynamical consistency, which potentially renders trajectories inexecutable. We address these shortcomings by proposing SAD-Flower, a novel framework for generating

2026-04-29

International Conference on Machine Learning (accepted)

doi.org

openreview.net

Search-R2: Enhancing Search-Integrated Reasoning via Actor-Refiner Collaboration

Bowei He

Minda Hu

Zenan Xu

Hongru Wang

Licheng Zong

Yankai Chen

Chen Ma

Xue Liu

Pluto Zhou

Irwin King

Search-integrated reasoning enables language agents to transcend static parametric knowledge by actively querying external sources. However,… (see more) training these agents via reinforcement learning is hindered by the *multi-scale credit assignment* problem: existing methods typically rely on sparse, trajectory-level rewards that fail to distinguish between high-quality reasoning and fortuitous guesses, leading to redundant or misleading search behaviors. To address this, we propose Search-R2, a novel Actor–Refiner collaboration framework that enhances reasoning through targeted intervention, with both components jointly optimized during training. Our approach decomposes the generation process into an Actor, which produces initial reasoning trajectories, and a Meta-Refiner, which selectively diagnoses and repairs flawed steps via a ``cut-and-regenerate'' mechanism. To provide fine-grained supervision, we introduce a hybrid reward design that couples outcome correctness with a dense process reward quantifying the information density of retrieved evidence. Theoretically, we formalize the Actor–Refiner interaction as a smoothed mixture policy, proving that selective correction yields strict performance gains over strong baselines. Extensive experiments across various general and multi-hop QA datasets demonstrate that Search-R2 consistently outperforms strong RAG and RL-based baselines across model scales, achieving superior reasoning accuracy with minimal overhead.

2026-04-29

International Conference on Machine Learning (accepted)

doi.org

openreview.net

Synthesizable Molecular Generation via Soft-constrained GFlowNets with Rich Chemical Priors

Louis Vaillancourt

Yves V. Brun

Yoshua Bengio

Alex Hernández-García

2026-04-29

International Conference on Machine Learning (accepted)

doi.org

openreview.net

$\tau^2$-Bench: Evaluating Conversational Agents in a Dual-Control Environment

Victor Barres

Honghua Dong

Soham Ray

Xujie Si

Karthik R Narasimhan

Existing benchmarks for conversational AI agents simulate _single-control_ environments, where only the AI agent can use tools to interact w… (see more)ith the world, while the user remains a passive information provider. This differs from real-world scenarios like technical support, where users need to actively participate in modifying the state of the (shared) world. In order to address this gap, we introduce

2026-04-29

International Conference on Machine Learning (spotlight)

openreview.net

The Cost of Commitment in Option-Based Hierarchical RL

Randy Lefebvre

Audrey Durand

Empirically, option-based hierarchical reinforcement (HRL) learning often produces longer and more diverse options when a deliberation cost … (see more)is charged at option boundaries. However, when options are executed for many steps under an approximate dynamics model, small model errors compound along the option, degrading the quality of the resulting plan. In this work, we introduce the commitment loss to formalize the tradeoff between deliberation cost and model error as a function of option duration. We characterize how optimal termination probabilities vary with this tradeoff under two model-error mechanisms. First, the model is learned from finite data via maximum-likelihood estimation, producing statistical error that interacts with option duration. Second, we consider an input-driven setting where an exogenous input is only observed at option boundaries and evolves unobserved between them, creating a drift-induced mismatch between planned and realized dynamics. In both cases, we solve for the optimal termination behavior as a function of deliberation cost and the error scale, clarifying the behavior of some popular HRL algorithms that approach the deliberation cost as a heuristic.

2026-04-29

International Conference on Machine Learning (accepted)

openreview.net