Yuejie Chi

Foundations and Applications of Generative AI

Generative AI or foundation models, such as large language models (LLMs) and diffusion models, are claiming major successes in the recent wave of AI developments. These models have demonstrated tremendous potentials in mastering complex tasks and generating new contents, exhibiting surprising emergent capabilities such as in-context learning. At the same time, the fundamental understandings of such models are yet again falling far behind, with their training and inference posing significant resource challenges in order to democratize their use; the sheer scale of state-of-the-art LLMs thwarts frugal entities from deploying them. My group is interested in developing the algorithmic foundations of generative AI models, and pushing their use in important application domains across science and engineering.

Diffusion Models

• A Sharp Convergence Theory for The Probability Flow ODEs of Diffusion Models [Arxiv]
  G. Li, Y. Wei, Y. Chi, and Y. Chen, preprint.

• Provably Robust Score-Based Diffusion Posterior Sampling for Plug-and-Play Image Reconstruction [Arxiv] [Code]
  X. Xu and Y. Chi, Conference on Neural Information Processing Systems (NeurIPS), 2024.

• Accelerating Convergence of Score-Based Diffusion Models, Provably [Arxiv] [Code]
  G. Li*, Y. Huang*, T. Efimov, Y. Wei, Y. Chi, and Y. Chen, International Conference on Machine Learning (ICML), 2024.

• Towards Non-Asymptotic Convergence for Diffusion-Based Generative Models [Arxiv]
  G. Li, Y. Wei, Y. Chen, and Y. Chi, International Conference on Learning Representations (ICLR), 2024.

Training Dynamics of Transformers

• A Theoretical Analysis of Self-Supervised Learning for Vision Transformers [Arxiv]
  Y. Huang*, Z. Wen*, Y. Chi, and Y. Liang, International Conference on Learning Representations (ICLR), 2025.

• In-Context Learning with Representations: Contextual Generalization of Trained Transformers [Arxiv]
  T. Yang, Y. Huang, Y. Liang, and Y. Chi, Conference on Neural Information Processing Systems (NeurIPS), 2024.

LLM Alignment

• LoRe: Personalizing LLMs via Low-Rank Reward Modeling [Arxiv]
  A. Bose, Z. Xiong, Y. Chi, S. Du, L. Xiao, and M. Fazel, preprint.

• Faster WIND: Accelerating Iterative Best-of-N Distillation for LLM Alignment [Arxiv]
  T. Yang, J. Mei, H. Dai, Z. Wen, S. Cen, D. Schuurmans, Y. Chi, and B. Dai, International Conference on Artificial Intelligence and Statistics (AISTATS), 2025.

• Value-Incentivized Preference Optimization: A Unified Approach to Online and Offline RLHF [Arxiv]
  S. Cen, J. Mei, K. Goshvadi, H. Dai, T. Yang, S. Yang, D. Schuurmans, Y. Chi, and B. Dai, International Conference on Learning Representations (ICLR), 2025.

LLM Efficiency

• Scalable LLM Math Reasoning Acceleration with Low-rank Distillation [Arxiv]
  H. Dong, B. Acun, B Chen, and Y. Chi, preprint.

• ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference [Arxiv] [Code]
  H. Sun, L.-W. Chang, W. Bao, S. Zheng, N. Zheng, X. Liu, H. Dong, Y. Chi, and B. Chen, International Conference on Machine Learning (ICML), 2025, spotlight presentation.

• Prompt-prompted Adaptive Structured Pruning for Efficient LLM Generation [Arxiv] [Code]
  H. Dong, B. Chen, and Y. Chi, Conference on Language Modeling (COLM), 2024.

• Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference [Arxiv] [Code]
  H. Dong, X. Yang, Z. Zhang, Z. Wang, Y. Chi, and B. Chen, International Conference on Machine Learning (ICML), 2024.

• Towards Structured Sparsity in Transformers for Efficient Inference
  H. Dong, B. Chen, and Y. Chi, ICML Workshop on Efficient Systems for Foundation Models, 2023.

Generative AI for Materials Science

• Leveraging Multimodal Diffusion Models to Accelerate Imaging with Side Information [Arxiv]
  T. Efimov, H. Dong, M. Shah, J. Simmons, S. Donegan, and Y. Chi, International Conference on Acoustics, Speech, and Signal Processing (ICASSP), 2025.

• A Lightweight Transformer for Faster and Robust EBSD Data Collection [Arxiv] [Code]
  H. Dong, S. Donegan, M. Shah, and Y. Chi, Scientific Reports, vol. 23, pp. 21253, 2023.