PaperlayerView exact revision on arXiv

Source-linked AI summary

MAST: Motif-Augmented Diffusion with Search Tree for Spectroscopic Molecular Structure Elucidation

Chenghao Jia, Mengdi Liu, Hong Chang, Shiguang Shan, Xilin Chen

arXiv:2610.12067v1physics.chem-phcs.AI

TL;DR

Spectra-to-structure elucidation is difficult because distinct molecules can have similar spectral signatures, and current methods face limitations in learning fine-grained evidence and sampling efficiently. MAST combines motif-augmented joint 2D–3D diffusion with reward-guided tree search. On QM9S, it reports state-of-the-art Top-1 recovery, improved exact recovery and 3D reconstruction, and a favorable accuracy–efficiency trade-off.

  • Problem

    Similar spectral signatures can correspond to distinct structures, while global spectral conditioning, repeated full-chain sampling, and scarce direct 3D recovery remain challenges.

    Source: §1 · pp. 1–2

  • Method

    MAST combines an interpretable probabilistic spectra-to-motif prior with joint 2D–3D diffusion and reward-guided tree search over denoising trajectories.

    Source: §1 · p. 2

  • Results

    On QM9S, MAST achieves state-of-the-art Top-1 recovery and substantially improves exact recovery and 3D reconstruction, with a favorable accuracy–efficiency trade-off.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The authors recommend using MAST as decision support that returns ranked candidates, reports uncertainty, and supports verification with orthogonal evidence.

    Source: Impact Statement · p. 9

    Experiments focus on simulated spectra; applying MAST to experimental data may require calibration, domain adaptation, and careful post-generation filtering and validation.

    Source: Impact Statement · p. 9

Abstract

from arXiv · show

Elucidating molecular structures from spectra is a foundational problem in chemical and materials characterization, yet remains challenging due to spectral ambiguity and the vast molecular space. Although recent diffusion-based generators show strong promise for spectra-conditioned elucidation, existing methods struggle to learn robust spectra-structure relationships from limited paired data when relying solely on global spectral representation. Moreover, the repeated full sampling inference strategy incurs substantial computation overhead. To address these limitations, we propose \textbf{MAST}, a \textbf{M}otif-\textbf{A}ugmented diffusion framework with \textbf{S}earch \textbf{T}ree, for joint 2D-3D spectroscopic molecular structure elucidation. MAST introduces explicit, interpretable \emph{motif priors} as intermediate evidences throughout denoising, reducing conditional ambiguity and facilitating spectra-conditioned optimization. We further cast diffusion sampling as \emph{reward-guided tree search} to prioritize high-reward denoising trajectories, yielding a compact set of spectra-consistent candidates under limited budgets. On the QM9S multi-spectra benchmark, MAST achieves \textbf{94.89\%} exact recovery and improves 3D fidelity, while preserving high chemical validity and stability. Code is available at https://github.com/Jia040223/MAST.

1. Introduction

Spectral ambiguity, limited paired data, and costly full-chain sampling complicate molecular elucidation, while direct 3D recovery remains scarce. MAST combines interpretable motif evidence with tree-guided diffusion and achieves strong QM9S recovery results.

1. Introduction

  • Existing methods often rely on a compressed global spectral representation and repeated full-chain sampling, while direct 3D geometry recovery remains comparatively scarce.The text identifies limited paired spectra–structure data, computational expense, and more complex spectra–geometry relationships as relevant challenges.Source: §1 · pp. 1–2
  • MAST learns a probabilistic spectra-to-motif prior and injects interpretable, composable evidence throughout joint 2D–3D denoising to shrink the hypothesis space.The framework combines this motif-augmented diffusion with tree search.Source: §1 · p. 2
  • MCTS searches checkpointed denoising trajectories, allocating computation to high-reward intermediates and pruning low-value branches early.Segment-wise adaptive guidance tunes guidance strength to improve spectral consistency while preserving the diffusion prior.Source: §1 · p. 2
  • MAST achieves state-of-the-art Top-1 recovery on QM9S, with high chemical validity and stability; MCTS-guided search improves exact recovery under the same denoising-step budget.The introduction also reports improved spectra-consistent recovery under limited sampling budgets without retraining or fine-tuning.Source: §1 · p. 2

