Kun Yang

I am a Machine Learning Engineer at Google in Sunnyvale, CA, focusing on LLM inference efficiency, self-improving multi-agent code optimization, and reinforcement learning environments for post-training and evaluation.

Previously, I served as a Senior Data Science Engineer at Juniper Networks and earned my Ph.D. in Electrical Engineering from the University of Virginia (UVA) advised by Prof. Cong Shen. Prior to UVA, I received my M.S. in Electrical Engineering from Texas A&M University (advised by Prof. Tie Liu) and my B.E. in Electrical and Information Science from Tsinghua University. My research spans multi-agent systems, bandit & prompt optimization, offline reinforcement learning, and federated learning.

Email Google Scholar GitHub LinkedIn DBLP

selected projects & focus

Google · LLM Systems

Recursive Self-Improvement (RSI) for Inference

Designing autonomous iterative search and self-improving optimization loops to improve LLM decoding and serving efficiency at scale.

Google · Post-Training

RL Environments for Gemini Training

Building scalable reinforcement learning environments and challenging coding-efficiency evaluation suites to support Gemini evaluation and post-training.

Google & NeurIPS 2024

Multi-Agent Code & Prompt Optimization

Developing self-evolving multi-agent code optimization harnesses with persistent memory, alongside bandit-based prompt optimization (TRIPLE).

news

Mar 10, 2025	Joined Google in Sunnyvale, CA as a Machine Learning Engineer working on LLM inference optimization, multi-agent systems, and RL post-training environments!
Sep 25, 2024	I’m thrilled to share that two of our papers have been accepted to NeurIPS 2024, with online version and videos coming soon. Looking forward to see you in Vancouver! Efficient Prompt Optimization Through the Lens of Best Arm Identification! (Co-first author): This is a work together with Dr. Chengshuai Shi (Bloomberg, UVA), Zihan Chen (UVA), Prof. Jundong Li (UVA), Prof. Cong Shen (UVA) and Prof. Jing Yang (PSU). Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Models: This is a work together with Dr. Chengshuai Shi (Bloomberg, UVA), Prof. Cong Shen (UVA) and Prof. Jing Yang (PSU).
Sep 18, 2024	Our paper Average Reward Reinforcement Learning for Wireless Radio Resource Management at IEEE Asilomar 2024 (with Prof. Cong Shen and Prof. Jing Yang) was selected as a Best Student Paper Award Finalist!
Aug 1, 2024	Successfully defended my Ph.D. dissertation in Electrical Engineering at the University of Virginia and joined Juniper Networks in Cupertino, CA as a Senior Data Science Engineer.
Jun 15, 2024	Our paper, Harnessing the Power of Federated Learning in Federated Contextual Bandits!, co-authored with Chengshuai Shi (UVA), Ruida Zhou (UCLA), and Cong Shen (UVA), has been accepted by Transactions on Machine Learning Research (TMLR).
May 17, 2024	Our work on utilizing offline RL for wireless communication has been accepted for publication in IEEE Transactions on Wireless Communications (TWC). Check out the manuscript here!

selected publications

2024

NeurIPS

Efficient Prompt Optimization Through the Lens of Best Arm Identification

Chengshuai Shi*, Kun Yang*, Zihan Chen, and 3 more authors

In Advances in Neural Information Processing Systems (NeurIPS), 2024

arXiv Bib

@inproceedings{shi2024efficient,
  title = {Efficient Prompt Optimization Through the Lens of Best Arm Identification},
  author = {Shi*, Chengshuai and Yang*, Kun and Chen, Zihan and Li, Jundong and Yang, Jing and Shen, Cong},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year = {2024},
}

NeurIPS

Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Models

Chengshuai Shi, Kun Yang, Jing Yang, and 1 more author

In Advances in Neural Information Processing Systems (NeurIPS), 2024

arXiv Bib

@inproceedings{shi2024transformers,
  title = {Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Models},
  author = {Shi, Chengshuai and Yang, Kun and Yang, Jing and Shen, Cong},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year = {2024},
}

IEEE TWC

Offline Reinforcement Learning for Wireless Network Optimization with Mixture Datasets

Kun Yang, Chengshuai Shi, Cong Shen, and 3 more authors

IEEE Transactions on Wireless Communications, 2024

Bib HTML

@article{yang2024offline_twc,
  title = {Offline Reinforcement Learning for Wireless Network Optimization with Mixture Datasets},
  author = {Yang, Kun and Shi, Chengshuai and Shen, Cong and Yang, Jing and Yeh, Shu-Ping and Sydir, Jaroslaw J.},
  journal = {IEEE Transactions on Wireless Communications},
  year = {2024},
  doi = {10.1109/TWC.2024.3395624},
}

TMLR

Harnessing the Power of Federated Learning in Federated Contextual Bandits

Chengshuai Shi, Ruida Zhou, Kun Yang, and 1 more author

Transactions on Machine Learning Research, 2024

Bib HTML

@article{shi2024harnessing,
  title = {Harnessing the Power of Federated Learning in Federated Contextual Bandits},
  author = {Shi, Chengshuai and Zhou, Ruida and Yang, Kun and Shen, Cong},
  journal = {Transactions on Machine Learning Research},
  issn = {2835-8856},
  year = {2024},
}

2023

IEEE JSAC

On the Convergence of Hybrid Server-Clients Collaborative Training

Kun Yang, Shengbo Chen, and Cong Shen

IEEE Journal on Selected Areas in Communications, 2023

Bib

@article{yang2023convergence_jsac,
  title = {On the Convergence of Hybrid Server-Clients Collaborative Training},
  author = {Yang, Kun and Chen, Shengbo and Shen, Cong},
  journal = {IEEE Journal on Selected Areas in Communications},
  volume = {41},
  number = {3},
  pages = {802--819},
  year = {2023},
  publisher = {IEEE},
}