ReMedI: Inference-Time Refinement and Self-Distillation for Controllable Molecular Optimization

1 INTRODUCTION

• The central challenge is improving multiple properties while preserving satisfactory properties and lead-relative similarity, which current CLMs do not reliably achieve without scaffold drift.Task-specific RL adapters and editing data can be costly and limit transfer to new objective combinations.Source: §1 · p. 1
• ReMedI uses a pretrained CLM to propose local edits and Metropolis–Hastings refinement to guide search under an objective balancing property improvement, preservation, and lead similarity.The priorities can be changed at inference time without retraining the model.Source: §1 · p. 2
• A shared proposal adapter trained on 500 lead–task instances achieves performance comparable to C-MORAL’s five task-specific policies trained on 3,883 prompts.Search traces include successful endpoints and intermediate states, avoiding a separate specialized editing dataset.Source: §1 · p. 2
• Search-discovered endpoints, intermediate states, and improving transitions become supervised pairs that are distilled into the base CLM to improve zero-shot generation and later refinement.The refinement trajectories are conditioned on the lead and search state.Source: §1 · p. 2, Figure 1 · p. 2
• ReMedI achieves approximately 2.8× the strict-success coverage of C-MORAL on matched GeLLMO-Mistral leads, while MH raises frozen-model SOR by 66–143% on C-MuMOInstruct and 12–50% on SMDD.The C-MuMOInstruct gains span three backbones and IND/OOD splits.Source: §1 · p. 2

2 RELATED WORKS

• Post-training methods, including C-MORAL and RePO, adapt molecular models through specialized training, whereas ReMedI keeps the CLM fixed during refinement and sets the improvement–similarity trade-off at inference time.The related work also situates these methods among supervised fine-tuning, preference optimization, and reinforcement learning approaches.Source: §2 · p. 3
• ReMedI uses a CLM as a state-dependent proposal in a Metropolis–Hastings loop, then trains a shared proposal adapter on successful endpoints and intermediate states to test transfer to unseen objective combinations.This combines inference-time refinement with proposal improvement.Source: §2 · p. 3

3 PRELIMINARY

• Metropolis–Hastings proposes a candidate from the current state and accepts it probabilistically; the Hastings correction accounts for asymmetric state-dependent proposals, while the target’s normalizing constant cancels in the ratio.The proposal controls exploration and the target determines the distribution being sampled.Source: §3 · p. 3
• The Nash bargaining solution maximizes the product of objective gains over a reference point, favoring balanced gains rather than a weighted-sum objective.If any objective remains at its reference value, the product is zero regardless of gains on other objectives.Source: §3 · pp. 3–4
• Focused refinement iterates over chains from the incumbent, ranks candidates, updates the incumbent, and shifts focus to the least-satisfied unmet requirement or to similarity when none remain.The procedure returns both the incumbent and a transition log.Source: §3 · p. 4

• The improvement score multiplies per-property satisfactions, so gains on one objective cannot fully mask poor satisfaction on another.Penalties activate when guarded properties decline beyond their allowed margins or lead similarity falls below τs.Source: §4 · p. 4
• Metropolis–Hastings refinement accepts or rejects proposed edits using the step-specific target and proposal probabilities in both directions.Target ratios require oracle scores for the current and proposed states.Source: §4 · p. 5
• Search traces train the proposal model on lead-conditioned endpoints, chain-conditioned endpoints, and some accepted transitions that do not meet every requirement.Seed- and chain-conditioned examples support direct generation and refinement; weighted examples favor higher-scoring targets according to βdist.Source: §4 · p. 5
• For each state, ReMedI focuses on the least-satisfied unmet improvement requirement, or on similarity when none remain.The focused task clause guides which edits the model proposes.Source: §4 · p. 5

5.1 SETUP

• Evaluation covers ten C-MuMOInstruct property-combination tasks, split into five IND and five OOD, plus 314 SMDD-Bench Task 4 instances; success metrics assess requested improvements and near-optimal joint satisfaction.For SMDD, SOR requires every requested objective to improve by its instance-specific threshold, while SSOR captures simultaneous near-optimal satisfaction.Source: §5.1 · pp. 5–6

5.2 MAIN RESULTS

