Recent Publications

2025

1. ICML

   Controlled Generation with Equivariant Variational Flow Matching

   Eijkelboom, Floor, Zimmermann, Heiko, Bekkers, Erik, Welling, Max, Naesseth, Christian, and van de Meent, Jan-Willem

   In International Conference on Machine Learning 2025
2. ICML

   Exponential Family Variational Flow Matching for Tabular Data Generation

   Guzmán-Cordero*, Andrés, Eijkelboom*, Floor, and van de Meent, Jan-Willem

   In International Conference on Machine Learning 2025
3. ICML

   Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems

   Zhdanov, Maksim, Welling, Max, and van de Meent, Jan-Willem

   In International Conference on Machine Learning 2025

   Abs HTML PDF

   Large-scale physical systems defined on irregular grids pose significant scalability challenges for deep learning methods, especially in the presence of long-range interactions and multi-scale coupling. Traditional approaches that compute all pairwise interactions, such as attention, become computationally prohibitive as they scale quadratically with the number of nodes. We present Erwin, a hierarchical transformer inspired by methods from computational many-body physics, which combines the efficiency of tree-based algorithms with the expressivity of attention mechanisms. Erwin employs ball tree partitioning to organize computation, which enables linear-time attention by processing nodes in parallel within local neighborhoods of fixed size. Through progressive coarsening and refinement of the ball tree structure, complemented by a novel cross-ball interaction mechanism, it captures both fine-grained local details and global features. We demonstrate Erwin’s effectiveness across multiple domains, including cosmology, molecular dynamics, and particle fluid dynamics, where it consistently outperforms baseline methods both in accuracy and computational efficiency.

4. PRL

   Learning Neural Free-Energy Functionals with Pair-Correlation Matching

   Dijkman, Jacobus, Dijkstra, Marjolein, van Roij, René, Welling, Max, van de Meent, Jan-Willem, and Ensing, Bernd

   Physical Review Letters Feb 2025

   Abs HTML PDF

   The intrinsic Helmholtz free-energy functional, the centerpiece of classical density functional theory, is at best only known approximately for 3D systems. Here we introduce a method for learning a neural-network approximation of this functional by exclusively training on a dataset of radial distribution functions, circumventing the need to sample costly heterogeneous density profiles in a wide variety of external potentials. For a supercritical Lennard-Jones system with planar symmetry, we demonstrate that the learned neural free-energy functional accurately predicts inhomogeneous density profiles under various complex external potentials obtained from simulations.

5. ICRA

   On-Robot Reinforcement Learning with Goal-Contrastive Rewards

   Biza, Ondrej, Weng, Thomas, Sun, Lingfeng, Schmeckpeper, Karl, Kelestemur, Tarik, Ma, Yecheng Jason, Platt, Robert, van de Meent, Jan-Willem, and Wong, Lawson L. S.

   In Proceedings of the 2025 IEEE International Conference on Robotics and Automation, ICRA’25 Feb 2025

   HTML PDF

2024

1. NeurIPS

   Practical Shuffle Coding

   Kunze, Julius, Severo, Daniel, van de Meent, Jan-Willem, and Townsend, James

   In Advances in Neural Information Processing Systems 2024

   HTML PDF
2. NeurIPS

   Variational Flow Matching for Graph Generation

   Eijkelboom, Floor, Bartosh, Grigory, Naesseth, Christian A., Welling, Max, and van de Meent, Jan-Willem

   In Advances in Neural Information Processing Systems 2024

   HTML PDF
3. NeurIPS

   VISA: Variational Inference with Sequential Sample-Average Approximations

   Zimmermann, Heiko, Naesseth, Christian A., and van de Meent, Jan-Willem

   In Advances in Neural Information Processing Systems 2024

   HTML PDF
