About me

I am a third-year Ph.D. student in Computer Science at the University of Massachusetts Amherst, advised by Prof. Justin Domke.

My research lies at the intersection of probabilistic machine learning and stochastic optimization. I develop scalable methods for approximate inference and large-scale learning, with interests spanning Bayesian inference, generative modeling, and optimization for modern deep learning systems.

I am currently a Guest Researcher at the Flatiron Institute, where I work on optimization for large-scale model training, including LLM pre-training. Previously, I interned at the Flatiron Institute, working on optimization methods for deep learning and Mixture-of-Experts models, and at Flagship Pioneering, where I worked on sampling-based protein design.

Prior to joining UMass, I completed my Master’s at KAIST and earned my undergraduate degree at Yonsei University.

Selected Publications (full list)

SoftServe: A Scalable Quasi-Newton Method for Deep Learning [pdf, code, blog]
Joohwan Ko, Tetiana Parshakova, Diana Cai, Robert M. Gower
Preprint, 2026

Amortized Factor Inference Networks for Posterior Inference [pdf, code]
Joohwan Ko, Justin Domke
Preprint, 2026

Model Informed Flows for Bayesian Inference [pdf, code]
Joohwan Ko, Justin Domke
NeurIPS 2025

Latent Target Score Matching, with an Application to Simulation-Based Inference [pdf]
Joohwan Ko, Tomas Geffner
NeurIPS MLPS Workshop, 2025

Learning to Scale Logits for Temperature-Conditional GFlowNets [pdf, code]
Minsu Kim*, Joohwan Ko*, Taeyoung Yun*, Dinghuai Zhang, Ling Pan, Woo Chang Kim, Jinkyoo Park, Emmanuel Bengio, Yoshua Bengio
ICML 2024

Provably Scalable Black-Box Variational Inference with Structured Variational Families [pdf]
Joohwan Ko*, Kyurae Kim*, Woo Chang Kim, Jacob R. Gardner
ICML 2024

(* denotes equal contribution)