(*) denotes equal contribution.
Preprints
Demystifying Diffusion Objectives: Reweighted Losses are Better Variational Bounds
Jiaxin Shi, Michalis K. Titsias
Refereed Publications
Variational Learning for Insertion-based Generation
Yangtian Zhang, Zhe Wang, Arthur Gretton, Rex Ying, David van Dijk, Michalis K. Titsias, Jiaxin Shi
Spotlight Presentation (top 2.2%).
CANDI: Hybrid Discrete-Continuous Diffusion Models
Patrick Pynadath, Jiaxin Shi, Ruqi Zhang
The Efficiency Gap in Byte Modeling
Celine Lee, Jing Nathan Yan, Chen Liang, Jiaxin Shi, Yin Zhang, Jeremiah Liu, Pengcheng Yin, Fernando Pereira, Ed Chi, Derek Cheng, Alexander M. Rush, Ruoxi Wang
Test-time Regression: A Unifying Framework for Designing Sequence Models with Associative Memory
Ke Alexander Wang, Jiaxin Shi, Emily B. Fox
Self-Speculative Masked Diffusions
Andrew Campbell, Valentin De Bortoli, Jiaxin Shi, Arnaud Doucet
Neural Eigenfunctions Are Structured Representation Learners
Zhijie Deng*, Jiaxin Shi*, Hao Zhang, Peng Cui, Cewu Lu, Jun Zhu.
IEEE Transaction on Pattern Analysis and Machine Intelligence (TPAMI). [pdf] [abs]
Informed Correctors for Discrete Diffusion Models
Yixiu Zhao, Jiaxin Shi, Feng Chen, Shaul Druckmann, Lester Mackey, Scott Linderman
NeurIPS 2025 [pdf] [abs] [code]
Generating Creative Chess Puzzles
Xidong Feng, Vivek Veeriah, Marcus Chiam, Michael D Dennis, Federico Barbero, Johan Obando-Ceron, Jiaxin Shi, Satinder Singh, Shaobo Hou, Nenad Tomasev, Tom Zahavy
NeurIPS 2025 [pdf] [abs] [booklet & review] [media]
Learning-Order Autoregressive Models with Application to Molecular Graph Generation
Zhe Wang, Jiaxin Shi, Nicolas Heess, Arthur Gretton, Michalis K. Titsias
Simplified and Generalized Masked Diffusion for Discrete Data
Jiaxin Shi*, Kehang Han*, Zhe Wang, Arnaud Doucet, Michalis K. Titsias
NeurIPS 2024 [pdf] [abs] [code] [slides]
A Finite-Particle Convergence Rate for Stein Variational Gradient Descent
Jiaxin Shi, Lester Mackey
Sequence Modeling with Multiresolution Convolutional Memory
Jiaxin Shi, Ke Alexander Wang, Emily B. Fox
Gradient Estimation with Discrete Stein Operators
Jiaxin Shi, Yuhao Zhou, Jessica Hwang, Michalis K. Titsias, Lester Mackey
NeurIPS 2022 [pdf] [abs] [code]
NeurIPS 2022 Outstanding Paper Award
NeuralEF: Deconstructing Kernels by Deep Neural Networks
Zhijie Deng, Jiaxin Shi, Jun Zhu
Double Control Variates for Gradient Estimation in Discrete Latent Variable Models
Michalis K. Titsias, Jiaxin Shi
AISTATS 2022 [pdf] [abs] [code]
Sampling with Mirrored Stein Operators
Jiaxin Shi, Chang Liu, Lester Mackey
ICLR 2022 [pdf] [abs] [code] [slides]
Spotlight Presentation (top 5.1%).
Scalable Variational Gaussian Processes via Harmonic Kernel Decomposition
Shengyang Sun, Jiaxin Shi, Andrew Gordon Wilson, Roger Grosse
Nonparametric Score Estimators
Yuhao Zhou, Jiaxin Shi, Jun Zhu
ICML 2020 [pdf] [abs] [code] [slides]
Sparse Orthogonal Variational Inference for Gaussian Processes
Jiaxin Shi, Michalis K. Titsias, Andriy Mnih
AISTATS 2020 [pdf] [abs] [code] [slides]
Best Student Paper Runner-Up at AABI Symposium, 2019
Sliced Score Matching: A Scalable Approach to Density and Score Estimation
Yang Song*, Sahaj Garg*, Jiaxin Shi, Stefano Ermon
UAI 2019 [pdf] [abs] [code] [video] [blog]
Oral Presentation (top 8.7%)
Scalable Training of Inference Networks for Gaussian-Process Models
Jiaxin Shi, Mohammad Emtiyaz Khan, Jun Zhu
ICML 2019 [pdf] [abs] [code] [slides]
Functional Variational Bayesian Neural Networks
Shengyang Sun*, Guodong Zhang*, Jiaxin Shi*, Roger Grosse
ICLR 2019 [pdf] [abs] [code] [video]
Semi-crowdsourced Clustering with Deep Generative Models
Yucen Luo, Tian Tian, Jiaxin Shi, Jun Zhu, Bo Zhang
NeurIPS 2018 [pdf] [abs] [code]
A Spectral Approach to Gradient Estimation for Implicit Distributions
Jiaxin Shi, Shengyang Sun, Jun Zhu
ICML 2018. [pdf] [abs] [code] [slides]
Message Passing Stein Variational Gradient Descent
Jingwei Zhuo, Chang Liu, Jiaxin Shi, Jun Zhu, Ning Chen, Bo Zhang
Kernel Implicit Variational Inference
Jiaxin Shi*, Shengyang Sun*, Jun Zhu
Workshop Papers
Learning Absorption Rates in Glucose-Insulin Dynamics from Meal Covariates
Ke Alexander Wang*, Matthew E. Levine*, Jiaxin Shi, Emily B. Fox [pdf]
NeurIPS Workshop: Learning from Time Series for Health, 2022
Neural Networks as Inter-domain Inducing Points
Shengyang Sun*, Jiaxin Shi*, Roger Grosse [pdf] [slides] [video]
Symposium on Advances in Approximate Bayesian Inference, 2020
Spectral Estimators for Gradient Fields of Log-Densities
Yuhao Zhou, Jiaxin Shi, Jun Zhu
ICML Workshop on Stein’s Method, 2019
Visualization & Graphics
Analyzing the Training Processes of Deep Generative Models
Mengchen Liu, Jiaxin Shi, Kelei Cao, Jun Zhu, Shixia Liu
IEEE Transactions on Visualization and Computer Graphics, 2018 [pdf]
Towards Better Analysis of Deep Convolutional Neural Networks
Mengchen Liu, Jiaxin Shi, Zhen Li, Chongxuan Li, Jun Zhu, Shixia Liu
IEEE Transactions on Visualization and Computer Graphics, 2017 [pdf]
Most cited paper of TVCG in 2017
Plenopatch: Patch-based Plenoptic Image Manipulation
Fanglue Zhang, Jue Wang, Eli Shechtman, Ziye Zhou, Jiaxin Shi, Shimin Hu
IEEE Transactions on Visualization and Computer Graphics, 2017 [paper]