2. Related Work

Prior spectra-to-structure work spans library matching, interpretable cues, direct generation, and diffusion-based approaches. Related research also motivates combining diffusion guidance with test-time tree search to improve spectra-consistent generation under limited budgets.

Spectra-to-structure molecular elucidation

  • Spectra-to-structure methods have progressed from expert heuristics and library matching to cue-based ranking, reinforcement-learning graph construction, Transformer translation, language models, and direct generation.Retrieval methods depend on reference-library coverage and curation, which the text says are often insufficient in practice.Source: §2 · p. 2

Diffusion-based molecular modeling

  • Molecular diffusion work includes discrete 2D graph denoising, equivariant 3D coordinate generation, and joint 2D–3D models for coherent topology and geometry.The cited joint models are described as enabling coherent bond structures and feasible conformations.Source: §2 · pp. 2–3

Test-time scaling and guided search

  • Test-time scaling includes sampling, self-consistency, and tree search; diffusion guidance offers a complementary way to steer conditional generation.The paper combines MCTS with diffusion guidance to improve spectra-consistent generation under limited budgets.Source: §2 · p. 3

3. Preliminaries

The task is to learn a conditional generative model from multimodal spectra to joint molecular structures that include both a graph and a 3D conformation. The section also outlines diffusion’s forward noising process and score-based reverse generation.

3.1. Task Definition

  • The task is to learn pθ(M | S) from paired spectra–structure data, with M combining molecular graph attributes and 3D coordinates.S contains observed spectra from multiple modalities; the model captures the one-to-many mapping from spectra to structures.Source: §3.1 · p. 3

3.2. Diffusion Models

  • Diffusion gradually noises a clean representation toward a tractable prior, with the terminal marginal typically chosen to approximate a standard Gaussian.The forward process uses time-dependent αt and σt, with q(G1) ≈ N(0, I).Source: §3.2 · p. 3
  • Reverse-time generation follows a stochastic differential equation whose drift depends on the score, typically approximated by a neural network trained with mean squared error.The reverse dynamics include a Wiener-process term and schedule-dependent coefficients.Source: §3.2 · p. 3

4. Method

MAST conditions joint 2D–3D diffusion on both global spectral context and interpretable motif evidence, then uses reward-guided tree search to find spectra-consistent structures under limited sampling budgets.

4.1. Conditioning Evidence from Spectra

  • Multi-modal spectra produce a global embedding and a probabilistic motif prior, which is injected throughout denoising as complementary conditioning evidence.The motif vocabulary includes reusable local chemical patterns beyond functional groups.Source: §4.1 · pp. 3–4

4.2. Motif-Augmented Diffusion

  • An SE(3)-equivariant denoiser jointly refines 3D geometry and graph attributes, using both global spectral context and motif evidence at each denoising layer.Training uses a variance-preserving diffusion process and weighted reconstruction loss over geometry and topology.Source: §4.2 · pp. 4–5
  • Reverse sampling is discretized into checkpoints, forming a tree whose branches represent alternative segment-level sampling actions and trajectories.Each action deterministically advances between checkpoint states under the chosen noise trajectory and guidance configuration.Source: §4.2 · p. 5

4.3. Reward-guided Tree Search

  • MCTS searches this tree by balancing exploration and exploitation, expanding selected leaves into alternative segment-level sampling actions.This allocates computation to promising prefixes rather than repeatedly executing complete diffusion trajectories.Source: §4.3 · p. 5
  • Fast partial rollouts score candidate molecules using spectral alignment and motif consistency, with invalid candidates downweighted.The combined reward includes a motif term weighted by λs.Source: §4.3 · p. 6
  • Softmax backup propagates rollout rewards to node values while emphasizing high-reward trajectories, which the method describes as improving robustness to noisy rollouts and search efficiency.The node values and visit counts are updated after child evaluations.Source: §4.3 · p. 6

5. Experiments