• MH improves q0 SOR by ≈80% and SSOR by ≈90% on average across three model families and IND/OOD splits.After distillation, q1-MH outperforms q1-base on every backbone and split.Source: §5.2 · pp. 6–7, Table 1 · p. 6
• ReMedI maintains higher strict-success coverage than C-MORAL across evaluated lead-similarity requirements, including pooled and held-out OOD tasks.Figure 2’s SSOR–similarity frontier indicates the advantage is not simply from larger structural changes.Source: Figure 2 · p. 7, §5.2 · p. 7
• On GeLLMO-Mistral, one distillation round raises zero-shot base SOR by 33% on IND (0.354 →0.472) and 23% on OOD (0.271 →0.332).It also raises MH SOR by 20% on IND (0.590 →0.706) and 18% on OOD (0.467 →0.552).Source: Table 1 · p. 6, §5.2 · p. 7
• On 314 SMDD-Bench instances, ReMedI-MH improves SOR by 0.13–0.29 and clearance-target success on every backbone versus pooled C-MORAL, while C-MORAL retains higher similarity.The comparison uses a shared adapter per backbone for pooled C-MORAL at the same 20-candidate inference budget.Source: §5.2 · p. 7, Table 2 · p. 7
• ReMedI’s shared adapter uses trajectories from 500 lead–task instances, about 7.8× fewer than C-MORAL’s 3,883 RL training prompts.Source: §5.2 · p. 8
• For Evobrutinib, ReMedI reduces AMES (0.858 →0.050), DILI (0.826 →0.546), and PGP (0.648 →0.474) together while retaining similarity 0.738.C-MORAL improves PGP but worsens DILI and AMES, with similarity 0.707; ReMedI finds one candidate improving all three liabilities.Source: §5.2 · p. 8, Figure 3 · p. 8

Effect of search strategy

• At K = 20, MH reaches SOR 0.41 versus 0.39 for unguided editing and 0.38 for greedy feedback, while preserving similarity 0.44 versus 0.39 for unguided editing.Sequential strategies outperform independent best-of-K sampling, and MH has the highest SOR at every tested candidate budget.Source: §6 · p. 8, Figure 4 · p. 8

Choosing the distillation round

• From q0 to q1, SOR rises from approximately 0.45 to 0.63 and SSOR from 0.12 to 0.16, with similarity decreasing from 0.47 to 0.43.Later rounds provide diminishing returns: q2 improves SOR but not SSOR, while q3 raises SSOR as similarity falls to approximately 0.36; q1 is used as the primary proposal.Source: §6 · p. 8, Figure 5 · p. 9

Effect of distillation sources

• Removing successful-endpoint supervision causes the largest SSOR drop: q1-Base falls from 0.115 to 0.070, and q1-MH from 0.160 to 0.090.Without endpoints, similarity rises from 0.30 to 0.505 for q1-Base and from 0.23 to 0.414 after MH refinement. Removing intermediate states or improving transitions also lowers refined SSOR while reducing similarity; the full mixture gives the highest refined SSOR, 0.160.Source: §6 · p. 9, Figure 6 · p. 9

7 LIMITATIONS & CONCLUSION

• ReMedI uses a CLM to propose sequential edits, MH to select them under a lead-relative objective, and successful search trajectories to train the proposal.Across C-MuMOInstruct and SMDD-Bench, this improves frozen proposals and yields a stronger strict-success–similarity tradeoff; one distillation round further improves direct generation and subsequent search.Source: §7 · p. 9

AI USE STATEMENT

• Generative AI tools assisted with prose editing and routine code and analysis; the authors say they verified claims, figures, and references before submission.Source: AI USE STATEMENT · p. 10

SOR

• SOR requires every targeted property to improve by its prescribed margin and every protected property to stay within its prescribed tolerance.The metric compares output properties with the lead, using the preferred direction for each property.Source: §A.1 · p. 13

SSOR

• SSOR checks whether all relevant properties meet their prescribed near-optimal thresholds.Unlike SOR's lead-relative test, SSOR uses absolute property thresholds.Source: §A.1 · p. 13

Metric example

• An output with permeability 0.6 and unchanged toxicity passes SOR but fails SSOR, while permeability 0.85 with unchanged toxicity passes both.The example uses a lead with permeability 0.4 and toxicity 0.1, requiring at least a 0.1 permeability increase and preservation of toxicity.Source: §A.1 · p. 13

C-MuMOInstruct results

• q1-MH increases both SOR and SSOR over q1-base for every backbone and split.Distillation further strengthens proposals on the GeLLMO backbones; combining one distillation round with refinement yields the highest SOR and SSOR in nearly all settings.Source: §A.2 · p. 13

OOD generalization

• Gains extend to OOD property combinations, where trajectories distilled from IND tasks improve optimization beyond the training combinations.Source: §A.2 · p. 13

Interpretation and scope

• Similarity generally decreases as optimization becomes more aggressive, so success and similarity are interpreted together through the main-text similarity-constrained coverage analysis.RI is included only as a within-method diagnostic because its scale is not comparable to the C-MORAL implementation.Source: §A.2 · p. 13

A.3 TRAJECTORY DEPTH ABLATION

• Overall SSOR doubles from 0.070 to 0.140 when 80 independent proposals are replaced by 10 trajectories of eight steps.The comparison uses the same 80 proposals per lead in each configuration.Source: §A.3 · p. 13
• Sequential conditioning lets multiple edits compose into solutions satisfying stringent multi-property constraints, unlike independent proposals, which rarely reach the target.Because each configuration evaluates the same number of proposals, the gain is not attributable to additional sampling.Source: §A.3 · p. 13
• All trajectory-depth configurations use 80 proposals per lead and aggregate IND and OOD results on frozen GeLLMO-Mistral.Source: Table 4 · p. 14