Skip to content
Levitating Labs
Go back

Physics-Guided Lattice Energy Prediction with Target-Quantile Scaffold Cross-Validation

15 min read

Molecular machine learning models often boast high accuracy on paper, but when deployed to discover novel drug candidates or materials, performance drops sharply. Why? Standard evaluation relies on random train/test splits, where the model accidentally memorizes chemical core frameworks it has already seen during training.

In this work, we benchmark LatticePhysicsTower—a physics-guided neural network combining 2,243-dimensional multimodal descriptors (2D topology + 3D Boltzmann-weighted conformer physics) with thermodynamic scaling on 23,496 organic compounds from the Bradley Open Melting Point Dataset. We systematically evaluate standard random splits against out-of-distribution (OOD) Bemis-Murcko scaffold splits, quantify the real-world generalization penalty, and investigate why Retrieval-Augmented Generation (RAG) memory banks fail under scaffold shift.

Executive Summary (TL;DR)
  • The 14.1% OOD Gap: Standard random splits yield an inflated 0.7983 R². Holding out unseen molecular scaffolds drops performance to 0.6570 R² (ΔR2=−0.1413\Delta R^2 = -0.1413).
  • Method Fix: Target-quantile scaffold stratification combined with Normalized MAE (nMAE=MAE/σy\text{nMAE} = \text{MAE}/\sigma_y) stabilizes 5-fold cross-validation at 0.4729 ± 0.0136 nMAE.
  • Why Naive RAG Failed: 2D fingerprint retrieval returns structural neighbors with wildly different 3D crystal packing energies, introducing noise that degrades performance below the standalone physics baseline.

1. Core Hypotheses & Key Findings

Hypothesis 1: Random Split Performance Inflation

Question: Do standard random cross-validation splits overestimate model performance for molecular lattice energy and melting point prediction?

Outcome (Confirmed): Yes. Random splits yield an inflated 0.7983 R² (dH MAE = 1.723 kJ/mol, Tm MAE = 30.50 K). When evaluated on a Bemis-Murcko scaffold-disjoint test set, performance drops to 0.6570 R²—exposing a 14.1% structural generalization gap (ΔR² = −0.1413).

Hypothesis 2: Stratified Scaffold Shift & Scale-Invariant Metric Rescue

Question: Can target-quantile scaffold stratification stabilize 5-fold cross-validation, and how should we evaluate error across chemical families with varying energy distributions?

Outcome (Confirmed): Yes. Target-quantile scaffold bin packing stabilizes 5-fold cross-validation at 0.6215 ± 0.0253 R². Introducing Normalized MAE (nMAE = MAE / σy,fold) resolves data geometry artifacts, yielding a scale-invariant error of 0.4729 ± 0.0136 across all scaffold folds.

Why it worked: Scaffold groups differ in baseline energy variance (σy). Normalizing MAE by fold standard deviation removes scale distortion, proving that relative model precision remains stable across diverse chemical scaffold topologies.


2. Problem Statement & Mathematical Formulation

Standard molecular property prediction models yield inflated performance under random cross-validation splits—where our baseline achieves 0.7983 R² and a melting point MAE of 30.50 K. However, when evaluated on out-of-distribution (OOD) test sets where core molecular frameworks are completely unseen (Scaffold(xi)∉Vtrain\text{Scaffold}(x_i) \notin \mathcal{V}_{\text{train}}), performance degrades to 0.6570 R², increasing normalized error from 0.3315 to 0.4518.

Thermodynamic Proxy & Physics Mapping

Following empirical crystal thermo-kinetics (Walden’s and Carnelley’s rules), we establish the thermodynamic relationship between melting point TmT_m (in Kelvin) and solid-state lattice enthalpy proxy ΔHlat\Delta H_{\text{lat}} (in kJ/mol):

ΔHlat=α⋅Tmwhere α=0.0565 kJ/(mol⋅K)\Delta H_{\text{lat}} = \alpha \cdot T_m \quad \text{where } \alpha = 0.0565 \text{ kJ/(mol}\cdot\text{K)}

