Steering GFlowNets for Drug Discovery: Which Directions Control Molecule Generation
Abstract
Lead optimization asks a generative model to improve certain properties while preserving others, and steering, adding a fixed direction to the model's internal state during generation, is an appealing alternative to retraining a scorer for every objective. Prior work shows chemical properties are linearly decodable from the internal states of SynFlowNet, a GFlowNet that builds molecules from reaction templates and purchasable blocks, so we asked whether those directions can also steer it. They cannot: sparse autoencoder features succeed at 6 of 40 trials and linear probes at 2 of 8, both indistinguishable from the rate a random direction of the same size achieves by construction, and all of them move the model less than that random direction does. Directions derived in closed form from the model's own weights do steer, raising a chosen synthesis action from about 1% of selections to 56% and moving all eight physicochemical properties tested on each of three independently trained policies, with 90 to 97% of molecules remaining distinct and the targeted property moving about three times as far as untargeted ones. What can be decoded from a representation is therefore not what steers it; what controls generation is the discrete action space the model chooses among.