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.

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.

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