Experiments on QM9S evaluate joint 2D–3D elucidation, motif-prior effects, MCTS-guided inference, and robustness to spectral perturbations. MAST improves exact recovery and 3D fidelity, while analyses examine its motif evidence and inference efficiency.

5.1. Datasets and Experimental Setup

  • QM9S supplies simulated IR, Raman, and UV/Vis spectra for joint 2D–3D prediction, evaluated with exact-match, structural-similarity, geometry, validity, and distribution metrics.Baselines include spectra-to-SMILES systems and diffusion generators; RMSD permits mirror inversion because the spectra are enantiomer-invariant.Source: §5.1 · pp. 6–7

5.2. Benchmarking Molecular Structure Elucidation

  • 94.89% Acc@1 with MCTS, alongside 0.1048 MCES and improved 3D fidelity, follows a 30.4-percentage-point Acc@1 gain over the strongest baseline without tree search.MAST also maintains near-perfect AtomStable and MolStable; valid, unique, and novel molecules reach 97.1%, versus 90.5% for DiffSpectra.Source: §5.2 · pp. 7–8

5.3. Analysis: Motif Prior

  • Removing the motif prior degrades exact recovery and structural similarities, while Figure 3 illustrates missing motifs without the prior and spectra-consistent samples with it.The motif predictor provides reliable motif-level evidence; conditioning raises overall macro-F1 from 0.9087 to 0.9430, and 65% of corrected cases coincide with motif differences across 100 of 103 motif types.Source: Figure 3 · p. 7, §5.3 · p. 8

5.4. Analysis: MCTS-guided Inference

  • At equal denoising-step budgets, tree search achieves higher exact recovery than Best-of-N and Beam Search, and its gains increase faster with budget.Around 2000 steps, MCTS reaches 94.34% Acc@1 at 85.40s per sample, versus Best-of-N at 86.13% and 79.68s; GPU memory is similar, though MCTS uses moderately more CPU memory.Source: §5.4 · p. 8, Figure 4 · p. 8
  • Reward improves with deeper search but saturates after moderate depth, while greater tree width mainly helps at shallow depths; search can return multiple high-reward, spectra-consistent structures.The illustrated alternatives include a near-ground-truth structure and a mirror-symmetric configuration.Source: §5.4 · pp. 8–9, Figure 4 · p. 8

5.5. Robustness and Sensitivity Analyses

  • MAST without MCTS degrades gracefully under Gaussian noise and baseline drift and outperforms DiffSpectra across tested combined perturbations; motif evidence is sensitive to corruption, while moderate motif-reward weighting works best.Softmax backup performs best among tested strategies, and the final gain reflects both informative motif conditioning and reward-guided search.Source: §5.5 · p. 9

6. Conclusion

The conclusion reports improved exact recovery and 3D reconstruction on QM9S, with robustness evaluations and runtime–memory profiling supporting MAST’s performance and efficiency trade-off. It identifies broader chemical spaces, additional modalities, and real experimental spectra as future evaluation targets.

6. Conclusion

  • MAST substantially improves exact recovery and 3D reconstruction on the QM9S benchmark.The framework combines motif-augmented diffusion with a search tree for 3D molecular elucidation from spectra.Source: §6 · p. 9
  • Robustness evaluations find MAST remains effective under noisy spectral perturbations, while runtime–memory profiling indicates a favorable accuracy–efficiency trade-off against repeated-sampling baselines.Source: §6 · p. 9
  • Future work is to scale MAST to broader chemical spaces, incorporate more spectral modalities, and evaluate more diverse datasets, especially real-world experimental spectra.Source: §6 · p. 9

Impact Statement

MAST is presented as a spectra-conditioned molecular elucidation framework with interpretable motif evidence and reward-guided search. The accompanying material describes its spectral encoder, motif predictor, conditioning fusion, equivariant denoiser, diffusion objective, and motif vocabulary.

Impact Statement

  • MAST’s probabilistic motif prior provides interpretable midlevel evidence that may improve transparency and human-in-the-loop review.Source: Impact Statement · p. 9

