Publications

* indicates equal contribution

** indicates corresponding author

Preprints

• Task-Agnostic Continual Prompt Tuning with Gradient-Based Selection and Decoding
  Anushka Tiwari, Sayantan Pal, Rohini K Srihari, Kaiyi Ji**. Submitted, 2025.

• LDC-MTL: Balancing Multi-Task Learning through Scalable Loss Discrepancy Control
  Peiyao Xiao, Chaosheng Dong, Shaofeng Zou, Kaiyi Ji**. Submitted, 2025.

• ReGNet: Reciprocal Space-Aware Long-Range Modeling for Crystalline Property Prediction
  Jianan Nie, Peiyao Xiao, Kaiyi Ji, Peng Gao. Submitted, 2025.

• Enhancing Diffusion Posterior Sampling for Inverse Problems by Integrating Crafted Measurements
  Shijie Zhou, Huaisheng Zhu, Rohan Sharma, Jiayi Chen, Ruiyi Zhang, Kaiyi Ji, Changyou Chen. Submitted, 2025.

2025

• (New!) SAMO: A Lightweight Sharpness-Aware Approach for Multi-Task Optimization with Joint Global-Local Perturbation <Code>
  Hao Ban, Gokul Ram Subramani and Kaiyi Ji**.
  International Conference on Computer Vision (ICCV) 2025.

• (New!) Tuning-Free Bilevel Optimization: New Algorithms and Convergence Analysis <Code>
  Yifan Yang, Hao Ban, Minhui Huang, Shiqian Ma, Kaiyi Ji**.
  International Conference on Learning Representations (ICLR) 2025.

• (New!) MGDA Converges under Generalized Smoothness, Provably
  Qi Zhang, Peiyao Xiao, Shaofeng Zou, Kaiyi Ji**.
  International Conference on Learning Representations (ICLR) 2025.

• (New!) Theoretical Study of Conflict-Avoidant Multi-Objective Reinforcement Learning
  Yudan Wang, Peiyao Xiao, Hao Ban, Kaiyi Ji**, Shaofeng Zou.
  IEEE Transactions on Information Theory (TIT) 2025.

• (New!) Efficiently Escaping Saddle Points in Bilevel Optimization
  Minhui Huang, Kaiyi Ji, Shiqian Ma, Lifeng Lai.
  Journal of Machine Learning Research (JMLR) 2025.

• (New!) Imperative Learning: A Self-supervised Neural-Symbolic Learning Framework for Robot Autonomy <Open-Source Library>
  Chen Wang, Kaiyi Ji, Junyi Geng, Zhongqiang Ren, Taimeng Fu, Fan Yang, Yifan Guo, Haonan He, Xiangyu Chen,
  Zitong Zhan, Qiwei Du, Shaoshu Su, Bowen Li, Yuheng Qiu, Yi Du, Qihang Li, Yifan Yang, Xiao Lin, Zhipeng Zhao.
  The International Journal of Robotics Research (IJRR) 2025.

• (New!) First-Order Federated Bilevel Learning
  Yifan Yang, Peiyao Xiao, Shiqian Ma, Kaiyi Ji**
  AAAI Conference on Artificial Intelligence (AAAI) 2025.

• (New!) Meta-Learning with Heterogeneous Tasks
  Zhaofeng Is, She Hu, Kaiyi Ji and Siwei Lyu
  Generalizing from Limited Resources in the Open World Workshop (GLOW) (Oral Presentation)
  International Joint Conferences on Artificial Intelligence Organization (IJCAI) 2025.

2024

• Understanding Forgetting in Continual Learning with Linear Regression
  Meng Ding, Kaiyi Ji**, Di Wang, Jinhui Xu.
  International Conference on Machine Learning (ICML) 2024.

• Fair Resource Allocation in Multi-Task Learning <Code>
  Hao Ban and Kaiyi Ji**.
  International Conference on Machine Learning (ICML) 2024.

• AUC-CL: A Batchsize-Robust Framework for Self-Supervised Contrastive Representation Learning <Code>
  Rohan Sharma, Kaiyi Ji, Zhiqiang Xu, Changyou Chen.
  International Conference on Learning Representations (ICLR) 2024.

