Post-training as inference
Adapting pretrained diffusion and language models to rewards as exact sampling of reward-tilted distributions, not just mode-seeking (DDPP, RTB).
PhD Student · Mila · Université de Montréal · Caltech
Jarrid
Rector-Brooks
generative models · probabilistic inference · AI for science
I am a PhD student at Mila and the Université de Montréal, supervised by Yoshua Bengio, winner of the 2018 A.M. Turing Award. I’m currently a visiting researcher in the lab of Frances Arnold, winner of the 2018 Nobel Prize in Chemistry, at Caltech.
I work on algorithms for training, adapting, and sampling from generative models in both discrete and continuous spaces — especially RL post-training viewed as probabilistic inference. Much of my work applies these methods to protein and molecular design, with designs tested in the lab.
During my PhD I co-founded Dreamfold, a drug discovery startup, where as Chief Scientific Officer I directed design research and experimental campaigns. I also co-founded the GEM workshop on biomolecular design. I studied computer science at the University of Michigan during undergrad.
If you’re interested in working together, feel free to reach out by email.
01 — Focus
What I work on
Sampling & amortized inference
Neural samplers for unnormalized densities — Boltzmann distributions and Bayesian posteriors — that find rare modes and get their probabilities right (iDEM, diffusion samplers).
Efficient generative models
Flow matching, diffusion, and discrete diffusion models built for fast generation, usable likelihoods, and cheap adaptation (OT-CFM, PAPL, P2).
Experimentally validated molecular design
Models that co-design protein sequence and structure for the de novo design of proteins, enzymes, and binders (DISCO, FoldFlow).
DISCO
General Multimodal Protein Design Enables DNA-Encoding of Chemistry
PAPL
Planner-Aware Path Learning in Diffusion Language Model Training
OXtal
OXtal: An All-Atom Diffusion Model for Organic Crystal Structure Prediction
Disc2Cont
From Discrete-Time Policies to Continuous-Time Diffusion Samplers: Asymptotic Equivalences and Faster Training
INCO
Inverse-Confidence Sampling for Continuous Diffusion Language Models
DDPP
Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction
Adversarial Exploration
Adaptive Teachers for Amortized Samplers
BIP-RL
Solving Bayesian Inverse Problems with Diffusion Priors and Off-Policy RL
P2
Path Planning for Masked Diffusion Model Sampling
iDEM
Iterated Denoising Energy Matching for Sampling from Boltzmann Densities
DiffSampler
Improved Off-Policy Training of Diffusion Samplers
FoldFlow-2
Sequence-Augmented SE(3)-Flow Matching for Conditional Protein Backbone Generation
RTB
Amortizing Intractable Inference in Diffusion Models for Vision, Language, and Control
FoldFlow
SE(3)-Stochastic Flow Matching for Protein Backbone Generation
RL Crystals
Learning Conditional Policies for Crystal Design Using Offline Reinforcement Learning
LambdaZero
Generative Active Learning for the Search of Small-Molecule Protein Binders
OT-CFM
Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport
SubTB
Learning GFlowNets from Partial Episodes for Improved Convergence and Stability
MOGFN
Multi-Objective GFlowNets
TS-GFN
Thompson Sampling for Improved Exploration in GFlowNets
FragDiff
Molecular Fragment-Based Diffusion Model for Drug Discovery
RECOVER
RECOVER: Sequential Model Optimization Platform for Combination Drug Repurposing Identifies Novel Synergistic Compounds in vitro
DEUP
DEUP: Direct Epistemic Uncertainty Prediction
ProtGFN
Biological Sequence Design with GFlowNets
Frank-Wolfe
Revisiting Projection-Free Optimization for Strongly Convex Constraint Sets
03 — Beyond the lab
Other stuff
When I’m not building models, you’ll usually find me hiking or biking. I grew up in small-town Michigan and played snare drum in the Michigan Marching Band along the way. Most of all, I’m lucky to share my days with my brilliant partner Alana Valko, our dog Ella, and cat Clementine.