A.1. Multi-spectra Encoder

  • The multi-spectra encoder patches modality signals, concatenates their embeddings, applies a Transformer, and pools them into a global spectrum representation.Source: §A.1 · p. 13

A.2. Spectra-to-Motif Prior

  • A multi-label predictor maps spectra to motif probabilities, using a vocabulary of ring and chain motifs extracted and filtered to retain frequent patterns.The vocabulary has 103 motifs spanning local topology, rings, and functional-group-bearing fragments; the predictor uses a hierarchical Transformer and weighted binary cross-entropy.Source: §A.2 · p. 13, B. Motif Vocabulary Construction · p. 15

A. Model Details

  • The denoiser conditions joint 2D–3D molecular structure updates on multi-spectra embeddings, motif probabilities, and diffusion time.Figure 5 depicts an SE(3)-equivariant denoiser with stacked blocks updating node, edge, and coordinate features.Source: Figure 5 · p. 14, §A.3 · p. 14, §A.4 · p. 14

A.3. Conditioning Embedding

  • Two successive vector-valued gates fuse spectral and timestep embeddings, then motif probabilities, into the conditioning vector used to modulate denoiser blocks.Source: §A.3 · p. 14

A.4. Joint 2D–3D Denoiser with Conditional Equivariant Blocks

  • Stacked equivariant blocks jointly update node, edge, and coordinate streams, with relational attention using edge features and invariant geometry.Coordinate updates aggregate relative displacement vectors, followed by mean-centering to remove global translation.Source: §A.4 · pp. 14–15

A.5. Diffusion Training Objective

  • The diffusion objective samples noisy structures from a VP marginal, predicts the clean structure, and minimizes a weighted reconstruction loss.Source: §A.5 · p. 15

B. Motif Vocabulary Construction

  • The 103-motif vocabulary is built from canonicalized ring and acyclic-chain patterns and spans reusable local chemistry between atom types and full-graph templates.Motifs appearing in more than 1% of QM9S molecules are retained; representative motifs are grouped by coarse chemical semantics.Source: B. Motif Vocabulary Construction · p. 15, Table 5 · p. 16

C. Implementation Details

The implementation uses a pretrained-and-fine-tuned spectral encoder, a separately trained frozen motif predictor, and an SE(3)-equivariant denoiser. MCTS operates over checkpointed reverse-diffusion segments with bounded expansions and sampling steps.

C. Implementation Details

  • The three-modality Transformer encoder is pretrained with masked autoencoding and then fine-tuned jointly with the denoiser; the 103-motif predictor is trained separately and frozen during diffusion training.The encoder uses depth 4, width 256, 8 heads, and output dimension 128.Source: C. Implementation Details · p. 16
  • The denoiser uses a cosine VP-SDE schedule with 1000 steps and an 8-block, 16-head SE(3)-equivariant Transformer.Training uses AdamW with learning rate 2 × 10^-4 and node, edge, and coordinate loss weights of 1.0, 0.25, and 0.1.Source: C. Implementation Details · p. 16
  • MCTS searches 50 reverse-diffusion checkpoints, expands up to four children per node, and caps sampling at 3000 denoising steps per molecule.Each segment uses a 10-step rollout; invalid 2D structures receive zero reward, and invalid 3D geometry downweights reward by 0.5.Source: C. Implementation Details · p. 16

D. Evaluation Metrics

The evaluation covers molecular generation quality and spectra-conditioned structure elucidation, using validity, stability, distributional, exact-recovery, structural-similarity, and 3D-geometry measures.

D. Evaluation Metrics

  • The evaluation separates molecular generation quality from spectra-conditioned elucidation, which measures exact recovery and graded structural or geometry similarity.Metric implementations follow DiffSpectra unless otherwise stated.Source: D. Evaluation Metrics · p. 16

D.1. Molecular Generation Quality

  • Generation quality is assessed through validity, uniqueness, and novelty, atom- and molecule-level stability, and distributional fidelity.FCD measures distance between ChemNet feature distributions; SNN, Frag, and Scaf assess similarity, fragment statistics, and scaffold statistics.Source: §D.1 · pp. 16–17

