Welcome!

Hi! I am a Ph.D. candidate in the department of Electrical and Computer Engineering (ECE) at University of Wisconsin-Madison, advised by Prof. Dimitris Papailiopoulos and Prof. Kangwook Lee. I received my M.S. in 2018 from Seoul National University, where I was advised by Prof. Jungwoo Lee and learned about communication systems. I also received my B.S. in ECE from Seoul National University. I am a recipient of the Korean Government Scholarship Program for Study Overseas.

I am open to research collaboration and internship / full time opportunities.

Research Interest

My research focuses on understanding and improving the generalization and reasoning capabilities of foundation models. I study how language models generalize beyond their training distributions—such as compositional, length-based, and easy-to-hard generalization—and how algorithmic and abstract reasoning emerge. I explore tasks where reasoning and knowledge can be disentangled, including algorithmic, mathematical, and logical problems. My work explores the interplay between data, inductive biases, and learning dynamics, leveraging self-improvement frameworks, prompting strategies, and optimization techniques.

News

• (Jun. 2025) Our paper on Extrapolation by Association: Length Generalization Transfer in Transformers is on Arxiv!
• (May. 2025) Excited to start an internship at Amazon AWS as an Applied Scientist Intern! See you in NYC!
• (May. 2025) Our paper on Self-Improving Transformers Overcome Easy-to-Hard and Length Generalization Challenges is accepted at ICML 2025!
• (Mar. 2025) Our paper on Self-Improving Transformers Overcome Easy-to-Hard and Length Generalization Challenges is accepted at ICLR 2025 SSI-FM Workshop (Oral Spotlight)! - Paper | Poster
• (May. 2024) Our paper on Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks is accepted at ICML 2024!
• (Jan. 2024) Our paper on Teaching Arithmetic to Small Transformers is accepted at ICLR 2024! - Paper | Poster
• (Aug. 2023) Gave a Short Talk at the Simons Institute LLM workshop! - Video | Slides
• (Jun. 2022) Our paper on Super Seeds: extreme model compression by trading off storage with computation is accepted at ICML 2022 UpML Workshop (Oral Spotlight)! - Paper | Poster

Selected Publications

• Extrapolation by Association: Length Generalization Transfer in Transformers, Unders submission 2025
  Z. Cai, N. Lee, A. Schwarzschild, S. Oymak, D. Papailiopoulos
  

• Self-Improving Transformers Overcome Easy-to-Hard and Length Generalization Challenges, ICML 2025
  N. Lee, Z. Cai, A. Schwarzschild, K. Lee, D. Papailiopoulos
  

• Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks, ICML 2024
  J. Park, J. Park, Z. Xiong, N. Lee, J. Cho, S. Oymak, K. Lee, D. Papailiopoulos
  

• Teaching Arithmetic to Small Transformers, ICLR 2024
  N. Lee, K. Sreenivasan, J. D. Lee, K. Lee, D. Papailiopoulos
  

• Super Seeds: extreme model compression by trading off storage with computation, ICML 2022 Workshop (oral spotlight)
  N. Lee, S. Rajput, J. Sohn, H. Wang, A. Nagle, E. Xing, K. Lee, D. Papailiopoulos
  

Work Experience

Amazon AWS (2025.05 - 2025.08)

• Applied Scientist Intern
• Mentor: Dr. Dionysis Manousakas

Communications and Machine Learning Lab, SNU (2018.03 - 2020.02)

• Graduate Research Assistant
• Advisor: Prof. Jungwoo Lee
• Topics: Channel Estimation, Non-Orthogonal Multiple Access

Nanyang Technological University, SNU (2018.03 - 2020.02)

• Undergraduate Research Intern
• Mentor: Prof. Junsong Yuan
• Projects: Hand Gesture Recognition Task Using Deep Learning Based on Google Soli

Teaching Experience

Graduate Teaching Assistant

• (UW-Madison CS/ECE 561) Probability and Information Theory in Machine Learning, 2022 Fall
• (UW-Madison CS/ECE/ME 532) Matrix Methods in Machine Learning, 2021 Fall
• (SNU) Information Theory, 2019-1
• (SNU) Introduction to Communications, 2018-2

Undergraduate Teaching Assistant

• Foundation of Physics 1, 2015-1, 2016-1,2017-1
• Foundation of Physics 2, 2015-2, 2016-2, 2017-2

Talks

• (02.24.2025) Talk at the FLaNN Seminars
• (09.11.2023) Talk at the IFDS Ideas Forum

• (08.18.2023) Short Talk at the Simons Institute LLM Workshop - Video

  Slides

---