Publications

Efficient AI

Reasoning

• TERMINATOR: Learning Optimal Exit Points for Early Stopping in Chain-of-Thought Reasoning
  with Alliot Nagle, Jakhongir Saydaliev, Dhia Garbaya, Michael C. Gastpar, and Hyeji Kim
  TTU Workshop, ICLR 2026
  

Prompt Compression

• Fundamental Limits of Prompt Compression: A Rate-Distortion Framework for Black-Box Language Models
  with Adway Girish, Alliot Nagle, Marco Bondaschi, Michael C. Gastpar, and Hyeji Kim
  NeurIPS, 2024
  

Optimal Transport & Mixture of Experts

• Optimal transport mapping via input convex neural networks [Slides]
  with Amirhossein Taghvaei, Jason Lee, and Sewoong Oh
  ICML, 2020
  

• Learning in Gated Neural Networks [Slides]
  with Sreeram Kannan, Sewoong Oh, and Pramod Viswanath
  AISTATS, 2020
  

• Breaking the gridlock in Mixture-of-Experts: Consistent and Efficient Algorithms
  with Sewoong Oh, Sreeram Kannan, and Pramod Viswanath
  ICML, 2019
  

• Learning One-hidden-layer Neural Networks under General Input Distributions
  with Weihao Gao, Sewoong Oh, and Pramod Viswanath
  AISTATS, 2019
  

Interpretable AI

A Markovian Tale of Transformers and SSMs

[Slides]

• From Markov to Laplace: How Mamba In-Context Learns Markov Chains
  with Marco Bondaschi, Nived Rajaraman, Xiuying Wei, Kannan Ramchandran, Razvan Pascanu, Caglar Gulcehre, and Michael C. Gastpar
  ICLR 2026, Oral (1% out of 18,949 papers)
  

• What One Cannot, Two Can: Two-Layer Transformers Provably Represent Induction Heads on Any-Order Markov Chains
  with Chanakya Ekbote, Marco Bondaschi, Nived Rajaraman, Jason D. Lee, Paul Pu Liang, and Michael C. Gastpar
  NeurIPS 2025, Spotlight (3% out of 12,575 papers)
  

• Attention with Markov: A Framework for Principled Analysis of Transformers via Markov Chains
  with Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael C. Gastpar
  ICLR 2025, Spotlight (5% out of 11,670 papers)
  

• Transformers on Markov Data: Constant Depth Suffices
  with Nived Rajaraman, Marco Bondaschi, Kannan Ramchandran, and Michael C. Gastpar
  NeurIPS, 2024
  

• Local to Global: Learning Dynamics and Effect of Initialization for Transformers
  with Marco Bondaschi, Chanakya Ekbote, Adway Girish, Alliot Nagle, Hyeji Kim, and Michael C. Gastpar
  NeurIPS, 2024
  

Learning Codes via Deep Learning

• LASER: Linear Compression in Wireless Distributed Optimization [Slides]
  with Marco Bondaschi, Thijs Vogels, Martin Jaggi, Hyeji Kim, and Michael C. Gastpar
  NeurIPS Federated Learning Workshop, 2023
  ICML, 2024
  

• CRISP: Curriculum based Sequential Neural Decoders for Polar Code Family
  with S Ashwin Hebbar, Viraj Nadkarni, Suma Bhat, Sewoong Oh, and Pramod Viswanath
  ICML, 2023
  

• Machine Learning-Aided Efficient Decoding of Reed-Muller Subcodes
  with Mohammad Vahid Jamali, Xiyang Liu, Hessam Mahdavifar, Sewoong Oh, and Pramod Viswanath
  IEEE Journal on Selected Topics in Information Theory (JSAIT), 2023
  

• TinyTurbo: Efficient Turbo Decoders on Edge
  with S Ashwin Hebbar, Rajesh K Mishra, Sravan Kumar Ankireddy, Hyeji Kim, and Pramod Viswanath
  ISIT, 2022
  

• KO codes: Inventing Nonlinear Encoding and Decoding for Reliable Wireless Communication via Deep-learning [Slides]
  with Xiyang Liu, Mohammad Vahid Jamali, Hessam Mahdavifar, Sewoong Oh, and Pramod Viswanath
  ICML, 2021
  

• Reed-Muller Subcodes: Machine Learning-Aided Design of Efficient Soft Recursive Decoding
  with Mohammad Vahid Jamali, Xiyang Liu, Hessam Mahdavifar, Sewoong Oh, and Pramod Viswanath
  ISIT, 2021
  

Blockchains

• Barracuda: The Power of l-polling in Proof-of-Stake Blockchains
  with Ranvir Rana, Jiantao Jao, Sewoong Oh, Giulia Fanti, and Pramod Viswanath
  ACM MobiHoc, 2019 (Best paper award)

Information theory

• Equivalence of additive-combinatorial linear inequalities for Shannon entropy and differential entropy
  with Yihong Wu
  IEEE Transactions on Information Theory, 2018

• On additive-combinatorial affine inequalities for Shannon entropy and differential entropy
  with Yihong Wu
  IEEE International Symposium on Information Theory (ISIT), 2016

Stochastic approximations

• Event-driven stochastic approximation
  with Vivek S. Borkar and Neeraja Sahasrabudhe
  Indian Journal of Pure and Applied Mathematics, June 2016, Volume 47, Issue 2, pp 291–299