RxDiff: Discrete Diffusion for Medication Recommendation with Inference-Time Safety Control
Abstract
Medication recommendation requires generating drug combinations that are therapeutically effective while minimizing harmful drug-drug interactions (DDIs). Both objectives are rooted in the combinatorial nature of prescriptions. Existing discriminative methods predict drugs independently, neglecting inter-drug dependencies; autoregressive methods introduce sequential dependencies but impose arbitrary generation orders and accumulate errors. Both paradigms confine DDI mitigation to training-time penalties, offering limited capacity to regulate interactions during inference for individual prescriptions. We reformulate medication recommendation as a discrete diffusion process and propose RxDiff, which leverages iterative denoising for bidirectional and revisable generation and provides a natural interface for dynamic DDI regulation at each step. A clinically grounded dual-guided forward process and an inference-time interaction-aware modulation mechanism are designed to jointly enable RxDiff to achieve consistent accuracy improvements, lower and more individually controlled DDI rates, and robustness to evolving DDI knowledge across MIMIC-III, MIMIC-IV, and a real-world inpatient dataset. Code is available at https://anonymous.4open.science/r/RxDiff-9157.