4. NAACL

   Towards Reducing Diagnostic Errors with Interpretable Risk Prediction

   McInerney, Denis Jered, Dickinson, William, Flynn, Lucy C., Young, Andrea C., Young, Geoffrey S., van de Meent, Jan-Willem, and Wallace, Byron C.

   In 2024 Annual Conference of the North American Chapter of the Association for Computational Linguistics (NAACL) 2024

   Abs HTML PDF

   Many diagnostic errors occur because clinicians cannot easily access relevant information in patient Electronic Health Records (EHRs). In this work we propose a method to use LLMs to identify pieces of evidence in patient EHR data that indicate increased or decreased risk of specific diagnoses; our ultimate aim is to increase access to evidence and reduce diagnostic errors. In particular, we propose a Neural Additive Model to make predictions backed by evidence with individualized risk estimates at time-points where clinicians are still uncertain, aiming to specifically mitigate delays in diagnosis and errors stemming from an incomplete differential. To train such a model, it is necessary to infer temporally fine-grained retrospective labels of eventual "true" diagnoses. We do so with LLMs, to ensure that the input text is from before a confident diagnosis can be made. We use an LLM to retrieve an initial pool of evidence, but then refine this set of evidence according to correlations learned by the model. We conduct an in-depth evaluation of the usefulness of our approach by simulating how it might be used by a clinician to decide between a pre-defined list of differential diagnoses.

5. ICLR

   Entropy Coding of Unordered Data Structures

   Kunze, Julius, Severo, Daniel, Zani, Giulio, van de Meent, Jan-Willem, and Townsend, James

   In International Conference on Learning Representations (ICLR) 2024

   HTML PDF

2023

1. EMNLP

   CHiLL: Zero-shot Custom Interpretable Feature Extraction from Clinical Notes with Large Language Models

   McInerney, Denis Jered, Young, Geoffrey, Meent, Jan-Willem, and Wallace, Byron

   In The 2023 Conference on Empirical Methods in Natural Language Processing 2023

   HTML PDF
2. NeurIPS

   Topological Obstructions and How to Avoid Them

   Esmaeili, Babak, Walters, Robin, Zimmermann, Heiko, and van de Meent, Jan-Willem

   In Thirty-seventh Conference on Neural Information Processing Systems (to appear) Dec 2023

   HTML PDF
3. CoRL

   One-shot Imitation Learning via Interaction Warping

   Biza, Ondrej, Thompson, Skye, Pagidi, Kishore Reddy, Kumar, Abhinav, Pol, Elise, Walters, Robin, Kipf, Thomas, Meent, Jan-Willem, Wong, Lawson L.S., and Platt, Robert

   In 7th Annual Conference on Robot Learning Nov 2023

   HTML PDF
4. ACT

   String Diagrams with Factorized Densities

   Sennesh, Eli, and van de Meent, Jan-Willem

   In Applied Category Theory Jul 2023

   Abs HTML PDF

   A growing body of research on probabilistic programs and causal models has highlighted the need to reason compositionally about model classes that extend directed graphical models. Both probabilistic programs and causal models define a joint probability density over a set of random variables, and exhibit sparse structure that can be used to reason about causation and conditional independence. This work builds on recent work on Markov categories of probabilistic mappings to define a category whose morphisms combine a joint density, factorized over each sample space, with a deterministic mapping from samples to return values. This is a step towards closing the gap between recent category-theoretic descriptions of probability measures, and the operational definitions of factorized densities that are commonly employed in probabilistic programming and causal inference.

5. ICML WNC

   Entropy Coding of Unordered Data Structures

   Kunze, Julius, Severo, Daniel, Zani, Giulio, Meent, Jan-Willem, and Townsend, James

   In ICML 2023 Workshop Neural Compression: From Information Theory to Applications Jul 2023
6. TMLR

   A Variational Perspective on Generative Flow Networks

   Zimmermann, Heiko, Lindsten, Fredrik, Meent, Jan-Willem, and Naesseth, Christian A

   Transactions on Machine Learning Research Apr 2023

   Abs HTML PDF Code

   Generative flow networks (GFNs) are a class of probabilistic models for sequential sampling of composite objects, proportional to a target distribution that is defined in terms of an energy function or a reward. GFNs are typically trained using a flow matching or trajectory balance objective, which matches forward and backward transition models over trajectories. In this work we introduce a variational objective for training GFNs, which is a convex combination of the reverse- and forward KL divergences, and compare it to the trajectory balance objective when sampling from the forward- and backward model, respectively. We show that, in certain settings, variational inference for GFNs is equivalent to minimizing the trajectory balance objective, in the sense that both methods compute the same score-function gradient. This insight suggests that in these settings, control variates, which are commonly used to reduce the variance of score-function gradient estimates, can also be used with the trajectory balance objective. We evaluate our findings and the performance of the proposed variational objective numerically by comparing it to the trajectory balance objective on two synthetic tasks.

