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
Joohwan Ko, Tetiana Parshakova, Diana Cai, Robert M. Gower
Preprint, 2026

Amortized Factor Inference Networks for Posterior Inference
Joohwan Ko, Justin Domke
Preprint, 2026

Model Informed Flows for Bayesian Inference
Joohwan Ko, Justin Domke
NeurIPS 2025

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

Learning to Scale Logits for Temperature-Conditional GFlowNets
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
Joohwan Ko*, Kyurae Kim*, Woo Chang Kim, Jacob R. Gardner
ICML 2024

(* denotes equal contribution)