Steering GFlowNets for Drug Discovery: Which Directions Control Molecule Generation
Amirtha Varshini Anbuchezhiyan Sindhanai ⋅ Hok Hei Tam ⋅ Matthew Parks
Abstract
Once a drug discovery program has a lead scaffold, the goal becomes optimization: improving certain properties while maintaining other desired ones. The standard approach, enumerating and filtering analogues, requires a scorer for each criterion and struggles as objectives grow and interact. An alternative is to steer a generative model directly. SynFlowNet, a GFlowNet trained to sample molecules in proportion to a multi-objective RDKit reward, builds molecules from reaction templates and purchasable blocks, representing each partial molecule as a vector it reads to choose its next step. Steering adds a fixed direction to this vector during generation, nudging the model toward a target property. Prior work shows such properties are linearly decodable from these vectors. But a direction that can be decoded is not necessarily one that can steer; we test which directions control what kinds of molecules the model builds. We compare two sources for the direction. Derived directions come from the model's own weights. Because the components picking the next block and template are linear, the direction that most favors any action follows in closed form directly from the weights, with no data and no fitting. Fitted directions are learned from the embeddings of generated molecules, using linear probes, difference in class means, and sparse autoencoder features. To separate a genuine steering effect from chance, we compare every direction against the largest of five random directions of equal size, since even a random push moves properties across roughly $25\%$ percent of the reachable range, and we require the model to keep generating diverse molecules rather than collapse to one. All runs score molecules with an exact cheminformatics calculation, not a learned surrogate, so the comparisons are reliable. Fitted directions do not work. Sparse autoencoder features succeed on $6$ of $40$ tests ($8$ substructures $\times$ $5$ autoencoders), indistinguishable from the $1$-in-$6$ rate expected by chance ($p = 0.68$); linear probes succeed on $2$ of $8$ properties ($p = 0.40$) and difference in means on $1$ of $8$ ($p = 0.77$), and the one probe that looked positive fails to replicate across four independently trained models. These directions move the model less than a random direction of equal size, so the problem is no effect to begin with, not too strong a push. What can be read from the representation, then, is not what controls it. Directions derived from the weights work at a nudge strength that keeps the molecules diverse. A single chosen action can be made much more likely while molecules stay varied: a building block from about $1\%$ of selections to as much as $56\%$, and every reaction template tested from its usual $6$ to $12\%$ of reactions to roughly half. For properties that no single action controls, we compose a direction from the weights of all the blocks, still without fitting; this controls all eight physicochemical properties we tested, moving each one $1.62$ to $3.91\times$ further than a random nudge of the same size would, again without sacrificing diversity. Because the direction comes from the weights, we can predict each logit shift in advance; the shift we then measure matches that prediction to six significant figures. How hard the direction is applied matters: as strength rises the effect grows but diversity falls, so the usable range is bounded only from above. A gentle nudge already beats the null; past a point the model repeats one molecule. It transfers, holding for six properties after tuning on two, but the diversity headroom varies up to twentyfold across properties, so effect and diversity should be checked together for each target. In short, weight-derived directions steer generation where representation-fitted ones fail, provided the strength stays low enough to keep outputs diverse. In practice this lets a medicinal chemist push a property such as polarity or $\log P$ toward a target during generation, obtaining novel, synthesizable molecules without training a new scorer for every objective.
Chat is not available.
Successful Page Load