Deep Sets GFlowNets for Subset Selection, with an Application to Combinatorial Chemistry

DNA-encoded libraries (DELs) are combinatorial chemical libraries screened against a protein target in a single pooled experiment, avoiding the high costs associated with evaluating molecules one-by-one. We formulate DEL design as a combinatorial optimization problem which involves choosing sets of building blocks which produce many high-quality combinations. Generative Flow Networks (GFlowNets; Bengio et al. (2023)), which learn to sample compositional objects proportional to a reward, are well-suited to this problem, generating diverse, high-reward libraries rather than converging to a single optimum. However, applying GFlowNets to real-world DELs presents two open challenges: 1) DELs must be represented in a way that preserves the chemical structure of the underlying building blocks for efficient learning; 2) the cost of scoring candidate libraries must remain bearable at real-world scales. To address these obstacles, we leverage DeepSets representations to amortize the computation cost of combinatorial rewards and obtain a structure-aware GFlowNet architecture suitable for the optimization of combinatorial set functions on large structured environments. Applied to DEL design, our approach scales to libraries containing 125 billion elements, three orders of magnitude beyond prior work.