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

• Provably Robust Score-Based Diffusion Posterior Sampling for Plug-and-Play Image Reconstruction [Arxiv]
  X. Xu and Y. Chi, preprint.

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

• Towards Faster Non-Asymptotic Convergence for Diffusion-Based Generative Models [Arxiv]
  G. Li, Y. Wei, Y. Chen, and Y. Chi, preprint. Short version at ICLR 2024.

Hardware-aware Algorithms for LLM Efficiency

• Prompt-prompted Mixture of Experts for Efficient LLM Generation [Arxiv]
  H. Dong, B. Chen, and Y. Chi, preprint.

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

Theory of Transformers

• Transformers Provably Learn Feature-Position Correlations in Masked Image Modeling [Arxiv]
  Y. Huang*, Z. Wen*, Y. Chi, and Y. Liang, preprint.

Generative AI for Materials Science

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