(*) denotes equal contribution.

Preprints

Learning-Order Autoregressive Models with Application to Molecular Graph Generation

Zhe Wang, Jiaxin Shi, Nicolas Heess, Arthur Gretton, Michalis K. Titsias

Preprint, 2025 [pdf] [abs]

Test-time Regression: A Unifying Framework for Designing Sequence Models with Associative Memory

Ke Alexander Wang, Jiaxin Shi, Emily B. Fox

Preprint, 2025 [pdf] [abs]

Informed Correctors for Discrete Diffusion Models

Yixiu Zhao, Jiaxin Shi, Lester Mackey, Scott Linderman

Preprint, 2024 [pdf] [abs]

Neural Eigenfunctions Are Structured Representation Learners

Zhijie Deng*, Jiaxin Shi*, Hao Zhang, Peng Cui, Cewu Lu, Jun Zhu.

Preprint, 2022 [pdf] [abs]

Refereed Conference Publications

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

NeurIPS 2023 [pdf] [abs]

Sequence Modeling with Multiresolution Convolutional Memory

Jiaxin Shi, Ke Alexander Wang, Emily B. Fox

ICML 2023 [pdf] [abs] [code]

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

ICML 2022 [pdf] [abs] [code]

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

ICML 2021 [pdf] [abs] [code]

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

ICML 2018 [pdf] [abs]

Kernel Implicit Variational Inference

Jiaxin Shi*, Shengyang Sun*, Jun Zhu

ICLR 2018 [pdf] [abs]

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]