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: Sampling, Alignment and Inverse Problems

• Diffusion Controller: Framework, Algorithms and Parameterization [Arxiv]
  T. Yang, M. Ryu, C.-W. Hsu, G. Tennenholtz, Y. Chi, C. Boutilier, and B. Dai, International Conference on Machine Learning (ICML), 2026.

• Polynomial Convergence of Riemannian Diffusion Models [Arxiv]
  X. Xu, Z. Zhang, Y. Nakahira, G. Qu, and Y. Chi, International Conference on Learning Representations (ICLR), 2026.

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

Learning Dynamics of Transformers

• On the Learning Dynamics of RLVR at the Edge of Competence [Arxiv]
  Y. Huang*, Z. Wen*, Y. Chi, Y. Wei, A. Singh, Y. Liang, and Y. Chen, International Conference on Machine Learning (ICML), 2026.

• Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization [Arxiv]
  Y. Huang*, Z. Wen*, A. Singh, Y. Chi, and Y. Chen, Conference on Neural Information Processing Systems (NeurIPS), 2025.

• Multi-Head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent [Arxiv]
  T. Yang, Y. Huang, Y. Liang, and Y. Chi, Conference on Neural Information Processing Systems (NeurIPS), 2025.

• 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, Inference and Scaling

• Generalized Parallel Scaling with Interdependent Generations [Arxiv]
  H. Dong, D. Brandfonbrener, E. Helenowski, Y. He, M. Kumar, H. Fang, Y. Chi, and K. Sankararaman, International Conference on Learning Representations (ICLR), 2026.
  NeurIPS 2025 Workshop on Efficient Reasoning Best Paper Nomination

• Scalable LLM Reasoning Acceleration with Low-rank Distillation [Arxiv]
  H. Dong, B. Acun, B Chen, and Y. Chi, Conference on Parsimony and Learning (CPAL), 2026.

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

• LoRe: Personalizing LLMs via Low-Rank Reward Modeling [Arxiv]
  A. Bose, Z. Xiong, Y. Chi, S. Du, L. Xiao, and M. Fazel, Conference on Language Modeling (COLM), 2025.

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

• 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

• Multimodal Diffusion to Mutually Enhance Polarized Light and Low Resolution EBSD Data [Arxiv]
  H. Dong, T. Efimov, M. Shah, J. Simmons, S. Donegan, M. De Graef, and Y. Chi, preprint. Short version at 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.