Target-Quantile Stratified Bemis-Murcko Split Formulation

Let D={(xi,yi)}i=1N\mathcal{D} = \{(x_i, y_i)\}_{i=1}^N represent the dataset of 23,496 organic molecules. We partition compounds by core Bemis-Murcko ring framework Scaffold(xi)\text{Scaffold}(x_i) into 5 disjoint folds using target-quantile greedy bin packing:

Scaffold(Dtest,k)∩Scaffold(Dtrain,k)=∅,∀k∈{1,…,5}\text{Scaffold}(\mathcal{D}_{\text{test}, k}) \cap \text{Scaffold}(\mathcal{D}_{\text{train}, k}) = \emptyset, \quad \forall k \in \{1, \dots, 5\}

We evaluate scale-invariant error using Normalized MAE (nMAE):

nMAEk=MAEkσy,k+ϵ=1∣Dk∣∑i∈Dk∣yi−y^i∣1∣Dk∣∑i∈Dk(yi−yˉk)2+10−8\text{nMAE}_k = \frac{\text{MAE}_k}{\sigma_{y,k} + \epsilon} = \frac{\frac{1}{|\mathcal{D}_k|} \sum_{i \in \mathcal{D}_k} |y_i - \hat{y}_i|}{\sqrt{\frac{1}{|\mathcal{D}_k|} \sum_{i \in \mathcal{D}_k} (y_i - \bar{y}_k)^2} + 10^{-8}}

3. The Setup & Model Architecture

Preprocessing & Dataset Details

Built from 23,496 valid, standardized compounds in the Bradley Open Melting Point Dataset.

  • Input Feature Dimension: 2,243 multimodal features
  • Target Enthalpy Mean: 21.46 kJ/mol21.46\text{ kJ/mol} (379.8 K379.8\text{ K})
  • Loss Function: Smooth Huber Loss (δ=5.0\delta = 5.0)
  • Optimizer & Scheduler: AdamW (lr=10−3\text{lr}=10^{-3}, wd=10−3\text{wd}=10^{-3}), Cosine Annealing, EMA (β=0.999\beta=0.999)
Residual Target Mean Initialization Rationale

The network explicitly incorporates a parameter-initialized target global mean (μy=21.46 kJ/mol\mu_y = 21.46\text{ kJ/mol}), transforming regression into residual offset learning (y^=μy+Δy\hat{y} = \mu_y + \Delta y). This prevents Epoch-1 gradient spikes and enables instant convergence by allowing early epochs to focus directly on learning complex chemical feature interactions rather than fitting baseline energy scales.

Evaluation Protocols
  1. RANDOM SPLIT: Standard 70/15/15 (train/val/test) uniform random partition (In-Distribution baseline).
  2. BEMIS-MURCKO SCAFFOLD SPLIT: 70/15/15 (train/val/test) core framework disjoint split (Single-fold OOD check).
  3. 5-FOLD STRATIFIED SCAFFOLD CV: 5-fold cross-validation with target-quantile stratification across scaffold clusters.

Multimodal Feature Extractor Pipeline (2,243 Dimensions)

