PaperlayerView exact revision on arXiv

Source-linked AI summary

Elite-Weighted Supervised Fine-tuning for Goal-Directed Molecular Optimization

Shiyun Wa, Yifei Wang, Anna G. Green, Simone Sciabola, Ye Wang

arXiv:2609.00189v1cs.LG

TL;DR

Goal-directed molecular optimization often relies on architecture-dependent trajectory-level RL, while fixed-dataset supervised fine-tuning does not use oracle reward. EW-SFT injects reward through elite selection and native-loss updates, achieving consistent improvements across generators, constraints, references, and oracles. Its computational gains do not establish biological or therapeutic properties without experimental validation.

  • Problem

    Trajectory-level RL requires architecture- and procedure-dependent generation log-probabilities, whereas fixed-dataset supervised fine-tuning does not incorporate oracle reward.

  • Method

    EW-SFT selects high-scoring molecules with a rolling elite buffer and updates each generator using its native pretraining loss.

  • Results

    EW-SFT consistently outperforms corresponding native optimizers across autoregressive, masked-diffusion, and discrete-flow generators and de novo, motif-extension, and linker-design tasks.

  • Takeaways & Limitations

    Elite selection is the primary route by which reward enters the update, enabling a unified optimizer without trajectory-level RL.

  • Takeaways & Limitations

    3D- and 2D-similarity improvements are computational evidence and do not establish binding, selectivity, synthesis feasibility, safety, or biological activity without experimental validation.

Abstract

from arXiv · show

Goal-directed optimization is essential for steering molecular generators to propose candidates with desired properties. However, it is often implemented with policy-gradient reinforcement learning, which requires a generation-trajectory log-probability whose form depends on the model architecture and generation procedure. This makes an optimizer difficult to reuse across architectures and conditional generative designs. Supervised fine-tuning needs none of that machinery, but its update is driven by a fixed dataset, so the reward never enters the update. We introduce Elite-Weighted Supervised Fine-tuning (EW-SFT), which uses reward to guide elite selection of high-scoring molecules, and updates the model by its own pretraining loss on that set. Ablations show that reward information is passed primarily through elite selection, rather than through continuous weighting within the selected set. Because the update consumes only scored molecules and the model's native loss, the same rule applies across autoregressive, masked-diffusion, and discrete-flow generators, and across de novo, motif-extension, and linker-design tasks. Under a fixed budget of 3D shape alignment oracle calls on two kinase reference compounds, EW-SFT consistently outperforms the corresponding native optimizers. It further improves goal-directed optimization under a 2D similarity oracle on four held-out references and achieves comparable performance on a sample-efficiency benchmark without a trajectory-level RL formulation. These results demonstrate that EW-SFT is a unified and effective optimizer across molecular generators, design constraints, references, and oracles.

1 INTRODUCTION

EW-SFT addresses the difficulty of applying trajectory-level RL across heterogeneous molecular generators by injecting oracle reward through elite selection into native supervised fine-tuning. It transfers across architectures and constrained design tasks while improving goal-directed optimization under molecular similarity oracles.

  • Policy-gradient RL requires generation-trajectory log-probabilities whose form depends on the generator architecture and generation procedure.
  • EW-SFT selects high-scoring molecules into its training set so reward guides a supervised update using each generator’s native pretraining loss.The workflow combines oracle scoring, elite selection, and native-loss updating.
  • EW-SFT consistently improves Top-1k similarity over corresponding native optimizers under a fixed 3D similarity-oracle budget on two kinase references.
  • EW-SFT applies across autoregressive, masked-diffusion, and discrete-flow generators and across de novo, motif-extension, and linker-design tasks.
  • Elite selection, rather than continuous within-set weighting, is the effective channel through which reward reaches the update.

2 PRELIMINARIES

Molecular generators differ in architecture, representation, and native objective, while existing optimization approaches expose a tension between trajectory-free supervised learning and reward-sensitive updates. Reward-guided supervised fine-tuning motivates a shared interface based on scored completed molecules and native losses.

  • Molecular generation and optimization: Generalist molecular generators support de novo, motif-extension, and linker-design tasks through heterogeneous architectures and representations.SAFE-GPT is autoregressive, GenMol uses masked diffusion, and InVirtuoGen uses discrete flow.
  • Policy-gradient RL: Trajectory-level policy-gradient RL scales generation log-probability gradients by reward or advantage, with factorization depending on the generator’s sampling procedure.
  • Supervised fine-tuning: Supervised fine-tuning avoids trajectory log-probabilities but gives higher- and lower-scoring molecules identical contributions when trained on a fixed dataset.
  • Reward-guided supervised fine-tuning: Reward-guided supervised fine-tuning can use completed molecules, oracle scores, and each generator’s native pretraining loss without relying on the generation trajectory.The remaining design choices concern data source, molecule selection, and within-set weighting.

3 METHOD

