Deep Scientific Reasoning under Physical Constraints:
Structure-Aware Spectrum Prediction

DeepSciReasoner

A crystal \(c\) is a graph of atoms with interatomic edges (each atom keeps at most \(n_{\max}\) neighbors), and the target is a spectrum \(y \in \mathbb{R}^{l_y}\) on a fixed grid of \(l_y\) spectral bins. The formulation is generic: the only target-specific choices are the spectral axis, its discretization, and optional conditioning (for XANES, the absorbing element). We describe it with eDOS as the running example.

Stage I: structure-aware spectrum decoding

Edge-aware crystal graph Transformer. Each layer attends over an atom's neighbors with keys and values modulated by the edge embedding \(g_{ij}\) of the interatomic distance, followed by a gated residual connection:

\[ \begin{gathered} \alpha_{ij} = \operatorname{softmax}_{j \in \mathcal{N}(i)}\!\Big(\frac{q_i^\top (k_j + g_{ij})}{\sqrt{d}}\Big), \\ \hat h_i = \textstyle\sum_{j \in \mathcal{N}(i)} \alpha_{ij}\, (v_j + g_{ij}). \end{gathered} \]

Spectrum decoding with source attention. Instead of decoding a single pooled embedding with an MLP, a Transformer decoder holds one query per spectral bin, initialized from the bin's index embedding and the global crystal embedding. Self-attention across bins models correlations along the spectral axis, and source attention from each bin to the atom embeddings lets every bin aggregate the atoms that contribute at that energy, mirroring the decomposition of the eDOS into atom-projected contributions. A Softplus head keeps the output nonnegative.

Stage I is trained with a shape-and-mass objective. With \(\tilde y\) and \(\tilde{\hat y}\) the sum-normalized target and prediction, a KL divergence supervises where spectral mass lies, and a second term supervises the predicted total \(\hat Z = \sum_i \hat y_i\):

\[ \mathcal{L} = D_{\mathrm{KL}}\big(\tilde y \,\|\, \tilde{\hat y}\big) + \lambda \,\big\|\hat Z - Z\big\|_2^2 . \]

For eDOS and phDOS the total (the number of states or modes) has an analytic form and can be supplied directly; for XANES and Raman it is learned through this term.

Stage II: mass-conserving iterative refinement

The Stage I prediction can still blur band edges, leave noise in the gap, or misplace mass. Stage II corrects it directly in probability space. The normalized prediction \(p_0\) is evolved toward the target \(p_T\) by pure transport along the spectral axis, governed by the continuity equation with a learned velocity field \(f_\theta\):

\[ \begin{gathered} \partial_t p(x,t) + \partial_x \big(p(x,t)\, f_\theta(x,t)\big) = 0, \\ \frac{d}{dt}\int_a^b p(x,t)\,dx = -\big[\,p f_\theta\,\big]_a^b = 0 . \end{gathered} \]

With no flux through the boundaries, mass can move but can never be created or destroyed. As a physically principled reference flow, we use the 2-Wasserstein displacement interpolation between \(p_0\) and \(p_T\). In one dimension it moves every quantile \(s\) on a straight line at constant speed:

\[ \begin{gathered} x_t(s) = (1-t)\,F_0^{-1}(s) + t\,F_T^{-1}(s), \\ f_t\big(x_t(s)\big) = F_T^{-1}(s) - F_0^{-1}(s), \end{gathered} \]

where \(F_0\) and \(F_T\) are the CDFs of \(p_0\) and \(p_T\). This pair exactly satisfies the continuity equation.

Inference pipeline (Algorithm 1 in the paper)

function INFER(G):
    {h_i}, h_x ← Enc(G);  ŷ₀ ← Dec({h_i}, h_x)       # Stage I: structure-aware decoding
    p ← ŷ₀ / Σ_i ŷ₀,i;  F_i ← Σ_{m≤i} p_m Δx         # normalize; initial CDF
    for k = 0, …, N_t − 2:                           # Stage II: mass-conserving refinement
        v_i ← f_θ(x_i, kΔt; h_x)                     # velocity on the spectral grid
        F_i ← F_i − Δt · v_i · D_x^up F_i            # upwind CDF update
        p_i ← (F_i − F_{i−1}) / Δx;  renormalize p   # recover the density
    return p                                         # optionally rescaled by the predicted total Ẑ

\(D_x^{\mathrm{up}}\) is the backward difference where \(v_i \ge 0\) and the forward difference where \(v_i < 0\), and a CFL-like condition \(\max_i |v_i|\,\Delta t / \Delta x \le 1\) keeps the CDF monotone. Derivations and proofs of the end-to-end validity guarantee are in the paper's appendix.

Code and Data

The repository contains Stage I and Stage II training, test-set evaluation, and launch scripts for every benchmark in the paper. The six benchmarks are released on Hugging Face as prepared crystal graphs with spectrum targets and fixed 80/10/10 splits.

Samples are materials, except for xanes_full, where they are (material, absorbing element) pairs. Curation details are in the paper's appendix and on the dataset card, which also lists the licenses of the source data and the works to cite.

git clone https://github.com/yingheng-wang/DeepSciReasoner && cd DeepSciReasoner
CONDA_CHANNEL_PRIORITY=flexible conda env create -f environment.yml && conda activate dsr

bash scripts/download_data.sh edos_caltech          # no argument: all six benchmarks
for seed in 1 4; do
  bash scripts/train_stage1.sh edos_caltech $seed   # Stage I: structure-aware decoding
  bash scripts/train_stage2.sh edos_caltech $seed   # Stage II: refinement on the frozen Stage I model
done
bash scripts/test.sh edos_caltech 1 4               # test-set metrics of the two-seed ensemble