About
I am a Ph.D. student in Computer Science at the University of Central Florida (Institute of Artificial Intelligence), advised by Dr. Mubarak Shah.
My research focuses on inference-time scaling and alignment for generative models. I specialize in developing training-free methods to enhance the safety, consistency, and personalization of Diffusion Models and Diffusion-LLMs.
Previously, I worked as a Computer Engineer building simulation systems. I hold a B.Sc. in Computer Science, Physics, and Mathematics (Triple Major) from Vanderbilt University.
Selected Manuscripts
Test-Time Scaling in Diffusion LLMs via Hidden Semi-Autoregressive Experts
J. Lee, H. Moon, K. Zhai, A.K. Chithanar, A.K. Sahu, S. Kar, C. Lee, S. Chakraborty, A.S. Bedi
[ACCEPTED] ICLR 2026
Proposing a training-free inference scaling method (HEX) that outperforms fine-tuned baselines by leveraging hidden semi-autoregressive experts.
Mitigating Reward Hacking in Inference-Time Noise Optimization for Text-to-Image Diffusion Models
K. Zhai, U. Singh, A. Thatipelli, S. Chakraborty, A.K. Sahu, F. Huang, A.S. Bedi, M. Shah
Submitted to ICML 2026
CORE: Context-Robust Remasking for Diffusion Language Models
K. Zhai, S. Mollah, Z. Wang, M. Shah
Submitted to ICML 2026
Consensus-to-Personal: Inference-Time Alignment for Text-to-Image Personalization
K. Zhai, U. Singh, S. Chakraborty, C. Rastogi, F. Huang, A.S. Bedi, M. Shah
Submitted to CVPR 2026
Repair-Aware Forgetting: An Iterative Approach to Unlearning in T2I Diffusion Models
S. Ghosh, K. Zhai, U. Singh, S. Chakraborty, A.S. Bedi, M. Shah
Submitted to ICML 2026