About Me

I am a researcher in machine learning and artificial intelligence, interested in understanding the relationship between training data and the properties of trained models. I completed my PhD at the University of Toronto in 2025, studying under Roger Grosse. I received my HBSc in Computer Science and Statistics from the same university in 2019.

Selected Publications

Studying Large Language Model Generalization with Influence Functions

Roger Grosse*, Juhan Bae*, Cem Anil*, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamilė Lukošiūtė, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, Samuel Bowman

Paper Blog Code Video Slides

---

Training Data Attribution via Approximate Unrolled Differentiation

Juhan Bae, Wu Lin, Jonathan Lorraine, Roger Grosse

Paper Code Video Slides Poster

---

What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions

Sang Keun Choe, Hwijeen Ahn*, Juhan Bae*, Kewen Zhao*, Minsoo Kang, Youngseog Chung, Adithya Pratapa, Willie Neiswanger, Emma Strubell, Teruko Mitamura, Jeff Schneider, Eduard Hovy, Roger Grosse, Eric Xing

Paper Code Slides

---

If Influence Functions are the Answer, Then What is the Question?

Juhan Bae, Nathan Ng, Alston Lo, Marzyeh Ghassemi, Roger Grosse

Paper Code (PyTorch) Code (Jax) Video Poster

---

Procedural Knowledge in Pretraining Drives Reasoning in Large Language Models

Laura Ruis, Maximilian Mozes, Juhan Bae, Siddhartha Rao Kamalakara, Dwaraknath Gnaneshwar, Acyr Locatelli, Robert Kirk, Tim Rocktaschel, Edward Grefenstette, Max Bartolo

Paper Blog Code Demo

---

Influence Functions for Scalable Data Attribution in Diffusion Models

Bruno Mlodozeniec, Runa Eschenhagen, Juhan Bae, Alexander Immer, David Krueger, Richard Turner

Paper Code

---

Benchmarking Neural Network Training Algorithms

George Dahl*, Frank Schneider*, Zachary Nado*, Naman Agarwal*, Chandramouli Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, Juhan Bae, Justin Gilmer and 13 more authors

Paper Blog Code

---

Accelerating Neural Network Training: An Analysis of the AlgoPerf Competition

Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, Boyuan Feng, Less Wright, Edward Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig, Michael Rabbat, George Dahl

Paper Blog Code

---

Amortized Proximal Optimization

Juhan Bae*, Paul Vicol*, Jeff HaoChen, Roger Grosse

Paper Code Video Slides Poster

Teaching

Instructor:

TUSK: Machine Learning Software Foundations, Fall 2022

CSC311: Introduction to Machine Learning, Fall 2020

Teaching Assistant:

CSC2547: AI Alignment, Winter 2023

CSC110: Foundations of Computer Science I, Fall 2023

CSC2702: Technical Entrepreneurship, Summer 2022

HLP101: Undergraduate CS Course Help Centre, Summer 2022

CSC2541: Neural Net Training Dynamics, Winter 2021

STA314: Statistical Methods for Machine Learning I, Fall 2021

CSC110: Foundations of Computer Science I, Fall 2021

CSC320: Introduction to Visual Computing, Winter 2020

Vector Institute AI Certificate Course: Deep Learning II, Winter 2020

CSC412/2506: Probabilistic Learning and Reasoning, Fall 2020

CSC165: Mathematical Expression and Reasoning for CS, Fall 2019

Vector Institute AI Certificate Course: Deep Learning I, Fall 2019

CSC165: Mathematical Expression and Reasoning for CS, Winter 2017