7. LAFI

   Verified Reversible Programming for Verified Lossless Compression

   Townsend, James, and van de Meent, Jan-Willem

   In POPL Workshop on Languages for Inference (LAFI) Apr 2023

   HTML PDF
8. LAFI

   A convenient category of tracing measure kernels

   Sennesh, Eli, and van de Meent, Jan-Willem

   In POPL Workshop on Languages for Inference (LAFI) Apr 2023

   HTML PDF
9. LAFI

   Semantics of Probabilistic Program Traces

   Lew, Alexander, Sennesh, Eli, van de Meent, Jan-Willem, and Mansinghka, Vikash

   In POPL Workshop on Languages for Inference (LAFI) Apr 2023

   HTML PDF

2022

1. EMNLP

   That’s the Wrong Lung! Evaluating and Improving the Interpretability of Unsupervised Multimodal Encoders for Medical Data

   McInerney, Denis Jered, Young, Geoffrey, van de Meent, Jan-Willem, and Wallace, Byron C.

   Proceedings of the Conference on Empirical Methods in Natural Language Processing. Conference on Empirical Methods in Natural Language Processing Dec 2022

   Abs

   Pretraining multimodal models on Electronic Health Records (EHRs) provides a means of learning representations that can transfer to downstream tasks with minimal supervision. Recent multimodal models induce soft local alignments between image regions and sentences. This is of particular interest in the medical domain, where alignments might highlight regions in an image relevant to specific phenomena described in free-text. While past work has suggested that attention “heatmaps” can be interpreted in this manner, there has been little evaluation of such alignments. We compare alignments from a state-of-the-art multimodal (image and text) model for EHR with human annotations that link image regions to sentences. Our main finding is that the text has an often weak or unintuitive influence on attention; alignments do not consistently reflect basic anatomical information. Moreover, synthetic modifications — such as substituting “left” for “right” — do not substantially influence highlights. Simple techniques such as allowing the model to opt out of attending to the image and few-shot finetuning show promise in terms of their ability to improve alignments with very little or no supervision. We make our code and checkpoints open-source.

2. NeurIPS WSGNR

   Understanding Optimization Challenges when Encoding to Geometric Structures

   Esmaeili, Babak, Walters, Robin, Zimmermann, Heiko, and Meent, Jan-Willem

   In NeurIPS 2022 Workshop on Symmetry and Geometry in Neural Representations Dec 2022
3. PLOS Comp Bio

   Probabilistic Program Inference in Network-Based Epidemiological Simulations

   Smedemark-Margulies, Niklas, Walters, Robin, Zimmermann, Heiko, Laird, Lucas, Loo, Christian, Kaushik, Neela, Caceres, Rajmonda, and Meent, Jan-Willem

   PLOS Computational Biology Nov 2022

   Abs HTML PDF

   Accurate epidemiological models require parameter estimates that account for mobility patterns and social network structure. We demonstrate the effectiveness of probabilistic programming for parameter inference in these models. We consider an agent-based simulation that represents mobility networks as degree-corrected stochastic block models, whose parameters we estimate from cell phone co-location data. We then use probabilistic program inference methods to approximate the distribution over disease transmission parameters conditioned on reported cases and deaths. Our experiments demonstrate that the resulting models improve the quality of fit in multiple geographies relative to baselines that do not model network topology.

4. Bio. Pysch.

   Interoception as Modeling, Allostasis as Control

   Sennesh, Eli, Theriault, Jordan, Brooks, Dana, Meent, Jan-Willem, Barrett, Lisa Feldman, and Quigley, Karen S.

   Biological Psychology Nov 2022

   Abs HTML PDF

   The brain regulates the body by anticipating its needs and attempting to meet them before they arise – a process called allostasis. Allostasis requires a model of the changing sensory conditions within the body, a process called interoception. In this paper, we examine how interoception may provide performance feedback for allostasis. We suggest studying allostasis in terms of control theory, reviewing control theory’s applications to related issues in physiology, motor control, and decision making. We synthesize these by relating them to the important properties of allostatic regulation as a control problem. We then sketch a novel formalism for how the brain might perform allostatic control of the viscera by analogy to skeletomotor control, including a mathematical view on how interoception acts as performance feedback for allostasis. Finally, we suggest ways to test implications of our hypotheses.