• Boosting One-Point Derivative-Free Online Optimization via Residual Feedback
  Yan Zhang, Yi Zhou, Kaiyi Ji, Yi Shen, and Michael M. Zavlanos. Accepted by IEEE Transactions on Automatic Control, 2024.

• First-Order Minimax Bilevel Optimization.
  Yifan Yang, Zhaofeng Si, Siwei Lyu, Kaiyi Ji**
  Conference on Neural Information Processing Systems (NeurIPS) 2024.

2023

• Direction-oriented Multi-objective Learning: Simple and Provable Stochastic Algorithms <Code>
  Peiyao Xiao*, Hao Ban*, Kaiyi Ji**. Accepted by Conference on Neural Information Processing Systems (NeurIPS) 2023.

• SimFBO: Towards Simple, Flexible and Communication-efficient Federated Bilevel Learning
  Yifan Yang, Peiyao Xiao, Kaiyi Ji**. Accepted by Conference on Neural Information Processing Systems (NeurIPS) 2023. (Spotlight)

• Achieving O(\epsilon^{-1.5}) Complexity in Hessian/Jacobian-free Stochastic Bilevel Optimization
  Yifan Yang, Peiyao Xiao, Kaiyi Ji**.
  Conference on Neural Information Processing Systems (NeurIPS) 2023.

• Non-Convex Bilevel Optimization with Time-Varying Objective Functions
  Sen Lin, Daouda Sow, Kaiyi Ji, Yingbin Liang, Ness Shroff. Accepted by Conference on Neural Information Processing Systems (NeurIPS) 2023.

• Bilevel Coreset Selection in Continual Learning: A New Formulation and Algorithm <Code>
  Jie Hao, Kaiyi Ji, Mingrui Liu.
  Conference on Neural Information Processing Systems (NeurIPS) 2023.

• Network Utility Maximization with Unknown Utility Functions: A Distributed, Data-Driven Bilevel Optimization Approach
  Kaiyi Ji and Lei Ying. ACM MobiHoc 2023! (acceptance rate: 21.9%)

• Communication-Efficient Federated Hypergradient Computation via Aggregated Iterative Differentiation.
  Peiyao Xiao and Kaiyi Ji**. International Conference on Machine Learning (ICML) 2023.

• Achieving Linear Speedup in Non-IID Federated Bilevel Learning
  Minhui Huang*, Dewei Zhang*, Kaiyi Ji**. International Conference on Machine Learning (ICML) 2023.

2022

• On the Convergence Theory for Hessian-Free Bilevel Algorithms <Code>
  Daouda Sow, Kaiyi Ji, Yingbin Liang
  Conference on Neural Information Processing Systems (NeurIPS) 2022.

• Will Bilevel Optimizers Benefit from Loops
  Kaiyi Ji, Mingrui Liu, Yingbin Liang, Lei Ying
  Conference on Neural Information Processing Systems (NeurIPS) 2022. (Spotlight, 5% acceptance rate)

• Lower Bounds and Accelerated Algorithms for Bilevel Optimization
  Kaiyi Ji, Yingbin Liang
  Journal of Machine Learning Research (JMLR) 2022.

• Data Sampling Affects the Complexity of Online SGD over Dependent Data
  Shaocong Ma, Ziyi Chen, Yi Zhou, Kaiyi Ji, Yingbin Liang.
  Conference on Uncertainty in Artificial Intelligence (UAI) 2022.

2021

• Provably Faster Algorithms for Bilevel Optimization <Code>
  Junjie Yang, Kaiyi Ji, Yingbin Liang
  Conference on Neural Information Processing Systems (NeurIPS) 2021. (Spotlight, 3% acceptance rate)

• Bilevel Optimization: Nonasymptotic Analysis and Faster Algorithms <Code>
  Kaiyi Ji, Junjie Yang, Yingbin Liang
  International Conference on Machine Learning (ICML) 2021 (21.4% acceptance rate)
  Short version appears at ICML 2021 Workshop: Beyond First-Order Methods in ML Systems.

