Ankit Singh Rawat

Google Research NY 111 8th Avenue New York, NY 10011, USA E-mail: ankitsrawat AT google(dot)com

About Me

I am a Research Scientist at Google Research, New York.

Previously, I held postdoctoral appointments at Massachusetts Institute of Technology, Carnegie Mellon University, and University of Massachusetts Amherst, where I worked with Prof. Gregory W. Wornell, Prof. Venkatesan Guruswami, and Prof. Arya Mazumdar, respectively. I received a Ph.D. in Electrical and Computer Engineering from The University of Texas at Austin in 2015, where I was advised by Prof. Sriram Vishwanath. At UT Austin, I also had the pleasure of working with Prof. Alex Dimakis. Before attending UT Austin, I received a B.Tech. degree in Electrical Engineering from Indian Institute of Technology (IIT) Kanpur in 2010.

Research Interests

• Large Scale Machine Learning

• Coding Theory

• Information Theory

News

• Three papers accepted at ICML 2024:

  • From Self-attention to Markov Models: Unveiling the Dynamics of Generative Transformers

  • USTAD: Unified Single-model Training Achieving Diverse Scores for Information Retrival

  • A Statistical Framework for Data-dependent Retrieval-augmented Models

• Gave a tutorial on Fundamentals of Transformers with Samet Oymak, Christos Thrampoulidis, and Mahdi Soltanolkotabi at ICASSP 2024.

• New preprint:

  • From Self-attention to Markov Models: Unveiling the Dynamics of Generative Transformers

• One paper accepted at AISTATS 2024:

  • Mechanics of Next Token Prediction with Self-attention

• Four papers accepted at ICLR 2024:

  • DistillSpec: Improving Speculative Decoding via Knowledge Distillation

  • Think Before You Speak: Training Language Models with Pause Tokens

  • Language Model Cascades: Token-level Uncertainty and Beyond

  • Efficacy of Dual-encoders for Extreme Multi-label Classification

• Two papers accepted at NeurIPS 2023:

  • When Does Confidence-based Cascade Deferral Suffice?

  • ResMem: Learn What You Can and Memorize the Rest

• Two papers accepted at ICML 2023:

  • A Statistical Perspective on Retrieval-based Models

  • On the Role of Attention in Prompt-tuning

• Research featured in Google AI Blog.

• New preprints:

  • EmbedDistill: A Geometric Knowledge Distillation for Information Retrieval.

  • ResMem: Learn What You Can and Memorize the Rest.

• Four papers accepted at ICLR 2023:

  • Teacher Guided Training: An Efficient Framework for Knowledge Transfer.

  • Supervision Complexity and Its Role in Knowledge Distillation.

  • Serving Graph Compression for Graph Neural Networks.

  • On Emergence of Activation Sparsity in Trained Transformers.

• New preprint:

  • Generalization Properties of Retrieval-based Models.

  • Large Language Models with Controllable Working Memory.

• Two papers accepted at NeurIPS 2022:

  • A Fourier Approach to Mixture Learning.

  • Post-hoc Estimators for Learning to Defer to an Expert.

• New preprint:

  • Teacher Guided Training: An Efficient Framework for Knowledge Transfer.

• One paper accepted at ICML 2022:

  • In Defense of Dua-encoders for Neural Ranking.

• New preprints:

  • ELM: Embedding and Logit Margins for Long-Tail Learning.

  • FedLite: A Scalable Approach for Federated Learning on Resource-constrained Clients.

• New preprint:

  • When in Doubt, Summon the Titans: Efficient Inference with Large Models.

• Two papers accepted at ICML 2021:

  • Disentangling Sampling and Labeling Bias for Learning in Large-output Spaces.

  • A Statistical Perspective on Distillation.

• One paper accepted at ISIT 2021:

  • Byzantine Resilient Distributed Clustering with Redundant Data Assignment.

• New preprints:

  • Modifying Memories in Transformer Models.

  • Distilling Double Descent.

  • On the Reproducibility of Neural Network Predictions.

• Two papers accepted at ICLR 2021:

  • Overparameterisation and Worst-case Generalisation: Friend or Foe?

  • Long-tail Learning via Logit Adjustment.

• One paper accepted at AISTATS 2021:

  • RankDistil: Knowledge Distillation for Ranking.

• Three papers accepted at NeurIPS 2020:

  • O(n) Connections are Expressive Enough: Universal Approximability of Sparse Transformers.

  • Adversarial Robustness via Robust Low Rank Representations.

  • Robust Large-Margin Learning in Hyperbolic Space.

• Two papers accepted at ICML 2020:

  • Federated Learning with Only Positive Labels.

  • Low-Rank Bottleneck in Multi-head Attention Models.

• 2020 EURASIP JASP Best Paper Award for this paper coauthored with Arya Mazumdar and Sriram Vishwanath.

• New preprints:

  • Why distillation helps: a statistical perspective.

  • Doubly-stochastic mining for heterogeneous retrieval.

  • Robust Large-Margin Learning in Hyperbolic Space.

