Balaji Lakshminarayanan

Hello! I am a Staff Research Scientist at Google Brain in Mountain View (USA), working on Machine Learning and its applications.

My research interests are in scalable, probabilistic machine learning. My PhD thesis [17] was focused on exploring (and exploiting :) connections between neat mathematical ideas in (non-parametric) Bayesian land and computationally efficient tricks in decision tree land, to get the best of both worlds.

More recently, I have focused on probabilistic deep learning:

Quick Links

Publications

Here's my Google Scholar page.

  1. Soft Calibration Objectives for Neural Networks
    Archit Karandikar, Nicholas Cain, Dustin Tran, Balaji Lakshminarayanan, Jonathon Shlens, Michael C. Mozer and Becca Roelofs
    [arXiv]

  2. BEDS-Bench: Behavior of EHR-models under Distributional Shift–A Benchmark
    Anand Avati, Martin Seneviratne, Emily Xue, Zhen Xu, Balaji Lakshminarayanan and Andrew M. Dai
    [arXiv]

  3. Exploring the Limits of Out-of-Distribution Detection
    Stanislav Fort, Jie Ren and Balaji Lakshminarayanan
    [arXiv]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2021.

  4. A Simple Fix to Mahalanobis Distance for Improving Near-OOD Detection
    Jie Ren, Stanislav Fort, Jeremiah Liu, Abhijit Guha Roy, Shreyas Padhy and Balaji Lakshminarayanan
    [arXiv]
    Note: Presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2021.

  5. Predicting Unreliable Predictions by Shattering a Neural Network
    Xu Ji, Razvan Pascanu, Devon Hjelm, Andrea Vedaldi, Balaji Lakshminarayanan, Yoshua Bengio
    [arXiv]

  6. What are effective labels for augmented data? Improving robustness with AutoLabel
    Yao Qin, Xuezhi Wang, Balaji Lakshminarayanan, Ed Chi, Alex Beutel
    [link]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2021.

  7. Uncertainty Baselines: Benchmarks for Uncertainty & Robustness in Deep Learning
    Zachary Nado, Neil Band, Mark Collier, Josip Djolonga, Michael Dusenberry, Sebastian Farquhar, Angelos Filos, Marton Havasi, Rodolphe Jenatton, Ghassen Jerfel, Jeremiah Liu, Zelda Mariet, Jeremy Nixon, Shreyas Padhy, Jie Ren, Tim Rudner, Yeming Wen, Florian Wenzel, Kevin Murphy, D. Sculley, Balaji Lakshminarayanan, Jasper Snoek, Yarin Gal, Dustin Tran
    [arXiv]

  8. A Realistic Simulation Framework for Learning with Label Noise
    Keren Gu, Xander Masotto, Vandana Bachani, Balaji Lakshminarayanan, Jack Nikodem, Dong Yin
    [link]

  9. Task-agnostic Continual Learning with Hybrid Probabilistic Models
    Polina Kirichenko, Mehrdad Farajtabar, Dushyant Rao, Balaji Lakshminarayanan, Nir Levine, Ang Li, Huiyi Hu, Andrew Gordon Wilson, Razvan Pascanu
    [link]
    Note: A short version of this work was presented at the ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021.

  10. Does Your Dermatology Classifier Know What It Doesn't Know? Detecting the Long-Tail of Unseen Conditions
    Abhijit Guha Roy, Jie Ren, Shekoofeh Azizi, Aaron Loh, Vivek Natarajan, Basil Mustafa, Nick Pawlowski, Jan Freyberg, Yuan Liu, Zach Beaver, Nam Vo, Peggy Bui, Samantha Winter, Patricia MacWilliams, Greg S. Corrado, Umesh Telang, Yun Liu, Taylan Cemgil, Alan Karthikesalingam, Balaji Lakshminarayanan, Jim Winkens
    [arXiv]

  11. Normalizing flows for probabilistic modeling and inference
    George Papamakarios, Eric Nalisnick, Danilo J. Rezende, Shakir Mohamed and Balaji Lakshminarayanan
    [arXiv]
    JMLR, 2021

  12. Density of States Estimation for Out-of-Distribution Detection
    Warren R. Morningstar, Cusuh Ham, Andrew G. Gallagher, Balaji Lakshminarayanan, Alexander A. Alemi and Josh Dillon
    AISTATS, 2021
    [arXiv]

  13. Training independent subnetworks for robust prediction
    Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Liu, Jasper Snoek, Balaji Lakshminarayanan, Andrew Dai, Dustin Tran
    ICLR, 2021
    [arXiv]

  14. Combining Ensembles and Data Augmentation can Harm your Calibration
    Yeming Wen, Ghassen Jerfel, Rafael Muller, Michael Dusenberry, Jasper Snoek, Balaji Lakshminarayanan and Dustin Tran
    ICLR, 2021
    [arXiv]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2020.

  15. Why Aren’t Bootstrapped Neural Networks Better?
    Jeremy Nixon, Dustin Tran and Balaji Lakshminarayanan
    [pdf]
    Note: This work was presented at the “I Can’t Believe It’s Not Better!” workshop at NeurIPS 2020.

  16. Bayesian Deep Ensembles via the Neural Tangent Kernel
    Bobby He, Balaji Lakshminarayanan and Yee Whye Teh
    NeurIPS, 2020
    [arXiv]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2020.

  17. Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness
    Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax-Weiss and Balaji Lakshminarayanan
    NeurIPS, 2020
    [arXiv]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2020.

  18. Revisiting One-vs-All Classifiers for Predictive Uncertainty and Out-of-Distribution Detection in Neural Networks
    Shreyas Padhy, Zachary Nado, Jie Ren, Jeremiah Liu, Jasper Snoek and Balaji Lakshminarayanan
    [arXiv]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2020.

  19. Evaluating Prediction-Time Batch Normalization for Robustness under Covariate Shift
    Zachary Nado, Shreyas Padhy, D Sculley, Alexander D'Amour, Balaji Lakshminarayanan and Jasper Snoek
    [arXiv]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2020.

  20. Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors
    Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yi-an Ma, Jasper Snoek, Katherine Heller, Balaji Lakshminarayanan and Dustin Tran
    ICML, 2020
    [arXiv]

  21. AugMix: A simple data processing method to improve robustness and uncertainty
    Dan Hendrycks, Norman Mu, Ekin D. Cubuk, Barret Zoph, Justin Gilmer and Balaji Lakshminarayanan
    ICLR, 2020
    [arXiv] [code]

  22. Deep ensembles: A loss landscape perspective
    Stanislav Fort, Clara Huiyi Hu and Balaji Lakshminarayanan
    [arXiv] [poster] [slides]
    Note: A short version of this work was presented as a contributed oral talk at the NeurIPS workshop on Bayesian deep learning, 2019.

  23. Detecting out-of-distribution inputs to deep generative models using a test for typicality
    Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh and Balaji Lakshminarayanan
    [arXiv] [poster]
    Note: A short version of this work was presented at the NeurIPS workshop on Bayesian deep learning, 2019.

  24. Likelihood ratios for out-of-distribution detection
    Jie Ren, Peter Liu, Emily Fertig, Jasper Snoek, Ryan Poplin, Mark DePristo, Josh Dillon and Balaji Lakshminarayanan
    NeurIPS, 2019
    [arXiv] [code] [poster] [3-minute video] [blog]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2019.

  25. Can you trust your model's uncertainty? Evaluating predictive uncertainty under dataset shift
    Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Josh Dillon, Balaji Lakshminarayanan and Jasper Snoek
    NeurIPS, 2019
    [arXiv] [code] [poster] [blog]
    Note: A short version of this work was presented at the ICML workshop on Uncertainty and Robustness in Deep Learning, 2019.

  26. Learning from delayed outcomes via proxies with applications to recommender systems
    Timothy Mann*, Sven Gowal*, Andras Gyorgy, Ray Jiang, Clara Huiyi Hu, Balaji Lakshminarayanan and Prav Srinivasan
    ICML, 2019
    [link]

  27. Hybrid models with deep and invertible features
    Eric Nalisnick*, Akihiro Matsukawa*, Yee Whye Teh, Dilan Gorur and Balaji Lakshminarayanan
    ICML, 2019
    [arXiv] [poster]
    Note: A short version of this work was presented at the NeurIPS workshop on Bayesian deep learning, 2018.

  28. Do deep generative models know what they don't know?
    Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur and Balaji Lakshminarayanan
    [arXiv] [poster]
    ICLR, 2019
    Note: A short version of this work was presented at the NeurIPS workshop on Bayesian deep learning, 2018.

  29. Adapting auxiliary losses using gradient similarity
    Yunshu Du, Wojciech Czarnecki, Siddhant Jayakumar, Razvan Pascanu and Balaji Lakshminarayanan
    [arXiv] [poster]
    Note: A short version of this work was presented at the NeurIPS workshop on continual learning, 2018.

  30. Clinically applicable deep learning for diagnosis and referral in retinal disease
    Jeffrey De Fauw, Joseph R Ledsam, Bernardino Romera-Paredes, Stanislav Nikolov, Nenad Tomasev, Sam Blackwell, Harry Askham, Xavier Glorot, Brendan O'Donoghue, Daniel Visentin, George van den Driessche, Balaji Lakshminarayanan, Clemens Meyer, Faith Mackinder, Simon Bouton, Kareem Ayoub, Reena Chopra, Dominic King, Alan Karthikesalingam, Cían O Hughes, Rosalind Raine, Julian Hughes, Dawn A Sim, Catherine Egan, Adnan Tufail, Hugh Montgomery, Demis Hassabis, Geraint Rees, Trevor Back, Peng T. Khaw, Mustafa Suleyman, Julien Cornebise, Pearse A. Keane, Olaf Ronneberger
    Nature Medicine, 2018
    [link] [blog]

  31. Distribution matching in variational inference
    Mihaela Rosca, Balaji Lakshminarayanan and Shakir Mohamed
    [arXiv]

  32. Many paths to equilibrium: GANs do not need to decrease a divergence at every step
    William Fedus*, Mihaela Rosca*, Balaji Lakshminarayanan, Andrew Dai, Shakir Mohamed and Ian Goodfellow
    ICLR, 2018
    [arXiv]

  33. Variational approaches for auto-encoding generative adversarial networks
    Mihaela Rosca*, Balaji Lakshminarayanan*, David Warde-Farley and Shakir Mohamed
    * denotes equal contribution
    [arXiv]

  34. The Cramer distance as a solution to biased Wasserstein gradients
    Marc G. Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer and Remi Munos
    [arXiv]

  35. Comparison of maximum likelihood and GAN-based training of Real NVPs
    Ivo Danihelka, Balaji Lakshminarayanan, Benigno Uria, Daan Wierstra and Peter Dayan
    [arXiv]

  36. Simple and scalable predictive uncertainty estimation using deep ensembles
    Balaji Lakshminarayanan, Alexander Pritzel and Charles Blundell
    NeurIPS, 2017
    [arXiv] [slides] [poster]
    Note: A version of this work was presented at the NeurIPS workshop on Bayesian deep learning, 2016.

  37. Learning in implicit generative models
    Shakir Mohamed* and Balaji Lakshminarayanan*
    * denotes equal contribution
    [arXiv]
    Note: A version of this work was presented at the NeurIPS workshop on adversarial training, 2016.

  38. Learning deep nearest neighbor representations using differentiable boundary trees
    Daniel Zoran, Balaji Lakshminarayanan and Charles Blundell
    [arXiv]

  39. Distributed Bayesian learning with stochastic natural-gradient expectation propagation and the posterior server
    Leonard Hasenclever, Stefan Webb, Thibaut Lienart, Sebastian Vollmer, Balaji Lakshminarayanan, Charles Blundell and Yee Whye Teh
    JMLR, 2017
    [arXiv] [code]

  40. Decision trees and forests: a probabilistic perspective
    Balaji Lakshminarayanan
    Ph.D. thesis, University College London, 2016
    [pdf]

  41. The Mondrian kernel
    Matej Balog, Balaji Lakshminarayanan, Zoubin Ghahramani, Daniel M. Roy and Yee Whye Teh
    UAI, 2016
    [arXiv] [code] [slides] [poster]

  42. Approximate inference with the variational Holder bound
    Guillaume Bouchard and Balaji Lakshminarayanan
    [arXiv]

  43. Mondrian forests for large-scale regression when uncertainty matters
    Balaji Lakshminarayanan, Daniel M. Roy and Yee Whye Teh
    AISTATS, 2016
    [pdf] [code] [slides]

  44. Kernel-based just-in-time learning for passing expectation propagation messages
    Wittawat Jitkrittum, Arthur Gretton, Nicolas Heess, Ali Eslami, Balaji Lakshminarayanan, Dino Sejdinovic and Zoltan Szabo
    UAI, 2015
    [arXiv] [code]

  45. Particle Gibbs for Bayesian additive regression trees
    Balaji Lakshminarayanan, Daniel M. Roy and Yee Whye Teh
    AISTATS, 2015
    [pdf] [code]

  46. Mondrian forests: Efficient online random forests
    Balaji Lakshminarayanan, Daniel M. Roy and Yee Whye Teh
    NeurIPS, 2014
    [pdf] [code] [slides]
    Note: See here for a much faster, scikit-learn compatible re-implementation of Mondrian forests.

  47. Distributed Bayesian posterior sampling via moment sharing
    Minjie Xu, Balaji Lakshminarayanan, Yee Whye Teh, Jun Zhu and Bo Zhang
    NeurIPS, 2014
    [pdf] [code]

  48. Latent IBP compound Dirichlet allocation
    Cedric Archambeau, Balaji Lakshminarayanan, and Guillaume Bouchard
    TPAMI special issue on Bayesian Nonparametrics, 2015
    [pdf] [code available upon request] [IEEE link]
    Note: A short version of this work appeared at the NeurIPS workshop on Bayesian nonparametrics, 2011.

  49. Inferring ground truth from multi-annotator ordinal data: a probabilistic approach
    Balaji Lakshminarayanan and Yee Whye Teh
    [arXiv] [code]

  50. Top-down particle filtering for Bayesian decision trees
    Balaji Lakshminarayanan, Daniel M. Roy and Yee Whye Teh
    ICML, 2013
    [pdf] [code] [slides]

  51. Acoustic classification of multiple simultaneous bird species: A multi-instance multi-label approach
    Forrest Briggs, Balaji Lakshminarayanan, Lawrence Neal, Xiaoli Z. Fern, Raviv Raich, Sarah J. K. Hadley, Adam S. Hadley, and Matthew G. Betts
    JASA, 2012
    [link]

  52. Inference in supervised latent Dirichlet allocation
    Balaji Lakshminarayanan and Raviv Raich
    MLSP, 2011
    [pdf] [code available upon request] [IEEE link]

  53. Robust Bayesian matrix factorization
    Balaji Lakshminarayanan, Guillaume Bouchard, and Cedric Archambeau
    AISTATS, 2011
    [revised pdf] [code available upon request]
    Note: the updates for a_n, c_m were wrong in the original version of the pdf. The increments ought to be ell_n and ell_m respectively instead of 1.

  54. Probabilistic models for classification of bioacoustic data
    Balaji Lakshminarayanan
    M.S. thesis, Oregon State University, 2010
    [pdf]

  55. Non-negative matrix factorization for parameter estimation in Hidden Markov models
    Balaji Lakshminarayanan and Raviv Raich
    MLSP, 2010
    [pdf] [IEEE link]

  56. A syllable-level probabilistic framework for bird species identification
    Balaji Lakshminarayanan, Raviv Raich and Xiaoli Z. Fern
    ICMLA, 2009
    [pdf] [IEEE link]

Copyright notice: This material is presented to ensure timely dissemination of scholarly and technical work. Copyright and all rights therein are retained by authors or by other copyright holders. Personal use of this material is permitted. Permission must be obtained from the copyright holder for all other uses.

Selected Talks

About me

Education

Industry Experience

Academic Activities