To capture both 2D chemical topology and 3D solid-state crystal energetics, the feature engine converts SMILES strings into a unified 2,243-dimensional L2-normalized vector composed of 6 distinct feature categories:

  1. 2,048-D Morgan Fingerprints (2048-D): 2D circular topological bit-vectors with radius r=2r=2 capturing local atomic environments.
  2. 167-D MACCS Structural Keys (167-D): Pre-defined 167-bit dictionary of chemical functional groups and substructural keys.
  3. 12-D Scaled Physicochemical Descriptors (12-D):
    • Molecular Weight (MolWt/500\text{MolWt} / 500)
    • Topological Polar Surface Area (TPSA/200\text{TPSA} / 200)
    • Rotatable Bond Count (NumRotatableBonds/10\text{NumRotatableBonds} / 10)
    • Hydrogen Bond Donors (NumHDonors/10\text{NumHDonors} / 10)
    • Hydrogen Bond Acceptors (NumHAcceptors/10\text{NumHAcceptors} / 10)
    • Formal Charge (FormalCharge/4\text{FormalCharge} / 4)
    • Molar Refractivity (MolMR/150\text{MolMR} / 150)
    • Ring Count (RingCount/5\text{RingCount} / 5)
    • Fraction of sp3sp^3 Carbons (FractionCSP3\text{FractionCSP3})
    • Labute Accessible Surface Area (LabuteASA/300\text{LabuteASA} / 300)
    • Heavy Atom Ratio (NumHeavyAtoms/50\text{NumHeavyAtoms} / 50)
    • Heteroatom Count (Non-C/Non-H atoms / 10)
  4. 4-D Topological Symmetry & Geometry Features (4-D):
    • symmetry_ratio: Graph automorphism rank over total atom count (CanonicalRankAtoms/NumAtoms\text{CanonicalRankAtoms} / \text{NumAtoms}) to measure graph symmetry.
    • aromatic_atom_ratio: Fraction of heavy atoms that are aromatic (AromaticAtoms/NumHeavyAtoms\text{AromaticAtoms} / \text{NumHeavyAtoms}).
    • aromatic_ring_ratio: Fraction of total rings that are aromatic (AromaticRings/TotalRings\text{AromaticRings} / \text{TotalRings}).
    • planarity_score: Fraction of heavy atoms with sp2sp^2 or spsp hybridization to capture π\pi-stacking propensity.
  5. 4-D Crystal Packing Motif Descriptors (4-D): Pre-compiled SMARTS pattern matching for strong directional hydrogen-bonding functional groups:
    • Amides: [NX3][CX3](=[OX1])
    • Carboxylic Acids: [CX3](=O)[OX2H1]
    • Sulfonamides: [SX4](=[OX1])(=[OX1])-[NX3]
    • Phenol −OH-\text{OH} Groups: [OX2H]
  6. 8-D Boltzmann-Weighted 3D Conformer Physics (8-D):
    • Generates 3D conformers using RDKit ETKDGv3 distance geometry (up to 5 conformers).
    • Minimizes conformer energies via MMFF94 / UFF force fields.
    • Computes Boltzmann thermal weighting (wi∝e−Ei/RTw_i \propto e^{-E_i / RT}) at 298.15 K298.15\text{ K} (RT=2.479 kJ/molRT = 2.479\text{ kJ/mol}).
    • Extracts 8 Boltzmann-averaged 3D descriptors: Principal Moments of Inertia (PMI1,PMI2,PMI3/1000\text{PMI1}, \text{PMI2}, \text{PMI3} / 1000), Asphericity, Eccentricity, Spherocity, Radius of Gyration, and Inertial Shape Factor.

Final Vector Normalization: All 6 feature groups are concatenated (2048+167+12+4+4+8=2,2432048 + 167 + 12 + 4 + 4 + 8 = 2,243) and passed through L2 normalization (xnorm=x/(∥x∥2+10−8)\mathbf{x}_{\text{norm}} = \mathbf{x} / (\|\mathbf{x}\|_2 + 10^{-8})).

Model Architecture & Subsystem Pipelines

Figure 1: Overall System Pipeline & Multimodal Feature Engine Figure 1: Overall system pipeline featurizing 2,243-D multimodal molecular vectors for lattice enthalpy and melting point prediction.

Figure 2: LatticePhysicsTower Neural Network Architecture Zoom-In

Figure 2: Zoom-in view of the LatticePhysicsTower MLP architecture (~148k parameters) with global target mean initialization (μy=21.46 kJ/mol\mu_y = 21.46\text{ kJ/mol}).


4. Training Strategy: From Baseline Calibration to Weight Smoothing

