Skip to content
Nicholas Bai

Nicholas Bai

Computer Science @ Caltech

Hi! I'm an undergraduate at Caltech studying computer science. I'm interested in building machine learning systems and in understanding why they work.

My research spans interpretability, generative modeling, and machine learning for scientific inverse problems. Currently, I am a researcher in Anima Anandkumar's AI + Science Lab, designing diffusion-based sampling methods for Bayesian inverse PDE problems.

I am inspired by methods that make powerful AI models more reliable, interpretable, and applicable in real-world settings.

Portrait of Nicholas Bai

Recent News

  1. Award

    Named a Carl F. Braun SURF Fellow for summer research on Bayesian inverse problems.

  2. Update

    Attending ICML 2026 in Seoul, Korea. See you there!

  3. Update

    Attending Y Combinator Startup School in San Francisco, CA.

  4. Publication

    “Count Me If You Can: Geometric Failure Modes in Language Model Counting” accepted to the ICML 2026 Workshop on Compositional Learning.

What I Work On

  • Interpretability and Reliable AI

    Understanding internal representations, explaining model behavior, and building methods that make systems more transparent and more trustworthy.

  • Generative Models

    Diffusion models, controllable generation, representation steering, and the behavior of modern generative systems.

  • AI for Science and Inverse Problems

    Machine learning for scientific computing: PDE inverse problems, Bayesian inference, neural operators, and function-space methods.

  • Foundations of Learning and Inference

    Particle methods, sampling, transport, geometry, and the mathematical behavior of high-dimensional learning systems.

Selected Publications

  1. ICML 2026 Workshop

    Count Me If You Can: Geometric Failure Modes in Language Model Counting

    Nicholas Bai, Ayushi Mehrotra

    Traces language-model counting errors to the geometry of their internal number representations, and improves high-count accuracy by up to 20%.

    Details
  2. NeurIPS 2025 Workshop

    CAT: Curvature-Adaptive Transformers for Geometry-Aware Learning

    R. Y. Lin, S. Ojha, Nicholas Bai

    Transformers that route tokens across three geometric attention branches with a mixture-of-experts gate.

    Details
  3. TMLR 2025ICML 2024 Workshop Spotlight

    Interpreting Neurons in Deep Vision Networks with Language Models

    Nicholas Bai, et al.

    A language-model-based framework that produces natural-language descriptions of individual neurons in vision networks.

    Details
View all research