• Two papers appeared at ICLR 2020:

  • Can Gradient Clipping Mitigate Label Noise?

  • Are Transformers Universal Approximators of Sequence-to-sequence Functions?

• One paper accepted at ISIT 2020:

  • Reliable Distributed Clustering with Redundant Data Assignment.

Selected Publications (Full List)

• USTAD: Unified Single-model Training Achieving Diverse Scores for Information Retrival
  Seungyeon Kim, Ankit Singh Rawat, Manzil Zaheer, Sadeep Jayasumana, Veeranjaneyulu Sadhanala, Wittawat Jitkrittum, Aditya Krishna Menon, Rob Fergus, Sanjiv Kumar
  To appear in International Conference on Machine Learning (ICML), 2024.

• DistillSpec: Improving Speculative Decoding via Knowledge Distillation
  Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, Rishabh Agarwal
  To appear in International Conference on Learning Representations (ICLR), 2024.

• Efficacy of Dual-encoders for Extreme Multi-label Classification
  Nilesh Gupta, Fnu Devvrit, Ankit Singh Rawat, Srinadh Bhojanapalli, Prateek Jain, Inderjit S Dhillon
  To appear in International Conference on Learning Representations (ICLR), 2024.

• Mechanics of Next Token Prediction with Self-attention
  Yingcong Li, Yixiao Huang, M. Emrullah Ildiz, Ankit Singh Rawat, Samet Oymak
  To appear in International Conference on Artificial Intelligence and Statistics (AISTATS), 2024.

• A Statistical Perspective on Retrieval-based Models
  Soumya Basu, Ankit Singh Rawat, Manzil Zaheer
  International Conference on Machine Learning (ICML), 2023.

• On the Role of Attention in Prompt-tuning
  Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, Christos Thrampoulidis
  International Conference on Machine Learning (ICML), 2023.

• Teacher Guided Training: An Efficient Framework for Knowledge Transfer
  Manzil Zaheer, Ankit Singh Rawat, Seungyeon Kim, Chong You, Himanshu Jain, Andreas Veit, Rob Fergus, Sanjiv Kumar
  International Conference on Learning Representations (ICLR), 2023.

• Supervision Complexity and Its Role in Knowledge Distillation
  Hrayr Harutyunyan, Ankit Singh Rawat, Aditya K Menon, Seungyeon Kim, Sanjiv Kumar
  International Conference on Learning Representations (ICLR), 2023.

• A Fourier Approach to Mixture Learning
  Mingda Qiao, Guru Guruganesh, Ankit Singh Rawat, Kumar Avinava Dubey, Manzil Zaheer
  Advances in Neural Information Processing Systems (NeurIPS), 2022.

• A Statistical Perspective on Distillation
  Aditya Krishna Menon, Ankit Singh Rawat, Sashank J. Reddi, Seungyeon Kim, Sanjiv Kumar
  International Conference on Machine Learning (ICML), 2021.

• Long-tail Learning via Logit Adjustment
  Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, Sanjiv Kumar
  Spotlight, International Conference on Learning Representations (ICLR), 2021.

• Federated Learning with Only Positive Labels
  Felix X. Yu, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar
  International Conference on Machine Learning (ICML), 2020.

• Can Gradient Clipping Mitigate Label Noise?
  Aditya Krishna Menon, Ankit Singh Rawat, Sashank J. Reddi, Sanjiv Kumar
  International Conference on Learning Representations (ICLR), 2020.

• Are Transformers Universal Approximators of Sequence-to-sequence Functions?
  Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, Sanjiv Kumar
  International Conference on Learning Representations (ICLR), 2020.

• Sampled Softmax with Random Fourier Features
  Ankit Singh Rawat, Jiecao Chen, Felix Yu, Ananda Theertha Suresh, Sanjiv Kumar
  Advances in Neural Information Processing Systems (NeurIPS), 2019.

• Robust Gradient Descent via Moment Encoding with LDPC Codes
  Raj Kumar Maity, Ankit Singh Rawat, Arya Mazumdar
  IEEE International Symposium on Information Theory (ISIT), 2019.

• MDS Code Constructions with Small Sub-packetization and Near-optimal Repair Bandwidth
  Ankit Singh Rawat, Itzhak Tamo, Venkatesan Guruswami, Klim Efremenko
  IEEE Transactions on Information Theory, October 2018.

• Locality and Availability in Distributed Storage
  Ankit Singh Rawat, Dimitris S. Papailiopoulos, Alexandros G. Dimakis, Sriram Vishwanath
  IEEE Transactions on Information Theory, August 2016.

• Batch Codes through Dense Graphs with High Girth
  Ankit Singh Rawat, Zhao Song, Alexandros G. Dimakis, Anna Gal
  IEEE Transactions on Information Theory, April 2016.

• Optimal Locally Repairable and Secure Codes for Distributed Storage Systems
  Ankit Singh Rawat, O. Ozan Koyluoglu, Natalia Silberstein, Sriram Vishwanath
  IEEE Transactions on Information Theory, January 2014.

Click here for the full list of publications.