Understanding how LatticePhysicsTower optimizes across training (up to 30–40 epochs with early stopping, patience=10) provides key insight into model stability and structural convergence:

Phase 1: Early Epochs (Epochs 1–5) — Baseline Anchoring & Coarse Feature Mapping

Phase 2: Later Epochs (Epochs 6–30+) — High-Order Fine-Tuning & Weight Smoothing


5. Benchmark Performance Results

The performance below summarizes the Stratified 5-Fold Cross-Validation and Random vs. Scaffold Split Comparison obtained from the experiment execution.

Summary Table

Experiment / Evaluation SchemeR² ScoredH MAE (kJ/mol)nMAEdH RMSE (kJ/mol)Tm MAE (K)
5-Fold Stratified Scaffold CV (Mean ± Std)0.6215 ± 0.02532.212 ± 0.4230.4729 ± 0.01362.874 ± 0.52939.15 ± 7.49 K
Random Split (In-Distribution)0.79831.7230.33152.33430.50 K
Bemis-Murcko Scaffold Split (OOD)0.65701.7930.45182.32431.73 K

Random vs. Scaffold Generalization Gap

Splitting SchemeR² ScoredH MAE (kJ/mol)nMAETm MAE (K)
Random Split (In-Distribution)0.79831.7230.331530.50 K
Bemis-Murcko Scaffold Split (OOD)0.65701.7930.451831.73 K
Scaffold Shift Generalization Drop (Δ\Delta)-0.1413+0.070+0.1203+1.23 K

6. Benchmark Figures

Figure 1: Scaffold Shift Generalization Drop Banner Figure 1: Generalization penalty under Bemis-Murcko Scaffold Shift comparing R² score and Normalized MAE.

Figure 2: Random vs. Scaffold Parity & Residual Scatter Plots Figure 2: Parity comparison between In-Distribution Random Split and OOD Scaffold Split predictions.

Figure 3: Stratified 5-Fold Scaffold CV Metrics Summary Figure 3: Summary of R², MAE, nMAE, and RMSE metrics across 5 Stratified Bemis-Murcko Scaffold Folds.

Figure 4: Scale-Invariant nMAE Stability Profile Figure 4: Demonstration of Normalized MAE (nMAE) stability across heterogeneous target standard deviations.


7. Findings & Analysis

  1. Quantifying the 14.1% OOD Generalization Penalty: Random train/test splits yield an inflated 0.7983 R², whereas holding out core scaffolds drops performance to 0.6570 R². Over 14% of random split accuracy stems from memorizing core scaffold topologies rather than learning transferrable physical principles.
  2. Scale-Invariant Metric Stabilization via nMAE: Raw MAE varies from 1.78 kJ/mol1.78\text{ kJ/mol} to 2.78 kJ/mol2.78\text{ kJ/mol} across folds due to baseline target variance differences (σy\sigma_y). Evaluating via Normalized MAE (nMAE=0.4729±0.0136\text{nMAE} = 0.4729 \pm 0.0136) confirms that relative model error is stable across chemical families.
  3. Capacity Regularization Prevents Scaffold Overfitting: Bottleneck architectures (2243→64→64→12243 \to 64 \to 64 \to 1) provide optimal structural regularization, preventing wide networks from memorizing complex fragment combinations tied strictly to training scaffolds.

8. Lab Notebook: Architectural & RAG Memory Gating Experiments

Our original objective was to develop a Retrieval-Augmented Generation (RAG) memory-guided model for lattice energy prediction, leveraging non-parametric lookup over known molecular structures. However, across multiple architectural iterations, the retrieval gating mechanisms repeatedly failed or completely suppressed the memory signals when encountering novel scaffolds. This caused the model to fall back to (or be degraded by) the underlying neural backbone—resulting in virtually no performance gain between the memory-guided RAG model and the standalone non-guided physics model.

Below is an honest record of our architectural, data-leakage, and RAG memory gating experiments, documenting how these retrieval failures led us to focus on standardizing the standalone physics baseline.