EW-SFT combines architecture-specific genetic proposal mechanisms with a shared elite-selection and native-loss update. A rolling buffer retains the best molecules found across rounds, and the selected set drives fine-tuning under structural constraints.

  • EW-SFT adapts to each generator through its native genetic operator and pretraining loss while retaining the same reward-guided selection and weighting formulation.
  • 3.1 PROBLEM SETTING: The optimization evaluates compounds under a finite oracle-call budget across de novo, motif-extension, and linker-design constraints.Motif extension retains one fragment, whereas linker design joins two retained fragments.
  • Each round generates valid, unique molecules from the current population, scores them with the oracle, and updates the population and elite buffer.
  • Elite-weighted native-loss update: EW-SFT performs one weighted native-loss gradient step over the scored update pool, using each model’s own pretraining loss.
  • Elite selection: Elite selection maintains a rolling buffer of the K highest-scoring molecules seen so far, with a threshold set by the lowest score in that buffer.The threshold rises monotonically as better molecules are found, so samples compete against the run’s history.
  • Weighting: Under elite+binary weighting, the update reduces to the mean pretraining loss over the best K molecules found so far.

4 EXPERIMENTS

Experiments evaluate EW-SFT across molecular generators, structural constraints, references, and oracle types. EW-SFT consistently improves 3D optimization over native optimizers, while ablations identify elite selection and genetic sampling as key components and show maintained diversity.

  • 3D shape similarity-directed optimization: EW-SFT improves Top-1k 3D shape similarity over native optimizers across all evaluated model–task settings.At 20k oracle calls, the improvement is nearly flat across de novo (+0.064), motif (+0.051), and linker (+0.055) tasks.
  • 3D shape similarity-directed optimization: EW-SFT adds significantly beyond native optimization, with gains of +0.079 over the strongest Genetic-PPO baseline and +0.091 over the weakest SAFE-GPT PPO baseline.Holm-adjusted Wilcoxon tests report pSAFE-GPT = 9.8 × 10^-4, pGenMol = 4.1 × 10^-2, and pInVirtuoGen = 2.3 × 10^-5.
  • 3D shape similarity-directed optimization: EW-SFT improves 3D scores while often moving away from the reference in 2D similarity, scoring higher in 3D in all 16 cells and lower in 2D in 9 of 16.This pattern is consistent with finding chemically distinct molecules that satisfy the shape objective rather than reconstructing the reference.
  • Ablation study: Genetic sampling is necessary for improvement beyond the pretrained model: direct policy-generated samples fail to enhance it, while SFT-only behaves essentially as distillation.The genetic operator mutates or refines samples, breaking the closed loop around the generator’s prior output.
  • Ablation study: Elite selection is the effective reward channel, whereas continuous emphasis within the elite buffer is redundant and can weaken the update.Positive emphasis helps when using the full batch, but adds nothing on top of elite filtering; signed emphasis reduces average improvement by 0.073 on InVirtuoGen.
  • Ablation study: EW-SFT remains above the expected score–diversity tradeoff, with diversity plateauing near 0.65–0.68 for SAFE-GPT and GenMol and 0.77–0.79 for InVirtuoGen.The moving-average diversity stabilizes while batch maxima continue improving in selected settings, consistent with continued exploration rather than complete collapse.

5 CONCLUSION

EW-SFT is presented as a unified optimizer across molecular generators, design tasks, references, and oracle objectives. It improves native optimizers under 3D and 2D objectives and matches Genetic-PPO on PMO.

  • EW-SFT applies across autoregressive, masked-diffusion, and discrete-flow generators.
  • EW-SFT covers de novo, motif-extension, and linker-design tasks.
  • EW-SFT improves native optimizers’ Top-1k 3D shape similarity across settings.
  • EW-SFT generalizes to held-out references under a 2D oracle.
  • EW-SFT achieves comparable performance to Genetic-PPO on PMO.

ETHICS STATEMENT

The paper situates EW-SFT within molecular generative optimization and emphasizes that computational oracle improvements require oversight and downstream experimental validation. It also contrasts policy-gradient approaches with supervised updates and identifies open transfer questions.

  • EW-SFT targets computational molecular design, while generated molecules require institutional, safety, and regulatory oversight.
  • 3D-shape and 2D-similarity oracle gains do not establish binding, selectivity, synthesis feasibility, safety, or biological activity.
  • Policy-gradient molecular optimization requires architecture-specific generation-trajectory machinery, whereas supervised fine-tuning uses the model’s native loss but does not inherently incorporate reward.
  • Open questions include transfer across likelihood-bound objectives, two-pass constrained generation, and whether weighting remains necessary after elite selection.

B IMPLEMENTATION DETAILS