• A New One-Point Residual-Feedback Oracle For Black-Box Learning and Control
  Yan Zhang, Yi Zhou, Kaiyi Ji, Michael M. Zavlanos
  Automatica 2021.

• Theoretical Convergence of Multi-Step Model-Agnostic Meta-Learning
  Kaiyi Ji, Junjie Yang, Yingbin Liang
  Journal of Machine Learning Research (JMLR) 2021
  Short version appears at ICML 2021 Workshop: Beyond First-Order Methods in ML Systems

• Understanding Estimation and Generalization Error of Generative Adversarial Networks
  Kaiyi Ji, Yi Zhou, Yingbin Liang
  IEEE Trans. on Information Theory (TIT), 2021.

• When Will Gradient Methods Converge to Max-margin Classifier under ReLU Models?
  Tengyu Xu, Yi Zhou, Kaiyi Ji, Yingbin Liang
  Stat, Special Issue of Deep Learning from Statistical Perspective, 2021

2020

• Convergence of Meta-Learning with Task-Specific Adaptation over Partial Parameters
  Kaiyi Ji, Jason D. Lee, Yingbin Liang, H. Vincent Poor
  Conference on Neural Information Processing Systems (NeurIPS) 2020 (20.1% acceptance rate).

• History-Gradient Aided Batch Size Adaptation for Variance Reduced Algorithms
  Kaiyi Ji, Zhe Wang, Bowen Weng, Yi Zhou, Wei Zhang, Yingbin Liang,
  International Conference on Machine Learning (ICML) 2020 (21.8% acceptance rate).

• Learning Latent Features with Pairwise Penalties in Matrix Completion
  Kaiyi Ji, Jian Tan, Yuejie Chi, Jinfeng Xu
  IEEE Trans. on Signal Processing (TSP), 2020
  Short version as an invited paper to a special session at the IEEE SAM 2020

• Proximal Gradient Algorithm with Momentum and Flexible Parameter Restart for Nonconvex Optimization
  Yi Zhou, Zhe Wang, Kaiyi Ji, Yingbin Liang, Vahid Tarokh
  International Joint Conferences on Artificial Intelligence Organization (IJCAI) 2020 (12.6% acceptance rate).

• Robust Stochastic Bandit Algorithms under Probabilistic Unbounded Adversarial Attack
  Ziwei Guan, Kaiyi Ji, Donald J. Bucci Jr, Timothy Hu, Joseph Palombo, Michael Liston, Yingbin Liang
  AAAI Conference on Artificial Intelligence (AAAI) 2020 (20.6% acceptance rate).

2019

• SpiderBoost and Momentum: Faster Stochastic Variance Reduction Algorithms
  Zhe Wang, Kaiyi Ji, Yi Zhou, Yingbin Liang, Vahid Tarokh
  Conference on Neural Information Processing Systems (NeurIPS) 2019 (21% acceptance rate).

• Improved Zeroth-Order Variance Reduced Algorithms and Analysis for Nonconvex Optimization
  Kaiyi Ji, Zhe Wang, Yi Zhou, Yingbin Liang
  International Conference on Machine Learning (ICML) 2019 (22.5% acceptance rate).

• Minimax Estimation of Neural Net Distance
  Kaiyi Ji, Yingbin Liang
  Conference on Neural Information Processing Systems (NeurIPS) 2018 (21% acceptance rate).

2018

• Asymptotic Miss Ratio of LRU Caching with Consistent Hashing
  Kaiyi Ji, Guocong Quan, Jian Tan
  IEEE INFOCOM 2018 (19% acceptance rate).

• LRU Caching with Dependent Competing Requests
  Guocong Quan, Kaiyi Ji, Jian Tan
  IEEE INFOCOM 2018 (19% acceptance rate).

• On Resource Pooling and Separation for LRU Caching
  Jian Tan, Guocong Quan, Kaiyi Ji, Ness Shroff
  ACM SIGMETRICS 2018 (20% acceptance rate).

• On Resource Pooling and Separation for LRU Caching
  Jian Tan, Guocong Quan, Kaiyi Ji, Ness Shroff
  Proceedings of the ACM on Measurement and Analysis of Computing Systems (POMACS), 2018