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.

• Learning Discrete Concepts in Latent Hierarchical Models [Arxiv]
  L. Kong, G. Chen, B. Huang, E. Xing, Y. Chi, and K. Zhang, 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.

Mechanisms of Transformers

• A Theoretical Analysis of Self-Supervised Learning for Vision Transformers [Arxiv]
  Y. Huang*, Z. Wen*, Y. Chi, and Y. Liang, preprint.

• 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

• 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, preprint.

• 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, preprint.

LLM Efficiency

• 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, preprint.

• 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.