← Search

Nikolay Malkin

38 accepted papers

2026

Discrete Compositional Generation via General Soft Operators and Robust Reinforcement Learning

ICLR 2026poster

A major bottleneck in scientific discovery consists of narrowing an exponentially large set of objects, such as proteins or molecules, to a small set of promising candidates with desirable properties. While this process can rely on expert knowledge, recent methods leverage reinforcement learning (RL…

Cited by 0SourceScholar
2026

Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

ICLR 2026poster

Learning the dynamics of a process given sampled observations at several time points is an important but difficult task in many scientific applications. When no ground-truth trajectories are available, but one has only snapshots of data taken at discrete time steps, the problem of modelling the dyna…

Cited by 0SourcecodeScholar
2025

Action abstractions for amortized sampling

ICLR 2025poster

As trajectories sampled by policies used by reinforcement learning (RL) and generative flow networks (GFlowNets) grow longer, credit assignment and exploration become more challenging, and the long planning horizon hinders mode discovery and generalization. The challenge is particularly pronounced i…

Cited by 0SourcePDFScholar
2025

Adaptive teachers for amortized samplers

ICLR 2025poster

Amortized inference is the task of training a parametric model, such as a neural network, to approximate a distribution with a given unnormalized density where exact sampling is intractable. When sampling is modeled as a sequential decision-making process, reinforcement learning (RL) methods, such a…

2025

Can a Bayesian Oracle Prevent Harm from an Agent?

UAI 2025

Is there a way to design powerful AI systems based on machine learning methods that would satisfy probabilistic safety guarantees? With the long-term goal of obtaining a probabilistic guarantee that would apply in every context, we consider estimating a context-dependent bound on the probability of

2025

Fast Flow-based Visuomotor Policies via Conditional Optimal Transport Couplings

CoRL 2025poster

Diffusion and flow matching policies have recently demonstrated remarkable performance in robotic applications by accurately capturing multimodal robot trajectory distributions. However, their computationally expensive inference, due to the numerical integration of an ODE or SDE, limits their applic…

Cited by 0SourceScholar
2025

Learning Diverse Attacks on Large Language Models for Robust Red-Teaming and Safety Tuning

ICLR 2025poster

Red-teaming, or identifying prompts that elicit harmful responses, is a critical step in ensuring the safe and responsible deployment of large language models (LLMs). Developing effective protection against many modes of attack prompts requires discovering diverse attacks. Automated red-teaming typi…

2025

Mixtures of In-Context Learners

ACL 2025long

In-context learning (ICL) adapts LLMs by providing demonstrations without fine-tuning the model parameters; however, it is very sensitive to the choice of in-context demonstrations, and processing many demonstrations can be computationally demanding. We propose Mixtures of In-Context Learners (MoICL…

2025

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models

ICML 2025poster

Any well-behaved generative model over a variable $\mathbf{x}$ can be expressed as a deterministic transformation of an exogenous (‘*outsourced'*) Gaussian noise variable $\mathbf{z}$: $\mathbf{x}=f_\theta(\mathbf{z})$. In such a model (*eg*, a VAE, GAN, or continuous-time flow-based model), sampli…

Cited by 0SourcePDFScholar
2025

PQMass: Probabilistic Assessment of the Quality of Generative Models using Probability Mass Estimation

ICLR 2025poster

We propose a likelihood-free method for comparing two distributions given samples from each, with the goal of assessing the quality of generative models. The proposed approach, PQMass, provides a statistically rigorous method for assessing the performance of a single generative model or the comparis…

Cited by 3SourcePDFScholar
2024

Amortizing intractable inference in diffusion models for vision, language, and control

NeurIPS 2024poster

Diffusion models have emerged as effective distribution estimators in vision, language, and reinforcement learning, but their use as priors in downstream tasks poses an intractable posterior inference problem. This paper studies *amortized* sampling of the posterior over data, $\mathbf{x}\sim p^{\rm…

2024

Amortizing intractable inference in large language models

ICLR 2024oral

Autoregressive large language models (LLMs) compress knowledge from their training data through next-token conditional distributions. This limits tractable querying of this knowledge to start-to-end autoregressive sampling. However, many tasks of interest---including sequence continuation, infilling…

2024

Delta-AI: Local objectives for amortized inference in sparse graphical models

ICLR 2024poster

We present a new algorithm for amortized inference in sparse probabilistic graphical models (PGMs), which we call $\Delta$-amortized inference ($\Delta$-AI). Our approach is based on the observation that when the sampling of variables in a PGM is seen as a sequence of actions taken by an agent, spar…

2024

Discrete Probabilistic Inference as Control in Multi-path Environments

UAI 2024poster

We consider the problem of sampling from a discrete and structured distribution as a sequential decision problem, where the objective is to find a stochastic policy such that objects are sampled at the end of this sequential process proportionally to some predefined reward. While we could use maximu…

2024

Expected flow networks in stochastic environments and two-player zero-sum games

ICLR 2024poster

Generative flow networks (GFlowNets) are sequential sampling models trained to match a given distribution. GFlowNets have been successfully applied to various structured object generation tasks, sampling a diverse set of high-reward objects quickly. We propose expected flow networks (EFlowNets), whi…

2024

Improved off-policy training of diffusion samplers

NeurIPS 2024poster

We study the problem of training diffusion models to sample from a distribution with a given unnormalized density or energy function. We benchmark several diffusion-structured inference methods, including simulation-based variational approaches and off-policy methods (continuous generative flow netw…

2024

Improving Gradient-Guided Nested Sampling for Posterior Inference

ICML 2024poster

We present a performant, general-purpose gradient-guided nested sampling (GGNS) algorithm, combining the state of the art in differentiable programming, Hamiltonian slice sampling, clustering, mode separation, dynamic nested sampling, and parallelization. This unique combination allows GGNS to scale…

2024

Iterated Denoising Energy Matching for Sampling from Boltzmann Densities

ICML 2024poster

Efficiently generating statistically independent samples from an unnormalized probability distribution, such as equilibrium samples of many-body systems, is a foundational problem in science. In this paper, we propose Iterated Denoising Energy Matching (iDEM), an iterative algorithm that uses a nove…

2024

PhyloGFN: Phylogenetic inference with generative flow networks

ICLR 2024poster

Phylogenetics is a branch of computational biology that studies the evolutionary relationships among biological entities. Its long history and numerous applications notwithstanding, inference of phylogenetic trees from sequence data remains challenging: the high complexity of tree space poses a sign…

2024

Simulation-Free Schrödinger Bridges via Score and Flow Matching

AISTATS 2024poster

We present simulation-free score and flow matching ([SF]$^2$M), a simulation-free objective for inferring stochastic dynamics given unpaired samples drawn from arbitrary source and target distributions. Our method generalizes both the score-matching loss used in the training of diffusion models and…

2023

A theory of continuous generative flow networks

ICML 2023poster

Generative flow networks (GFlowNets) are amortized variational inference algorithms that are trained to sample from unnormalized target distributions over compositional objects. A key limitation of GFlowNets until this time has been that they are restricted to discrete spaces. We present a theory fo…

2023

Better Training of GFlowNets with Local Credit and Incomplete Trajectories

ICML 2023poster

Generative Flow Networks or GFlowNets are related to Monte-Carlo Markov chain methods (as they sample from a distribution specified by an energy function), reinforcement learning (as they learn a policy to sample composed objects through a sequence of steps), generative models (as they learn to repr…

2023

GFlowNet-EM for Learning Compositional Latent Variable Models

ICML 2023poster

Latent variable models (LVMs) with discrete compositional latents are an important but challenging setting due to a combinatorially large number of possible configurations of the latents. A key tradeoff in modeling the posteriors over latents is between expressivity and tractable optimization. For a…

2023

GFlowNets and variational inference

ICLR 2023poster

This paper builds bridges between two families of probabilistic algorithms: (hierarchical) variational inference (VI), which is typically used to model distributions over continuous spaces, and generative flow networks (GFlowNets), which have been used for distributions over discrete structures such…

2023

GFlowOut: Dropout with Generative Flow Networks

ICML 2023poster

Bayesian inference offers principled tools to tackle many critical problems with modern neural networks such as poor calibration and generalization, and data inefficiency. However, scaling Bayesian inference to large architectures is challenging and requires restrictive approximations. Monte Carlo D…

Cited by 24SourcePDFScholar
2023

Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network

NeurIPS 2023poster

Generative Flow Networks (GFlowNets), a class of generative models over discrete and structured sample spaces, have been previously applied to the problem of inferring the marginal posterior distribution over the directed acyclic graph (DAG) of a Bayesian Network, given a dataset of observations. Ba…

2023

Learning GFlowNets From Partial Episodes For Improved Convergence And Stability

ICML 2023oral

Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks. Existing training objectives for GFlowNets are either local to states o…

2023

Let the Flows Tell: Solving Graph Combinatorial Problems with GFlowNets

NeurIPS 2023spotlight

Combinatorial optimization (CO) problems are often NP-hard and thus out of reach for exact algorithms, making them a tempting domain to apply machine learning methods. The highly structured constraints in these problems can hinder either optimization or sampling directly in the solution space. On th…

2023

ThinkSum: Probabilistic reasoning over sets using large language models

ACL 2023long

Large language models (LLMs) have a substantial capacity for high-level analogical reasoning: reproducing patterns in linear text that occur in their training data (zero-shot evaluation) or in the provided context (few-shot in-context learning). However, recent studies show that even the more advanc…

Cited by 31SourcePDFScholar
2022

Coherence boosting: When your pretrained language model is not paying enough attention

ACL 2022long

Long-range semantic coherence remains a challenge in automatic language generation and understanding. We demonstrate that large language models have insufficiently learned the effect of distant words on next-token prediction. We present coherence boosting, an inference procedure that increases a LM’…

2022

Diffusion Models as Plug-and-Play Priors

NeurIPS 2022accept

We consider the problem of inferring high-dimensional data $x$ in a model that consists of a prior $p(x)$ and an auxiliary differentiable constraint $c(x,y)$ on $x$ given some additional information $y$. In this paper, the prior is an independently trained denoising diffusion generative model. The a…

2022

Generative Flow Networks for Discrete Probabilistic Modeling

ICML 2022spotlight

We present energy-based generative flow networks (EB-GFN), a novel probabilistic modeling algorithm for high-dimensional discrete data. Building upon the theory of generative flow networks (GFlowNets), we model the generation process by a stochastic data construction policy and thus amortize expensi…

2022

Resolving label uncertainty with implicit posterior models

UAI 2022poster

We propose a method for jointly inferring labels across a collection of data samples, where each sample consists of an observation and a prior belief about the label. By implicitly assuming the existence of a generative model for which a differentiable predictor is the posterior, we derive a trainin…

2022

Trajectory balance: Improved credit assignment in GFlowNets

NeurIPS 2022accept

Generative flow networks (GFlowNets) are a method for learning a stochastic policy for generating compositional objects, such as graphs or strings, from a given unnormalized density by sequences of actions, where many possible action sequences may lead to the same object. We find previously proposed…

2021

GPT Perdetry Test: Generating new meanings for new words

NAACL 2021long

Human innovation in language, such as inventing new words, is a challenge for pretrained language models. We assess the ability of one large model, GPT-3, to process new words and decide on their meaning. We create a set of nonce words and prompt GPT-3 to generate their dictionary definitions. We fi…

2021

Studying word order through iterative shuffling

EMNLP 2021main

As neural language models approach human performance on NLP benchmark tasks, their advances are widely seen as evidence of an increasingly complex understanding of syntax. This view rests upon a hypothesis that has not yet been empirically tested: that word order encodes meaning essential to perform…

2020

Mining self-similarity: Label super-resolution with epitomic representations

ECCV 2020poster

We show that simple patch-based models, such as epitomes (Jojic et al., 2003), can have superior performance to the current state of the art in semantic segmentation and label super-resolution, which uses deep convolutional neural networks. We derive a new training algorithm for epitomes which allow…