Lee Lab @ UW Madison

Research focus

• Deep Learning with Large Pretrained Models (Large Language Models, Diffusion Models, …)
  • LIFT: Solving non-language ML tasks with language models. NeurIPS’22 [Github] [Summary]
  • Designing a CPU with a looped Transformer. ICML’23 [Summary]
  • Teaching arithmetic to a small Transformer. Arxiv’23 [Github] [Summary]
  • A compute-latency trade-off for language model decoding. Arxiv’23
  • Training CLIP models with mini-batches. Arxiv’23 [Github]
  • Designing a chatbot for long open-domain conversation with a inter-connected language models. ACL’23 (Findings) [Github] [Summary]
  • Reinforcment learning for finetuning text-to-image model. NeurIPS’23 [Summary]
  • Reinforcment learning for short-cut DDPM sampling. ICML’23 [Summary]
  • Diffusion models minimize the Wasserstein distance. NeurIPS’22 [Summary]
  • WALIP: Connecting two CLIP models for word translation. EMNLP’22 (Findings) [Github] [Summary]
  • Invited talks on LLMs/GPT/Transformers The Future Frontier of Large Language Models — A UW-Madison Panel Discussion (June 2023), BarryFest (June 2023), Innovation in Data Seminar @ Early Warning (June 2023), KSEA Distinguished Guest Series (May 2023), The second annual Wisconsin Digital Symposium (May 2023), University of Wisconsin-Madison Law School (May 2023), Midwest Regional Conference (Mar. 2023), Samsung Advanced Institute of Technology (Aug. 2022), EIRIC Korea (July 2022)
  • Invited talks on Diffusion Models The Machine Learning for Medical Imaging (ML4MI) at UW-Madison (Sept. 2023), The Coordinated Science Laboratory Student Conference (Feb. 2023), Information Theory and Applications Workshop (Feb. 2023), Information Theory and Data Science Workshop (Feb. 2023), UW Madison SILO (Jan. 2023), UW Madison Physics Department (Jan. 2022)
• ML Fairness
  • Equal Improvability ICLR’23 (Github) [Summary] [Montreal AI Ethics Institute]
  • FairBatch ICLR’21, its robust variant NeurIPS’21, and its decentralized variant MLSys-CrossFL 2022
  • Fundamental limits of local fair training in federated learning. [ISIT’23]
  • Fair training under distribution shift. ICML’23
  • Invited talks The Fairness and Ethics in ML Seminars at AmFam (Aug. 2023), The US-Mexico Workshop on Optimization and its Applications (Jan. 2023), KAIST CS Colloquium (Sept. 2022), USC EE (Apr. 2022), UCSB CCDC Seminar Series (Apr. 2022), KAIST EE Colloquium Lecture Series (Mar. 2022), UC Berkeley BLISS (Feb. 2022), IOS’22, KAIST AI International Symposium (Dec. 2021), AI+Society seminar @ University of Wisconsin (Oct. 2021), POSTECH (June 2021), Shannon meets Turing Colloquium @ Seoul National University (May 2021), Machine Learning Ideas @ Microsoft Research New England (Dec. 2020), UW Madison SILO (Nov. 2020), UC Berkeley BLISS (Nov. 2020)

New preprints

• Mini-Batch Optimization of Contrastive Loss
  Jaewoong Cho, Kartik Sreenivasan, Keon Lee, Kyunghoo Mun, Soheun Yi, Jeong-Gwan Lee, Anna Lee, Jy-yong Sohn, Dimitris Papailiopoulos, Kangwook Lee
  Arxiv 2023
• Predictive Pipelined Decoding: A Compute-Latency Trade-off for Exact LLM Decoding
  Seongjun Yang, Gibbeum Lee, Jaewoong Cho, Dimitris Papailiopoulos, Kangwook Lee
  Arxiv 2023
• Teaching Arithmetic to Small Transformers
  Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, Dimitris Papailiopoulos
  Arxiv 2023
• Vector-Valued Variation Spaces and Width Bounds for DNNs: Insights on Weight Decay Regularization
  Joseph Shenouda, Rahul Parhi, Kangwook Lee, and Robert D. Nowak
  Arxiv 2023

News

• (Sep. 2023) [NeurIPS’23]
  • DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models
• (June 2023) [ICML’23 Workshop]
  • Predictive Pipelined Decoding: A Compute-Latency Trade-off for Exact LLM Decoding
  • Looped Transformers are Better at Learning Learning Algorithms
  • A Representer Theorem for Vector-Valued Neural Networks: Insights on Weight Decay Training and Widths of Deep Neural Networks
  • Teaching Arithmetic to Small Transformers
• (June 2023) [The Federated Learning Systems (FLSys) Workshop @ MLSys 2023]
  • FedGP: Buffer-based Gradient Projection for Continual Federated Learning
  • This paper received the best paper award of the workshop! Congrats, Shenghong and Bryce!
• (May 2023) [ACL’23 (Findings)]
  • Prompted LLMs as Chatbot Modules for Long Open-domain Conversation
• (Apr. 2023) [ICML’23]
  • Optimizing DDPM Sampling with Shortcut Fine-Tuning
  • Looped Transformers as Programmable Computers
  • Improving Fair Training under Correlation Shifts
• (Apr. 2023) Bryce Chen received the 2023–2024 Wisconsin Hilldale Undergraduate/Faculty Research Fellowship! Congratulations!
• (Apr. 2023) [ISIT’23]
  • Federated Learning with Local Fairness Constraints
• (Mar. 2023) Dr. Dinh will join UCSF as a postdoc! Congratulations!
• (Mar. 2023) [ICLR’23 Workshop on Mathematical and Empirical Understanding of Foundation Models]
  • Mini-Batch Optimization of Contrastive Loss
  • Looped Transformers as Programmable Computers
• (Jan. 2023) [ICLR’23]
  • Equal Improvability: A New Fairness Notion Considering the Long-Term Impact
• (Dec. 2022) Dr. Sohn will join Yonsei University as an assistant professor! Congratulations!