About Me

Hi! I am Kai Yang (杨铠), a final-year undergraduate student in Turing Class, School of EECS, Peking University. I am advised by Prof. Liwei Wang and Prof. Di He. During the summer of 2024, I worked as an student intern advised by Prof. Pradeep Ravikumar in Carnegie Mellon University.

My research interest lies in machine learning theory and algorithms. I study the foundations of Transformers and other contemporary models, understanding their power and limitations. Recently, I am also working on developing new representation learning algorithms with theoretical guarantees.

I’m applying for a Ph.D. position in the area of machine learning theory starting from Fall, 2025. Please contact me through email if you are interested in my work!

Publications and Preprints

* means equal contribution.

  1. Do Efficient Transformers Really Save Computation?
    Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, Liwei Wang. In ICML 2024.

  2. Two Stones Hit One Bird: Bilevel Positional Encoding for Better Length Extrapolation
    Zhenyu He*, Guhao Feng*, Shengjie Luo*, Kai Yang, Liwei Wang, Jingjing Xu, Zhi Zhang, Hongxia Yang, Di He. In ICML 2024. Code.

  3. How Numerical Precision Affects Mathematical Reasoning Capabilities of LLMs
    Guhao Feng*, Kai Yang*, Yuntian Gu, Xinyue Ai, Shengjie Luo, Jiacheng Sun, Di He, Zhenguo Li, Liwei Wang. arXiv Preprint.

Professional Services

Peer Review:

  • ICLR: 2025.
  • NeurIPS workshop: M3L 2024.

Teaching Assistant:

  • Peking University: Mathematical Foundations for the Information Age, Fall 2023 & Fall 2024.
  • Peking University: Randomized Algorithms, Spring 2024.

Selected Awards and Honors

  • Chinese National Scholarship, The highest level of honor for university students, 2021 - 2022.
  • May Fourth Scholarship, The highest level scholarship at Peking University, 2023 - 2024.
  • First Prize, 14th & 15th National University Mathematical Contest, 2022 & 2023.
  • Silver Medal, 36th Chinese Mathematical Olympiad, 2020.