The implementation evaluates three pretrained generators with native and genetic optimization procedures, alongside EW-SFT, using model-specific generation interfaces and shared oracle-budget accounting. EW-SFT applies a common AdamW-based update while adapting sampling and elite-buffer details to each model.

  • The experiments use pretrained GenMol, InVirtuoGen, and SAFE-GPT checkpoints as the three generative baselines.
  • Native optimization uses SAFE-GPT PPO, InVirtuoGen Genetic-PPO, and GenMol’s fragment-remasking genetic algorithm with model-specific configurations.
  • GenMol and SAFE-GPT attach population fragments before regenerating a region, whereas InVirtuoGen recombines fragment-level molecules and samples a flow-model completion.
  • EW-SFT uses AdamW at 1×10^-5 with gradient clipping at 1.0, sampling 128 molecules for GenMol and SAFE-GPT or 100 for InVirtuoGen.
  • The 3D oracle scores molecules by their best shape-and-color overlap across up to 4 stereoisomers and 25 conformers per stereoisomer.
  • Early stopping is required for InVirtuoGen Genetic-PPO and reimplemented SAFE-GPT PPO, but not for GenMol’s GA or EW-SFT.

C.1 FRAGMENT-CONSTRAINED ANCHORS

The fragment-constrained evaluation distinguishes motif-extension and linker-design anchors and uses four held-out references alongside two kinase references. Statistical comparisons are paired by seed or oracle, with Holm-adjusted tests and a separate correlation analysis for quality–diversity claims.

  • C.1 FRAGMENT-CONSTRAINED ANCHORS: Motif-extension anchors have one attachment point, while linker-design anchors contain two fixed end fragments connected by a designed core.
  • C.1 FRAGMENT-CONSTRAINED ANCHORS: The anchor set includes two kinase references for the main experiments and four fragment-constrained references for held-out evaluation.
  • Hypothesis tests use paired differences with one-sided tests for directional optimizer comparisons, two-sided tests for PMO comparisons, and Holm adjustment within hypothesis families.
  • Table 1 uses n = 18 per model and n = 12 for SAFE-GPT, while Table 3 uses n = 24 across four references and two tasks.
  • Quality–diversity claims instead test Pearson correlations across 12 model-by-reference-by-task cells using one-sample t-tests against zero.

C.3 FROZEN GENERATION

Frozen sampling is reliable for unconstrained de novo design but degrades under fragment and linker constraints, motivating optimization comparisons against these baselines. The reported optimization views use Top-1k gains and AUC Top-10 convergence, with some native runs stopping early.

  • Frozen sampling: All four frozen models produce near-perfect validity and uniqueness (> 0.926) with high diversity (> 0.870) for de novo design.
  • Frozen sampling: Fragment constraints largely preserve validity but sharply reduce uniqueness and diversity by narrowing the feasible chemical design space.
  • Frozen sampling: On linker design, SAFE-GPT and InVirtuoGen have lower BTK validity (0.461, 0.181) than TYK2 validity (0.998, 0.256).
  • Optimization comparison: Table 8 separates native-optimizer gains from EW-SFT-to-native and EW-SFT-to-frozen improvements using Top-1k mean 3D similarity.
  • Optimization comparison: AUC Top-10 convergence is normalized by the full 20k-call budget, while SAFE-GPT and InVirtuoGen retain their native early-stopping behavior.
  • Optimization comparison: SAFE-GPT’s linker column has no native PPO curve, whereas GenMol’s native optimizer runs to the full budget.

C.6 2D NOVELTY

The 2D-novelty analysis evaluates chemical similarity to references among molecules selected by 3D similarity, while convergence figures provide context on native optimizer stopping behavior. High 3D similarity can therefore correspond to chemically distinct molecules.

  • Convergence context: InVirtuoGen’s native Genetic-PPO exhausts its novelty budget and early-stops, while EW-SFT continues producing novel candidates through the full budget.
  • 2D novelty: Table 10 measures ECFP4 Tanimoto similarity to the reference over the Top-1k molecules ranked by 3D shape similarity.
  • 2D novelty: A molecule with high 3D similarity but low 2D similarity is shape-matched while chemically distinct from the reference.
  • 2D novelty: Motif and linker columns have higher 2D similarity because the supplied anchor is already a fragment of the reference.

C.7 ABLATION SETTINGS

The ablations vary how molecules are selected and reweighted, showing that elite selection is the central reward pathway while signed emphasis is mathematically unstable. Exploration analysis associates EW-SFT’s gain with longer productive search without significant diversity loss.

  • Ablation design: Each architecture uses one molecule-level elite stage, but the elite buffer differs because GenMol and SAFE-GPT lack complete winners while InVirtuoGen ranks whole molecules.
  • Ablation design: The ablation varies the selection threshold and emphasis reweighting, with elite defined as a rolling buffer of capacity K = 64.
  • Signed emphasis: Signed emphasis makes the loss unbounded below when negative weights multiply non-negative, unbounded native losses.
  • Signed emphasis: The signed per-round training loss crosses zero and keeps descending, unlike other settings that remain bounded below.
  • Exploration and exploitation: EW-SFT’s optimization gain is associated with longer productive search without a significant reduction in diversity.
  • Cost accounting: Wall-clock comparisons reflect the oracle budget actually consumed because native optimizers often terminate before exhausting it.
Loading 2609.00189v1…