Research

Statistics and AI

At the interface of statistics and AI, I study how statistical principles can inform the construction of learning systems. How should a model account for differences across individuals? Which relationships should a graph representation capture? What information should be shared across different measurement types? These questions connect AI model design to statistical ideas about heterogeneity, dependence, and generalization.

In reinforcement learning, our work develops individualized policies from previously collected data on heterogeneous populations. By representing individual differences through latent variables, we learn decision rules that account for variation in how people respond to actions. This framework combines individualized policy estimation with theoretical guarantees on policy performance.

In graph learning, our MaGNET framework integrates information from local neighborhoods and more distant graph relationships. It identifies influential nodes, edges, and features, connecting predictions to interpretable graph structures. In multimodal learning, our work on variational autoencoders investigates how distinct measurement types can contribute to useful representations. Our recent Meta Fusion framework extends this direction through mutual learning across models, unifying early, intermediate, and late fusion.

Together, these projects investigate how the treatment of variation and structure within a model affects what it can learn and how its results can be interpreted.

• Miao, R., Shahbaba, B., & Qu, A. (2025). Reinforcement Learning for Individual Optimal Policy from Heterogeneous Data. Annals of Statistics, 53(4), 1513–1534.
• Zhou, W., Qu, A., Cooper, K., Fortin, N., & Shahbaba, B. (2024). A Model-Agnostic Graph Neural Network for Integrating Local and Global Information. Journal of the American Statistical Association, 120(550), 1225–1238.
• Sutter, T. M., Meng, Y., Fortin, N., Vogt, J. E., Shahbaba, B., & Mandt, S. (2024). Unity by Diversity: Improved Representation Learning for Multimodal VAEs. NeurIPS.
• Liang, Z., Qu, A., & Shahbaba, B. (2026). Meta Fusion: A Unified Framework for Multimodality Fusion with Mutual Learning. Journal of the Royal Statistical Society, Series B (to appear).

Bayesian modeling and computation

My Bayesian research develops flexible probability models and the computational methods needed to use them. A central question is how to represent complex relationships in data while keeping inference reliable and computationally feasible. My early work on Dirichlet process mixtures developed nonlinear prediction models that adapt their complexity to the data. Related work uses Gaussian processes and latent-variable models to describe dependence in neural and other biomedical measurements.

A sustained thread develops sampling methods that exploit the structure and geometry of posterior distributions. Split Hamiltonian Monte Carlo separates the Hamiltonian into components so that much of the simulation can be performed at lower computational cost. Spherical Hamiltonian Monte Carlo addresses constrained distributions, while wormhole Hamiltonian Monte Carlo facilitates movement between separated modes. Our work on distributed stochastic-gradient MCMC extends Bayesian computation to large datasets by coordinating sampling across workers.

More recently, I have investigated how neural networks can support Bayesian modeling and inference. This includes methods that approximate expensive computations within sampling workflows, a Calibration–Emulation–Sampling strategy for Bayesian neural networks, and fully Bayesian autoencoders with sparse Gaussian-process priors. These developments extend my broader interest in designing models and algorithms together, so that computational advances support richer statistical inference.

• Shahbaba, B., & Neal, R. M. (2009). Nonlinear models using Dirichlet process mixtures. Journal of Machine Learning Research, 10, 1829–1850.
• Shahbaba, B., Lan, S., Johnson, W. O., & Neal, R. M. (2014). Split Hamiltonian Monte Carlo. Statistics and Computing, 24(3), 339–349.
• Lan, S., Streets, J., & Shahbaba, B. (2014). Wormhole Hamiltonian Monte Carlo. AAAI.
• Tran, B., Shahbaba, B., Mandt, S., & Filippone, M. (2023). Fully Bayesian Autoencoders with Latent Sparse Gaussian Processes. ICML.
• Moslemi, Z., Meng, Y., Lan, S., & Shahbaba, B. (2024). Scaling Up Bayesian Neural Networks with Neural Networks. Transactions on Machine Learning Research.

Scientific discovery in neuroscience and health

My scientific collaborations investigate how biological systems represent information, change over time, and relate to health outcomes. I contribute to question formulation, study design, method development, and interpretation, connecting statistical advances to the scientific problems that motivate them.

In neuroscience, we study how populations of neurons represent sequences and support memory and decision-making. Our work on hippocampal ensembles showed that neural representations of sequential relationships extend to discrete, nonspatial events. Related projects use latent-factor Gaussian processes to model changing functional connectivity and optimal transport to align neuronal activity across heterogeneous datasets. Current work investigates multi-step planning and goal-directed choice, connecting learned representations to temporal and decision models.

My biomedical collaborations also span Alzheimer’s disease, stroke recovery, circadian biology, and maternal and infant health. Recent projects examine how low-burden clinical measurements can help identify dementia neuropathology and how transcriptomic data can reveal circadian time. Ongoing research on wearable and health data develops methods for individualized risk prediction and intervention. Across these applications, the aim is to connect statistical evidence to specific scientific questions and decisions.

• Li, L., Pluta, D., Shahbaba, B., Fortin, N., Ombao, H., & Baldi, P. (2019). Modeling Dynamic Functional Connectivity with Latent Factor Gaussian Processes. NeurIPS.
• Shahbaba, B., et al. (2022). Hippocampal ensembles represent sequential relationships among discrete nonspatial events. Nature Communications, 13, 787.
• Duan, J., et al. (2024). tauFisher accurately predicts circadian time from a single sample of bulk and single-cell transcriptomic data. Nature Communications, 15, 3840.
• Yuan, Y., et al. (2025). Optimal Transport based Cross-Domain Integration for Heterogeneous Data. Journal of the American Statistical Association, 120(551), 1449–1462.
• Ren, Y., Shahbaba, B., & Stark, C. (2025). Identifying dementia neuropathology using low-burden clinical data. Alzheimer’s & Dementia, 21(8), e70539.