D.2. Spectra-conditioned Structure Elucidation

  • Acc@K measures the fraction of test spectra with at least one candidate matching the ground-truth molecule exactly.The exact match is checked using canonical SMILES; MCES and fingerprint or functional-group similarities capture partial structural overlap, while RMSD and MapAcc assess 3D consistency.Source: §D.2 · p. 17

E. Robustness to Spectral Perturbations

The robustness tests evaluate MAST on noisy and baseline-shifted spectra without retraining. MAST outperforms DiffSpectra under combined perturbations, while motif prediction remains relatively stable.

Individual perturbations

  • Without MCTS or retraining, MAST degrades gracefully: Acc@1 falls from 79.12% on clean spectra to 75.98% at SNR = 30 dB and 74.61% at drift = 5.Structural similarity remains high under these perturbations.Source: E. Robustness to Spectral Perturbations · p. 18

Combined perturbations

  • 73.83% Acc@1 for MAST versus 9.38% for DiffSpectra at SNR = 30 dB and drift = 5; at SNR = 20 dB and drift = 5, MAST retains 74.41% Acc@1.MAST outperforms DiffSpectra across the combined perturbation settings.Source: E. Robustness to Spectral Perturbations · p. 18

Motif predictor robustness

  • The motif predictor's micro-F1 and macro-F1 scores drop only slightly under noisy and shifted spectra, supporting its relative stability as intermediate chemical evidence.The passage identifies this stability as an important reason MAST preserves spectra–structure consistency under corrupted inputs.Source: E. Robustness to Spectral Perturbations · p. 18

F. Additional Sensitivity Analyses

Sensitivity analyses examine reward backup, scorer perturbations, motif-reward weighting, and motif-prior quality in MCTS-guided diffusion.

Perturbing the reward scorer

Motif-consistency reward weight

  • A motif-consistency reward weight of 0.10 gives the best result; increasing it to 0.20 reduces exact recovery.The authors suggest that excessive motif weighting may weaken the global spectra–structure reward's role in exact recovery.Source: Table 12 · p. 19, F. Additional Sensitivity Analyses · p. 19

Reward backup strategy

  • Softmax backup achieves the best Acc@1 among softmax, max, and mean strategies, balancing emphasis on high-reward rollouts against sensitivity to noisy maxima.Mean backup can dilute strong signals from promising branches.Source: F. Additional Sensitivity Analyses · p. 19

Perturbing the motif prior

  • Corrupting the motif prior lowers exact recovery: Acc@1 falls to 51.37% at flip rate r = 0.3 and 29.49% with a random prior.Small Gaussian perturbations cause moderate degradation, while random bit flips and fully random priors cause larger drops.Source: F. Additional Sensitivity Analyses · p. 19

G. Efficiency and Resource Profiling

The profiling compares inference accuracy, runtime, and memory across denoising-step budgets. MCTS-guided diffusion achieves higher exact recovery than Best-of-N and Beam Search under comparable budgets, with moderately higher CPU RAM use.

Search strategy comparison

  • MCTS adaptively expands high-reward denoising prefixes and prunes less promising branches, achieving better accuracy under comparable runtime.Best-of-N repeats full-chain sampling without reusing partially successful prefixes, while Beam Search continues multiple candidates at each stage.Source: G. Efficiency and Resource Profiling · pp. 19–20

Accuracy–runtime–memory trade-off

  • At around 3000 denoising steps, MCTS-guided diffusion reaches 94.81% Acc@1 versus 89.45% for Best-of-N; at around 4000 steps, it reaches 95.11%.The runtime at around 3000 steps is 124.63s per sample for MCTS-guided diffusion and 118.94s for Best-of-N.Source: Table 14 · p. 20, G. Efficiency and Resource Profiling · p. 20

Resource use

  • Peak GPU memory is similar across methods, but MCTS uses moderately more CPU RAM to store tree states and search statistics.Its runtime is slightly higher at the same nominal step budget because of tree operations and reward evaluations.Source: G. Efficiency and Resource Profiling · p. 20
Loading 2610.12067v1…