Part A: Model Architecture & Evaluation Failures

Experiment A1 — Over-Parameterized Deep Funnel MLP (>2M Parameters)

What we tried: A wide deep funnel MLP (2243 → 512 → 256 → 128 → 64 → 1) with over 2 million trainable parameters, expecting high capacity to capture complex non-linear interactions.

Why it failed: Under random splits, the model reached high R² scores. However, on scaffold-disjoint test sets, performance collapsed due to severe scaffold memorization. With ~2M parameters and 23,496 training samples, the network memorized fragment combinations instead of physical principles.

Lesson: Model capacity must be calibrated against the count of independent scaffold families. A 148k-parameter bottleneck tower (2243 → 64 → 64 → 1) proved optimal.

Experiment A2 — Naive Scaffold CV Without Target Stratification

What we tried: Unstratified Bemis-Murcko scaffold cross-validation, assigning scaffold groups to folds arbitrarily.

Why it failed: Target energy distributions were wildly imbalanced across folds (e.g., ionic-character aromatic folds vs. aliphatic chain folds), causing severe cross-fold R² variance (±0.08+).

Lesson: Target-quantile greedy bin packing is mandatory to ensure every fold spans the complete thermodynamic range of the dataset.

Experiment A3 — Evaluation with Raw MAE Only

What we tried: Using raw MAE (kJ/mol) as the primary cross-fold performance indicator.

Why it failed: Scaffold families possess different baseline target variances (σy). Raw MAE penalizes high-variance folds unnaturally, conflating data geometry with prediction accuracy.

Lesson: Normalized MAE (nMAE = MAE / σy,fold) provides a scale-invariant metric necessary for benchmarking heterogeneous chemical series.

Part B: Retrieval-Augmented Generation (RAG) & Memory Gating Exploration

We benchmarked non-parametric memory retrieval (RAG) over molecular feature banks to evaluate if retrieving structural neighbors could improve OOD predictions. Across our experiments, gating either completely ignored/shut off retrieval or retrieved noisy neighbors, leaving the final accuracy virtually indistinguishable from (or worse than) the non-guided model:

Scientific Rationale: Why Naive RAG Fails in Crystal Thermochemistry
  1. Topological Fragment vs. 3D Lattice Packing Disconnect: Standard molecular retrieval relies on 2D fragment overlap (e.g., 2048-bit Morgan Fingerprints). However, lattice energy (ΔHlat\Delta H_{\text{lat}}) is determined by long-range 3D crystal packing forces (intermolecular hydrogen bonding networks, π\pi-π\pi stacking, and steric close-packing). Minor 2D structural substitutions (such as adding a methyl or halogen group) frequently cause complete crystal space-group rearrangement—drastically altering melting point (TmT_m) and lattice enthalpy despite >85%>85\% 2D fingerprint similarity.
  2. Scaffold-Disjoint Cosine Space Collapse: Under Bemis-Murcko scaffold shift, query molecules possess unseen core frameworks absent from the memory bank. Cosine similarities between query features and training memory entries collapse into an uninformative narrow band (0.65−0.720.65 - 0.72), preventing similarity-based gates from distinguishing true thermodynamic analogues from irrelevant structural matches.
  3. Joint Gradient Interference: In single-stage joint training (Y=αyrag+(1−α)yphysicsY = \alpha y_{\text{rag}} + (1-\alpha) y_{\text{physics}}), backpropagated loss gradients force the parametric physics backbone to distort its parameter weights to compensate for noisy non-parametric retrieval priors (yragy_{\text{rag}}), degrading the model’s underlying physics learning.
Experiment B1 — Static FAISS Index & Data Leakage Discovery

What we tried: Built an initial static FAISS memory index over 2,243-D multimodal vectors with a heuristic memory trust gate g.

