Discrete Compositional Generation via General Soft Operators and Robust Reinforcement Learning

TL;DR

This paper introduces a new RL-based method with a unified operator and robustness perspective to improve the quality and diversity of candidate generation in large search spaces, such as proteins or molecules.

Contribution

It proposes a novel unified operator for regularized RL that better targets peakier distributions and introduces a robust RL framework for filtering candidates.

Findings

TGM outperforms baselines in synthetic tasks

TGM identifies higher quality candidates in real-world applications

The method enhances diversity and robustness of candidate generation

Abstract

A major bottleneck in scientific discovery consists of narrowing an exponentially large set of objects, such as proteins or molecules, to a small set of promising candidates with desirable properties. While this process can rely on expert knowledge, recent methods leverage reinforcement learning (RL) guided by a proxy reward function to enable this filtering. By employing various forms of entropy regularization, these methods aim to learn samplers that generate diverse candidates that are highly rated by the proxy function. In this work, we make two main contributions. First, we show that these methods are liable to generate overly diverse, suboptimal candidates in large search spaces. To address this issue, we introduce a novel unified operator that combines several regularized RL operators into a general framework that better targets peakier sampling distributions. Secondly, we offer a novel, robust RL perspective of this filtering process. The regularization can be interpreted as robustness to a compositional form of uncertainty in the proxy function (i.e., the true evaluation of a candidate differs from the proxy's evaluation). Our analysis leads us to a novel, easy-to-use algorithm we name trajectory general mellowmax (TGM): we show it identifies higher quality, diverse candidates than baselines in both synthetic and real-world tasks. Code: https://github.com/marcojira/tgm.