5. Neuroinf.

   A Computational Neural Model for Mapping Degenerate Neural Architectures

   Khan, Zulqarnain, Wang, Yiyu, Sennesh, Eli, Dy, Jennifer, Ostadabbas, Sarah, van de Meent, Jan-Willem, Hutchinson, J. Benjamin, and Satpute, Ajay B.

   Neuroinformatics Mar 2022

   Abs HTML PDF

   Degeneracy in biological systems refers to a many-to-one mapping between physical structures and their functional (including psychological) outcomes. Despite the ubiquity of the phenomenon, traditional analytical tools for modeling degeneracy in neuroscience are extremely limited. In this study, we generated synthetic datasets to describe three situations of degeneracy in fMRI data to demonstrate the limitations of the current univariate approach. We describe a novel computational approach for the analysis referred to as neural topographic factor analysis (NTFA). NTFA is designed to capture variations in neural activity across task conditions and participants. The advantage of this discovery-oriented approach is to reveal whether and how experimental trials and participants cluster into task conditions and participant groups. We applied NTFA on simulated data, revealing the appropriate degeneracy assumption in all three situations and demonstrating NTFA’s utility in uncovering degeneracy. Lastly, we discussed the importance of testing degeneracy in fMRI data and the implications of applying NTFA to do so.

6. WACV

   Enhancing Few-Shot Image Classification With Unlabelled Examples

   Bateni, Peyman, Barber, Jarred, Meent, Jan-Willem, and Wood, Frank

   In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV) Jan 2022

   Abs HTML PDF Code

   We develop a transductive meta-learning method that uses unlabelled instances to improve few-shot image classification performance. Our approach combines a regularized Mahalanobis-distance-based soft k-means clustering procedure with a modified state of the art neural adaptive feature extractor to achieve improved test-time classification accuracy using unlabelled data. We evaluate our method on transductive few-shot learning tasks, in which the goal is to jointly predict labels for query (test) examples given a set of support (training) examples. We achieve state of the art performance on the Meta-Dataset, mini-ImageNet and tiered-ImageNet benchmarks. All trained models and code have been made publicly available

2021

1. AISTATS

   Rate-Regularization and Generalization in Variational Autoencoders

   Bozkurt, Alican, Esmaeili, Babak, Tristan, Jean-Baptiste, Brooks, Dana, Dy, Jennifer, and Meent, Jan-Willem

   In International Conference on Artificial Intelligence and Statistics Mar 2021

   Abs HTML PDF

   Variational autoencoders (VAEs) optimize an objective that comprises a reconstruction loss (the distortion) and a KL term (the rate). The rate is an upper bound on the mutual information, which is...

2. UAI

   Learning proposals for probabilistic programs with inference combinators

   Stites, Sam, Zimmermann, Heiko, Wu, Hao, Sennesh, Eli, and Meent, Jan-Willem

   In Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence 27–30 jul 2021

   Abs HTML PDF Code

   We develop operators for construction of proposals in probabilistic programs, which we refer to as inference combinators. Inference combinators define a grammar over importance samplers that compose primitive operations such as application of a transition kernel and importance resampling. Proposals in these samplers can be parameterized using neural networks, which in turn can be trained by optimizing variational objectives. The result is a framework for user-programmable variational methods that are correct by construction and can be tailored to specific models. We demonstrate the flexibility of this framework by implementing advanced variational methods based on amortized Gibbs sampling and annealing.