Outcome: Reported an initial R² ≈ 0.8081 (vs. 0.7961 standalone). However, audit revealed test-set data leakage: test molecules resided in the un-isolated index. Enforcing strict Bemis-Murcko scaffold isolation dropped true test performance to R² ≈ 0.65.

Lesson: RAG memory banks must be strictly isolated to training scaffolds to prevent subtle data leakage in molecular retrieval benchmarks.

Experiment B2 — Additive Gate Boosting & Engine Collapse

What we tried: Replaced FAISS with a trainable PyTorch GPU retriever and boosted memory trust (g = glearned + 0.85 · gmatch).

Outcome: Memory trust locked at g = 1.0 (100% memory override). This starved the physics engine of target gradients, causing internal engine performance to collapse (R² = 0.2162). Overall model score dropped to R² = 0.4922 on novel test scaffolds.

Lesson: Over-privileging retrieval gates starves the parametric backbone during joint training, destroying fallback capability on novel scaffolds.

Experiment B3 — Soft MoE Mixture (α) & Naive Structural Noise Bottleneck

What we tried: Implemented a Soft Mixture-of-Experts gate (Yfinal = α · yrag + (1−α) · yphysics) with binary entropy regularization −H(α) and a tri-partite multi-task loss.

Outcome: Entropy regularization stabilized blending (α ≈ 0.43). However, RAG model (R² = 0.6078) performed worse than standalone physics (R² = 0.6560). Naive 2D structural similarity retrieves neighbors whose thermodynamic properties differ wildly across scaffold boundaries.

Lesson: Structural similarity (e.g. Morgan fingerprints) does not equal property similarity. Uncalibrated structural retrieval acts as additive noise in lattice energy estimation.

Solution — 2-Stage Property-Guided Contrastive RAG Pipeline

The Fix: Decoupled metric learning from downstream regression into a modular 2-stage pipeline:

  1. Stage 1 (Supervised Contrastive Metric Learning): Pre-train MetricProjectionEncoder using Continuous Supervised Contrastive Loss (LSupCon-Reg) to align metric cosine distance with continuous target property Gaussian distance (wij = exp(−|yi − yj|² / 2σy²)).
  2. Stage 2 (Thermodynamic Prototype Attention): Extract K = 64 cluster prototypes (μc, Ȳc) from property-projected memory space to query dual k-NN property neighbors.

Key Takeaway: Decoupling contrastive property-metric pre-training eliminates gradient interference and guarantees retrieved memory neighbors are thermodynamically aligned, turning RAG into a reliable prior.


9. Limitations

  1. Single-Component Neutral Solid Limit: The dataset contains single-component organic compounds. Co-crystals, solvates, and ionic salts exhibit non-covalent lattice energies not covered by standard molecular descriptors.
  2. Crystal Polymorphism: Experimental melting points correspond to the thermodynamic room-pressure polymorph; alternative polymorphic forms are not explicitly differentiated in 2D SMILES representations.

10. Conclusions & Takeaways

  1. Random CV severely underestimates real-world error: Evaluating on scaffold-disjoint splits is essential for realistic deployment in molecular design.
  2. Target-Quantile Stratification is mandatory: Stratifying scaffold bin packing prevents variance imbalance across cross-validation folds.
  3. Use Normalized MAE (nMAE) for scale invariance: nMAE removes data geometry distortion when benchmarking across diverse chemical series.

11. Reproducibility & Code

ResourceDescriptionReference Link
Executable Kaggle NotebookFull executable PyTorch pipeline for feature extraction, model training, and 5-fold scaffold CVKaggle Notebook
Melting Point DatasetStandardized dataset of 23,496 organic compounds with experimental melting pointsKaggle Dataset


Previous Article
From 29% to 96.7% ROC-AUC: How ESM-2 Language Model Fixes Topological Memorization in Drug-Target Graph Networks
Next Article
Calibrated Uncertainty for DFT HOMO-LUMO Gap Prediction, and What Active Learning Did (and Didn't) Add

Comments & Discussions