Jae Hyun Lim

I am a Member of Technical Staff at Microsoft AI (London office). My research during graduate studies centered on deep generative modeling–particularly diffusion- and flow-based approaches–as well as approximate inference and deep reinforcement learning. I completed my Ph.D. in Computer Science at Mila, Université de Montréal, advised by Chris Pal. Before that, I received a B.S. in Bio and Brain Engineering and an M.S. in Electrical Engineering from KAIST.

email | google scholar | github | linkedin | x

news

Dec 1, 2025

I’ve started my new role at Microsoft AI!

Sep 8, 2025

Our paper Score-Based Diffusion Models in Function Space is published in JMLR!

Oct 15, 2021

I’m selected to receive a NeurIPS 2021 Outstanding Reviewer Award (the top 8% of reviewers)!

Sep 29, 2021

Our paper A Variational Perspective on Diffusion-Based Generative Models and Score Matching is accepted to NeurIPS 2021 and selected as a spotlight!

Jan 1, 2021

Our paper Bijective-Contrastive Estimation is accepted to AABI symposium 2021! (selected for a contributed talk)

projects

Score-based Diffusion Models in Function Space Jae Hyun Lim*, Nikola B Kovachki*, Ricardo Baptista*, Christopher Beckham, Kamyar Azizzadenesheli, Jean Kossaifi, Vikram Voleti, Jiaming Song, Karsten Kreis, Jan Kautz, Christopher Pal, Arash Vahdat, and Anima Anandkumar Journal of Machine Learning Research 2025 [code] [pdf]
Diffusion models have recently emerged as a powerful framework for generative modeling. They consist of a forward process that perturbs input data with Gaussian white noise and a reverse process that learns a score function to generate samples by denoising. Despite their tremendous success, they are mostly formulated on finite-dimensional spaces, e.g. Euclidean, limiting their applications to many domains where the data has a functional form such as in scientific computing and 3D geometric data analysis. In this work, we introduce a mathematically rigorous framework called Denoising Diffusion Operators (DDOs) for training diffusion models in function space. In DDOs, the forward process perturbs input functions gradually using a Gaussian process. The generative process is formulated by integrating a function-valued Langevin dynamic. Our approach requires an appropriate notion of the score for the perturbed data distribution, which we obtain by generalizing denoising score matching to function spaces that can be infinite-dimensional. We show that the corresponding discretized algorithm generates accurate samples at a fixed cost that is independent of the data resolution. We theoretically and numerically verify the applicability of our approach on a set of problems, including generating solutions to the Navier-Stokes equation viewed as the push-forward distribution of forcings from a Gaussian Random Field (GRF).

A variational perspective on diffusion-based generative models and score matching Chin-Wei Huang, Jae Hyun Lim, and Aaron Courville In Advances in Neural Information Processing Systems 2021 [code] [colab] [pdf]
Discrete-time diffusion-based generative models and score matching methods have shown promising results in modeling high-dimensional image data. Recently, Song et al. (2021) show that diffusion processes that transform data into noise can be reversed via learning the score function, i.e. the gradient of the log-density of the perturbed data. They propose to plug the learned score function into an inverse formula to define a generative diffusion process. Despite the empirical success, a theoretical underpinning of this procedure is still lacking. In this work, we approach the (continuous-time) generative diffusion directly and derive a variational framework for likelihood estimation, which includes continuous-time normalizing flows as a special case, and can be seen as an infinitely deep variational autoencoder. Under this framework, we show that minimizing the score-matching loss is equivalent to maximizing a lower bound of the likelihood of the plug-in reverse SDE proposed by Song et al. (2021), bridging the theoretical gap.

Bijective-Contrastive Estimation Jae Hyun Lim*, Chin-Wei Huang*, Aaron Courville, and Christopher Pal 3rd Symposium on Advances in Approximate Bayesian Inference 2021 [pdf] [video]
In this work, we propose Bijective-Contrastive Estimation (BCE), a classification-based learning criterion for energy-based models. We generate a collection of contrasting distributions using bijections, and solve all the classification problems between the original data distribution and the distributions induced by the bijections using a classifier parameterized by an energy model. We show that if the classification objective is minimized, the energy function will uniquely recover the data density up to a normalizing constant. This has the benefit of not having to explicitly specify a contrasting distribution, like noise contrastive estimation. Experimentally, we demonstrate that the proposed method works well on 2D synthetic datasets. We discuss the difficulty in high dimensional cases, and propose potential directions to explore for future work.

AR-DAE: Towards Unbiased Neural Entropy Gradient Estimation Jae Hyun Lim, Aaron Courville, Christopher Pal, and Chin-Wei Huang In International Conference on Machine Learning 2020 [vae code] [rl code] [colab] [arxiv]
Entropy is ubiquitous in machine learning, but it is in general intractable to compute the entropy of the distribution of an arbitrary continuous random variable. In this paper, we propose the amortized residual denoising autoencoder (AR-DAE) to approximate the gradient of the log density function, which can be used to estimate the gradient of entropy. Amortization allows us to significantly reduce the error of the gradient approximator by approaching asymptotic optimality of a regular DAE, in which case the estimation is in theory unbiased. We conduct theoretical and experimental analyses on the approximation error of the proposed method, as well as extensive studies on heuristics to ensure its robustness. Finally, using the proposed gradient approximator to estimate the gradient of entropy, we demonstrate state-of-the-art performance on density estimation with variational autoencoders and continuous control with soft actor-critic.

Neural Multisensory Scene Inference Jae Hyun Lim, Pedro O Pinheiro, Negar Rostamzadeh, Christopher Pal, and Sungjin Ahn In Advances in Neural Information Processing Systems 2019 [project page] [code] [pdf]
For embodied agents to infer representations of the underlying 3D physical world they inhabit, they should efficiently combine multisensory cues from numerous trials, e.g., by looking at and touching objects. Despite its importance, multisensory 3D scene representation learning has received less attention compared to the unimodal setting. In this paper, we propose the Generative Multisensory Network (GMN) for learning latent representations of 3D scenes which are partially observable through multiple sensory modalities. We also introduce a novel method, called the Amortized Product-of-Experts, to improve the computational efficiency and the robustness to unseen combinations of modalities at test time. Experimental results demonstrate that the proposed model can efficiently infer robust modality-invariant 3D-scene representations from arbitrary combinations of modalities and perform accurate cross-modal generation. To perform this exploration we have also developed a novel multi-sensory simulation environment for embodied agents.

Geometric Gan Jae Hyun Lim, and Jong Chul Ye arXiv preprint arXiv:1705.02894 2017 [code] [arxiv]
Generative Adversarial Nets (GANs) represent an important milestone for effective generative models, which has inspired numerous variants seemingly different from each other. One of the main contributions of this paper is to reveal a unified geometric structure in GAN and its variants. Specifically, we show that the adversarial generative model training can be decomposed into three geometric steps: separating hyperplane search, discriminator parameter update away from the separating hyperplane, and the generator update along the normal vector direction of the separating hyperplane. This geometric intuition reveals the limitations of the existing approaches and leads us to propose a new formulation called geometric GAN using SVM separating hyperplane that maximizes the margin. Our theoretical analysis shows that the geometric GAN converges to a Nash equilibrium between the discriminator and generator. In addition, extensive numerical results show that the superior performance of geometric GAN.