Thomas Zeng, Shuibai Zhang, Shutong Wu, Christian Classen, Daewon Chae, Ethan Ewer, Minjae Lee, Heeju Kim, Wonjun Kang, Jackson Kunde, Ying Fan, Jungtaek Kim, Hyung Il Koo, Kannan Ramchandran, Dimitris Papailiopoulos, Kangwook Lee
Process Reward Models (PRMs) have proven effective at enhancing mathematical reasoning for Large Language Models (LLMs) by leveraging increased inference-time computation. However, they are predominantly trained on mathematical data and their generalizability to non-mathematical domains has not been rigorously studied. In response, this work first shows that current PRMs have poor performance in other domains. To address this limitation, we introduce VersaPRM, a multi-domain PRM trained on synthetic reasoning data generated using our novel data generation and annotation method. VersaPRM achieves consistent performance gains across diverse domains. For instance, in the MMLU-Pro category of Law, VersaPRM via weighted majority voting, achieves a 7.9% performance gain over the majority voting baseline -- surpassing Qwen2.5-Math-PRM's gain of 1.3%. We further contribute to the community by open-sourcing all data, code and models for VersaPRM.
Ethan Ewer, Daewon Chae, Thomas Zeng, Jinkyu Kim, Kangwook Lee
Next-token prediction is conventionally done using decoder-only Transformers with causal attention, as this approach allows for efficient reuse of keys and values. What if we were not compute-limited, should we still use decoder-only Transformers? We introduce Encoder-only Next Token Prediction (ENTP). We use small scale experiments to explore the differences between ENTP and decoders, highlighting potential advantages of ENTP in setting with unbounded compute.
KRAFTON AI Fellowship, June 2025– August 2025 (Upcoming)
Selected through a highly competitive multi-stage evaluation process
Conducting research under guidance of leading AI researchers and professors
University of Wisconsin Madison, April 2024 – Present
Explored using encoder-only transformers for next-token prediction (ENTP)
Investigated search algorithms to improve large language model reasoning performance, by scaling inference-time compute (VersaPRM)
Supervisor: Kangwook Lee
University of Wisconsin Madison, February 2024 – Present
Researched relationships between brain connectivity and consciousness using machine learning and signal processing techniques
Supervisor: Matthew Banks
Tutor, UW Madison Undergraduate Learning Center, February 2023 - Present
Led group sessions and worked one-on-one with students to develop engineering problem-solving skills.
Tutored physics, differential equations, and linear algebra.
Tutor and Course Assistant, UW Madison Math Learning Center, October 2022 - May 2024
Led group sessions and worked one-on-one with students to develop creative problem-solving skills.
Tutored pre-calculus, calculus, linear algebra, and combinatorics.
Wisconsin Hilldale Undergraduate/Faculty Research Fellowship, 2025 - 2026
The Hilldale Undergraduate/Faculty Research Fellowship provides research training and support to undergraduates. Students have the opportunity to undertake their own research project in collaboration with UW–Madison faculty or research/instructional academic staff.
Summer Music Clinic Tuition Waiver, 2022 - 2026
One of 10 Wisconsin high school juniors selected for a competitive tuition waiver to UW–Madison, awarded for outstanding musical achievement and potential.