Nikunj Saunshi

Research Scientist, Google

LinkTwitter

I am a researcher at Google Research NYC. My interests lie in using a mix of theory and empirics to demystify the incredible success of deep learning and AI, and to design efficient and reliable learning algorithms. My research has spanned topics like language models and reasoning, self-supervised representation learning, meta-learning, natural language processing, interpretability of deep learning models.

I received my PhD in Computer Science from Princeton University, where I was advised by Sanjeev Arora. Before joining Princeton, I worked in the R&D center of Samsung Electronics in Suwon, South Korea. I completed my B.Tech. with Honors in Computer Science and Engineering and Minor in Mathematics from the Indian Institute of Technology Bombay.

Updates


[Sept 2024] New paper On the Inductive Bias of Stacking Towards Improving Reasoning accepted at NeurIPS 2024

[May 2024] Our paper Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? was accepted at ICML 2024

[Feb 2024] Paper on Efficient Stagewise Pretraining via Progressive Subnetworks with interns is out now!

[Aug 2023] Had the pleasure of hosting two summer interns at Google Research - Abhishek Panigrahi and Kaifeng Lyu

[May 2023] Reasoning in Large Language Models Through Symbolic Math Word Problems accepted at Findings of ACL 2023 and Natural Language Reasoning and Structured Explanations (NLRSE) workshop

[Apr 2023] Task-Specific Skill Localization in Fine-tuned Language Models accepted at ICML 2023!

[Apr 2023] Understanding Influence Functions and Datamodels via Harmonic Analysis presented at ICLR 2023

[Apr 2023] Selected as Notable Reviewer for ICLR 2023

[Jan 2023] Started as a Research Scientist at Google

[Jan 2023] Presented a talk on Towards Understanding Self-Supervised Learning at CMU

[July 2022] Defended my PhD thesis!

Publications by year