3. EMNLP

   Disentangling Representations of Text by Masking Transformers

   Zhang, Xiongyi, van de Meent, Jan-Willem, and Wallace, Byron

   In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing Nov 2021

   Abs HTML PDF

   Representations from large pretrained models such as BERT encode a range of features into monolithic vectors, affording strong predictive accuracy across a range of downstream tasks. In this paper we explore whether it is possible to learn disentangled representations by identifying existing subnetworks within pretrained models that encode distinct, complementary aspects. Concretely, we learn binary masks over transformer weights or hidden units to uncover subsets of features that correlate with a specific factor of variation; this eliminates the need to train a disentangled model from scratch for a particular task. We evaluate this method with respect to its ability to disentangle representations of sentiment from genre in movie reviews, toxicity from dialect in Tweets, and syntax from semantics. By combining masking with magnitude pruning we find that we can identify sparse subnetworks within BERT that strongly encode particular aspects (e.g., semantics) while only weakly encoding others (e.g., syntax). Moreover, despite only learning masks, disentanglement-via-masking performs as well as — and often better than —previously proposed methods based on variational autoencoders and adversarial training.

4. NAACL

   On the Impact of Random Seeds on the Fairness of Clinical Classifiers

   Amir, Silvio, van de Meent, Jan-Willem, and Wallace, Byron

   In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies Jun 2021

   Abs HTML PDF

   Recent work has shown that fine-tuning large networks is surprisingly sensitive to changes in random seed(s). We explore the implications of this phenomenon for model fairness across demographic groups in clinical prediction tasks over electronic health records (EHR) in MIMIC-III —— the standard dataset in clinical NLP research. Apparent subgroup performance varies substantially for seeds that yield similar overall performance, although there is no evidence of a trade-off between overall and subgroup performance. However, we also find that the small sample sizes inherent to looking at intersections of minority groups and somewhat rare conditions limit our ability to accurately estimate disparities. Further, we find that jointly optimizing for high overall performance and low disparities does not yield statistically significant improvements. Our results suggest that fairness work using MIMIC-III should carefully account for variations in apparent differences that may arise from stochasticity and small sample sizes.

5. NeurIPS

   Nested Variational Inference

   Zimmermann, Heiko, Wu, Hao, Esmaeili, Babak, and Meent, Jan-Willem

   In Advances in Neural Information Processing Systems Jun 2021

   Abs HTML PDF

   We develop nested variational inference (NVI), a family of methods that learn proposals for nested importance samplers by minimizing an forward or reverse KL divergence at each level of nesting. NVI is applicable to many commonly-used importance sampling strategies and provides a mechanism for learning intermediate densities, which can serve as heuristics to guide the sampler. Our experiments apply NVI to (a) sample from a multimodal distribution using a learned annealing path (b) learn heuristics that approximate the likelihood of future observations in a hidden Markov model and (c) to perform amortized inference in hierarchical deep generative models. We observe that optimizing nested objectives leads to improved sample quality in terms of log average weight and effective sample size.

6. AAMAS

   Action Priors for Large Action Spaces in Robotics

   Biza, Ondrej, Wang, Dian, Platt, Robert, Meent, Jan-Willem, and Wong, Lawson L.S.

   In Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems Jun 2021

   Abs HTML PDF Code

   In robotics, it is often not possible to learn useful policies using pure model-free reinforcement learning without significant reward shaping or curriculum learning. As a consequence, many researchers rely on expert demonstrations to guide learning. However, acquiring expert demonstrations can be expensive. This paper proposes an alternative approach where the solutions of previously solved tasks are used to produce an action prior that can facilitate exploration in future tasks. The action prior is a probability distribution over actions that summarizes the set of policies found solving previous tasks. Our results indicate that this approach can be used to solve robotic manipulation problems that would otherwise be infeasible without expert demonstrations. Source code is available at https://github.com/ondrejba/action_priors.

7. ICML

   Conjugate Energy-Based Models

   Wu*, Hao, Esmaeili*, Babak, Wick, Michael L, Tristan, Jean-Baptiste, and van de Meent, Jan-Willem

   In International Conference on Machine Learning Jun 2021

   Abs

   In this paper, we propose conjugate energy-based models (CEBMs), a new class of energy-based models that define a joint density over data and latent variables. The joint density of a CEBM decomposes into an intractable distribution over data and a tractable posterior over latent variables. CEBMs have similar use cases as variational autoencoders, in the sense that they learn an unsupervised mapping from data to latent variables. However, these models omit a generator network, which allows them to learn more flexible notions of similarity between data points. Our experiments demonstrate that conjugate EBMs achieve competitive results in terms of image modelling, predictive power of latent space, and out-of-domain detection on a variety of datasets.