Shangzhen Zhu

I am a PhD student at the University of Illinois Urbana-Champaign, advised by Prof. Tomasz Kozlowski.

My current research asks how much numerical precision softmax attention actually needs. I study which parts of softmax a pretrained LLM depends on, how to turn that into faster attention kernels (for example, a calibration-free variant inside FlashAttention-4), and what changes when a transformer is pretrained with a coarse, quantized softmax from the start. I am broadly interested in LLM efficiency, GPU kernels, and AI infrastructure.

Before moving to machine learning, I worked on PDE-based inverse problems, Monte Carlo radiation transport (FLUKA, MCNP, Serpent), and 3D radiation-field reconstruction from sparse measurements. I received my master’s degree from Tsinghua University and my bachelor’s degree from Southeast University.

A short explainer of my recent work: How Much Softmax Do You Need in LLMs?

latest posts

Oct 18, 2026	How Much Softmax Do You Need in LLMs?

selected publications

arXiv

Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration

Shangzhen Zhu, Muyan Hu, and Tomasz Kozlowski

Sep 2026

arXiv Bib Video Blog

@misc{zhu2026approximating,
  title = {Approximating Softmax in Pretrained {LLMs}: Model Sensitivity and Kernel Acceleration},
  author = {Zhu, Shangzhen and Hu, Muyan and Kozlowski, Tomasz},
  year = {2026},
  month = sep,
  archiveprefix = {arXiv},
}

arXiv

Pretraining Transformers with Quantized Softmax in Attention

Shangzhen Zhu, Muyan Hu, and Tomasz Kozlowski

Sep 2026

arXiv Bib Video Blog

@misc{zhu2026pretraining,
  title = {Pretraining Transformers with Quantized Softmax in Attention},
  author = {Zhu, Shangzhen and Hu, Muyan and Kozlowski, Tomasz},
  year = {2026},
  month = sep,
  archiveprefix = {arXiv},
}