(#) denotes alphabetical order

On the Inductive Bias of Stacking Towards Improving Reasoning

Nikunj Saunshi, Stefani Karp, Shankar Krishnan, Sobhan Miryoosefi, Sashank Reddi, Sanjiv Kumar

To Appear at NeurIPS 2024

Landscape-Aware Growing: The Power of a Little LAG

Nikunj Saunshi*, Stefani Karp*, Sobhan Miryoosefi, Sashank Reddi, Sanjiv Kumar

Arxiv

Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?

Khashayar Gatmiry, Nikunj Saunshi, Sashank Reddi, Stefanie Jegelka, Sanjiv Kumar

ICML 2024

Efficient Stagewise Pretraining via Progressive Subnetworks

Abhishek Panigrahi*, Nikunj Saunshi*, Kaifeng Lyu, Sobhan Miryoosefi, Sashank Reddi, Satyen Kale, Sanjiv Kumar

Arxiv

Task-Specific Skill Localization in Fine-tuned Language Models

Abhishek Panigrahi*, Nikunj Saunshi*, Haoyu Zhang, Sanjeev Arora

ICML 2023

Understanding Influence Functions and Datamodels via Harmonic Analysis

Nikunj Saunshi, Arushi Gupta, Mark Braverman, Sanjeev Arora

ICLR 2023

New Definitions and Evaluations for Saliency Methods: Staying Intrinsic, Complete and Sound

Arushi Gupta*, Nikunj Saunshi*, Dingli Yu*, Kaifeng Lyu, Sanjeev Arora

NeurIPS 2022 (Oral)

Understanding Contrastive Learning Requires Incorporating Inductive Biases

Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, Akshay Krishnamurthy

ICML 2022

On Predicting Generalization using GANs

Yi Zhang, Arushi Gupta, Nikunj Saunshi, Sanjeev Arora

ICLR 2022 (Oral)

Predicting What You Already Know Helps: Provable Self-Supervised Learning

(#) Jason D. Lee, Qi Lei, Nikunj Saunshi, Jiacheng Zhuo

NeurIPS 2021

A Mathematical Exploration of Why Language Models Help Solve Downstream Tasks

Nikunj Saunshi, Sadhika Malladi, Sanjeev Arora

ICLR 2021

[Talk]

A Sample Complexity Separation between Non-Convex and Convex Meta-Learning

Nikunj Saunshi, Yi Zhang, Mikhail Khodak, Sanjeev Arora

ICML 2020

[Talk]

Provable Representation Learning for Imitation Learning via Bi-level Optimization

(#) Sanjeev Arora, Simon S. Du, Sham Kakade, Yuping Luo, Nikunj Saunshi

ICML 2020

A Theoretical Analysis of Contrastive Unsupervised Representation Learning

(#) Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, Nikunj Saunshi

ICML 2019

[Blog] [Talk]

A La Carte Embedding: Cheap but Effective Induction of Semantic Feature Vectors

Mikhail Khodak*, Nikunj Saunshi*, Yingyu Liang, Tengyu Ma, Brandon Stewart and Sanjeev Arora

ACL 2018

[Blog] [Talk]

A Compressed Sensing View of Unsupervised Text Embeddings, Bag-of-n-Grams, and LSTMs

(#) Sanjeev Arora, Mikhail Khodak, Nikunj Saunshi, Kiran Vodrahalli

ICLR 2018

[Blog]

A Large Self-Annotated Corpus for Sarcasm

Mikhail Khodak, Nikunj Saunshi, Kiran Vodrahalli

LREC 2018

Publications by topic

Language models and reasoning

On the Inductive Bias of Stacking Towards Improving Reasoning

Nikunj Saunshi, Stefani Karp, Shankar Krishnan, Sobhan Miryoosefi, Sashank Reddi, Sanjiv Kumar

To Appear at NeurIPS 2024

Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?

Khashayar Gatmiry, Nikunj Saunshi, Sashank Reddi, Stefanie Jegelka, Sanjiv Kumar

ICML 2024

Efficient Stagewise Pretraining via Progressive Subnetworks

Abhishek Panigrahi*, Nikunj Saunshi*, Kaifeng Lyu, Sobhan Miryoosefi, Sashank Reddi, Satyen Kale, Sanjiv Kumar

Arxiv

Task-Specific Skill Localization in Fine-tuned Language Models

Abhishek Panigrahi*, Nikunj Saunshi*, Haoyu Zhang, Sanjeev Arora

ICML 2023

A Mathematical Exploration of Why Language Models Help Solve Downstream Tasks

Nikunj Saunshi, Sadhika Malladi, Sanjeev Arora

ICLR 2021

[Talk]

Efficiency for large models

On the Inductive Bias of Stacking Towards Improving Reasoning

Nikunj Saunshi, Stefani Karp, Shankar Krishnan, Sobhan Miryoosefi, Sashank Reddi, Sanjiv Kumar

To Appear at NeurIPS 2024

Landscape-Aware Growing: The Power of a Little LAG

Nikunj Saunshi*, Stefani Karp*, Sobhan Miryoosefi, Sashank Reddi, Sanjiv Kumar

Arxiv

Efficient Stagewise Pretraining via Progressive Subnetworks

Abhishek Panigrahi*, Nikunj Saunshi*, Kaifeng Lyu, Sobhan Miryoosefi, Sashank Reddi, Satyen Kale, Sanjiv Kumar

Arxiv

Self-supervised learning

Understanding Contrastive Learning Requires Incorporating Inductive Biases

Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, Akshay Krishnamurthy

ICML 2022

Predicting What You Already Know Helps: Provable Self-Supervised Learning

(#) Jason D. Lee, Qi Lei, Nikunj Saunshi, Jiacheng Zhuo

NeurIPS 2021

A Mathematical Exploration of Why Language Models Help Solve Downstream Tasks

Nikunj Saunshi, Sadhika Malladi, Sanjeev Arora

ICLR 2021

[Talk]

A Theoretical Analysis of Contrastive Unsupervised Representation Learning

(#) Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, Nikunj Saunshi

ICML 2019

[Blog] [Talk]

Interpretability

Understanding Influence Functions and Datamodels via Harmonic Analysis

Nikunj Saunshi, Arushi Gupta, Mark Braverman, Sanjeev Arora

ICLR 2023

New Definitions and Evaluations for Saliency Methods: Staying Intrinsic, Complete and Sound

Arushi Gupta*, Nikunj Saunshi*, Dingli Yu*, Kaifeng Lyu, Sanjeev Arora

NeurIPS 2022 (Oral)

Meta learning

A Sample Complexity Separation between Non-Convex and Convex Meta-Learning

Nikunj Saunshi, Yi Zhang, Mikhail Khodak, Sanjeev Arora

ICML 2020

[Talk]

Provable Representation Learning for Imitation Learning via Bi-level Optimization

(#) Sanjeev Arora, Simon S. Du, Sham Kakade, Yuping Luo, Nikunj Saunshi

ICML 2020

Representation learning & NLP

A La Carte Embedding: Cheap but Effective Induction of Semantic Feature Vectors

Mikhail Khodak*, Nikunj Saunshi*, Yingyu Liang, Tengyu Ma, Brandon Stewart and Sanjeev Arora

ACL 2018

[Blog] [Talk]

A Compressed Sensing View of Unsupervised Text Embeddings, Bag-of-n-Grams, and LSTMs

(#) Sanjeev Arora, Mikhail Khodak, Nikunj Saunshi, Kiran Vodrahalli

ICLR 2018

[Blog]

A Large Self-Annotated Corpus for Sarcasm

Mikhail Khodak, Nikunj Saunshi, Kiran Vodrahalli

LREC 2018

Other topics

On Predicting Generalization using GANs

Yi Zhang, Arushi Gupta, Nikunj Saunshi, Sanjeev Arora

ICLR 2022 (Oral)

Peer review