AuditRxn: Auditable Fast-Weight Memories for Multi-Step Chemical Reaction Prefix-Prediction

Multi-step organic synthesis planning requires tracking dynamic molecular transformations, protecting groups, and reaction conditions across sequential chemical stages. Existing chemical language models treat multi-step reaction sequences as flattened token streams within standard autoregressive Transformers. Under this paradigm, context lengths grow quadratically, low-information solvents and auxiliary reagents dilute attention away from reactive transformation centers, and hidden state transitions remain uninterpretable black boxes that hallucinate physically impossible steps violating valence rules and stoichiometric conservation. In this work, we propose AuditRxn, an auditable fast-weight memory framework for multi-step chemical reaction prefix-prediction and synthetic pathway planning. AuditRxn introduces a Chemical Fast-Weight Memory (CFWM) that dynamically binds reacting functional groups and catalytic states into a recurrent associative matrix state via an online Hebbian update rule, achieving 𝑂 (1) recurrent step complexity. To guarantee physical validity, a Constraint- Audited Recurrent Alignment (CARA) module inspects memory readouts against hard chemical constraints—including atomic valence, stoichiometric mass balance, and solvent-reagent hazard compatibility—projecting unconstrained latent trajectories onto an admissible reaction manifold while issuing step-wise audit certificates. Furthermore, Counterfactual Prefix-Aware Calibration (CPAC) trains the memory dynamics against perturbed reaction prefixes to suppress error cascades in autonomous multi-step rollouts. Comprehensive evaluations across four benchmark datasets demonstrate the superiority of AuditRxn: on USPTO-50K retrosynthesis, AuditRxn attains a Top-1 accuracy of 51.82% (+1.87% over competitive baselines) with 99.92% chemical validity; on SMolInstruct forward synthesis, it reaches 66.8% Top-1 exact match; on USPTO- Condition, it achieves 0.2985 overall condition accuracy; and on PaRoutes multi-step route planning, it improves Solve@10 route completion to 75.4% while reducing 10-step trajectory inference latency by 62.3%.