Single-Source Bottleneck Path Algorithm Faster than Sorting for Sparse Graphs. - Ran Duan; Kaifeng Lyu; Yuanhang Xie (ICALP, 2018)

Single-Source Bottleneck Path Algorithm Faster than Sorting for Sparse Graphs. - Ran Duan; Kaifeng Lyu; Hongxun Wu; Yuanhang Xie (CoRR, 2018)

Theoretical Analysis of Auto Rate-Tuning by Batch Normalization. - Sanjeev Arora; Zhiyuan Li 005; Kaifeng Lyu (ICLR (Poster), 2019)

Fine-grained Complexity Meets IP = PSPACE. - Lijie Chen 001; Shafi Goldwasser; Kaifeng Lyu; Guy N. Rothblum; Aviad Rubinstein (SODA, 2019)

Gradient Descent Maximizes the Margin of Homogeneous Neural Networks. - Kaifeng Lyu; Jian Li 015 (ICLR, 2020)

Reconciling Modern Deep Learning with Traditional Optimization Analyses: The Intrinsic Learning Rate. - Zhiyuan Li 005; Kaifeng Lyu; Sanjeev Arora (NeurIPS, 2020)

Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learning. - Zhiyuan Li 005; Yuping Luo; Kaifeng Lyu (ICLR, 2021)

Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias. - Kaifeng Lyu; Zhiyuan Li 005; Runzhe Wang; Sanjeev Arora (NeurIPS, 2021)

New Definitions and Evaluations for Saliency Methods: Staying Intrinsic, Complete and Sound. - Arushi Gupta; Nikunj Saunshi; Dingli Yu; Kaifeng Lyu; Sanjeev Arora (NeurIPS, 2022)

Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction. - Kaifeng Lyu; Zhiyuan Li 005; Sanjeev Arora (NeurIPS, 2022)

On the SDEs and Scaling Rules for Adaptive Gradient Algorithms. - Sadhika Malladi; Kaifeng Lyu; Abhishek Panigrahi; Sanjeev Arora (NeurIPS, 2022)

Why (and When) does Local SGD Generalize Better than SGD? - Xinran Gu; Kaifeng Lyu; Longbo Huang; Sanjeev Arora (ICLR, 2023)

Understanding Incremental Learning of Gradient Descent: A Fine-grained Analysis of Matrix Sensing. - Jikai Jin; Zhiyuan Li 005; Kaifeng Lyu; Simon Shaolei Du; Jason D. Lee (ICML, 2023)

A Quadratic Synchronization Rule for Distributed Deep Learning. - Xinran Gu; Kaifeng Lyu; Sanjeev Arora; Jingzhao Zhang; Longbo Huang (ICLR, 2024)

Dichotomy of Early and Late Phase Implicit Biases Can Provably Induce Grokking. - Kaifeng Lyu; Jikai Jin; Zhiyuan Li 005; Simon Shaolei Du; Jason D. Lee; Wei Hu 014 (ICLR, 2024)

The Marginal Value of Momentum for Small Learning Rate SGD. - Runzhe Wang; Sadhika Malladi; Tianhao Wang 017; Kaifeng Lyu; Zhiyuan Li 005 (ICLR, 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 (ICLR, 2024)

Keeping LLMs Aligned After Fine-tuning: The Crucial Role of Prompt Templates. - Kaifeng Lyu; Haoyu Zhao; Xinran Gu; Dingli Yu; Anirudh Goyal; Sanjeev Arora (NeurIPS, 2024)

AI-Assisted Generation of Difficult Math Questions. - Vedant Shah; Dingli Yu; Kaifeng Lyu; Simon Park 002; Nan Rosemary Ke; Michael Mozer; Yoshua Bengio; Sanjeev Arora; Anirudh Goyal (CoRR, 2024)

Feature Averaging: An Implicit Bias of Gradient Descent Leading to Non-Robustness in Neural Networks. - Binghui Li; Zhixuan Pan; Kaifeng Lyu; Jian Li (ICLR, 2025)

Towards Understanding Text Hallucination of Diffusion Models via Local Generation Bias. - Rui Lu; Runzhe Wang; Kaifeng Lyu; Xitai Jiang; Gao Huang 001; Mengdi Wang (ICLR, 2025)

A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules. - Kairong Luo; Haodong Wen; Shengding Hu; Zhenbo Sun; Zhiyuan Liu 001; Maosong Sun 001; Kaifeng Lyu; Wenguang Chen (ICLR, 2025)

Efficient stagewise pretraining via progressive subnetworks. - Abhishek Panigrahi; Nikunj Saunshi; Kaifeng Lyu; Sobhan Miryoosefi; Sashank J. Reddi; Satyen Kale; Sanjiv Kumar (ICLR, 2025)

Safety Alignment Should be Made More Than Just a Few Tokens Deep. - Xiangyu Qi; Ashwinee Panda; Kaifeng Lyu; Xiao Ma 010; Subhrajit Roy; Ahmad Beirami; Prateek Mittal; Peter Henderson 002 (ICLR, 2025)

RNNs are not Transformers (Yet): The Key Bottleneck on In-Context Retrieval. - Kaiyue Wen; Xingyu Dang; Kaifeng Lyu (ICLR, 2025)

Weak-to-Strong Generalization Even in Random Feature Networks, Provably. - Marko Medvedev; Kaifeng Lyu; Dingli Yu; Sanjeev Arora; Zhiyuan Li 005; Nathan Srebro (CoRR, 2025)

LEGO-Puzzles: How Good Are MLLMs at Multi-Step Spatial Reasoning? - Kexian Tang; Junyao Gao 002; Yanhong Zeng; Haodong Duan; Yanan Sun 005; Zhening Xing; Wenran Liu; Kaifeng Lyu; Kai Chen 026 (CoRR, 2025)

Data Mixing Can Induce Phase Transitions in Knowledge Acquisition. - Xinran Gu; Kaifeng Lyu; Jiazheng Li 015; Jingzhao Zhang (CoRR, 2025)

How Far Are We from Optimal Reasoning Efficiency? - Jiaxuan Gao; Shu Yan; Qixin Tan; Lu Yang; Shusheng Xu; Wei Fu; Zhiyu Mei; Kaifeng Lyu; Yi Wu 013 (CoRR, 2025)