Sanghyeok Choi

Diffusion Alignment as Variataional Expectation-Maximization

Inhyuck Song

Hyeongyu Kang

Woocheol Shin

Kiyoung Om

Jinkyoo Park

Diffusion alignment aims to optimize diffusion models for the downstream objective. While existing methods based on reinforcement learning o… (see more)r direct backpropagation achieve considerable success in maximizing rewards, they often suffer from reward over-optimization and mode collapse. We introduce Diffusion Alignment as Variational Expectation-Maximization (DAV), a framework that formulates diffusion alignment as an iterative process alternating between two complementary phases: the E-step and the M-step. In the E-step, we employ test-time search to generate diverse and reward-aligned samples. In the M-step, we refine the diffusion model using samples discovered by the E-step. We demonstrate that DAV can optimize reward while preserving diversity for both continuous and discrete tasks: text-to-image synthesis and DNA sequence design.

2025-12-31

International Conference on Learning Representations (Accept (Poster))

openreview.net

Reinforced Sequential Monte Carlo for Amortised Sampling

Sanghyeok Choi

Sarthak Mittal

Víctor Elvira

Jinkyoo Park

Esmeralda S. Whitammer

This paper proposes a synergy of amortised and particle-based methods for sampling from distributions defined by unnormalised density functi… (see more)ons. We state a connection between sequential Monte Carlo (SMC) and neural sequential samplers trained by maximum-entropy reinforcement learning (MaxEnt RL), wherein learnt sampling policies and value functions define proposal kernels and twist functions. Exploiting this connection, we introduce an off-policy RL training procedure for the sampler that uses samples from SMC -- using the learnt sampler as a proposal -- as a behaviour policy that better explores the target distribution. We describe techniques for stable joint training of proposals and twist functions and an adaptive weight tempering scheme to reduce training signal variance. Furthermore, building upon past attempts to use experience replay to guide the training of neural samplers, we derive a way to combine historical samples with annealed importance sampling weights within a replay buffer. On synthetic multi-modal targets (in both continuous and discrete spaces) and the Boltzmann distribution of alanine dipeptide conformations, we demonstrate improvements in approximating the true distribution as well as training stability compared to both amortised and Monte Carlo methods.

2025-12-31

International Conference on Machine Learning (Accept (spotlight))

openreview.net

Improved Off-policy Reinforcement Learning in Biological Sequence Design

Alex Hernández-García

Jinkyoo Park

Designing biological sequences with desired properties is challenging due to vast search spaces and limited evaluation budgets. Although rei… (see more)nforcement learning methods use proxy models for rapid reward evaluation, insufficient training data can cause proxy misspecification on out-of-distribution inputs. To address this, we propose a novel off-policy search,

2025-10-05

Proceedings of the 42nd International Conference on Machine Learning (published)

doi.org

proceedings.mlr.press

Ant Colony Sampling with GFlowNets for Combinatorial Optimization

Jiwoo Son

Jinkyoo Park

We present the Generative Flow Ant Colony Sampler (GFACS), a novel meta-heuristic method that hierarchically combines amortized inference an… (see more)d parallel stochastic search. Our method first leverages Generative Flow Networks (GFlowNets) to amortize a \emph{multi-modal} prior distribution over combinatorial solution space that encompasses both high-reward and diversified solutions. This prior is iteratively updated via parallel stochastic search in the spirit of Ant Colony Optimization (ACO), leading to the posterior distribution that generates near-optimal solutions. Extensive experiments across seven combinatorial optimization problems demonstrate GFACS's promising performances.

2025-05-02

International Conference on Artificial Intelligence and Statistics (Accept (Poster))

doi.org

proceedings.mlr.press

Adaptive Teachers for Amortized Samplers

Sungsoo Ahn

Jinkyoo Park

Nikolay Malkin

Yoshua Bengio

Amortized inference is the task of training a parametric model, such as a neural network, to approximate a distribution with a given unnorma… (see more)lized density where exact sampling is intractable. When sampling is implemented as a sequential decision-making process, reinforcement learning (RL) methods, such as generative flow networks, can be used to train the sampling policy. Off-policy RL training facilitates the discovery of diverse, high-reward candidates, but existing methods still face challenges in efficient exploration. We propose to use an adaptive training distribution (the \teacher) to guide the training of the primary amortized sampler (the \student). The \teacher, an auxiliary behavior model, is trained to sample high-loss regions of the \student and can generalize across unexplored modes, thereby enhancing mode coverage by providing an efficient training curriculum. We validate the effectiveness of this approach in a synthetic environment designed to present an exploration challenge, two diffusion-based sampling tasks, and four biochemical discovery tasks demonstrating its ability to improve sample efficiency and mode coverage. Source code is available at https://github.com/alstn12088/adaptive-teacher.

2025-04-22

International Conference on Learning Representations (Accept (Poster))

doi.org

openreview.net

Mila's Community of Practice: Agentic AI

Venture Sprint

AI Advantage: Productivity in Public Service

Sanghyeok Choi

Publications

Mila's Community of Practice: Agentic AI

Venture Sprint

AI Advantage: Productivity in Public Service

Popular keywords:

Sanghyeok Choi

Publications